allows toggling between the old as_chain to the new as_graph methods.…

… Tests run with as_graph.
vintasoftware · Sep 11, 2024 · 5c441e5 · 5c441e5
1 parent 8acb2e7
commit 5c441e5
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/django_ai_assistant/conf.py b/django_ai_assistant/conf.py
@@ -10,6 +10,7 @@
 
 DEFAULTS = {
     "INIT_API_FN": "django_ai_assistant.api.views.init_api",
+    "USE_LANGGRAPH": False,
     "CAN_CREATE_THREAD_FN": "django_ai_assistant.permissions.allow_all",
     "CAN_VIEW_THREAD_FN": "django_ai_assistant.permissions.owns_thread",
     "CAN_UPDATE_THREAD_FN": "django_ai_assistant.permissions.owns_thread",

diff --git a/django_ai_assistant/helpers/assistants.py b/django_ai_assistant/helpers/assistants.py
@@ -50,6 +50,7 @@
 from langgraph.graph.message import add_messages
 from langgraph.prebuilt import ToolNode
 
+from django_ai_assistant.conf import app_settings
 from django_ai_assistant.decorators import with_cast_id
 from django_ai_assistant.exceptions import (
     AIAssistantMisconfiguredError,
@@ -659,8 +660,10 @@ def invoke(self, *args: Any, thread_id: Any | None, **kwargs: Any) -> dict:
             dict: The output of the assistant chain,
                 structured like `{"output": "assistant response", "history": ...}`.
         """
-        # chain = self.as_chain(thread_id)
-        chain = self.as_graph(thread_id)
+        if app_settings.USE_LANGGRAPH:
+            chain = self.as_graph(thread_id)
+        else:
+            chain = self.as_chain(thread_id)
         return chain.invoke(*args, **kwargs)
 
     @with_cast_id

diff --git a/tests/settings.py b/tests/settings.py
@@ -107,6 +107,7 @@
 # django-ai-assistant
 
 # NOTE: set a OPENAI_API_KEY on .env.tests file at root when updating the VCRs.
+AI_ASSISTANT_USE_LANGGRAPH = True
 AI_ASSISTANT_INIT_API_FN = "django_ai_assistant.api.views.init_api"
 AI_ASSISTANT_CAN_CREATE_THREAD_FN = "django_ai_assistant.permissions.allow_all"
 AI_ASSISTANT_CAN_VIEW_THREAD_FN = "django_ai_assistant.permissions.owns_thread"