agentos-python 0.1.0__py3-none-any.whl

agentos/integrations/crewai.py

@@ -0,0 +1,221 @@
+"""CrewAI integration — decorators and hooks for CrewAI agent instrumentation.
+
+Usage::
+
+    from agentos.integrations.crewai import instrument_crewai
+
+    # Patches CrewAI to auto-capture agent runs, task executions, and LLM calls
+    instrument_crewai(api_key="aos_...", agent_id="my-crew")
+
+Or use decorators directly::
+
+    from agentos.integrations.crewai import trace_crew, trace_task
+
+    @trace_crew("research-crew")
+    def run_crew():
+        crew = Crew(agents=[...], tasks=[...])
+        return crew.kickoff()
+"""
+
+from __future__ import annotations
+
+import functools
+import logging
+import time
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+logger = logging.getLogger("agentos.integrations.crewai")
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def trace_crew(
+    crew_name: str,
+    *,
+    api_key: str | None = None,
+    base_url: str = "https://api.agentos.dev",
+) -> Callable[[F], F]:
+    """Decorator that wraps a CrewAI crew execution in a trace.
+
+    Captures the full crew run as a trace with business_event on completion.
+
+    Usage::
+
+        @trace_crew("research-crew")
+        def run_my_crew():
+            crew = Crew(agents=[researcher, writer], tasks=[...])
+            return crew.kickoff()
+    """
+
+    def decorator(fn: F) -> F:
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            from agentos.client import AgentOS, get_client
+            from agentos.tracing import trace
+
+            client = get_client()
+            if client is None and api_key:
+                client = AgentOS(api_key=api_key, base_url=base_url, flush_interval=0)
+
+            start = time.monotonic()
+
+            with trace(agent_id=crew_name) as ctx:
+                try:
+                    result = fn(*args, **kwargs)
+                    duration_ms = (time.monotonic() - start) * 1000
+
+                    if client:
+                        client.business_event(
+                            crew_name,
+                            event_name="crew.completed",
+                            trace_id=ctx.trace_id,
+                            span_id=ctx.span_id,
+                            metadata={"duration_ms": duration_ms},
+                        )
+
+                    return result
+
+                except Exception as exc:
+                    duration_ms = (time.monotonic() - start) * 1000
+
+                    if client:
+                        client.business_event(
+                            crew_name,
+                            event_name="crew.failed",
+                            trace_id=ctx.trace_id,
+                            span_id=ctx.span_id,
+                            metadata={
+                                "duration_ms": duration_ms,
+                                "error": str(exc),
+                            },
+                        )
+                    raise
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def trace_task(
+    task_name: str,
+) -> Callable[[F], F]:
+    """Decorator that wraps a CrewAI task function as a tool_call span.
+
+    Usage::
+
+        @trace_task("research")
+        def research_task(topic: str) -> str:
+            ...
+    """
+
+    def decorator(fn: F) -> F:
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            from agentos.client import get_client
+            from agentos.tracing import span
+
+            client = get_client()
+            start = time.monotonic()
+
+            with span() as ctx:
+                try:
+                    result = fn(*args, **kwargs)
+                    duration_ms = (time.monotonic() - start) * 1000
+
+                    if client:
+                        client.tool_call(
+                            ctx.agent_id,
+                            tool_name=task_name,
+                            status="success",
+                            duration_ms=duration_ms,
+                            trace_id=ctx.trace_id,
+                            span_id=ctx.span_id,
+                            parent_span_id=ctx.parent_span_id,
+                        )
+
+                    return result
+
+                except Exception as exc:
+                    duration_ms = (time.monotonic() - start) * 1000
+
+                    if client:
+                        client.tool_call(
+                            ctx.agent_id,
+                            tool_name=task_name,
+                            status="error",
+                            duration_ms=duration_ms,
+                            error={"type": type(exc).__name__, "message": str(exc)},
+                            trace_id=ctx.trace_id,
+                            span_id=ctx.span_id,
+                            parent_span_id=ctx.parent_span_id,
+                        )
+                    raise
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def instrument_crewai(
+    *,
+    api_key: str | None = None,
+    agent_id: str = "crewai",
+    base_url: str = "https://api.agentos.dev",
+    capture_content: bool = True,
+) -> None:
+    """Monkey-patch CrewAI to auto-capture all agent and task activity.
+
+    Patches ``Crew.kickoff()`` and ``Task._execute_core()`` to emit events.
+    """
+    try:
+        from crewai import Crew
+    except ImportError:
+        logger.error("crewai not installed — cannot instrument")
+        return
+
+    from agentos.client import AgentOS, get_client
+
+    client = get_client()
+    if client is None and api_key:
+        client = AgentOS(api_key=api_key, base_url=base_url, flush_interval=0)
+
+    if client is None:
+        logger.warning("No AgentOS client — instrument_crewai is a no-op")
+        return
+
+    # Patch Crew.kickoff
+    original_kickoff = Crew.kickoff
+
+    @functools.wraps(original_kickoff)
+    def patched_kickoff(self: Any, *args: Any, **kwargs: Any) -> Any:
+        from agentos.tracing import trace
+
+        crew_name = getattr(self, "name", agent_id) or agent_id
+        start = time.monotonic()
+
+        with trace(agent_id=crew_name) as ctx:
+            try:
+                result = original_kickoff(self, *args, **kwargs)
+                duration_ms = (time.monotonic() - start) * 1000
+                client.business_event(
+                    crew_name,
+                    event_name="crew.completed",
+                    trace_id=ctx.trace_id,
+                    span_id=ctx.span_id,
+                    metadata={"duration_ms": duration_ms},
+                )
+                return result
+            except Exception as exc:
+                duration_ms = (time.monotonic() - start) * 1000
+                client.business_event(
+                    crew_name,
+                    event_name="crew.failed",
+                    trace_id=ctx.trace_id,
+                    span_id=ctx.span_id,
+                    metadata={"duration_ms": duration_ms, "error": str(exc)},
+                )
+                raise
+
+    Crew.kickoff = patched_kickoff  # type: ignore[assignment]
+    logger.info("CrewAI instrumented — Crew.kickoff patched")

agentos/integrations/langchain.py

@@ -0,0 +1,307 @@
+"""LangChain integration — CallbackHandler that auto-captures LLM and tool events.
+
+Usage::
+
+    from agentos.integrations.langchain import AgentOSCallbackHandler
+
+    handler = AgentOSCallbackHandler(api_key="aos_...", agent_id="my-agent")
+    chain.invoke({"input": "..."}, config={"callbacks": [handler]})
+
+    # Or set globally:
+    import langchain
+    langchain.callbacks.set_handler(handler)
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Sequence
+from typing import Any
+from uuid import UUID
+
+logger = logging.getLogger("agentos.integrations.langchain")
+
+
+class AgentOSCallbackHandler:
+    """LangChain-compatible callback handler.
+
+    Captures LLM calls, tool invocations, chain runs, and retriever queries
+    as Agent OS events.
+
+    Implements the LangChain BaseCallbackHandler interface without importing it,
+    so this module works even if langchain isn't installed (duck typing).
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        agent_id: str = "langchain-agent",
+        base_url: str = "https://api.agentos.dev",
+        capture_content: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        from agentos.client import AgentOS, get_client
+
+        if api_key:
+            self._client = AgentOS(api_key=api_key, base_url=base_url, flush_interval=0, **kwargs)
+        else:
+            self._client = get_client()
+
+        self._agent_id = agent_id
+        self._capture_content = capture_content
+        self._run_starts: dict[str, float] = {}  # run_id → start time
+
+    # --- LLM callbacks ---
+
+    def on_llm_start(
+        self,
+        serialized: dict[str, Any],
+        prompts: list[str],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when an LLM starts."""
+        self._run_starts[str(run_id)] = time.monotonic()
+
+    def on_chat_model_start(
+        self,
+        serialized: dict[str, Any],
+        messages: list[list[Any]],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when a chat model starts."""
+        self._run_starts[str(run_id)] = time.monotonic()
+
+    def on_llm_end(
+        self,
+        response: Any,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when an LLM finishes. Sends agent.llm_call event."""
+        if self._client is None:
+            return
+
+        start = self._run_starts.pop(str(run_id), None)
+        duration_ms = (time.monotonic() - start) * 1000 if start else None
+
+        # Extract model info from response
+        llm_output = getattr(response, "llm_output", {}) or {}
+        model = llm_output.get("model_name", "unknown")
+        model_id = llm_output.get("model_id", "")
+        system = model_id.split("/")[0] if "/" in model_id else "unknown"
+
+        # Token usage
+        token_usage = llm_output.get("token_usage", {})
+        data: dict[str, Any] = {
+            "model": model,
+            "system": system,
+        }
+
+        if duration_ms is not None:
+            data["duration_ms"] = duration_ms
+
+        if token_usage:
+            if "prompt_tokens" in token_usage:
+                data["input_tokens"] = token_usage["prompt_tokens"]
+            if "completion_tokens" in token_usage:
+                data["output_tokens"] = token_usage["completion_tokens"]
+            if "total_tokens" in token_usage:
+                data["total_tokens"] = token_usage["total_tokens"]
+
+        # Content
+        if self._capture_content:
+            generations = getattr(response, "generations", [])
+            if generations and generations[0]:
+                gen = generations[0][0]
+                text = getattr(gen, "text", None)
+                if text:
+                    data["output"] = [{"role": "assistant", "content": text}]
+                # Check for message-based output
+                message = getattr(gen, "message", None)
+                if message:
+                    content = getattr(message, "content", None)
+                    role = getattr(message, "type", "assistant")
+                    if content:
+                        data["output"] = [{"role": role, "content": content}]
+
+        try:
+            self._client.llm_call(self._agent_id, **data)
+        except Exception:
+            logger.exception("Failed to capture LangChain LLM call")
+
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when an LLM errors."""
+        if self._client is None:
+            return
+
+        start = self._run_starts.pop(str(run_id), None)
+        duration_ms = (time.monotonic() - start) * 1000 if start else None
+
+        try:
+            self._client.llm_call(
+                self._agent_id,
+                model="unknown",
+                system="unknown",
+                duration_ms=duration_ms,
+                error={"type": type(error).__name__, "message": str(error)},
+            )
+        except Exception:
+            logger.exception("Failed to capture LangChain LLM error")
+
+    # --- Tool callbacks ---
+
+    def on_tool_start(
+        self,
+        serialized: dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when a tool starts."""
+        self._run_starts[str(run_id)] = time.monotonic()
+
+    def on_tool_end(
+        self,
+        output: Any,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when a tool finishes. Sends agent.tool_call event."""
+        if self._client is None:
+            return
+
+        start = self._run_starts.pop(str(run_id), None)
+        duration_ms = (time.monotonic() - start) * 1000 if start else None
+
+        tool_name = kwargs.get("name", "unknown")
+
+        data: dict[str, Any] = {
+            "tool_name": tool_name,
+            "status": "success",
+        }
+        if duration_ms is not None:
+            data["duration_ms"] = duration_ms
+        if self._capture_content and output is not None:
+            data["output"] = str(output)
+
+        try:
+            self._client.tool_call(self._agent_id, **data)
+        except Exception:
+            logger.exception("Failed to capture LangChain tool call")
+
+    def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when a tool errors."""
+        if self._client is None:
+            return
+
+        start = self._run_starts.pop(str(run_id), None)
+        duration_ms = (time.monotonic() - start) * 1000 if start else None
+
+        try:
+            self._client.tool_call(
+                self._agent_id,
+                tool_name=kwargs.get("name", "unknown"),
+                status="error",
+                duration_ms=duration_ms,
+                error={"type": type(error).__name__, "message": str(error)},
+            )
+        except Exception:
+            logger.exception("Failed to capture LangChain tool error")
+
+    # --- Retriever callbacks ---
+
+    def on_retriever_start(
+        self,
+        serialized: dict[str, Any],
+        query: str,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when a retriever starts."""
+        self._run_starts[str(run_id)] = time.monotonic()
+
+    def on_retriever_end(
+        self,
+        documents: Sequence[Any],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when a retriever finishes. Sends agent.retrieval_query event."""
+        if self._client is None:
+            return
+
+        start = self._run_starts.pop(str(run_id), None)
+        duration_ms = (time.monotonic() - start) * 1000 if start else None
+
+        try:
+            self._client.retrieval_query(
+                self._agent_id,
+                source="langchain-retriever",
+                results_count=len(documents),
+                duration_ms=duration_ms,
+            )
+        except Exception:
+            logger.exception("Failed to capture LangChain retriever call")
+
+    # --- Chain callbacks (no-ops, used for context) ---
+
+    def on_chain_start(
+        self, serialized: dict[str, Any], inputs: dict[str, Any], **kwargs: Any
+    ) -> None:
+        pass
+
+    def on_chain_end(self, outputs: dict[str, Any], **kwargs: Any) -> None:
+        pass
+
+    def on_chain_error(self, error: BaseException, **kwargs: Any) -> None:
+        pass
+
+    # --- Flush ---
+
+    def flush(self) -> None:
+        """Flush pending events."""
+        if self._client:
+            self._client.flush()
+
+    def shutdown(self) -> None:
+        """Flush and close."""
+        if self._client:
+            self._client.shutdown()

agentos/integrations/litellm.py

@@ -0,0 +1,123 @@
+"""LiteLLM integration — patches litellm.completion() globally.
+
+Usage::
+
+    from agentos.integrations.litellm import patch_litellm
+    patch_litellm(api_key="aos_...", agent_id="my-agent")
+
+    import litellm
+    response = litellm.completion(model="gpt-4o", messages=[...])
+    # Auto-captured
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any
+
+logger = logging.getLogger("agentos.integrations.litellm")
+
+_original_completion: Any = None
+
+
+def patch_litellm(
+    *,
+    api_key: str | None = None,
+    agent_id: str = "default",
+    base_url: str = "https://api.agentos.dev",
+    capture_content: bool = True,
+    **kwargs: Any,
+) -> None:
+    """Patch ``litellm.completion`` to auto-capture LLM calls."""
+    global _original_completion
+
+    try:
+        import litellm
+    except ImportError:
+        logger.error("litellm not installed — cannot patch")
+        return
+
+    from agentos.client import AgentOS, get_client
+
+    aos_client = None
+    if api_key:
+        aos_client = AgentOS(api_key=api_key, base_url=base_url, flush_interval=0, **kwargs)
+    else:
+        aos_client = get_client()
+
+    if aos_client is None:
+        logger.warning("No AgentOS client configured — patch_litellm is a no-op")
+        return
+
+    _original_completion = litellm.completion
+
+    def patched_completion(*args: Any, **call_kwargs: Any) -> Any:
+        model_arg = call_kwargs.get("model", args[0] if args else "unknown")
+        messages = call_kwargs.get("messages")
+
+        start = time.monotonic()
+        try:
+            response = _original_completion(*args, **call_kwargs)
+            duration_ms = (time.monotonic() - start) * 1000
+
+            # LiteLLM responses follow the OpenAI format
+            data: dict[str, Any] = {
+                "model": getattr(response, "model", model_arg),
+                "system": "litellm",
+                "duration_ms": duration_ms,
+            }
+
+            usage = getattr(response, "usage", None)
+            if usage:
+                data["input_tokens"] = getattr(usage, "prompt_tokens", None)
+                data["output_tokens"] = getattr(usage, "completion_tokens", None)
+
+            choices = getattr(response, "choices", [])
+            if choices:
+                fr = getattr(choices[0], "finish_reason", None)
+                if fr:
+                    data["finish_reason"] = fr
+
+            if capture_content and messages:
+                data["input"] = [
+                    {"role": m.get("role", ""), "content": m.get("content")}
+                    for m in messages
+                    if isinstance(m, dict)
+                ]
+
+            try:
+                aos_client.llm_call(agent_id, **data)
+            except Exception:
+                logger.exception("Failed to capture LiteLLM call")
+
+            return response
+
+        except Exception as exc:
+            duration_ms = (time.monotonic() - start) * 1000
+            try:
+                aos_client.llm_call(
+                    agent_id,
+                    model=str(model_arg),
+                    system="litellm",
+                    duration_ms=duration_ms,
+                    error={"type": type(exc).__name__, "message": str(exc)},
+                )
+            except Exception:
+                logger.exception("Failed to capture LiteLLM error")
+            raise
+
+    litellm.completion = patched_completion
+
+
+def unpatch_litellm() -> None:
+    """Restore original ``litellm.completion``."""
+    global _original_completion
+    if _original_completion is not None:
+        try:
+            import litellm
+
+            litellm.completion = _original_completion
+            _original_completion = None
+        except ImportError:
+            pass