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

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