Added VertexAI and fixed a bug in the websocket

pyproject.toml

@@ -22,7 +22,7 @@ include = ["src/backend/langflow/*", "src/backend/langflow/**/*"]
 langflow = "langflow.__main__:main"
 
 [tool.poetry.dependencies]
-python = ">=3.9,<3.12"
+python = ">=3.9,<3.11"
 fastapi = "^0.98.0"
 uvicorn = "^0.22.0"
 beautifulsoup4 = "^4.12.2"
@@ -69,6 +69,7 @@ pinecone-client = "^2.2.2"
 supabase = "^1.0.3"
 pymongo = "^4.4.0"
 certifi = "^2023.5.7"
+google-cloud-aiplatform = "^1.26.1"
 psycopg = "^3.1.9"
 psycopg-binary = "^3.1.9"
 

src/backend/langflow/api/v1/chat.py

@@ -19,8 +19,11 @@ async def chat(client_id: str, websocket: WebSocket):
         if client_id in chat_manager.in_memory_cache:
             await chat_manager.handle_websocket(client_id, websocket)
         else:
+            # We accept the connection but close it immediately
+            # if the flow is not built yet
+            await websocket.accept()
             message = "Please, build the flow before sending messages"
-            await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason=message)
+            await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=message)
     except WebSocketException as exc:
         logger.error(exc)
         await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=str(exc))

src/backend/langflow/config.yaml

@@ -119,6 +119,15 @@ llms:
     documentation: "https://python.langchain.com/docs/modules/model_io/models/chat/integrations/anthropic"
   HuggingFaceHub:
     documentation: "https://python.langchain.com/docs/modules/model_io/models/llms/integrations/huggingface_hub"
+  VertexAI:
+    documentation: "https://python.langchain.com/docs/modules/model_io/models/llms/integrations/google_vertex_ai_palm"
+  ###
+  # There's a bug in this component deactivating until we get it sorted: _language_models.py", line 804, in send_message
+  #     is_blocked=safety_attributes.get("blocked", False),
+  # AttributeError: 'list' object has no attribute 'get'
+  # ChatVertexAI:
+  #   documentation: "https://python.langchain.com/docs/modules/model_io/models/chat/integrations/google_vertex_ai_palm"
+  ###
 memories:
   # https://github.com/supabase-community/supabase-py/issues/482
   # ZepChatMessageHistory:

src/backend/langflow/graph/edge/base.py

@@ -27,12 +27,7 @@ def validate_edge(self) -> None:
         # Get what type of input the target node is expecting
 
         self.matched_type = next(
-            (
-                output
-                for output in self.source_types
-                for target_req in self.target_reqs
-                if output in target_req
-            ),
+            (output for output in self.source_types if output in self.target_reqs),
             None,
         )
         no_matched_type = self.matched_type is None

src/backend/langflow/interface/custom_lists.py

@@ -10,8 +10,12 @@
     text_splitter,
 )
 from langchain.agents import agent_toolkits
-from langchain.chat_models import AzureChatOpenAI, ChatOpenAI
-from langchain.chat_models import ChatAnthropic
+from langchain.chat_models import (
+    AzureChatOpenAI,
+    ChatOpenAI,
+    ChatVertexAI,
+    ChatAnthropic,
+)
 
 from langflow.interface.importing.utils import import_class
 from langflow.interface.agents.custom import CUSTOM_AGENTS
@@ -22,6 +26,7 @@
 llm_type_to_cls_dict["anthropic-chat"] = ChatAnthropic  # type: ignore
 llm_type_to_cls_dict["azure-chat"] = AzureChatOpenAI  # type: ignore
 llm_type_to_cls_dict["openai-chat"] = ChatOpenAI  # type: ignore
+llm_type_to_cls_dict["vertexai-chat"] = ChatVertexAI  # type: ignore
 
 
 # Toolkits

src/backend/langflow/interface/initialize/llm.py

@@ -0,0 +1,9 @@
+def initialize_vertexai(class_object, params):
+    if credentials_path := params.get("credentials"):
+        from google.oauth2 import service_account  # type: ignore
+
+        credentials_object = service_account.Credentials.from_service_account_file(
+            filename=credentials_path
+        )
+        params["credentials"] = credentials_object
+    return class_object(**params)

src/backend/langflow/interface/initialize/loading.py

@@ -6,6 +6,7 @@
 from langchain.agents.agent import AgentExecutor
 from langchain.agents.agent_toolkits.base import BaseToolkit
 from langchain.agents.tools import BaseTool
+from langflow.interface.initialize.llm import initialize_vertexai
 
 from langflow.interface.initialize.vector_store import vecstore_initializer
 
@@ -88,6 +89,15 @@ def instantiate_based_on_type(class_object, base_type, node_type, params):
         return class_object(**params)
 
 
+def instantiate_llm(node_type, class_object, params: Dict):
+    # This is a workaround so JinaChat works until streaming is implemented
+    # if "openai_api_base" in params and "jina" in params["openai_api_base"]:
+    # False if condition is True
+    if node_type == "VertexAI":
+        return initialize_vertexai(class_object=class_object, params=params)
+    return class_object(**params)
+
+
 def instantiate_memory(node_type, class_object, params):
     try:
         return class_object(**params)
@@ -117,10 +127,6 @@ def instantiate_retriever(node_type, class_object, params):
     return class_object(**params)
 
 
-def instantiate_llm(node_type, class_object, params: Dict):
-    return class_object(**params)
-
-
 def instantiate_chains(node_type, class_object: Type[Chain], params: Dict):
     if "retriever" in params and hasattr(params["retriever"], "as_retriever"):
         params["retriever"] = params["retriever"].as_retriever()

src/backend/langflow/template/frontend_node/llms.py

@@ -8,6 +8,49 @@
 
 
 class LLMFrontendNode(FrontendNode):
+    def add_extra_fields(self) -> None:
+        if "VertexAI" in self.template.type_name:
+            # Add credentials field which should of type file.
+            self.template.add_field(
+                TemplateField(
+                    field_type="file",
+                    required=False,
+                    show=True,
+                    name="credentials",
+                    value="",
+                    suffixes=[".json"],
+                    fileTypes=["json"],
+                )
+            )
+
+    @staticmethod
+    def format_vertex_field(field: TemplateField, name: str):
+        if "VertexAI" in name:
+            advanced_fields = [
+                "tuned_model_name",
+                "verbose",
+                "top_p",
+                "top_k",
+                "max_output_tokens",
+            ]
+            if field.name in advanced_fields:
+                field.advanced = True
+            show_fields = [
+                "tuned_model_name",
+                "verbose",
+                "project",
+                "location",
+                "credentials",
+                "max_output_tokens",
+                "model_name",
+                "temperature",
+                "top_p",
+                "top_k",
+            ]
+
+            if field.name in show_fields:
+                field.show = True
+
     @staticmethod
     def format_openai_field(field: TemplateField):
         if "openai" in field.name.lower():
@@ -60,6 +103,8 @@ def format_field(field: TemplateField, name: Optional[str] = None) -> None:
             LLMFrontendNode.format_azure_field(field)
         if name and "llama" in name.lower():
             LLMFrontendNode.format_llama_field(field)
+        if name and "vertex" in name.lower():
+            LLMFrontendNode.format_vertex_field(field, name)
         SHOW_FIELDS = ["repo_id"]
         if field.name in SHOW_FIELDS:
             field.show = True
@@ -95,6 +140,17 @@ def format_field(field: TemplateField, name: Optional[str] = None) -> None:
             "model_file",
             "model_type",
             "deployment_name",
+            "credentials",
         ]:
             field.advanced = False
             field.show = True
+        if field.name == "credentials":
+            field.field_type = "file"
+        if name == "VertexAI" and field.name not in [
+            "callbacks",
+            "client",
+            "stop",
+            "tags",
+            "cache",
+        ]:
+            field.show = True

src/frontend/src/App.tsx

@@ -51,6 +51,9 @@ export default function App() {
   useEffect(() => {
     // If there is an error alert open with data, add it to the alertsList
     if (errorOpen && errorData) {
+      if(alertsList.length > 0 && JSON.stringify(alertsList[alertsList.length - 1].data)===JSON.stringify(errorData)){
+        return;
+      }
       setErrorOpen(false);
       setAlertsList((old) => {
         let newAlertsList = [
@@ -62,6 +65,9 @@ export default function App() {
     }
     // If there is a notice alert open with data, add it to the alertsList
     else if (noticeOpen && noticeData) {
+      if(alertsList.length > 0 && JSON.stringify(alertsList[alertsList.length - 1].data)===JSON.stringify(noticeData)){
+        return;
+      }
       setNoticeOpen(false);
       setAlertsList((old) => {
         let newAlertsList = [
@@ -73,6 +79,9 @@ export default function App() {
     }
     // If there is a success alert open with data, add it to the alertsList
     else if (successOpen && successData) {
+      if(alertsList.length > 0 && JSON.stringify(alertsList[alertsList.length - 1].data)===JSON.stringify(successData)){
+        return;
+      }
       setSuccessOpen(false);
       setAlertsList((old) => {
         let newAlertsList = [
@@ -152,4 +161,4 @@ export default function App() {
       </div>
     </div>
   );
-}
+}

src/frontend/src/contexts/alertContext.tsx

@@ -78,9 +78,9 @@ export function AlertProvider({ children }: { children: ReactNode }) {
    * @param newState An object containing the new error data, including title and optional list of error messages
    */
   function setErrorData(newState: { title: string; list?: Array<string> }) {
-    setErrorDataState(newState);
-    setErrorOpen(true);
     if (newState.title && newState.title !== "") {
+      setErrorDataState(newState);
+      setErrorOpen(true);
       setNotificationCenter(true);
       pushNotificationList({
         type: "error",
@@ -95,9 +95,9 @@ export function AlertProvider({ children }: { children: ReactNode }) {
    * @param newState An object containing the title of the notice and optionally a link.
    */
   function setNoticeData(newState: { title: string; link?: string }) {
-    setNoticeDataState(newState);
-    setNoticeOpen(true);
     if (newState.title && newState.title !== "") {
+      setNoticeDataState(newState);
+      setNoticeOpen(true);
       // Add new notice to notification center
       setNotificationCenter(true);
       pushNotificationList({
@@ -113,11 +113,10 @@ export function AlertProvider({ children }: { children: ReactNode }) {
    * @param newState - A state object with a "title" property to set in the success data state.
    */
   function setSuccessData(newState: { title: string }) {
-    setSuccessDataState(newState); // update the success data state with the provided new state
-    setSuccessOpen(true); // open the success alert
-
     // If the new state has a "title" property, add a new success notification to the list
     if (newState.title && newState.title !== "") {
+      setSuccessDataState(newState); // update the success data state with the provided new state
+      setSuccessOpen(true); // open the success alert
       setNotificationCenter(true); // show the notification center
       pushNotificationList({
         // add the new notification to the list

src/frontend/src/icons/VertexAI/index.tsx

@@ -0,0 +1,9 @@
+import React, { forwardRef } from "react";
+import { ReactComponent as VertexAISVG } from "./vertex_ai.svg";
+
+export const VertexAIIcon = forwardRef<
+  SVGSVGElement,
+  React.PropsWithChildren<{}>
+>((props, ref) => {
+  return <VertexAISVG ref={ref} {...props} />;
+});

src/frontend/src/modals/chatModal/index.tsx

@@ -11,6 +11,7 @@ import { ChatMessageType } from "../../types/chat";
 import ChatInput from "./chatInput";
 
 import _ from "lodash";
+import { getHealth } from "../../controllers/API";
 
 export default function ChatModal({
   flow,
@@ -204,25 +205,25 @@ export default function ChatModal({
         handleOnClose(event);
       };
       newWs.onerror = (ev) => {
-        console.log(ev, "error");
-        if (flow.id === "") {
-          connectWS();
-        } else {
+        getHealth().then((res) => {
+          if (res.status === 200) {
+            connectWS();
+          }
+        }).catch((err) => {
           setErrorData({
-            title: "There was an error on web connection, please: ",
+            // message when the backend failed
+            title: "The backend is not responding. Please try again later.",
+            // possible solution list
             list: [
-              "Refresh the page",
-              "Use a new flow tab",
-              "Check if the backend is up",
+              "Check your internet connection.",
+              "Check if the backend is running."
             ],
           });
-        }
+        })
       };
       ws.current = newWs;
     } catch (error) {
-      if (flow.id === "") {
         connectWS();
-      }
       console.log(error);
     }
   }

src/frontend/src/utils.ts

@@ -47,6 +47,7 @@ import {
 } from "lucide-react";
 import { SupabaseIcon } from "./icons/supabase";
 import { MongoDBIcon } from "./icons/MongoDB";
+import { VertexAIIcon } from "./icons/VertexAI";
 
 export function classNames(...classes: Array<string>) {
   return classes.filter(Boolean).join(" ");
@@ -241,6 +242,12 @@ export const nodeIconsLucide: {
   SupabaseVectorStore: SupabaseIcon as React.ForwardRefExoticComponent<
     ComponentType<SVGProps<SVGSVGElement>>
   >,
+  VertexAI: VertexAIIcon as React.ForwardRefExoticComponent<
+    ComponentType<SVGProps<SVGSVGElement>>
+  >,
+  ChatVertexAI: VertexAIIcon as React.ForwardRefExoticComponent<
+    ComponentType<SVGProps<SVGSVGElement>>
+  >,
   agents: Rocket as React.ForwardRefExoticComponent<
     ComponentType<SVGProps<SVGSVGElement>>
   >,