tatlock/src/core/tool_tracking.py

"""
Tool call tracking.

Tracks which tools are recommended by the Steward versus which tools
are actually used by Tatlock for debugging and analysis.
"""
from typing import Optional

from src.core.logging_config import get_logger

logger = get_logger(__name__)


class ToolCallTracker:
    """
    Tracks tool calls for accuracy analysis.

    Compares Steward's recommendations with Tatlock's actual tool usage
    to measure recommendation accuracy.
    """

    def __init__(
        self,
        recommended_capabilities: list[str],
        conversation_id: Optional[str] = None
    ):
        """
        Initialize tool call tracker.

        Args:
            recommended_capabilities: List of capability names recommended by Steward
            conversation_id: Optional conversation ID for tracking
        """
        self.recommended_capabilities = set(recommended_capabilities)
        self.actual_calls: dict[str, list[float]] = {}  # tool_name -> [durations]
        self.conversation_id = conversation_id

        logger.debug(
            "tool_tracker_initialized",
            recommended=list(self.recommended_capabilities),
            conversation_id=conversation_id,
        )

    def _extract_capability(self, tool_name: str) -> str:
        """
        Extract capability name from tool name.

        Tool names like 'delegate_to_librarian' map to capability 'librarian'.
        """
        if tool_name.startswith("delegate_to_"):
            return tool_name.replace("delegate_to_", "")
        return tool_name

    def log_call(self, message: str):
        """Log a tool call message (for UI display)."""
        logger.debug("tool_call_message", message=message)

    async def track_call(self, tool_name: str, duration: float):
        """
        Record a tool call with timing.

        Args:
            tool_name: Name of the tool that was called
            duration: Duration of the call in seconds
        """
        # Record the call
        if tool_name not in self.actual_calls:
            self.actual_calls[tool_name] = []
        self.actual_calls[tool_name].append(duration)

        # Check if tool was recommended (normalize tool name to capability)
        capability = self._extract_capability(tool_name)
        was_recommended = capability in self.recommended_capabilities

        if not was_recommended:
            logger.warning(
                "tool_call_not_recommended",
                tool_name=tool_name,
                duration=duration,
                recommended=list(self.recommended_capabilities),
            )

        logger.debug(
            "tool_call_tracked",
            tool_name=tool_name,
            duration=duration,
            was_recommended=was_recommended,
        )

    async def finalize(self):
        """
        Finalize tracking and log unused recommended tools.

        Called after Tatlock completes its response to identify
        tools that were recommended but never used.
        """
        # Normalize actual tool names to capabilities for comparison
        used_capabilities = {
            self._extract_capability(tool) for tool in self.actual_calls.keys()
        }
        # Find tools that were recommended but not used
        unused_tools = self.recommended_capabilities - used_capabilities

        if unused_tools:
            logger.info(
                "recommended_tools_unused",
                unused=list(unused_tools),
                used=list(self.actual_calls.keys()),
                conversation_id=self.conversation_id,
            )

        # Log summary
        total_calls = sum(len(durations) for durations in self.actual_calls.values())
        logger.info(
            "tool_tracking_finalized",
            total_calls=total_calls,
            unique_tools_used=len(self.actual_calls),
            recommended_count=len(self.recommended_capabilities),
            unused_count=len(unused_tools),
        )

    def get_summary(self) -> dict:
        """
        Get tracking summary for debugging.

        Returns:
            Dict with tracking statistics
        """
        total_calls = sum(len(durations) for durations in self.actual_calls.values())
        # Normalize actual tool names to capabilities for comparison
        used_capabilities = {
            self._extract_capability(tool) for tool in self.actual_calls.keys()
        }
        unused = self.recommended_capabilities - used_capabilities

        return {
            "recommended_capabilities": list(self.recommended_capabilities),
            "tools_used": list(self.actual_calls.keys()),
            "tools_unused": list(unused),
            "total_calls": total_calls,
            "accuracy": {
                "recommended_and_used": len(
                    self.recommended_capabilities & used_capabilities
                ),
                "recommended_but_unused": len(unused),
                "not_recommended_but_used": len(
                    used_capabilities - self.recommended_capabilities
                ),
            },
        }