Merge pull request #219 from an7oine/async

rvlb · web-flow · commit 02c3fa5e91d8 · 2026-03-11T14:24:52.000-03:00
Enable langgraph astream usage
diff --git a/django_ai_assistant/helpers/assistants.py b/django_ai_assistant/helpers/assistants.py
@@ -1,7 +1,20 @@
 import abc
 import inspect
 import re
-from typing import Annotated, Any, ClassVar, Dict, Sequence, Type, TypedDict, cast
+from typing import (
+    Annotated,
+    Any,
+    AsyncIterable,
+    AsyncIterator,
+    ClassVar,
+    Dict,
+    Literal,
+    Sequence,
+    Type,
+    TypedDict,
+    cast,
+    overload,
+)
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import (
@@ -415,16 +428,20 @@ def get_history_aware_retriever(self) -> Runnable[dict, RetrieverOutput]:
         )
 
     @with_cast_id
-    def as_graph(self, thread_id: Any | None = None) -> Runnable[dict, dict]:
+    def as_graph(
+        self, thread_id: Any | None = None, thread: Any | None = None
+    ) -> Runnable[dict, dict]:
         """Create the LangGraph graph for the assistant.\n
         This graph is an agent that supports chat history, tool calling, and RAG (if `has_rag=True`).\n
         `as_graph` uses many other methods to create the graph for the assistant.
         Prefer to override the other methods to customize the graph for the assistant.
         Only override this method if you need to customize the graph at a lower level.
 
+        If both arguments are `None`, an in-memory chat message history is used.
+
         Args:
             thread_id (Any | None): The thread ID for the chat message history.
-                If `None`, an in-memory chat message history is used.
+            thread (Any | None): The thread object for the chat message history.
 
         Returns:
             the compiled graph
@@ -434,21 +451,11 @@ def as_graph(self, thread_id: Any | None = None) -> Runnable[dict, dict]:
         llm = self.get_llm()
         tools = self.get_tools()
         llm_with_tools = llm.bind_tools(tools) if tools else llm
-        if thread_id:
+        if thread is None and thread_id is not None:
             thread = Thread.objects.get(id=thread_id)
-        else:
-            thread = None
-
-        def custom_add_messages(left: list[BaseMessage], right: list[BaseMessage]):
-            result = add_messages(left, right)  # type: ignore
-            if thread:
-                # Save all messages, except the initial system message:
-                thread_messages = [m for m in result if not isinstance(m, SystemMessage)]
-                save_django_messages(cast(list[BaseMessage], thread_messages), thread=thread)
-            return result
 
         class AgentState(TypedDict):
-            messages: Annotated[list[AnyMessage], custom_add_messages]
+            messages: Annotated[list[AnyMessage], add_messages]
             input: str | None  # noqa: A003
             output: Any
 
@@ -522,6 +529,10 @@ def record_response(state: AgentState):
             else:
                 response = state["messages"][-1].content
 
+            if thread:
+                # Save all messages, except the initial system message:
+                thread_messages = [m for m in state["messages"] if not isinstance(m, SystemMessage)]
+                save_django_messages(cast(list[BaseMessage], thread_messages), thread=thread)
             return {"output": response}
 
         workflow = StateGraph(AgentState)
@@ -550,28 +561,62 @@ def record_response(state: AgentState):
 
         return workflow.compile()
 
+    @overload
+    def invoke(
+        self,
+        *args: Any,
+        thread_id: Any | None,
+        thread: Any | None = None,
+        mode: Literal["invoke"] = "invoke",
+        **kwargs: Any,
+    ) -> dict:
+        ...  # pragma: no cover
+
+    @overload
+    def invoke(
+        self,
+        *args: Any,
+        thread_id: Any | None,
+        thread: Any | None = None,
+        mode: Literal["astream"],
+        **kwargs: Any,
+    ) -> AsyncIterator[dict]:
+        ...  # pragma: no cover
+
     @with_cast_id
-    def invoke(self, *args: Any, thread_id: Any | None, **kwargs: Any) -> dict:
+    def invoke(
+        self,
+        *args: Any,
+        thread_id: Any | None = None,
+        thread: Any | None = None,
+        mode: Literal["invoke", "astream"] = "invoke",
+        **kwargs: Any,
+    ) -> dict | AsyncIterator[dict]:
         """Invoke the assistant LangChain graph with the given arguments and keyword arguments.\n
         This is the lower-level method to run the assistant.\n
         The graph is created by the `as_graph` method.\n
 
+        If thread_id and thread are `None`, an in-memory chat message history is used.
+
         Args:
             *args: Positional arguments to pass to the graph.
                 To add a new message, use a dict like `{"input": "user message"}`.
                 If thread already has a `HumanMessage` in the end, you can invoke without args.
             thread_id (Any | None): The thread ID for the chat message history.
-                If `None`, an in-memory chat message history is used.
+            thread (Any | None): The thread object for the chat message history.
+            mode (invoke | astream): call named graph method
             **kwargs: Keyword arguments to pass to the graph.
 
         Returns:
             dict: The output of the assistant graph,
                 structured like `{"output": "assistant response", "history": ...}`.
         """
-        graph = self.as_graph(thread_id)
+        graph = self.as_graph(thread_id=thread_id, thread=thread)
         config = kwargs.pop("config", {})
         config["max_concurrency"] = config.pop("max_concurrency", self.tool_max_concurrency)
-        return graph.invoke(*args, config=config, **kwargs)
+        if mode not in ("invoke", "astream"):
+            raise NotImplementedError(f"mode={mode!r}")
+        return getattr(graph, mode)(*args, config=config, **kwargs)
 
     @with_cast_id
     def run(self, message: str, thread_id: Any | None = None, **kwargs: Any) -> Any:
@@ -595,6 +640,34 @@ def run(self, message: str, thread_id: Any | None = None, **kwargs: Any) -> Any:
             **kwargs,
         )["output"]
 
+    @with_cast_id
+    async def astream(
+        self, message: str, thread: Any | None = None, **kwargs: Any
+    ) -> AsyncIterable[Any]:
+        """Async-stream the assistant with the given message and thread.\n
+        This is the higher-level method to run the assistant.\n
+
+        Args:
+            message (str): The user message to pass to the assistant.
+            thread (Any | None): The thread object for the chat message history.
+                If `None`, an in-memory chat message history is used.
+            **kwargs: Additional keyword arguments to pass to the graph.
+
+        Yields:
+            Any: The assistant response to the user message.
+        """
+        async for output, metadata in self.invoke(
+            {
+                "input": message,
+            },
+            thread=thread,
+            mode="astream",
+            stream_mode="messages",
+            **kwargs,
+        ):
+            if metadata.get("langgraph_node") == "agent" and (content := output.content):
+                yield content
+
     def _run_as_tool(self, message: str, **kwargs: Any) -> Any:
         return self.run(message, thread_id=None, **kwargs)
 
diff --git a/tests/test_helpers/cassettes/test_assistants/test_AIAssistant_astream.yaml b/tests/test_helpers/cassettes/test_assistants/test_AIAssistant_astream.yaml
@@ -0,0 +1,200 @@
+interactions:
+- request:
+    body: '{"messages":[{"content":"You are a temperature bot. Today is 2024-06-09.","role":"system"},{"content":"What
+      is the temperature today in Recife?","role":"user"}],"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true},"temperature":1.0,"tools":[{"type":"function","function":{"name":"fetch_current_temperature","description":"Fetch
+      the current temperature data for a location","parameters":{"properties":{"location":{"type":"string"}},"required":["location"],"type":"object"}}},{"type":"function","function":{"name":"fetch_forecast_temperature","description":"Fetch
+      the forecast temperature data for a location","parameters":{"type":"object","properties":{"location":{"type":"string"},"dt_str":{"description":"Date
+      in the format ''YYYY-MM-DD''","type":"string"}},"required":["location","dt_str"]}}}]}'
+    headers:
+      Accept:
+      - application/json
+      Accept-Encoding:
+      - gzip, deflate, zstd
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '812'
+      Content-Type:
+      - application/json
+      Host:
+      - api.openai.com
+      X-Stainless-Raw-Response:
+      - 'true'
+      authorization:
+      - DUMMY
+      user-agent:
+      - OpenAI/Python
+      x-stainless-retry-count:
+      - '0'
+    method: POST
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    body:
+      string: 'data: {"id":"chatcmpl-DIDpzDJSqT4EW7UPZCL857JRcFTq6","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,"id":"call_zuq818gZWJJlDhdrLcI5IPMm","type":"function","function":{"name":"fetch_current_temperature","arguments":""}}],"refusal":null},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"NW7i9IqiUa8"}
+
+
+        data: {"id":"chatcmpl-DIDpzDJSqT4EW7UPZCL857JRcFTq6","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\""}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"2K6"}
+
+
+        data: {"id":"chatcmpl-DIDpzDJSqT4EW7UPZCL857JRcFTq6","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"location"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"sQRXefPeqyM0h3"}
+
+
+        data: {"id":"chatcmpl-DIDpzDJSqT4EW7UPZCL857JRcFTq6","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\":\""}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"1"}
+
+
+        data: {"id":"chatcmpl-DIDpzDJSqT4EW7UPZCL857JRcFTq6","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"Rec"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"bbt"}
+
+
+        data: {"id":"chatcmpl-DIDpzDJSqT4EW7UPZCL857JRcFTq6","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"ife"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"2tw"}
+
+
+        data: {"id":"chatcmpl-DIDpzDJSqT4EW7UPZCL857JRcFTq6","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"}"}}]},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"sf0"}
+
+
+        data: {"id":"chatcmpl-DIDpzDJSqT4EW7UPZCL857JRcFTq6","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"tool_calls"}],"usage":null,"obfuscation":"gBbl"}
+
+
+        data: {"id":"chatcmpl-DIDpzDJSqT4EW7UPZCL857JRcFTq6","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[],"usage":{"prompt_tokens":111,"completion_tokens":16,"total_tokens":127,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":0,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}},"obfuscation":"UEfMPsiTrLLjA"}
+
+
+        data: [DONE]
+
+
+        '
+    headers:
+      Connection:
+      - keep-alive
+      Content-Type:
+      - text/event-stream; charset=utf-8
+      Date: Sun, 09 Jun 2024 23:39:08 GMT
+      Server: DUMMY
+      Strict-Transport-Security:
+      - max-age=31536000; includeSubDomains; preload
+      Transfer-Encoding:
+      - chunked
+      X-Content-Type-Options:
+      - nosniff
+      access-control-expose-headers:
+      - X-Request-ID
+      cf-cache-status:
+      - DYNAMIC
+      openai-project:
+      - proj_dummy_openai_id
+      x-openai-proxy-wasm:
+      - v0.1
+    status:
+      code: 200
+      message: OK
+- request:
+    body: '{"messages":[{"content":"You are a temperature bot. Today is 2024-06-09.","role":"system"},{"content":"What
+      is the temperature today in Recife?","role":"user"},{"content":null,"role":"assistant","tool_calls":[{"type":"function","id":"call_zuq818gZWJJlDhdrLcI5IPMm","function":{"name":"fetch_current_temperature","arguments":"{\"location\":
+      \"Recife\"}"}}]},{"content":"32 degrees Celsius","role":"tool","tool_call_id":"call_zuq818gZWJJlDhdrLcI5IPMm"}],"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true},"temperature":1.0,"tools":[{"type":"function","function":{"name":"fetch_current_temperature","description":"Fetch
+      the current temperature data for a location","parameters":{"properties":{"location":{"type":"string"}},"required":["location"],"type":"object"}}},{"type":"function","function":{"name":"fetch_forecast_temperature","description":"Fetch
+      the forecast temperature data for a location","parameters":{"type":"object","properties":{"location":{"type":"string"},"dt_str":{"description":"Date
+      in the format ''YYYY-MM-DD''","type":"string"}},"required":["location","dt_str"]}}}]}'
+    headers:
+      Accept:
+      - application/json
+      Accept-Encoding:
+      - gzip, deflate, zstd
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '1103'
+      Content-Type:
+      - application/json
+      Host:
+      - api.openai.com
+      X-Stainless-Raw-Response:
+      - 'true'
+      authorization:
+      - DUMMY
+      cookie:
+      - DUMMY
+      user-agent:
+      - OpenAI/Python
+      x-stainless-retry-count:
+      - '0'
+    method: POST
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    body:
+      string: 'data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"role":"assistant","content":"","refusal":null},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"XpqH8r6lKhOKG6"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"The"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"cIlWkGBpC6Lml"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"
+        current"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"t2j4AdCr"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"
+        temperature"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"d4F5"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"
+        in"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"HMWCJkSoPRIze"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"
+        Recife"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"b21LgWyaZ"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"
+        is"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"si0Fdd77rKVjw"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"
+        "},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"sXqWNkI7XGaTUss"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"32"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"RrCp3ZZkst7zWH"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"
+        degrees"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"SsSscPaq"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"
+        Celsius"},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"SbS9VfIu"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{"content":"."},"logprobs":null,"finish_reason":null}],"usage":null,"obfuscation":"uKYplWNzsKkx9s3"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}],"usage":null,"obfuscation":"doOrarXmUb"}
+
+
+        data: {"id":"chatcmpl-DIDpzBVzIdaiHSEUtACeflIWDRaWo","object":"chat.completion.chunk","created":1773235467,"model":"gpt-4o-2024-08-06","service_tier":"default","system_fingerprint":"fp_c7a156cce7","choices":[],"usage":{"prompt_tokens":139,"completion_tokens":12,"total_tokens":151,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":0,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}},"obfuscation":"jW6QWGp9HDgt3"}
+
+
+        data: [DONE]
+
+
+        '
+    headers:
+      Connection:
+      - keep-alive
+      Content-Type:
+      - text/event-stream; charset=utf-8
+      Date: Sun, 09 Jun 2024 23:39:08 GMT
+      Server: DUMMY
+      Strict-Transport-Security:
+      - max-age=31536000; includeSubDomains; preload
+      Transfer-Encoding:
+      - chunked
+      X-Content-Type-Options:
+      - nosniff
+      access-control-expose-headers:
+      - X-Request-ID
+      cf-cache-status:
+      - DYNAMIC
+      openai-project:
+      - proj_dummy_openai_id
+      x-openai-proxy-wasm:
+      - v0.1
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/test_helpers/test_assistants.py b/tests/test_helpers/test_assistants.py
