dao-ai 0.1.20__py3-none-any.whl → 0.1.21__py3-none-any.whl

dao_ai/middleware/tool_call_observability.py

@@ -0,0 +1,227 @@
+"""
+Tool call observability middleware for DAO AI agents.
+
+This middleware logs detailed information about tool call patterns to help
+diagnose whether tools are being called in parallel or sequentially.
+
+Example YAML config::
+
+    middleware:
+      - name: dao_ai.middleware.tool_call_observability.\
+              create_tool_call_observability_middleware
+        args:
+          log_level: INFO
+          include_args: false
+          track_timing: true
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+from langchain.agents.middleware import AgentMiddleware
+from langchain_core.messages import AIMessage, BaseMessage
+from langgraph.runtime import Runtime
+from loguru import logger
+
+from dao_ai.state import AgentState, Context
+
+__all__ = [
+    "ToolCallObservabilityMiddleware",
+    "create_tool_call_observability_middleware",
+]
+
+
+class ToolCallObservabilityMiddleware(AgentMiddleware[AgentState, Context]):
+    """
+    Middleware that observes and logs tool call patterns.
+
+    Tracks:
+    - Number of tool calls per model response
+    - Whether tools are called in parallel (multiple per response) or sequentially
+    - Tool execution timing
+    - Cumulative statistics across the conversation
+    """
+
+    def __init__(
+        self,
+        log_level: str = "INFO",
+        include_args: bool = False,
+        track_timing: bool = True,
+    ):
+        """
+        Initialize the observability middleware.
+
+        Args:
+            log_level: Logging level ("DEBUG", "INFO", "WARNING")
+            include_args: Whether to log tool call arguments (may be verbose)
+            track_timing: Whether to track execution timing
+        """
+        self.log_level = log_level.upper()
+        self.include_args = include_args
+        self.track_timing = track_timing
+
+        # Statistics tracking (per-run, reset on before_agent)
+        self._total_model_calls = 0
+        self._total_tool_calls = 0
+        self._parallel_batches = 0  # Responses with 2+ tool calls
+        self._sequential_calls = 0  # Responses with exactly 1 tool call
+        self._tool_execution_times: dict[str, list[float]] = {}
+        self._run_start_time: float | None = None
+
+    def _log(self, message: str, **kwargs: Any) -> None:
+        """Log at the configured level."""
+        log_fn = getattr(logger, self.log_level.lower(), logger.info)
+        log_fn(message, **kwargs)
+
+    def before_agent(
+        self,
+        state: AgentState,
+        runtime: Runtime[Context],
+    ) -> dict[str, Any] | None:
+        """Reset statistics at the start of each agent run."""
+        self._total_model_calls = 0
+        self._total_tool_calls = 0
+        self._parallel_batches = 0
+        self._sequential_calls = 0
+        self._tool_execution_times = {}
+        self._run_start_time = time.time()
+
+        self._log(
+            "Tool call observability: Agent run started",
+            thread_id=runtime.context.thread_id if runtime.context else None,
+        )
+        return None
+
+    def after_model(
+        self,
+        state: AgentState,
+        runtime: Runtime[Context],
+    ) -> dict[str, Any] | None:
+        """Analyze tool calls in model response."""
+        self._total_model_calls += 1
+
+        # Extract tool calls from state messages (last message is model response)
+        messages: list[BaseMessage] = state.get("messages", [])
+
+        # Look at the last message which should be the model's response
+        for msg in messages[-1:]:
+            if isinstance(msg, AIMessage) and msg.tool_calls:
+                num_tool_calls = len(msg.tool_calls)
+                self._total_tool_calls += num_tool_calls
+
+                if num_tool_calls > 1:
+                    self._parallel_batches += 1
+                    self._log(
+                        "PARALLEL tool calls detected",
+                        num_tools=num_tool_calls,
+                        tool_names=[tc["name"] for tc in msg.tool_calls],
+                        model_call_number=self._total_model_calls,
+                    )
+
+                    if self.include_args:
+                        for tc in msg.tool_calls:
+                            self._log(
+                                f"  Tool: {tc['name']}",
+                                args=tc.get("args", {}),
+                            )
+                elif num_tool_calls == 1:
+                    self._sequential_calls += 1
+                    tc = msg.tool_calls[0]
+                    log_kwargs: dict[str, Any] = {
+                        "tool_name": tc["name"],
+                        "model_call_number": self._total_model_calls,
+                    }
+                    if self.include_args:
+                        log_kwargs["args"] = tc.get("args", {})
+                    self._log("Sequential tool call", **log_kwargs)
+
+        return None
+
+    def after_agent(
+        self,
+        state: AgentState,
+        runtime: Runtime[Context],
+    ) -> dict[str, Any] | None:
+        """Log final statistics."""
+        total_time = time.time() - self._run_start_time if self._run_start_time else 0
+
+        # Calculate parallelism ratio
+        total_responses_with_tools = self._parallel_batches + self._sequential_calls
+        parallelism_ratio = (
+            self._parallel_batches / total_responses_with_tools * 100
+            if total_responses_with_tools > 0
+            else 0
+        )
+
+        # Calculate average tool times
+        avg_times = {
+            name: round(sum(times) / len(times) * 1000, 2)
+            for name, times in self._tool_execution_times.items()
+        }
+
+        self._log(
+            "Tool Call Observability Summary",
+            total_model_calls=self._total_model_calls,
+            total_tool_calls=self._total_tool_calls,
+            parallel_batches=self._parallel_batches,
+            sequential_calls=self._sequential_calls,
+            parallelism_ratio=f"{parallelism_ratio:.1f}%",
+            total_time_ms=round(total_time * 1000, 2),
+            avg_tool_times_ms=avg_times,
+        )
+
+        # Log verdict
+        if self._parallel_batches > 0:
+            logger.success(
+                f"Parallel tool calling IS happening: "
+                f"{self._parallel_batches} batches with multiple tools"
+            )
+        elif self._sequential_calls > 0:
+            logger.warning(
+                f"All tool calls are SEQUENTIAL: "
+                f"{self._sequential_calls} single-tool responses. "
+                f"Consider prompt engineering to encourage parallel calls."
+            )
+
+        return None
+
+
+def create_tool_call_observability_middleware(
+    log_level: str = "INFO",
+    include_args: bool = False,
+    track_timing: bool = True,
+) -> ToolCallObservabilityMiddleware:
+    """
+    Factory function to create tool call observability middleware.
+
+    Args:
+        log_level: Logging level ("DEBUG", "INFO", "WARNING")
+        include_args: Whether to log tool call arguments
+        track_timing: Whether to track execution timing
+
+    Returns:
+        ToolCallObservabilityMiddleware instance
+
+    Example YAML config::
+
+        middleware:
+          - name: dao_ai.middleware.tool_call_observability.\
+                  create_tool_call_observability_middleware
+            args:
+              log_level: INFO
+              include_args: false
+              track_timing: true
+    """
+    logger.debug(
+        "Creating tool call observability middleware",
+        log_level=log_level,
+        include_args=include_args,
+        track_timing=track_timing,
+    )
+    return ToolCallObservabilityMiddleware(
+        log_level=log_level,
+        include_args=include_args,
+        track_timing=track_timing,
+    )

dao_ai/utils.py

@@ -152,10 +152,14 @@ def dao_ai_version() -> str:
             return "dev"
 
 
-def get_installed_packages() -> dict[str, str]:
-    """Get all installed packages with versions"""
+def get_installed_packages() -> list[str]:
+    """Get all installed packages with versions.
 
-    packages: Sequence[str] = [
+    Returns a list of pip requirement strings for packages used by dao-ai.
+    This is used for MLflow model logging to ensure all dependencies are captured.
+    """
+
+    packages: list[str] = [
         f"databricks-agents=={version('databricks-agents')}",
         f"databricks-langchain[memory]=={version('databricks-langchain')}",
         f"databricks-mcp=={version('databricks-mcp')}",

{dao_ai-0.1.20.dist-info → dao_ai-0.1.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.1.20
+Version: 0.1.21
 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.1.20.dist-info → dao_ai-0.1.21.dist-info}/RECORD

@@ -1,7 +1,7 @@
 dao_ai/__init__.py,sha256=18P98ExEgUaJ1Byw440Ct1ty59v6nxyWtc5S6Uq2m9Q,1062
 dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
 dao_ai/cli.py,sha256=7hVCC8mn9S3c4wW-eRt-WoFKzV1wPdJVAeNhkyhfUGc,53251
-dao_ai/config.py,sha256=08eHf7HzXNEsFY1_-lCESCtC-RhEZy9nikTiWPmc25A,156770
+dao_ai/config.py,sha256=alAPiAtglojEIBUm4dwrpnirkRbgDPBS-6DZAKCxyuE,159798
 dao_ai/evaluation.py,sha256=4dveWDwFnUxaybswr0gag3ydZ5RGVCTRaiE3eKLClD4,18161
 dao_ai/graph.py,sha256=1-uQlo7iXZQTT3uU8aYu0N5rnhw5_g_2YLwVsAs6M-U,1119
 dao_ai/logging.py,sha256=lYy4BmucCHvwW7aI3YQkQXKJtMvtTnPDu9Hnd7_O4oc,1556
@@ -11,7 +11,7 @@ dao_ai/nodes.py,sha256=H7_C0ev0TpS5KWkGZD6eE4Wn6ouBwnN5HgYUyBeKe0A,10881
 dao_ai/optimization.py,sha256=phK6t4wYmWPObCjGUBHdZzsaFXGhQOjhAek2bAEfwXo,22971
 dao_ai/state.py,sha256=ifDTAC7epdowk3Z1CP3Xqw4uH2dIxQEVF3C747dA8yI,6436
 dao_ai/types.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dao_ai/utils.py,sha256=ImgH0jnHPCK2AR7KcueG_Zb7kltcBzTw78ujDUpARIE,17184
+dao_ai/utils.py,sha256=nQ5U92IU-JWLIikBg04s4O5gHUNX3j_Dy3YfAnT_ccM,17340
 dao_ai/vector_search.py,sha256=PfmT2PDMymk-3dTm2uszlOZNyHyiDge--imxKpKJRsY,4440
 dao_ai/apps/__init__.py,sha256=RLuhZf4gQ4pemwKDz1183aXib8UfaRhwfKvRx68GRlM,661
 dao_ai/apps/handlers.py,sha256=6-IhhklHSPnS8aqKp155wPaSnYWTU1BSOPwbdWYBkFU,3594
@@ -20,16 +20,16 @@ dao_ai/apps/resources.py,sha256=5l6UxfMq6uspOql-HNDyUikfqRAa9eH_TiJHrGgMb6s,4002
 dao_ai/apps/server.py,sha256=neWbVnC2z9f-tJZBnho70FytNDEVOdOM1YngoGc5KHI,1264
 dao_ai/genie/__init__.py,sha256=UpSvP6gZO8H-eAPokYpkshvFxYD4ETYZHz-pRPoK2sI,2786
 dao_ai/genie/core.py,sha256=eKZo4pRagwI6QglAXqHYjUaC3AicmGMiy_9WIuZ-tzw,9119
-dao_ai/genie/cache/__init__.py,sha256=hMD2q7VRX8fQmkd4rXhxGJdRIVbDqLh2DyjBu0CdM_k,2129
+dao_ai/genie/cache/__init__.py,sha256=gUKBcbSMicalQj6sc_bvuBUeVUUTZ90xIMytf9Bk_m8,2242
 dao_ai/genie/cache/base.py,sha256=nbWl-KTstUPGagdUtO8xtVUSosuqkaNc_hx-PgT1ROo,7155
 dao_ai/genie/cache/core.py,sha256=48sDY7dbrsmflb96OFEE8DYarNB6zyiFxZQG-qfhXD4,2537
 dao_ai/genie/cache/lru.py,sha256=dWoNottME8y6y_OKnQZ1xH4NmQxk2PdXvUgKcdzjlxI,19935
-dao_ai/genie/cache/optimization.py,sha256=8tyMrthHCxdOdvqDwGnGAvRNCWk9e0nIb3dJkE05t-4,30167
-dao_ai/genie/cache/context_aware/__init__.py,sha256=EaW2CmWxb0CDNcm3pKL6Vyaw2KuwanKV1LiFVlT2M-s,1197
-dao_ai/genie/cache/context_aware/base.py,sha256=o3S3f0B0jCB7vjb2XKWm5hFcO6v-jNYc6ZELyBkeBrg,40452
-dao_ai/genie/cache/context_aware/in_memory.py,sha256=Z-ZrMTVD1pjC55irhodBWsrs6D2h3Sv-Iifk0Cr1mRo,22316
+dao_ai/genie/cache/context_aware/__init__.py,sha256=8uxG-0_9SIqjmiNKyPXjn_pOhOi82RrcziljQOYdrwc,1946
+dao_ai/genie/cache/context_aware/base.py,sha256=rxiXqJF5VZFTKEAr_MKqGH1hcqF5rxeqsC75Z7nlrWk,42819
+dao_ai/genie/cache/context_aware/in_memory.py,sha256=9V7uDyXYe3A8GU3rxJq4icbt3IxSTuHhoPN6sJxOEDU,26809
+dao_ai/genie/cache/context_aware/optimization.py,sha256=foUCjK5g0aRUH1wXWoPIVaGC6DqSstTrsP5naQGpw7U,31758
 dao_ai/genie/cache/context_aware/persistent.py,sha256=cpn25Go6ZyN65lY_vh5cWcuqr__nNH7RPWJA_LP7wcE,28154
-dao_ai/genie/cache/context_aware/postgres.py,sha256=cakJuba3dFjKoEHlpBYgmr3pEuM88dM8ambgL6Vq_54,46403
+dao_ai/genie/cache/context_aware/postgres.py,sha256=HGadN7FmSXbDXc8itq6vz1tSbhyklxOnWB3qDegrVKc,53087
 dao_ai/hooks/__init__.py,sha256=uA4DQdP9gDf4SyNjNx9mWPoI8UZOcTyFsCXV0NraFvQ,463
 dao_ai/hooks/core.py,sha256=yZAfRfB0MyMo--uwGr4STtVxxen5s4ZUrNTnR3a3qkA,1721
 dao_ai/memory/__init__.py,sha256=Us3wFehvug_h83m-UJ7OXdq2qZ0e9nHBQE7m5RwoAd8,559
@@ -37,7 +37,7 @@ dao_ai/memory/base.py,sha256=99nfr2UZJ4jmfTL_KrqUlRSCoRxzkZyWyx5WqeUoMdQ,338
 dao_ai/memory/core.py,sha256=38H-JLIyUrRDIECLvpXK3iJlWG35X97E-DTo_4c3Jzc,6317
 dao_ai/memory/databricks.py,sha256=SM6nwLjhSRJO4hLc3GUuht5YydYtTi3BAOae6jPwTm4,14377
 dao_ai/memory/postgres.py,sha256=bSjtvEht0h6jy2ADN2vqISVQDxm_DeM586VDdGaShJQ,23168
-dao_ai/middleware/__init__.py,sha256=Qy8wbvjXF7TrUzi3tWziOwxqsrUcT1rzE3UWd3x5CrU,5108
+dao_ai/middleware/__init__.py,sha256=cmXRtyY4rvyv1p2BJktw-SllI5qZLk_EEeiWJtOkbiM,5358
 dao_ai/middleware/assertions.py,sha256=C1K-TnNZfBEwWouioHCt6c48i1ux9QKfQaX6AzghhgE,27408
 dao_ai/middleware/base.py,sha256=uG2tpdnjL5xY5jCKvb_m3UTBtl4ZC6fJQUkDsQvV8S4,1279
 dao_ai/middleware/context_editing.py,sha256=5rNKqH1phFFQTVW-4nzlVH5cbqomD-HFEIy2Z841D4I,7687
@@ -50,6 +50,7 @@ dao_ai/middleware/model_retry.py,sha256=SlWjAcaEmvj6KBOkjUicChYjhlg7bAJM7-e6KLpH
 dao_ai/middleware/pii.py,sha256=zetfoz1WlJ-V0vjJp37v8NGimXB27EkZfetUHpGCXno,5137
 dao_ai/middleware/summarization.py,sha256=gp2s9uc4DEJat-mWjWEzMaR-zAAeUOXYvu5EEYtqae4,7143
 dao_ai/middleware/tool_call_limit.py,sha256=WQ3NmA3pLo-pNPBmwM7KwkYpT1segEnWqkhgW1xNkCE,6321
+dao_ai/middleware/tool_call_observability.py,sha256=mU52HwF82Yb05TIAYhu5cNkn3jcPGn6cKhsxT3rkS8M,7673
 dao_ai/middleware/tool_retry.py,sha256=QfJ7yTHneME8VtnA88QcmnjXIegSFeJztyngy49wTgM,5568
 dao_ai/middleware/tool_selector.py,sha256=POj72YdzZEiNGfW4AQXPBeVVS1RUBsiG7PBuSENEhe0,4516
 dao_ai/orchestration/__init__.py,sha256=i85CLfRR335NcCFhaXABcMkn6WZfXnJ8cHH4YZsZN0s,1622
@@ -82,8 +83,8 @@ dao_ai/tools/time.py,sha256=tufJniwivq29y0LIffbgeBTIDE6VgrLpmVf8Qr90qjw,9224
 dao_ai/tools/unity_catalog.py,sha256=oBlW6pH-Ne08g60QW9wVi_tyeVYDiecuNoxQbIIFmN8,16515
 dao_ai/tools/vector_search.py,sha256=34uhd58FKHzvcdgHHoACRdZAUJWTaUuPYiwIqBwvGqk,29061
 dao_ai/tools/verifier.py,sha256=ociBVsGkQNyhWS6F6G8x17V7zAQfSuTe4Xcd6Y-7lPE,4975
-dao_ai-0.1.20.dist-info/METADATA,sha256=v-u8GWz1AnJXLm1LfSMrD-cHZYT-nbyU57VNOxG9KWY,16954
-dao_ai-0.1.20.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-dao_ai-0.1.20.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.1.20.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.1.20.dist-info/RECORD,,
+dao_ai-0.1.21.dist-info/METADATA,sha256=v1iHbA1jzdcEUmkmOa0kyh16V7-BMhIx5HkAMWUnRlE,16954
+dao_ai-0.1.21.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dao_ai-0.1.21.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.1.21.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.1.21.dist-info/RECORD,,