kalibr 1.0.28__py3-none-any.whl → 1.1.2a0__py3-none-any.whl

kalibr_crewai/instrumentor.py

@@ -0,0 +1,513 @@
+"""Kalibr Auto-Instrumentation for CrewAI.
+
+This module provides automatic instrumentation for CrewAI applications
+by wrapping key methods to capture telemetry.
+"""
+
+import os
+import time
+import traceback
+import uuid
+from datetime import datetime, timezone
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional
+
+from .callbacks import EventBatcher, _count_tokens, _get_provider_from_model
+
+# Import Kalibr cost adapters
+try:
+    from kalibr.cost_adapter import CostAdapterFactory
+except ImportError:
+    CostAdapterFactory = None
+
+
+class KalibrCrewAIInstrumentor:
+    """Auto-instrumentation for CrewAI.
+
+    This instrumentor wraps CrewAI's Crew, Agent, and Task classes to
+    automatically capture telemetry without requiring manual callbacks.
+
+    Args:
+        api_key: Kalibr API key
+        endpoint: Backend endpoint URL
+        tenant_id: Tenant identifier
+        environment: Environment (prod/staging/dev)
+        service: Service name
+        workflow_id: Workflow identifier
+        capture_input: Whether to capture inputs (default: True)
+        capture_output: Whether to capture outputs (default: True)
+
+    Usage:
+        from kalibr_crewai import KalibrCrewAIInstrumentor
+
+        # Instrument before creating crews
+        instrumentor = KalibrCrewAIInstrumentor(tenant_id="my-tenant")
+        instrumentor.instrument()
+
+        # Use CrewAI normally - all operations are traced
+        crew = Crew(agents=[...], tasks=[...])
+        result = crew.kickoff()
+
+        # Optionally uninstrument
+        instrumentor.uninstrument()
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        endpoint: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        environment: Optional[str] = None,
+        service: Optional[str] = None,
+        workflow_id: Optional[str] = None,
+        capture_input: bool = True,
+        capture_output: bool = True,
+    ):
+        self.api_key = api_key or os.getenv("KALIBR_API_KEY", "")
+        self.endpoint = endpoint or os.getenv(
+            "KALIBR_ENDPOINT",
+            os.getenv("KALIBR_API_ENDPOINT", "http://localhost:8001/api/v1/traces")
+        )
+        self.tenant_id = tenant_id or os.getenv("KALIBR_TENANT_ID", "default")
+        self.environment = environment or os.getenv("KALIBR_ENVIRONMENT", "prod")
+        self.service = service or os.getenv("KALIBR_SERVICE", "crewai-app")
+        self.workflow_id = workflow_id or os.getenv("KALIBR_WORKFLOW_ID", "default-workflow")
+        self.capture_input = capture_input
+        self.capture_output = capture_output
+
+        # Batcher for sending events
+        self._batcher: Optional[EventBatcher] = None
+
+        # Original methods to restore on uninstrument
+        self._originals: Dict[str, Any] = {}
+
+        # Instrumentation state
+        self._is_instrumented = False
+
+    def instrument(self) -> bool:
+        """Instrument CrewAI classes.
+
+        Returns:
+            True if instrumentation was successful
+        """
+        if self._is_instrumented:
+            return True
+
+        try:
+            from crewai import Crew, Agent, Task
+        except ImportError:
+            print("[Kalibr] CrewAI not installed, skipping instrumentation")
+            return False
+
+        # Initialize batcher
+        self._batcher = EventBatcher.get_instance(
+            endpoint=self.endpoint,
+            api_key=self.api_key,
+        )
+
+        # Instrument Crew.kickoff
+        self._originals["Crew.kickoff"] = Crew.kickoff
+        Crew.kickoff = self._wrap_crew_kickoff(Crew.kickoff)
+
+        # Instrument Crew.kickoff_async if available
+        if hasattr(Crew, "kickoff_async"):
+            self._originals["Crew.kickoff_async"] = Crew.kickoff_async
+            Crew.kickoff_async = self._wrap_crew_kickoff_async(Crew.kickoff_async)
+
+        # Instrument Agent.execute_task
+        if hasattr(Agent, "execute_task"):
+            self._originals["Agent.execute_task"] = Agent.execute_task
+            Agent.execute_task = self._wrap_agent_execute(Agent.execute_task)
+
+        # Instrument Task.execute_sync if available
+        if hasattr(Task, "execute_sync"):
+            self._originals["Task.execute_sync"] = Task.execute_sync
+            Task.execute_sync = self._wrap_task_execute(Task.execute_sync)
+
+        self._is_instrumented = True
+        return True
+
+    def uninstrument(self) -> bool:
+        """Remove instrumentation and restore original methods.
+
+        Returns:
+            True if uninstrumentation was successful
+        """
+        if not self._is_instrumented:
+            return True
+
+        try:
+            from crewai import Crew, Agent, Task
+
+            # Restore original methods
+            if "Crew.kickoff" in self._originals:
+                Crew.kickoff = self._originals["Crew.kickoff"]
+
+            if "Crew.kickoff_async" in self._originals:
+                Crew.kickoff_async = self._originals["Crew.kickoff_async"]
+
+            if "Agent.execute_task" in self._originals:
+                Agent.execute_task = self._originals["Agent.execute_task"]
+
+            if "Task.execute_sync" in self._originals:
+                Task.execute_sync = self._originals["Task.execute_sync"]
+
+            self._originals.clear()
+            self._is_instrumented = False
+            return True
+
+        except ImportError:
+            return False
+
+    def _wrap_crew_kickoff(self, original_method: Callable) -> Callable:
+        """Wrap Crew.kickoff to capture crew execution."""
+        instrumentor = self
+
+        @wraps(original_method)
+        def wrapper(crew_self, *args, **kwargs):
+            trace_id = str(uuid.uuid4())
+            span_id = str(uuid.uuid4())
+            start_time = time.time()
+            ts_start = datetime.now(timezone.utc)
+
+            # Capture crew info
+            crew_name = getattr(crew_self, "name", None) or "unnamed_crew"
+            agent_count = len(getattr(crew_self, "agents", []))
+            task_count = len(getattr(crew_self, "tasks", []))
+
+            status = "success"
+            error_type = None
+            error_message = None
+            result = None
+
+            try:
+                result = original_method(crew_self, *args, **kwargs)
+                return result
+
+            except Exception as e:
+                status = "error"
+                error_type = type(e).__name__
+                error_message = str(e)[:512]
+                raise
+
+            finally:
+                duration_ms = int((time.time() - start_time) * 1000)
+                ts_end = datetime.now(timezone.utc)
+
+                # Build output info
+                output_preview = None
+                if instrumentor.capture_output and result is not None:
+                    output_preview = str(result)[:500]
+
+                # Create event
+                event = {
+                    "schema_version": "1.0",
+                    "trace_id": trace_id,
+                    "span_id": span_id,
+                    "parent_span_id": None,
+                    "tenant_id": instrumentor.tenant_id,
+                    "workflow_id": instrumentor.workflow_id,
+                    "provider": "crewai",
+                    "model_id": "crew",
+                    "model_name": crew_name,
+                    "operation": f"crew:{crew_name}",
+                    "endpoint": "crew.kickoff",
+                    "duration_ms": duration_ms,
+                    "latency_ms": duration_ms,
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "total_tokens": 0,
+                    "cost_usd": 0.0,
+                    "total_cost_usd": 0.0,
+                    "status": status,
+                    "error_type": error_type,
+                    "error_message": error_message,
+                    "timestamp": ts_start.isoformat(),
+                    "ts_start": ts_start.isoformat(),
+                    "ts_end": ts_end.isoformat(),
+                    "environment": instrumentor.environment,
+                    "service": instrumentor.service,
+                    "runtime_env": os.getenv("RUNTIME_ENV", "local"),
+                    "sandbox_id": os.getenv("SANDBOX_ID", "local"),
+                    "metadata": {
+                        "span_type": "crew",
+                        "crewai": True,
+                        "crew_name": crew_name,
+                        "agent_count": agent_count,
+                        "task_count": task_count,
+                        "output_preview": output_preview,
+                    },
+                }
+
+                if instrumentor._batcher:
+                    instrumentor._batcher.enqueue(event)
+
+        return wrapper
+
+    def _wrap_crew_kickoff_async(self, original_method: Callable) -> Callable:
+        """Wrap Crew.kickoff_async to capture async crew execution."""
+        instrumentor = self
+
+        @wraps(original_method)
+        async def wrapper(crew_self, *args, **kwargs):
+            trace_id = str(uuid.uuid4())
+            span_id = str(uuid.uuid4())
+            start_time = time.time()
+            ts_start = datetime.now(timezone.utc)
+
+            crew_name = getattr(crew_self, "name", None) or "unnamed_crew"
+            agent_count = len(getattr(crew_self, "agents", []))
+            task_count = len(getattr(crew_self, "tasks", []))
+
+            status = "success"
+            error_type = None
+            error_message = None
+            result = None
+
+            try:
+                result = await original_method(crew_self, *args, **kwargs)
+                return result
+
+            except Exception as e:
+                status = "error"
+                error_type = type(e).__name__
+                error_message = str(e)[:512]
+                raise
+
+            finally:
+                duration_ms = int((time.time() - start_time) * 1000)
+                ts_end = datetime.now(timezone.utc)
+
+                output_preview = None
+                if instrumentor.capture_output and result is not None:
+                    output_preview = str(result)[:500]
+
+                event = {
+                    "schema_version": "1.0",
+                    "trace_id": trace_id,
+                    "span_id": span_id,
+                    "parent_span_id": None,
+                    "tenant_id": instrumentor.tenant_id,
+                    "workflow_id": instrumentor.workflow_id,
+                    "provider": "crewai",
+                    "model_id": "crew",
+                    "model_name": crew_name,
+                    "operation": f"crew:{crew_name}",
+                    "endpoint": "crew.kickoff_async",
+                    "duration_ms": duration_ms,
+                    "latency_ms": duration_ms,
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "total_tokens": 0,
+                    "cost_usd": 0.0,
+                    "total_cost_usd": 0.0,
+                    "status": status,
+                    "error_type": error_type,
+                    "error_message": error_message,
+                    "timestamp": ts_start.isoformat(),
+                    "ts_start": ts_start.isoformat(),
+                    "ts_end": ts_end.isoformat(),
+                    "environment": instrumentor.environment,
+                    "service": instrumentor.service,
+                    "runtime_env": os.getenv("RUNTIME_ENV", "local"),
+                    "sandbox_id": os.getenv("SANDBOX_ID", "local"),
+                    "metadata": {
+                        "span_type": "crew",
+                        "crewai": True,
+                        "async": True,
+                        "crew_name": crew_name,
+                        "agent_count": agent_count,
+                        "task_count": task_count,
+                        "output_preview": output_preview,
+                    },
+                }
+
+                if instrumentor._batcher:
+                    instrumentor._batcher.enqueue(event)
+
+        return wrapper
+
+    def _wrap_agent_execute(self, original_method: Callable) -> Callable:
+        """Wrap Agent.execute_task to capture agent execution."""
+        instrumentor = self
+
+        @wraps(original_method)
+        def wrapper(agent_self, task, *args, **kwargs):
+            span_id = str(uuid.uuid4())
+            start_time = time.time()
+            ts_start = datetime.now(timezone.utc)
+
+            # Get agent info
+            role = getattr(agent_self, "role", "unknown")
+            goal = getattr(agent_self, "goal", "")
+
+            # Get task info
+            task_description = ""
+            if hasattr(task, "description"):
+                task_description = str(task.description)
+
+            status = "success"
+            error_type = None
+            error_message = None
+            result = None
+
+            try:
+                result = original_method(agent_self, task, *args, **kwargs)
+                return result
+
+            except Exception as e:
+                status = "error"
+                error_type = type(e).__name__
+                error_message = str(e)[:512]
+                raise
+
+            finally:
+                duration_ms = int((time.time() - start_time) * 1000)
+                ts_end = datetime.now(timezone.utc)
+
+                output_preview = None
+                if instrumentor.capture_output and result is not None:
+                    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")
+
+                event = {
+                    "schema_version": "1.0",
+                    "trace_id": str(uuid.uuid4()),  # TODO: Link to crew trace
+                    "span_id": span_id,
+                    "parent_span_id": None,
+                    "tenant_id": instrumentor.tenant_id,
+                    "workflow_id": instrumentor.workflow_id,
+                    "provider": "crewai",
+                    "model_id": "agent",
+                    "model_name": role,
+                    "operation": f"agent:{role}",
+                    "endpoint": "agent.execute_task",
+                    "duration_ms": duration_ms,
+                    "latency_ms": duration_ms,
+                    "input_tokens": input_tokens,
+                    "output_tokens": output_tokens,
+                    "total_tokens": input_tokens + output_tokens,
+                    "cost_usd": 0.0,
+                    "total_cost_usd": 0.0,
+                    "status": status,
+                    "error_type": error_type,
+                    "error_message": error_message,
+                    "timestamp": ts_start.isoformat(),
+                    "ts_start": ts_start.isoformat(),
+                    "ts_end": ts_end.isoformat(),
+                    "environment": instrumentor.environment,
+                    "service": instrumentor.service,
+                    "runtime_env": os.getenv("RUNTIME_ENV", "local"),
+                    "sandbox_id": os.getenv("SANDBOX_ID", "local"),
+                    "metadata": {
+                        "span_type": "agent",
+                        "crewai": True,
+                        "agent_role": role,
+                        "agent_goal": goal[:200] if goal else None,
+                        "task_description": task_description[:200] if task_description else None,
+                        "output_preview": output_preview,
+                    },
+                }
+
+                if instrumentor._batcher:
+                    instrumentor._batcher.enqueue(event)
+
+        return wrapper
+
+    def _wrap_task_execute(self, original_method: Callable) -> Callable:
+        """Wrap Task.execute_sync to capture task execution."""
+        instrumentor = self
+
+        @wraps(original_method)
+        def wrapper(task_self, *args, **kwargs):
+            span_id = str(uuid.uuid4())
+            start_time = time.time()
+            ts_start = datetime.now(timezone.utc)
+
+            description = getattr(task_self, "description", "")
+            expected_output = getattr(task_self, "expected_output", "")
+
+            status = "success"
+            error_type = None
+            error_message = None
+            result = None
+
+            try:
+                result = original_method(task_self, *args, **kwargs)
+                return result
+
+            except Exception as e:
+                status = "error"
+                error_type = type(e).__name__
+                error_message = str(e)[:512]
+                raise
+
+            finally:
+                duration_ms = int((time.time() - start_time) * 1000)
+                ts_end = datetime.now(timezone.utc)
+
+                output_preview = None
+                if instrumentor.capture_output and result is not None:
+                    if hasattr(result, "raw"):
+                        output_preview = str(result.raw)[:500]
+                    else:
+                        output_preview = str(result)[:500]
+
+                input_tokens = _count_tokens(description, "gpt-4")
+                output_tokens = _count_tokens(output_preview or "", "gpt-4")
+
+                event = {
+                    "schema_version": "1.0",
+                    "trace_id": str(uuid.uuid4()),
+                    "span_id": span_id,
+                    "parent_span_id": None,
+                    "tenant_id": instrumentor.tenant_id,
+                    "workflow_id": instrumentor.workflow_id,
+                    "provider": "crewai",
+                    "model_id": "task",
+                    "model_name": "crewai-task",
+                    "operation": f"task:{description[:30]}..." if len(description) > 30 else f"task:{description}",
+                    "endpoint": "task.execute_sync",
+                    "duration_ms": duration_ms,
+                    "latency_ms": duration_ms,
+                    "input_tokens": input_tokens,
+                    "output_tokens": output_tokens,
+                    "total_tokens": input_tokens + output_tokens,
+                    "cost_usd": 0.0,
+                    "total_cost_usd": 0.0,
+                    "status": status,
+                    "error_type": error_type,
+                    "error_message": error_message,
+                    "timestamp": ts_start.isoformat(),
+                    "ts_start": ts_start.isoformat(),
+                    "ts_end": ts_end.isoformat(),
+                    "environment": instrumentor.environment,
+                    "service": instrumentor.service,
+                    "runtime_env": os.getenv("RUNTIME_ENV", "local"),
+                    "sandbox_id": os.getenv("SANDBOX_ID", "local"),
+                    "metadata": {
+                        "span_type": "task",
+                        "crewai": True,
+                        "task_description": description[:200] if description else None,
+                        "expected_output": expected_output[:200] if expected_output else None,
+                        "output_preview": output_preview,
+                    },
+                }
+
+                if instrumentor._batcher:
+                    instrumentor._batcher.enqueue(event)
+
+        return wrapper
+
+    def flush(self):
+        """Force flush pending events."""
+        if self._batcher:
+            self._batcher.flush()
+
+    @property
+    def is_instrumented(self) -> bool:
+        """Check if CrewAI is currently instrumented."""
+        return self._is_instrumented

kalibr_langchain/__init__.py

@@ -0,0 +1,47 @@
+"""Kalibr LangChain Integration - Observability for LangChain applications.
+
+This package provides a callback handler that integrates LangChain with
+Kalibr's observability platform, capturing:
+- LLM calls with token usage and costs
+- Chain executions with timing
+- Tool/Agent invocations
+- Retrieval operations
+- Error tracking with stack traces
+
+Usage:
+    from kalibr_langchain import KalibrCallbackHandler
+    from langchain_openai import ChatOpenAI
+
+    # Create callback handler
+    handler = KalibrCallbackHandler(
+        api_key="your-api-key",
+        tenant_id="my-tenant",
+    )
+
+    # Use with LangChain
+    llm = ChatOpenAI(model="gpt-4", callbacks=[handler])
+    response = llm.invoke("Hello, world!")
+
+    # Or attach to all components
+    from langchain.globals import set_llm_cache
+    chain = prompt | llm | parser
+    chain.invoke({"input": "Hello"}, config={"callbacks": [handler]})
+
+Environment Variables:
+    KALIBR_API_KEY: API key for authentication
+    KALIBR_ENDPOINT: Backend endpoint URL
+    KALIBR_TENANT_ID: Tenant identifier
+    KALIBR_ENVIRONMENT: Environment (prod/staging/dev)
+    KALIBR_SERVICE: Service name
+"""
+
+__version__ = "0.1.0"
+
+from .callback import KalibrCallbackHandler
+from .async_callback import AsyncKalibrCallbackHandler
+
+__all__ = [
+    "KalibrCallbackHandler",
+    "AsyncKalibrCallbackHandler",
+    "__version__",
+]