dao-ai 0.0.16__py3-none-any.whl → 0.0.17__py3-none-any.whl

dao_ai/graph.py

@@ -221,7 +221,7 @@ def _create_swarm_graph(config: AppConfig) -> CompiledStateGraph:
 
     return swarm_node
 
-    #return workflow.compile(checkpointer=checkpointer, store=store)
+    # return workflow.compile(checkpointer=checkpointer, store=store)
 
 
 def create_dao_ai_graph(config: AppConfig) -> CompiledStateGraph:

dao_ai/models.py

@@ -65,9 +65,19 @@ class LanggraphChatModel(ChatModel):
         context: Context = self._convert_to_context(params)
         custom_inputs: dict[str, Any] = {"configurable": context.model_dump()}
 
-        response: dict[str, Sequence[BaseMessage]] = self.graph.invoke(
-            request, context=context, config=custom_inputs
+        # Use async ainvoke internally for parallel execution
+        import asyncio
+
+        async def _async_invoke():
+            return await self.graph.ainvoke(
+                request, context=context, config=custom_inputs
+            )
+
+        loop = asyncio.get_event_loop()
+        response: dict[str, Sequence[BaseMessage]] = loop.run_until_complete(
+            _async_invoke()
         )
+
         logger.trace(f"response: {response}")
 
         last_message: BaseMessage = response["messages"][-1]
@@ -114,33 +124,51 @@ class LanggraphChatModel(ChatModel):
         context: Context = self._convert_to_context(params)
         custom_inputs: dict[str, Any] = {"configurable": context.model_dump()}
 
-        for nodes, stream_mode, messages_batch in self.graph.stream(
-            request,
-            context=context,
-            config=custom_inputs,
-            stream_mode=["messages", "custom"],
-            subgraphs=True,
-        ):
-            nodes: tuple[str, ...]
-            stream_mode: str
-            messages_batch: Sequence[BaseMessage]
-            logger.trace(
-                f"nodes: {nodes}, stream_mode: {stream_mode}, messages: {messages_batch}"
-            )
-            for message in messages_batch:
-                if (
-                    isinstance(
-                        message,
-                        (
-                            AIMessageChunk,
-                            AIMessage,
-                        ),
-                    )
-                    and message.content
-                    and "summarization" not in nodes
-                ):
-                    content = message.content
-                    yield self._create_chat_completion_chunk(content)
+        # Use async astream internally for parallel execution
+        import asyncio
+
+        async def _async_stream():
+            async for nodes, stream_mode, messages_batch in self.graph.astream(
+                request,
+                context=context,
+                config=custom_inputs,
+                stream_mode=["messages", "custom"],
+                subgraphs=True,
+            ):
+                nodes: tuple[str, ...]
+                stream_mode: str
+                messages_batch: Sequence[BaseMessage]
+                logger.trace(
+                    f"nodes: {nodes}, stream_mode: {stream_mode}, messages: {messages_batch}"
+                )
+                for message in messages_batch:
+                    if (
+                        isinstance(
+                            message,
+                            (
+                                AIMessageChunk,
+                                AIMessage,
+                            ),
+                        )
+                        and message.content
+                        and "summarization" not in nodes
+                    ):
+                        content = message.content
+                        yield self._create_chat_completion_chunk(content)
+
+        # Convert async generator to sync generator
+        loop = asyncio.get_event_loop()
+        async_gen = _async_stream()
+
+        try:
+            while True:
+                try:
+                    item = loop.run_until_complete(async_gen.__anext__())
+                    yield item
+                except StopAsyncIteration:
+                    break
+        finally:
+            loop.run_until_complete(async_gen.aclose())
 
     def _create_chat_completion_chunk(self, content: str) -> ChatCompletionChunk:
         return ChatCompletionChunk(
@@ -178,9 +206,18 @@ def _process_langchain_messages(
     messages: Sequence[BaseMessage],
     custom_inputs: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any] | Any:
+    """Process LangChain messages using async LangGraph calls internally."""
+    import asyncio
+
     if isinstance(app, LanggraphChatModel):
         app = app.graph
-    return app.invoke({"messages": messages}, config=custom_inputs)
+
+    # Use async ainvoke internally for parallel execution
+    async def _async_invoke():
+        return await app.ainvoke({"messages": messages}, config=custom_inputs)
+
+    loop = asyncio.get_event_loop()
+    return loop.run_until_complete(_async_invoke())
 
 
 def _process_langchain_messages_stream(
@@ -188,6 +225,9 @@ def _process_langchain_messages_stream(
     messages: Sequence[BaseMessage],
     custom_inputs: Optional[dict[str, Any]] = None,
 ) -> Generator[AIMessageChunk, None, None]:
+    """Process LangChain messages in streaming mode using async LangGraph calls internally."""
+    import asyncio
+
     if isinstance(app, LanggraphChatModel):
         app = app.graph
 
@@ -196,32 +236,48 @@ def _process_langchain_messages_stream(
     custom_inputs = custom_inputs.get("configurable", custom_inputs or {})
     context: Context = Context(**custom_inputs)
 
-    for nodes, stream_mode, messages in app.stream(
-        {"messages": messages},
-        context=context,
-        config=custom_inputs,
-        stream_mode=["messages", "custom"],
-        subgraphs=True,
-    ):
-        nodes: tuple[str, ...]
-        stream_mode: str
-        messages: Sequence[BaseMessage]
-        logger.trace(
-            f"nodes: {nodes}, stream_mode: {stream_mode}, messages: {messages}"
-        )
-        for message in messages:
-            if (
-                isinstance(
-                    message,
-                    (
-                        AIMessageChunk,
-                        AIMessage,
-                    ),
-                )
-                and message.content
-                and "summarization" not in nodes
-            ):
-                yield message
+    # Use async astream internally for parallel execution
+    async def _async_stream():
+        async for nodes, stream_mode, stream_messages in app.astream(
+            {"messages": messages},
+            context=context,
+            config=custom_inputs,
+            stream_mode=["messages", "custom"],
+            subgraphs=True,
+        ):
+            nodes: tuple[str, ...]
+            stream_mode: str
+            stream_messages: Sequence[BaseMessage]
+            logger.trace(
+                f"nodes: {nodes}, stream_mode: {stream_mode}, messages: {stream_messages}"
+            )
+            for message in stream_messages:
+                if (
+                    isinstance(
+                        message,
+                        (
+                            AIMessageChunk,
+                            AIMessage,
+                        ),
+                    )
+                    and message.content
+                    and "summarization" not in nodes
+                ):
+                    yield message
+
+    # Convert async generator to sync generator
+    loop = asyncio.get_event_loop()
+    async_gen = _async_stream()
+
+    try:
+        while True:
+            try:
+                item = loop.run_until_complete(async_gen.__anext__())
+                yield item
+            except StopAsyncIteration:
+                break
+    finally:
+        loop.run_until_complete(async_gen.aclose())
 
 
 def _process_mlflow_messages(
@@ -272,7 +328,7 @@ def process_messages_stream(
     Process messages through a ChatAgent in streaming mode.
 
     Utility function that normalizes message input formats and
-    streams the agent's responses as they're generated.
+    streams the agent's responses as they're generated using async LangGraph calls internally.
 
     Args:
         app: The ChatAgent to process messages with
@@ -302,7 +358,7 @@ def process_messages(
     Process messages through a ChatAgent in batch mode.
 
     Utility function that normalizes message input formats and
-    returns the complete response from the agent.
+    returns the complete response from the agent using async LangGraph calls internally.
 
     Args:
         app: The ChatAgent to process messages with

dao_ai/nodes.py

@@ -69,7 +69,7 @@ def summarization_node(app_model: AppModel) -> RunnableLike:
 
 
 def call_agent_with_summarized_messages(agent: CompiledStateGraph) -> RunnableLike:
-    def call_agent(state: SharedState, runtime: Runtime[Context]) -> SharedState:
+    async def call_agent(state: SharedState, runtime: Runtime[Context]) -> SharedState:
         logger.debug(f"Calling agent {agent.name} with summarized messages")
 
         # Get the summarized messages from the summarization node
@@ -81,7 +81,9 @@ def call_agent_with_summarized_messages(agent: CompiledStateGraph) -> RunnableLi
             "messages": messages,
         }
 
-        response: dict[str, Any] = agent.invoke(input=input, context=runtime.context)
+        response: dict[str, Any] = await agent.ainvoke(
+            input=input, context=runtime.context
+        )
         response_messages = response.get("messages", [])
         logger.debug(f"Agent returned {len(response_messages)} messages")
 
@@ -193,7 +195,9 @@ def message_hook_node(config: AppConfig) -> RunnableLike:
     message_hooks: Sequence[Callable[..., Any]] = create_hooks(config.app.message_hooks)
 
     @mlflow.trace()
-    def message_hook(state: IncomingState, runtime: Runtime[Context]) -> SharedState:
+    async def message_hook(
+        state: IncomingState, runtime: Runtime[Context]
+    ) -> SharedState:
         logger.debug("Running message validation")
         response: dict[str, Any] = {"is_valid": True, "message_error": None}
 

dao_ai/tools/agent.py

@@ -40,16 +40,16 @@ def create_agent_endpoint_tool(
 
     doc: str = description + "\n" + doc_signature
 
-    def agent_endpoint(prompt: str) -> AIMessage:
+    async def agent_endpoint(prompt: str) -> AIMessage:
         model: LanguageModelLike = llm.as_chat_model()
         messages: Sequence[BaseMessage] = [HumanMessage(content=prompt)]
-        response: AIMessage = model.invoke(messages)
+        response: AIMessage = await model.ainvoke(messages)
         return response
 
     name: str = name if name else agent_endpoint.__name__
 
     structured_tool: StructuredTool = StructuredTool.from_function(
-        func=agent_endpoint, name=name, description=doc, parse_docstring=False
+        coroutine=agent_endpoint, name=name, description=doc, parse_docstring=False
     )
 
     return structured_tool

dao_ai/tools/genie.py

@@ -62,14 +62,16 @@ def create_genie_tool(
 
     doc: str = description + "\n" + doc_signature
 
-    def genie_tool(question: str) -> GenieResponse:
+    async def genie_tool(question: str) -> GenieResponse:
+        # Use sync API for now since Genie doesn't support async yet
+        # Can be easily updated to await when Genie gets async support
         response: GenieResponse = genie.ask_question(question)
         return response
 
     name: str = name if name else genie_tool.__name__
 
     structured_tool: StructuredTool = StructuredTool.from_function(
-        func=genie_tool, name=name, description=doc, parse_docstring=False
+        coroutine=genie_tool, name=name, description=doc, parse_docstring=False
     )
 
     return structured_tool

dao_ai/tools/human_in_the_loop.py

@@ -50,7 +50,7 @@ def add_human_in_the_loop(
     logger.debug(f"Wrapping tool {tool} with human-in-the-loop functionality")
 
     @create_tool(tool.name, description=tool.description, args_schema=tool.args_schema)
-    def call_tool_with_interrupt(config: RunnableConfig, **tool_input) -> Any:
+    async def call_tool_with_interrupt(config: RunnableConfig, **tool_input) -> Any:
         logger.debug(f"call_tool_with_interrupt: {tool.name} with input: {tool_input}")
         request: HumanInterrupt = {
             "action_request": {
@@ -66,10 +66,10 @@ def add_human_in_the_loop(
         logger.debug(f"Human interrupt response: {response}")
 
         if response["type"] == "accept":
-            tool_response = tool.invoke(tool_input, config=config)
+            tool_response = await tool.ainvoke(tool_input, config=config)
         elif response["type"] == "edit":
             tool_input = response["args"]["args"]
-            tool_response = tool.invoke(tool_input, config=config)
+            tool_response = await tool.ainvoke(tool_input, config=config)
         elif response["type"] == "response":
             user_feedback = response["args"]
             tool_response = user_feedback

dao_ai/tools/mcp.py

@@ -77,6 +77,8 @@ def create_mcp_tools(
             logger.error(f"Failed to list MCP tools: {e}")
             return []
 
+    # Note: This still needs to run sync during tool creation/registration
+    # The actual tool execution will be async
     try:
         mcp_tools: list | ListToolsResult = asyncio.run(_list_mcp_tools())
         if isinstance(mcp_tools, ListToolsResult):
@@ -96,22 +98,19 @@ def create_mcp_tools(
             description=mcp_tool.description or f"MCP tool: {mcp_tool.name}",
             args_schema=mcp_tool.inputSchema,
         )
-        def tool_wrapper(**kwargs):
+        async def tool_wrapper(**kwargs):
             """Execute MCP tool with fresh session and authentication."""
             logger.debug(f"Invoking MCP tool {mcp_tool.name} with fresh session")
 
-            async def _invoke():
-                connection = _create_fresh_connection()
-                client = MultiServerMCPClient({function.name: connection})
+            connection = _create_fresh_connection()
+            client = MultiServerMCPClient({function.name: connection})
 
-                try:
-                    async with client.session(function.name) as session:
-                        return await session.call_tool(mcp_tool.name, kwargs)
-                except Exception as e:
-                    logger.error(f"MCP tool {mcp_tool.name} failed: {e}")
-                    raise
-
-            return asyncio.run(_invoke())
+            try:
+                async with client.session(function.name) as session:
+                    return await session.call_tool(mcp_tool.name, kwargs)
+            except Exception as e:
+                logger.error(f"MCP tool {mcp_tool.name} failed: {e}")
+                raise
 
         return as_human_in_the_loop(tool_wrapper, function)
 

{dao_ai-0.0.16.dist-info → dao_ai-0.0.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.0.16
+Version: 0.0.17
 Summary: DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML.
 Project-URL: Homepage, https://github.com/natefleming/dao-ai
 Project-URL: Documentation, https://natefleming.github.io/dao-ai

{dao_ai-0.0.16.dist-info → dao_ai-0.0.17.dist-info}/RECORD

@@ -4,11 +4,11 @@ dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
 dao_ai/chat_models.py,sha256=uhwwOTeLyHWqoTTgHrs4n5iSyTwe4EQcLKnh3jRxPWI,8626
 dao_ai/cli.py,sha256=Aez2TQW3Q8Ho1IaIkRggt0NevDxAAVPjXkePC5GPJF0,20429
 dao_ai/config.py,sha256=VyHqkW1UMQJ0fzyme1fV_3fi_6wDmKRQeCrx881fDQ4,45173
-dao_ai/graph.py,sha256=8A-qbZAC9BgDitWjxdGFmldKgSHMeoX3z6j3fwmr5Y0,7830
+dao_ai/graph.py,sha256=kXaGLGFVekDWqm-AHzti6LmrXnyi99VQ-AdCGuNb_xM,7831
 dao_ai/guardrails.py,sha256=-Qh0f_2Db9t4Nbrrx9FM7tnpqShjMoyxepZ0HByItfU,4027
 dao_ai/messages.py,sha256=tRZQTeb5YFKu8cm1xeaCkKhidq-0tdzncNEzVePvits,6806
-dao_ai/models.py,sha256=wME6hS7w2BTS4SUlhEYuCExMOYStnWk-YWsyatAHYyY,11836
-dao_ai/nodes.py,sha256=nusBk8tBLY4JlkmzObCKwDe5JR11A8XzYcenC-yS43o,8406
+dao_ai/models.py,sha256=Zf5Rqus5xcdpxSvuLlDau4JM1j1fF9v_MnQ7HW4BXU4,13862
+dao_ai/nodes.py,sha256=SSuFNTXOdFaKg_aX-yUkQO7fM9wvNGu14lPXKDapU1U,8461
 dao_ai/prompts.py,sha256=vpmIbWs_szXUgNNDs5Gh2LcxKZti5pHDKSfoClUcgX0,1289
 dao_ai/state.py,sha256=GwbMbd1TWZx1T5iQrEOX6_rpxOitlmyeJ8dMr2o_pag,1031
 dao_ai/types.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -24,17 +24,17 @@ dao_ai/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 dao_ai/providers/base.py,sha256=-fjKypCOk28h6vioPfMj9YZSw_3Kcbi2nMuAyY7vX9k,1383
 dao_ai/providers/databricks.py,sha256=XRPOqwF5SeA9rPAOWMg2gSMC7lw31BI5VI_4K0KIOqo,27931
 dao_ai/tools/__init__.py,sha256=ye6MHaJY7tUnJ8336YJiLxuZr55zDPNdOw6gm7j5jlc,1103
-dao_ai/tools/agent.py,sha256=_XMz6HtrybpVthhRyStADechF6vXLFyK97i01XTBhtw,1868
+dao_ai/tools/agent.py,sha256=WbQnyziiT12TLMrA7xK0VuOU029tdmUBXbUl-R1VZ0Q,1886
 dao_ai/tools/core.py,sha256=Kei33S8vrmvPOAyrFNekaWmV2jqZ-IPS1QDSvU7RZF0,1984
-dao_ai/tools/genie.py,sha256=yG0Y6FHsYJBdhY7TJoKiUF7XCmCBMsPlbASnZhsUN84,2636
-dao_ai/tools/human_in_the_loop.py,sha256=9HpbFQ9NW6R6eu8eNz2hWPekrB0FFUOJyfa5nUzqqNo,3446
-dao_ai/tools/mcp.py,sha256=rUeodCfGFXSvW-byTW9gsO6QI-p5SZkAP8ZFCFzREY0,4400
+dao_ai/tools/genie.py,sha256=GzV5lfDYKmzW_lSLxAsPaTwnzX6GxQOB1UcLaTDqpfY,2787
+dao_ai/tools/human_in_the_loop.py,sha256=IBmQJmpxkdDxnBNyABc_-dZhhsQlTNTkPyUXgkHKIgY,3466
+dao_ai/tools/mcp.py,sha256=auEt_dwv4J26fr5AgLmwmnAsI894-cyuvkvjItzAUxs,4419
 dao_ai/tools/python.py,sha256=XcQiTMshZyLUTVR5peB3vqsoUoAAy8gol9_pcrhddfI,1831
 dao_ai/tools/time.py,sha256=Y-23qdnNHzwjvnfkWvYsE7PoWS1hfeKy44tA7sCnNac,8759
 dao_ai/tools/unity_catalog.py,sha256=PXfLj2EgyQgaXq4Qq3t25AmTC4KyVCF_-sCtg6enens,1404
 dao_ai/tools/vector_search.py,sha256=EDYQs51zIPaAP0ma1D81wJT77GQ-v-cjb2XrFVWfWdg,2621
-dao_ai-0.0.16.dist-info/METADATA,sha256=D7Bus2E6nvTL3BVMF1sKPViSS-6CZpzqw6qRRwNGxxI,41378
-dao_ai-0.0.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dao_ai-0.0.16.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.0.16.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.0.16.dist-info/RECORD,,
+dao_ai-0.0.17.dist-info/METADATA,sha256=G50pidDVsQt5j4T2NhygeeYJyMOdL7fSI_zyMaWIGUo,41378
+dao_ai-0.0.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dao_ai-0.0.17.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.0.17.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.0.17.dist-info/RECORD,,