kalibr 1.1.3a0__py3-none-any.whl → 1.2.1__py3-none-any.whl

kalibr_crewai/callbacks.py

@@ -60,6 +60,72 @@ def _get_provider_from_model(model: str) -> str:
         return "openai"
 
 
+def _extract_model_from_agent(agent) -> tuple[str, str]:
+    """Extract model name and provider from agent's LLM config.
+
+    Args:
+        agent: CrewAI agent instance
+
+    Returns:
+        Tuple of (model_name, provider)
+    """
+    model_name = "unknown"
+    provider = "openai"
+
+    if not hasattr(agent, "llm"):
+        return model_name, provider
+
+    llm = agent.llm
+
+    # Case 1: LLM is a string like "openai/gpt-4o-mini" or "gpt-4"
+    if isinstance(llm, str):
+        if "/" in llm:
+            parts = llm.split("/", 1)
+            provider = parts[0]
+            model_name = parts[1]
+        else:
+            model_name = llm
+            provider = _get_provider_from_model(llm)
+        return model_name, provider
+
+    # Case 2: LLM has model or model_name attribute
+    if hasattr(llm, "model"):
+        model_name = str(llm.model)
+    elif hasattr(llm, "model_name"):
+        model_name = str(llm.model_name)
+
+    # Parse provider from model string if it contains "/"
+    if "/" in model_name:
+        parts = model_name.split("/", 1)
+        provider = parts[0]
+        model_name = parts[1]
+    else:
+        provider = _get_provider_from_model(model_name)
+
+    return model_name, provider
+
+
+def _calculate_cost(provider: str, model: str, input_tokens: int, output_tokens: int) -> float:
+    """Calculate cost using CostAdapterFactory.
+
+    Args:
+        provider: Provider name (openai, anthropic, etc.)
+        model: Model name
+        input_tokens: Number of input tokens
+        output_tokens: Number of output tokens
+
+    Returns:
+        Cost in USD
+    """
+    if CostAdapterFactory is None:
+        return 0.0
+
+    try:
+        return CostAdapterFactory.compute_cost(provider, model, input_tokens, output_tokens)
+    except Exception:
+        return 0.0
+
+
 class EventBatcher:
     """Shared event batching for callbacks."""
 
@@ -198,6 +264,7 @@ class KalibrAgentCallback:
         service: Service name
         workflow_id: Workflow identifier
         metadata: Additional metadata for all events
+        agent: Optional agent reference for model extraction
 
     Usage:
         from kalibr_crewai import KalibrAgentCallback
@@ -210,6 +277,7 @@ class KalibrAgentCallback:
             goal="Find information",
             step_callback=callback,
         )
+        callback.set_agent(agent)  # Set agent reference for model extraction
     """
 
     def __init__(
@@ -221,6 +289,7 @@ class KalibrAgentCallback:
         service: Optional[str] = None,
         workflow_id: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
+        agent: Optional[Any] = None,
     ):
         self.api_key = api_key or os.getenv("KALIBR_API_KEY", "")
         self.endpoint = endpoint or os.getenv(
@@ -232,6 +301,7 @@ class KalibrAgentCallback:
         self.service = service or os.getenv("KALIBR_SERVICE", "crewai-app")
         self.workflow_id = workflow_id or os.getenv("KALIBR_WORKFLOW_ID", "default-workflow")
         self.default_metadata = metadata or {}
+        self._agent = agent
 
         # Get shared batcher
         self._batcher = EventBatcher.get_instance(
@@ -244,6 +314,14 @@ class KalibrAgentCallback:
         self._agent_span_id: Optional[str] = None
         self._step_count: int = 0
 
+    def set_agent(self, agent: Any) -> None:
+        """Set the agent reference for model extraction.
+
+        Args:
+            agent: CrewAI agent instance
+        """
+        self._agent = agent
+
     def __call__(self, step_output: Any) -> None:
         """Called after each agent step.
 
@@ -271,6 +349,12 @@ class KalibrAgentCallback:
 
         span_id = str(uuid.uuid4())
 
+        # Extract model from agent if available
+        model_name = "unknown"
+        provider = "openai"
+        if self._agent:
+            model_name, provider = _extract_model_from_agent(self._agent)
+
         # Extract step information
         step_type = "agent_step"
         operation = "agent_step"
@@ -307,8 +391,11 @@ class KalibrAgentCallback:
             output_text = str(step_output)
 
         # Count tokens
-        input_tokens = _count_tokens(tool_input or "", "gpt-4")
-        output_tokens = _count_tokens(output_text, "gpt-4")
+        input_tokens = _count_tokens(tool_input or "", model_name)
+        output_tokens = _count_tokens(output_text, model_name)
+
+        # Calculate cost using CostAdapterFactory
+        cost_usd = _calculate_cost(provider, model_name, input_tokens, output_tokens)
 
         # Build event
         event = {
@@ -318,9 +405,9 @@ class KalibrAgentCallback:
             "parent_span_id": self._agent_span_id,
             "tenant_id": self.tenant_id,
             "workflow_id": self.workflow_id,
-            "provider": "crewai",
-            "model_id": "agent",
-            "model_name": "crewai-agent",
+            "provider": provider,
+            "model_id": model_name,
+            "model_name": model_name,
             "operation": operation,
             "endpoint": operation,
             "duration_ms": 0,  # Step timing not available
@@ -328,8 +415,8 @@ class KalibrAgentCallback:
             "input_tokens": input_tokens,
             "output_tokens": output_tokens,
             "total_tokens": input_tokens + output_tokens,
-            "cost_usd": 0.0,  # Cost tracked at LLM level
-            "total_cost_usd": 0.0,
+            "cost_usd": cost_usd,
+            "total_cost_usd": cost_usd,
             "status": status,
             "timestamp": now.isoformat(),
             "ts_start": now.isoformat(),
@@ -376,6 +463,7 @@ class KalibrTaskCallback:
         service: Service name
         workflow_id: Workflow identifier
         metadata: Additional metadata for all events
+        agent: Optional agent reference for model extraction
 
     Usage:
         from kalibr_crewai import KalibrTaskCallback
@@ -388,6 +476,7 @@ class KalibrTaskCallback:
             agent=my_agent,
             callback=callback,
         )
+        callback.set_agent(my_agent)  # Set agent reference for model extraction
     """
 
     def __init__(
@@ -399,6 +488,7 @@ class KalibrTaskCallback:
         service: Optional[str] = None,
         workflow_id: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
+        agent: Optional[Any] = None,
     ):
         self.api_key = api_key or os.getenv("KALIBR_API_KEY", "")
         self.endpoint = endpoint or os.getenv(
@@ -410,6 +500,7 @@ class KalibrTaskCallback:
         self.service = service or os.getenv("KALIBR_SERVICE", "crewai-app")
         self.workflow_id = workflow_id or os.getenv("KALIBR_WORKFLOW_ID", "default-workflow")
         self.default_metadata = metadata or {}
+        self._agent = agent
 
         # Get shared batcher
         self._batcher = EventBatcher.get_instance(
@@ -421,6 +512,14 @@ class KalibrTaskCallback:
         self._trace_id: Optional[str] = None
         self._crew_span_id: Optional[str] = None
 
+    def set_agent(self, agent: Any) -> None:
+        """Set the agent reference for model extraction.
+
+        Args:
+            agent: CrewAI agent instance
+        """
+        self._agent = agent
+
     def __call__(self, task_output: Any) -> None:
         """Called when task completes.
 
@@ -467,9 +566,18 @@ class KalibrTaskCallback:
         if hasattr(task_output, "agent"):
             agent_role = str(task_output.agent)
 
+        # Extract model from agent if available
+        model_name = "unknown"
+        provider = "openai"
+        if self._agent:
+            model_name, provider = _extract_model_from_agent(self._agent)
+
         # Token counting
-        input_tokens = _count_tokens(description, "gpt-4")
-        output_tokens = _count_tokens(raw_output, "gpt-4")
+        input_tokens = _count_tokens(description, model_name)
+        output_tokens = _count_tokens(raw_output, model_name)
+
+        # Calculate cost using CostAdapterFactory
+        cost_usd = _calculate_cost(provider, model_name, input_tokens, output_tokens)
 
         # Build operation name from description
         operation = "task_complete"
@@ -486,9 +594,9 @@ class KalibrTaskCallback:
             "parent_span_id": self._crew_span_id,
             "tenant_id": self.tenant_id,
             "workflow_id": self.workflow_id,
-            "provider": "crewai",
-            "model_id": "task",
-            "model_name": agent_role,
+            "provider": provider,
+            "model_id": model_name,
+            "model_name": model_name,
             "operation": operation,
             "endpoint": "task_complete",
             "duration_ms": 0,  # Task timing not available in callback
@@ -496,8 +604,8 @@ class KalibrTaskCallback:
             "input_tokens": input_tokens,
             "output_tokens": output_tokens,
             "total_tokens": input_tokens + output_tokens,
-            "cost_usd": 0.0,  # Cost tracked at LLM level
-            "total_cost_usd": 0.0,
+            "cost_usd": cost_usd,
+            "total_cost_usd": cost_usd,
             "status": "success",
             "timestamp": now.isoformat(),
             "ts_start": now.isoformat(),

kalibr_crewai/instrumentor.py

@@ -21,6 +21,72 @@ except ImportError:
     CostAdapterFactory = None
 
 
+def _extract_model_from_agent(agent) -> tuple[str, str]:
+    """Extract model name and provider from agent's LLM config.
+
+    Args:
+        agent: CrewAI agent instance
+
+    Returns:
+        Tuple of (model_name, provider)
+    """
+    model_name = "unknown"
+    provider = "openai"
+
+    if not hasattr(agent, "llm"):
+        return model_name, provider
+
+    llm = agent.llm
+
+    # Case 1: LLM is a string like "openai/gpt-4o-mini" or "gpt-4"
+    if isinstance(llm, str):
+        if "/" in llm:
+            parts = llm.split("/", 1)
+            provider = parts[0]
+            model_name = parts[1]
+        else:
+            model_name = llm
+            provider = _get_provider_from_model(llm)
+        return model_name, provider
+
+    # Case 2: LLM has model or model_name attribute
+    if hasattr(llm, "model"):
+        model_name = str(llm.model)
+    elif hasattr(llm, "model_name"):
+        model_name = str(llm.model_name)
+
+    # Parse provider from model string if it contains "/"
+    if "/" in model_name:
+        parts = model_name.split("/", 1)
+        provider = parts[0]
+        model_name = parts[1]
+    else:
+        provider = _get_provider_from_model(model_name)
+
+    return model_name, provider
+
+
+def _calculate_cost(provider: str, model: str, input_tokens: int, output_tokens: int) -> float:
+    """Calculate cost using CostAdapterFactory.
+
+    Args:
+        provider: Provider name (openai, anthropic, etc.)
+        model: Model name
+        input_tokens: Number of input tokens
+        output_tokens: Number of output tokens
+
+    Returns:
+        Cost in USD
+    """
+    if CostAdapterFactory is None:
+        return 0.0
+
+    try:
+        return CostAdapterFactory.compute_cost(provider, model, input_tokens, output_tokens)
+    except Exception:
+        return 0.0
+
+
 class KalibrCrewAIInstrumentor:
     """Auto-instrumentation for CrewAI.
 
@@ -341,6 +407,9 @@ class KalibrCrewAIInstrumentor:
             role = getattr(agent_self, "role", "unknown")
             goal = getattr(agent_self, "goal", "")
 
+            # Extract model from agent's LLM config
+            model_name, provider = _extract_model_from_agent(agent_self)
+
             # Get task info
             task_description = ""
             if hasattr(task, "description"):
@@ -370,8 +439,11 @@ class KalibrCrewAIInstrumentor:
                     output_preview = str(result)[:500]
 
                 # Token estimation
-                input_tokens = _count_tokens(task_description + goal, "gpt-4")
-                output_tokens = _count_tokens(output_preview or "", "gpt-4")
+                input_tokens = _count_tokens(task_description + goal, model_name)
+                output_tokens = _count_tokens(output_preview or "", model_name)
+
+                # Calculate cost using CostAdapterFactory
+                cost_usd = _calculate_cost(provider, model_name, input_tokens, output_tokens)
 
                 event = {
                     "schema_version": "1.0",
@@ -380,9 +452,9 @@ class KalibrCrewAIInstrumentor:
                     "parent_span_id": None,
                     "tenant_id": instrumentor.tenant_id,
                     "workflow_id": instrumentor.workflow_id,
-                    "provider": "crewai",
-                    "model_id": "agent",
-                    "model_name": role,
+                    "provider": provider,
+                    "model_id": model_name,
+                    "model_name": model_name,
                     "operation": f"agent:{role}",
                     "endpoint": "agent.execute_task",
                     "duration_ms": duration_ms,
@@ -390,8 +462,8 @@ class KalibrCrewAIInstrumentor:
                     "input_tokens": input_tokens,
                     "output_tokens": output_tokens,
                     "total_tokens": input_tokens + output_tokens,
-                    "cost_usd": 0.0,
-                    "total_cost_usd": 0.0,
+                    "cost_usd": cost_usd,
+                    "total_cost_usd": cost_usd,
                     "status": status,
                     "error_type": error_type,
                     "error_message": error_message,
@@ -430,6 +502,13 @@ class KalibrCrewAIInstrumentor:
             description = getattr(task_self, "description", "")
             expected_output = getattr(task_self, "expected_output", "")
 
+            # Try to extract model from task's agent
+            model_name = "unknown"
+            provider = "openai"
+            agent = getattr(task_self, "agent", None)
+            if agent:
+                model_name, provider = _extract_model_from_agent(agent)
+
             status = "success"
             error_type = None
             error_message = None
@@ -456,8 +535,11 @@ class KalibrCrewAIInstrumentor:
                     else:
                         output_preview = str(result)[:500]
 
-                input_tokens = _count_tokens(description, "gpt-4")
-                output_tokens = _count_tokens(output_preview or "", "gpt-4")
+                input_tokens = _count_tokens(description, model_name)
+                output_tokens = _count_tokens(output_preview or "", model_name)
+
+                # Calculate cost using CostAdapterFactory
+                cost_usd = _calculate_cost(provider, model_name, input_tokens, output_tokens)
 
                 event = {
                     "schema_version": "1.0",
@@ -466,9 +548,9 @@ class KalibrCrewAIInstrumentor:
                     "parent_span_id": None,
                     "tenant_id": instrumentor.tenant_id,
                     "workflow_id": instrumentor.workflow_id,
-                    "provider": "crewai",
-                    "model_id": "task",
-                    "model_name": "crewai-task",
+                    "provider": provider,
+                    "model_id": model_name,
+                    "model_name": model_name,
                     "operation": f"task:{description[:30]}..." if len(description) > 30 else f"task:{description}",
                     "endpoint": "task.execute_sync",
                     "duration_ms": duration_ms,
@@ -476,8 +558,8 @@ class KalibrCrewAIInstrumentor:
                     "input_tokens": input_tokens,
                     "output_tokens": output_tokens,
                     "total_tokens": input_tokens + output_tokens,
-                    "cost_usd": 0.0,
-                    "total_cost_usd": 0.0,
+                    "cost_usd": cost_usd,
+                    "total_cost_usd": cost_usd,
                     "status": status,
                     "error_type": error_type,
                     "error_message": error_message,

kalibr_langchain/__init__.py

@@ -29,7 +29,7 @@ Usage:
 
 Environment Variables:
     KALIBR_API_KEY: API key for authentication
-    KALIBR_ENDPOINT: Backend endpoint URL
+    KALIBR_COLLECTOR_URL: Backend endpoint URL
     KALIBR_TENANT_ID: Tenant identifier
     KALIBR_ENVIRONMENT: Environment (prod/staging/dev)
     KALIBR_SERVICE: Service name

kalibr_langchain/callback.py

@@ -29,6 +29,8 @@ try:
 except ImportError:
     CostAdapterFactory = None
 
+from kalibr.context import get_goal
+
 # Import tiktoken for token counting
 try:
     import tiktoken
@@ -288,6 +290,25 @@ class KalibrCallbackHandler(BaseCallbackHandler):
         # Compute cost
         cost_usd = self._compute_cost(provider, model, input_tokens, output_tokens)
 
+        # Extract tool_id from operation if this is a tool span
+        tool_id = ""
+        tool_input = ""
+        tool_output = ""
+
+        if span.get("span_type") == "tool":
+            operation = span.get("operation", "")
+            if operation.startswith("tool:"):
+                tool_id = operation[5:]  # Extract "browserless" from "tool:browserless"
+
+            # Get tool input/output from span (truncate to 10KB)
+            if span.get("input"):
+                tool_input = str(span["input"])[:10000]
+            if metadata and metadata.get("output"):
+                tool_output = str(metadata["output"])[:10000]
+
+        # Get goal from context (thread-safe)
+        current_goal = get_goal() or ""
+
         # Build event
         event = {
             "schema_version": "1.0",
@@ -318,6 +339,11 @@ class KalibrCallbackHandler(BaseCallbackHandler):
             "service": self.service,
             "runtime_env": os.getenv("RUNTIME_ENV", "local"),
             "sandbox_id": os.getenv("SANDBOX_ID", "local"),
+            # New fields for tool/goal tracking
+            "tool_id": tool_id,
+            "tool_input": tool_input,
+            "tool_output": tool_output,
+            "goal": current_goal,
             "metadata": {
                 **self.default_metadata,
                 "span_type": span.get("span_type", "llm"),

kalibr_openai_agents/__init__.py

@@ -26,7 +26,7 @@ Usage:
 
 Environment Variables:
     KALIBR_API_KEY: API key for authentication
-    KALIBR_ENDPOINT: Backend endpoint URL
+    KALIBR_COLLECTOR_URL: Backend endpoint URL
     KALIBR_TENANT_ID: Tenant identifier
     KALIBR_ENVIRONMENT: Environment (prod/staging/dev)
     KALIBR_SERVICE: Service name