[Bug fix] Fix vertex ai integration issue (#1257)

embedchain/embedchain.py

@@ -429,11 +429,10 @@ class EmbedChain(JSONSerializable):
 
         if dry_run:
             return list(documents), metadatas, ids, 0
-        
+
         # Count before, to calculate a delta in the end.
         chunks_before_addition = self.db.count()
 
-        
         # Filter out empty documents and ensure they meet the API requirements
         valid_documents = [doc for doc in documents if doc and isinstance(doc, str)]
 
@@ -441,7 +440,7 @@ class EmbedChain(JSONSerializable):
 
         # Chunk documents into batches of 2048 and handle each batch
         # helps wigth large loads of embeddings  that hit OpenAI limits
-        document_batches = [documents[i:i+2048] for i in range(0, len(documents), 2048)]
+        document_batches = [documents[i : i + 2048] for i in range(0, len(documents), 2048)]
         for batch in document_batches:
             try:
                 # Add only valid batches
@@ -452,12 +451,10 @@ class EmbedChain(JSONSerializable):
                 # Handle the error, e.g., by logging, retrying, or skipping
                 pass
 
-
         count_new_chunks = self.db.count() - chunks_before_addition
         print(f"Successfully saved {src} ({chunker.data_type}). New chunks count: {count_new_chunks}")
-        
-        return list(documents), metadatas, ids, count_new_chunks
 
+        return list(documents), metadatas, ids, count_new_chunks
 
     @staticmethod
     def _format_result(results):
@@ -493,9 +490,7 @@ class EmbedChain(JSONSerializable):
         :return: List of contents of the document that matched your query
         :rtype: list[str]
         """
-        print("Query passed in config:", config)
         query_config = config or self.llm.config
-        print("Final config:", query_config)
         if where is not None:
             where = where
         else:
@@ -506,7 +501,6 @@ class EmbedChain(JSONSerializable):
             if self.config.id is not None:
                 where.update({"app_id": self.config.id})
 
-        print('Number documents', query_config)
         contexts = self.db.query(
             input_query=input_query,
             n_results=query_config.number_documents,

embedchain/llm/base.py

@@ -5,7 +5,9 @@ from typing import Any, Optional
 from langchain.schema import BaseMessage as LCBaseMessage
 
 from embedchain.config import BaseLlmConfig
-from embedchain.config.llm.base import DEFAULT_PROMPT, DEFAULT_PROMPT_WITH_HISTORY_TEMPLATE, DOCS_SITE_PROMPT_TEMPLATE
+from embedchain.config.llm.base import (DEFAULT_PROMPT,
+                                        DEFAULT_PROMPT_WITH_HISTORY_TEMPLATE,
+                                        DOCS_SITE_PROMPT_TEMPLATE)
 from embedchain.helpers.json_serializable import JSONSerializable
 from embedchain.memory.base import ChatHistory
 from embedchain.memory.message import ChatMessage

embedchain/llm/openai.py

@@ -2,6 +2,7 @@ import json
 import os
 from typing import Any, Callable, Dict, Optional, Type, Union
 
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain.schema import BaseMessage, HumanMessage, SystemMessage
 from langchain_core.tools import BaseTool
 from langchain_openai import ChatOpenAI
@@ -41,8 +42,6 @@ class OpenAILlm(BaseLlm):
         if config.top_p:
             kwargs["model_kwargs"]["top_p"] = config.top_p
         if config.stream:
-            from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-
             callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()]
             chat = ChatOpenAI(**kwargs, streaming=config.stream, callbacks=callbacks, api_key=api_key)
         else:

embedchain/llm/vertex_ai.py

@@ -2,6 +2,9 @@ import importlib
 import logging
 from typing import Optional
 
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain_google_vertexai import ChatVertexAI
+
 from embedchain.config import BaseLlmConfig
 from embedchain.helpers.json_serializable import register_deserializable
 from embedchain.llm.base import BaseLlm
@@ -24,13 +27,17 @@ class VertexAILlm(BaseLlm):
 
     @staticmethod
     def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
-        from langchain_community.chat_models import ChatVertexAI
-
-        chat = ChatVertexAI(temperature=config.temperature, model=config.model)
-
         if config.top_p and config.top_p != 1:
             logging.warning("Config option `top_p` is not supported by this model.")
 
         messages = BaseLlm._get_messages(prompt, system_prompt=config.system_prompt)
 
-        return chat(messages).content
+        if config.stream:
+            callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()]
+            llm = ChatVertexAI(
+                temperature=config.temperature, model=config.model, callbacks=callbacks, streaming=config.stream
+            )
+        else:
+            llm = ChatVertexAI(temperature=config.temperature, model=config.model)
+
+        return llm.invoke(messages).content