Astra Assistants Support (#2041)

* ruff

* revert config.yaml

* add build_config

* ruff

* add build_config

src/backend/base/langflow/components/astra_assistants/__init__.py

@@ -0,0 +1,13 @@
+from .create_assistant import AssistantsCreateAssistant
+from .get_assistant import AssistantsGetAssistantName
+from .list_assistants import AssistantsListAssistants
+from .run import AssistantsRun
+from .getenvvar import GetEnvVar
+
+__all__ = [
+    "AssistantsCreateAssistant",
+    "AssistantsGetAssistantName",
+    "AssistantsListAssistants",
+    "AssistantsRun",
+    "GetEnvVar",
+]

src/backend/base/langflow/components/astra_assistants/create_assistant.py

@@ -0,0 +1,48 @@
+from langflow.custom import CustomComponent
+from openai import OpenAI
+from astra_assistants import patch
+
+
+class AssistantsCreateAssistant(CustomComponent):
+    display_name = "Create Assistant"
+    description = "Creates an Assistant and returns it's id"
+
+    def build_config(self):
+        return {
+            "name": {
+                "display_name": "Assistant Name",
+                "advanced": False,
+                "info": "Name for the assistant being created",
+            },
+            "instructions": {
+                "display_name": "Instructions",
+                "info": "Instructions for the assistant, think of these as the system prompt.",
+                "advanced": False,
+            },
+            "model": {
+                "display_name": "Model name",
+                "advanced": False,
+                "info": (
+                    "Model for the assistant.\n\n"
+                    "Environment variables for provider credentials can be set with the Dotenv Component.\n\n"
+                    "Models are supported via LiteLLM, see (https://docs.litellm.ai/docs/providers) for supported model names and env vars."
+                ),
+            },
+            "env_set": {
+                "display_name": "Environment Set",
+                "advanced": False,
+                "info": "Dummy input to allow chaining with Dotenv Component.",
+            },
+        }
+
+    def build(self, name: str, instructions: str, model: str, env_set: str = None) -> str:
+        print(f"env_set is {env_set}")
+        if env_set is None:
+            raise Exception("Environment variables not set")
+        client = patch(OpenAI())
+        assistant = client.beta.assistants.create(
+            name=name,
+            instructions=instructions,
+            model=model,
+        )
+        return assistant.id

src/backend/base/langflow/components/astra_assistants/create_thread.py

@@ -0,0 +1,25 @@
+from langflow.custom import CustomComponent
+from openai import OpenAI
+from astra_assistants import patch
+
+
+class AssistantsCreateThread(CustomComponent):
+    display_name = "Create Assistant Thread"
+    description = "Creates a thread and returns the thread id"
+
+    def build_config(self):
+        return {
+            "env_set": {
+                "display_name": "Environment Set",
+                "advanced": False,
+                "info": "Dummy input to allow chaining with Dotenv Component.",
+            },
+        }
+
+    def build(self, env_set: str = None) -> str:
+        client = patch(OpenAI())
+
+        thread = client.beta.threads.create()
+        thread_id = thread.id
+
+        return thread_id

src/backend/base/langflow/components/astra_assistants/dotenv.py

@@ -0,0 +1,28 @@
+import io
+from dotenv import load_dotenv
+from langflow.custom import CustomComponent
+
+
+class Dotenv(CustomComponent):
+    display_name = "Dotenv"
+    description = "Load .env file into env vars"
+
+    def build_config(self):
+        return {
+            "dotenv_file_content": {
+                "display_name": "Dotenv file content",
+                "advanced": False,
+                "info": (
+                    "Paste the content of your .env file directly\n\n"
+                    "Since contents are sensitive, using a Global variable set as 'password' is recommended"
+                ),
+            },
+        }
+
+    def build(self, dotenv_file_content: str) -> str:
+        try:
+            fake_file = io.StringIO(dotenv_file_content)
+            result = load_dotenv(stream=fake_file, override=True)
+            return result
+        except Exception as e:
+            raise e

src/backend/base/langflow/components/astra_assistants/get_assistant.py

@@ -0,0 +1,28 @@
+from langflow.custom import CustomComponent
+from openai import OpenAI
+from astra_assistants import patch
+
+
+class AssistantsGetAssistantName(CustomComponent):
+    display_name = "Get Assistant name"
+    description = "Assistant by id"
+
+    def build_config(self):
+        return {
+            "assistant_id": {
+                "display_name": "Assistant ID",
+                "advanced": False,
+            },
+            "env_set": {
+                "display_name": "Environment Set",
+                "advanced": False,
+                "info": "Dummy input to allow chaining with Dotenv Component.",
+            },
+        }
+
+    def build(self, assistant_id: str, env_set: str = None) -> str:
+        client = patch(OpenAI())
+        assistant = client.beta.assistants.retrieve(
+            assistant_id=assistant_id,
+        )
+        return assistant.name

src/backend/base/langflow/components/astra_assistants/getenvvar.py

@@ -0,0 +1,14 @@
+import os
+from langflow.custom import CustomComponent
+
+
+class GetEnvVar(CustomComponent):
+    display_name = "Get env var"
+    description = "Get env var"
+    icon = "custom_components"
+
+    def build_config(self):
+        return {"env_var_name": {"display_name": "Env var name"}}
+
+    def build(self, env_var_name: str) -> str:
+        return os.environ[env_var_name]

src/backend/base/langflow/components/astra_assistants/list_assistants.py

@@ -0,0 +1,18 @@
+from typing import List
+from langflow.custom import CustomComponent
+from openai import OpenAI
+from astra_assistants import patch
+
+
+class AssistantsListAssistants(CustomComponent):
+    display_name = "List Assistants"
+    description = "Returns a list of assistant id's"
+
+    def build_config(self):
+        return {}
+
+    def build(self) -> List[str]:
+        client = patch(OpenAI())
+        assistants = client.beta.assistants.list()
+        id_list = [assistant.id for assistant in assistants]
+        return id_list

src/backend/base/langflow/components/astra_assistants/run.py

@@ -0,0 +1,63 @@
+from langflow.custom import CustomComponent
+from openai import OpenAI
+from openai.lib.streaming import AssistantEventHandler
+from astra_assistants import patch
+
+
+class AssistantsRun(CustomComponent):
+    display_name = "Run Assistant"
+    description = "Executes an Assistant Run against a thread"
+
+    def build_config(self):
+        return {
+            "assistant_id": {
+                "display_name": "Assistant ID",
+                "advanced": False,
+                "info": (
+                    "The ID of the assistant to run. \n\n"
+                    "Can be retrieved using the List Assistants component or created with the Create Assistant component."
+                ),
+            },
+            "user_message": {
+                "display_name": "User Message",
+                "info": "User message to pass to the run.",
+                "advanced": False,
+            },
+            "thread_id": {
+                "display_name": "Thread ID",
+                "advanced": False,
+                "info": "Thread ID to use with the run. If not provided, a new thread will be created.",
+            },
+            "env_set": {
+                "display_name": "Environment Set",
+                "advanced": False,
+                "info": "Dummy input to allow chaining with Dotenv Component.",
+            },
+        }
+
+    def build(self, assistant_id: str, user_message: str, thread_id: str = None, env_set: str = None) -> str:
+        text = ""
+        client = patch(OpenAI())
+
+        if thread_id is None:
+            thread = client.beta.threads.create()
+            thread_id = thread.id
+
+        # add the user message
+        client.beta.threads.messages.create(thread_id=thread_id, role="user", content=user_message)
+
+        class EventHandler(AssistantEventHandler):
+            def __init__(self):
+                super().__init__()
+
+        event_handler = EventHandler()
+        with client.beta.threads.runs.create_and_stream(
+            thread_id=thread_id,
+            assistant_id=assistant_id,
+            event_handler=event_handler,
+        ) as stream:
+            # return stream.text_deltas
+            for part in stream.text_deltas:
+                text += part
+                print(part)
+        return text