trace-agent-langchain 0.0.1__tar.gz

trace_agent_langchain-0.0.1/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: trace-agent-langchain
+Version: 0.0.1
+Summary: LangChain callback integration for TraceAgent.
+Requires-Python: >=3.11
+Requires-Dist: trace-agent-sdk
+Requires-Dist: langchain-core<0.4.0,>=0.3.0
+Provides-Extra: dev
+Requires-Dist: pytest<9.0.0,>=8.3.2; extra == "dev"
+Requires-Dist: langchain<0.4.0,>=0.3.0; extra == "dev"
+Requires-Dist: langchain-openai<0.3.0,>=0.2.0; extra == "dev"
+Requires-Dist: trace-agent-server; extra == "dev"

trace_agent_langchain-0.0.1/pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "trace-agent-langchain"
+version = "0.0.1"
+description = "LangChain callback integration for TraceAgent."
+requires-python = ">=3.11"
+dependencies = [
+  "trace-agent-sdk",
+  "langchain-core>=0.3.0,<0.4.0",
+]
+
+[project.optional-dependencies]
+dev = [
+  "pytest>=8.3.2,<9.0.0",
+  "langchain>=0.3.0,<0.4.0",
+  "langchain-openai>=0.2.0,<0.3.0",
+  "trace-agent-server",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

trace_agent_langchain-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

trace_agent_langchain-0.0.1/src/trace_agent_langchain/__init__.py

@@ -0,0 +1,5 @@
+__all__ = ["__version__", "TraceAgentLangChainCallback"]
+
+__version__ = "0.1.0"
+
+from trace_agent_langchain.callback import TraceAgentLangChainCallback

trace_agent_langchain-0.0.1/src/trace_agent_langchain/callback.py

@@ -0,0 +1,654 @@
+"""LangChain callback handler that streams execution events into TraceAgent.
+
+Usage:
+    from trace_agent_sdk import TraceAgentClient
+    from trace_agent_langchain import TraceAgentLangChainCallback
+
+    client = TraceAgentClient(base_url="http://127.0.0.1:8000")
+    run = client.start_run(agent_name="my-agent", goal="Do something")
+
+    callback = run.as_langchain_callback()
+    chain.invoke({"input": "..."}, config={"callbacks": [callback]})
+"""
+
+from __future__ import annotations
+
+import time
+import traceback
+from typing import Any
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+    from langchain_core.outputs import ChatResult, LLMResult
+except ImportError:  # pragma: no cover
+    BaseCallbackHandler = None  # type: ignore[misc,assignment]
+    LLMResult = None  # type: ignore[misc,assignment]
+    ChatResult = None  # type: ignore[misc,assignment]
+
+
+class TraceAgentLangChainCallback(BaseCallbackHandler):
+    """Record LangChain execution as TraceAgent observations.
+
+    The callback translates LangChain lifecycle events into HTTP POSTs to the
+    TraceAgent backend.  If the backend is unreachable the callback logs a warning
+    but **never** crashes the calling chain/agent.
+    """
+
+    def __init__(
+        self,
+        run_session: Any,
+        *,
+        record_chain_events: bool = True,
+        record_retriever_events: bool = True,
+    ) -> None:
+        if BaseCallbackHandler is None:
+            raise RuntimeError(
+                "langchain-core is not installed. "
+                "Install it with: uv pip install -e '.[langchain]'"
+            )
+        self.run_session = run_session
+        self.record_chain_events = record_chain_events
+        self.record_retriever_events = record_retriever_events
+        # langchain run_id -> start timestamp (for duration calculation)
+        self._starts: dict[str, float] = {}
+
+    # ── Helpers ─────────────────────────────────────────────────────────────
+
+    def _post_observations(self, observations: list[dict[str, Any]]) -> None:
+        """Send observations to the TraceAgent backend."""
+        if not observations:
+            return
+        try:
+            self.run_session.client._request(
+                "POST",
+                f"/api/runs/{self.run_session.run_id}/observations",
+                {"observations": observations},
+            )
+        except Exception:
+            # Silently ignore backend errors so the agent keeps running.
+            traceback.print_exc()
+
+    def _post_tool_results(self, results: list[dict[str, Any]]) -> None:
+        """Send tool results to the TraceAgent backend."""
+        if not results:
+            return
+        try:
+            self.run_session.client._request(
+                "POST",
+                f"/api/runs/{self.run_session.run_id}/tool-results",
+                {"results": results},
+            )
+        except Exception:
+            traceback.print_exc()
+
+    @staticmethod
+    def _extract_token_usage(result: LLMResult | ChatResult | None) -> dict[str, Any]:
+        """Try to extract token usage from an LLM / Chat result."""
+        if result is None:
+            return {}
+        llm_output = getattr(result, "llm_output", None) or {}
+        usage = llm_output.get("token_usage") or {}
+        # OpenAI format
+        if "prompt_tokens" in usage:
+            return {
+                "prompt_tokens": usage.get("prompt_tokens"),
+                "completion_tokens": usage.get("completion_tokens"),
+                "total_tokens": usage.get("total_tokens"),
+            }
+        # Anthropic format
+        if "input_tokens" in usage:
+            return {
+                "prompt_tokens": usage.get("input_tokens"),
+                "completion_tokens": usage.get("output_tokens"),
+                "total_tokens": (
+                    (usage.get("input_tokens") or 0)
+                    + (usage.get("output_tokens") or 0)
+                ),
+            }
+        return dict(usage)
+
+    # ── Chain lifecycle ─────────────────────────────────────────────────────
+
+    def on_chain_start(
+        self,
+        serialized: dict[str, Any] | None,
+        inputs: dict[str, Any],
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if not self.record_chain_events:
+            return
+        self._starts[run_id] = time.time()
+        chain_type = (serialized or {}).get("name", "unknown")
+        inputs_keys = list(inputs.keys()) if isinstance(inputs, dict) else []
+        self._post_observations(
+            [
+                {
+                    "kind": "chain_step",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "running",
+                    "input_summary": f"Chain '{chain_type}' started",
+                    "metadata": {
+                        "chain_type": chain_type,
+                        "inputs_keys": inputs_keys,
+                        "tags": tags,
+                        **(metadata or {}),
+                    },
+                }
+            ]
+        )
+
+    def on_chain_end(
+        self,
+        outputs: dict[str, Any],
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if not self.record_chain_events:
+            return
+        duration_ms = int((time.time() - self._starts.pop(run_id, time.time())) * 1000)
+        output_keys = list(outputs.keys()) if isinstance(outputs, dict) else []
+        self._post_observations(
+            [
+                {
+                    "kind": "chain_step",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "succeeded",
+                    "duration_ms": duration_ms,
+                    "input_summary": f"Chain completed",
+                    "metadata": {
+                        "output_keys": output_keys,
+                        "tags": tags,
+                        **(metadata or {}),
+                    },
+                }
+            ]
+        )
+
+    def on_chain_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if not self.record_chain_events:
+            return
+        duration_ms = int((time.time() - self._starts.pop(run_id, time.time())) * 1000)
+        self._post_observations(
+            [
+                {
+                    "kind": "chain_step",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "failed",
+                    "duration_ms": duration_ms,
+                    "error_code": error.__class__.__name__,
+                    "error_summary": str(error),
+                    "metadata": {"tags": tags, **(metadata or {})},
+                }
+            ]
+        )
+
+    # ── LLM / Chat model lifecycle ──────────────────────────────────────────
+
+    def _on_llm_start(
+        self,
+        *,
+        run_id: str,
+        parent_run_id: str | None,
+        provider: str | None,
+        model: str | None,
+        metadata: dict[str, Any] | None,
+    ) -> None:
+        self._starts[run_id] = time.time()
+        self._post_observations(
+            [
+                {
+                    "kind": "llm_call",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "running",
+                    "provider": provider,
+                    "model": model,
+                    "metadata": metadata or {},
+                }
+            ]
+        )
+
+    def _on_llm_end(
+        self,
+        result: LLMResult | ChatResult | None,
+        *,
+        run_id: str,
+        parent_run_id: str | None,
+        provider: str | None,
+        model: str | None,
+        metadata: dict[str, Any] | None,
+    ) -> None:
+        duration_ms = int((time.time() - self._starts.pop(run_id, time.time())) * 1000)
+        usage = self._extract_token_usage(result)
+        self._post_observations(
+            [
+                {
+                    "kind": "llm_call",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "succeeded",
+                    "duration_ms": duration_ms,
+                    "provider": provider,
+                    "model": model,
+                    "metadata": {**usage, **(metadata or {})},
+                }
+            ]
+        )
+
+    def _on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: str,
+        parent_run_id: str | None,
+        provider: str | None,
+        model: str | None,
+        metadata: dict[str, Any] | None,
+    ) -> None:
+        duration_ms = int((time.time() - self._starts.pop(run_id, time.time())) * 1000)
+        self._post_observations(
+            [
+                {
+                    "kind": "llm_call",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "failed",
+                    "duration_ms": duration_ms,
+                    "provider": provider,
+                    "model": model,
+                    "error_code": error.__class__.__name__,
+                    "error_summary": str(error),
+                    "metadata": metadata or {},
+                }
+            ]
+        )
+
+    def on_llm_start(
+        self,
+        serialized: dict[str, Any] | None,
+        prompts: list[str],
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        provider = (serialized or {}).get("name", "unknown")
+        self._on_llm_start(
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            provider=provider,
+            model=None,
+            metadata={"tags": tags, **(metadata or {})},
+        )
+
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        provider = (metadata or {}).get("name", "unknown")
+        self._on_llm_end(
+            response,
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            provider=provider,
+            model=None,
+            metadata={"tags": tags, **(metadata or {})},
+        )
+
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        provider = (metadata or {}).get("name", "unknown")
+        self._on_llm_error(
+            error,
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            provider=provider,
+            model=None,
+            metadata={"tags": tags, **(metadata or {})},
+        )
+
+    def on_chat_model_start(
+        self,
+        serialized: dict[str, Any] | None,
+        messages: list[list[Any]],
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        provider = (serialized or {}).get("name", "unknown")
+        model = (kwargs.get("invocation_params") or {}).get("model")
+        self._on_llm_start(
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            provider=provider,
+            model=model,
+            metadata={"tags": tags, **(metadata or {})},
+        )
+
+    def on_chat_model_end(
+        self,
+        response: ChatResult,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        provider = (metadata or {}).get("name", "unknown")
+        model = (kwargs.get("invocation_params") or {}).get("model")
+        self._on_llm_end(
+            response,
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            provider=provider,
+            model=model,
+            metadata={"tags": tags, **(metadata or {})},
+        )
+
+    # ── Tool lifecycle ──────────────────────────────────────────────────────
+
+    def on_tool_start(
+        self,
+        serialized: dict[str, Any] | None,
+        input_str: str,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self._starts[run_id] = time.time()
+        tool_name = (serialized or {}).get("name", "unknown")
+        self._post_observations(
+            [
+                {
+                    "kind": "tool_call",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "running",
+                    "tool_name": tool_name,
+                    "input_summary": f"Tool '{tool_name}' started",
+                    "metadata": {
+                        "input": input_str,
+                        "tags": tags,
+                        **(metadata or {}),
+                    },
+                }
+            ]
+        )
+
+    def on_tool_end(
+        self,
+        output: Any,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        duration_ms = int((time.time() - self._starts.pop(run_id, time.time())) * 1000)
+        tool_name = (metadata or {}).get("name", "unknown")
+        output_str = str(output) if output is not None else ""
+        self._post_tool_results(
+            [
+                {
+                    "call_id": run_id,
+                    "name": tool_name,
+                    "step_id": run_id,
+                    "status": "succeeded",
+                    "output_summary": output_str[:500],
+                    "output_hash": None,
+                    "duration_ms": duration_ms,
+                }
+            ]
+        )
+        self._post_observations(
+            [
+                {
+                    "kind": "tool_call",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "succeeded",
+                    "tool_name": tool_name,
+                    "duration_ms": duration_ms,
+                    "input_summary": f"Tool '{tool_name}' succeeded",
+                    "metadata": {"tags": tags, **(metadata or {})},
+                }
+            ]
+        )
+
+    def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        duration_ms = int((time.time() - self._starts.pop(run_id, time.time())) * 1000)
+        tool_name = (metadata or {}).get("name", "unknown")
+        error_summary = str(error)
+        self._post_tool_results(
+            [
+                {
+                    "call_id": run_id,
+                    "name": tool_name,
+                    "step_id": run_id,
+                    "status": "failed",
+                    "output_summary": error_summary[:500],
+                    "output_hash": None,
+                    "error_code": error.__class__.__name__,
+                    "error_summary": error_summary,
+                    "duration_ms": duration_ms,
+                }
+            ]
+        )
+        self._post_observations(
+            [
+                {
+                    "kind": "tool_call",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "failed",
+                    "tool_name": tool_name,
+                    "duration_ms": duration_ms,
+                    "error_code": error.__class__.__name__,
+                    "error_summary": error_summary,
+                    "metadata": {"tags": tags, **(metadata or {})},
+                }
+            ]
+        )
+
+    # ── Retriever lifecycle ─────────────────────────────────────────────────
+
+    def on_retriever_start(
+        self,
+        serialized: dict[str, Any] | None,
+        query: str,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if not self.record_retriever_events:
+            return
+        self._starts[run_id] = time.time()
+        self._post_observations(
+            [
+                {
+                    "kind": "retriever_query",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "running",
+                    "input_summary": f"Retriever query: {query[:200]}",
+                    "metadata": {"query": query, "tags": tags, **(metadata or {})},
+                }
+            ]
+        )
+
+    def on_retriever_end(
+        self,
+        documents: list[Any],
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if not self.record_retriever_events:
+            return
+        duration_ms = int((time.time() - self._starts.pop(run_id, time.time())) * 1000)
+        self._post_observations(
+            [
+                {
+                    "kind": "retriever_query",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "succeeded",
+                    "duration_ms": duration_ms,
+                    "input_summary": f"Retriever returned {len(documents)} documents",
+                    "metadata": {
+                        "document_count": len(documents),
+                        "tags": tags,
+                        **(metadata or {}),
+                    },
+                }
+            ]
+        )
+
+    def on_retriever_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if not self.record_retriever_events:
+            return
+        duration_ms = int((time.time() - self._starts.pop(run_id, time.time())) * 1000)
+        self._post_observations(
+            [
+                {
+                    "kind": "retriever_query",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "failed",
+                    "duration_ms": duration_ms,
+                    "error_code": error.__class__.__name__,
+                    "error_summary": str(error),
+                    "metadata": {"tags": tags, **(metadata or {})},
+                }
+            ]
+        )
+
+    # ── Agent lifecycle ─────────────────────────────────────────────────────
+
+    def on_agent_action(
+        self,
+        action: Any,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        tool = getattr(action, "tool", None) or "unknown"
+        tool_input = getattr(action, "tool_input", None) or {}
+        self._post_observations(
+            [
+                {
+                    "kind": "agent_action",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "running",
+                    "tool_name": tool,
+                    "input_summary": f"Agent action: {tool}",
+                    "metadata": {
+                        "tool": tool,
+                        "tool_input": tool_input,
+                        "tags": tags,
+                        **(metadata or {}),
+                    },
+                }
+            ]
+        )
+
+    def on_agent_finish(
+        self,
+        finish: Any,
+        *,
+        run_id: str,
+        parent_run_id: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        output = getattr(finish, "return_values", {})
+        self._post_observations(
+            [
+                {
+                    "kind": "agent_action",
+                    "step_id": str(run_id),
+                    "parent_step_id": str(parent_run_id) if parent_run_id else None,
+                    "status": "succeeded",
+                    "input_summary": "Agent finished",
+                    "metadata": {
+                        "output_keys": list(output.keys()),
+                        "tags": tags,
+                        **(metadata or {}),
+                    },
+                }
+            ]
+        )

trace_agent_langchain-0.0.1/src/trace_agent_langchain.egg-info/PKG-INFO

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: trace-agent-langchain
+Version: 0.0.1
+Summary: LangChain callback integration for TraceAgent.
+Requires-Python: >=3.11
+Requires-Dist: trace-agent-sdk
+Requires-Dist: langchain-core<0.4.0,>=0.3.0
+Provides-Extra: dev
+Requires-Dist: pytest<9.0.0,>=8.3.2; extra == "dev"
+Requires-Dist: langchain<0.4.0,>=0.3.0; extra == "dev"
+Requires-Dist: langchain-openai<0.3.0,>=0.2.0; extra == "dev"
+Requires-Dist: trace-agent-server; extra == "dev"

trace_agent_langchain-0.0.1/src/trace_agent_langchain.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+pyproject.toml
+src/trace_agent_langchain/__init__.py
+src/trace_agent_langchain/callback.py
+src/trace_agent_langchain.egg-info/PKG-INFO
+src/trace_agent_langchain.egg-info/SOURCES.txt
+src/trace_agent_langchain.egg-info/dependency_links.txt
+src/trace_agent_langchain.egg-info/requires.txt
+src/trace_agent_langchain.egg-info/top_level.txt
+tests/test_langchain_callback.py
+tests/test_langchain_integration.py

trace_agent_langchain-0.0.1/src/trace_agent_langchain.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

trace_agent_langchain-0.0.1/src/trace_agent_langchain.egg-info/requires.txt

@@ -0,0 +1,8 @@
+trace-agent-sdk
+langchain-core<0.4.0,>=0.3.0
+
+[dev]
+pytest<9.0.0,>=8.3.2
+langchain<0.4.0,>=0.3.0
+langchain-openai<0.3.0,>=0.2.0
+trace-agent-server

trace_agent_langchain-0.0.1/src/trace_agent_langchain.egg-info/top_level.txt

@@ -0,0 +1 @@
+trace_agent_langchain

trace_agent_langchain-0.0.1/tests/test_langchain_callback.py

@@ -0,0 +1,277 @@
+"""Unit tests for TraceAgentLangChainCallback.
+
+These tests mock the HTTP layer and verify that LangChain callbacks are
+correctly translated into TraceAgent observations.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+langchain_core = pytest.importorskip("langchain_core")
+from trace_agent_langchain import TraceAgentLangChainCallback
+
+
+@pytest.fixture
+def mock_run_session() -> MagicMock:
+    session = MagicMock()
+    session.run_id = "run-123"
+    session.client._request = MagicMock(return_value={"status": "ok"})
+    return session
+
+
+@pytest.fixture
+def callback(mock_run_session) -> TraceAgentLangChainCallback:
+    return TraceAgentLangChainCallback(mock_run_session)
+
+
+class TestChainCallbacks:
+    def test_on_chain_start_posts_observation(self, callback, mock_run_session) -> None:
+        callback.on_chain_start(
+            serialized={"name": "LLMChain"},
+            inputs={"input": "hello"},
+            run_id="chain-1",
+            parent_run_id=None,
+        )
+        mock_run_session.client._request.assert_called_once()
+        args = mock_run_session.client._request.call_args
+        assert args[0][0] == "POST"
+        assert "/observations" in args[0][1]
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "chain_step"
+        assert obs["status"] == "running"
+        assert obs["step_id"] == "chain-1"
+        assert obs["metadata"]["chain_type"] == "LLMChain"
+
+    def test_on_chain_end_posts_observation(self, callback, mock_run_session) -> None:
+        callback._starts["chain-1"] = 0.0
+        callback.on_chain_end(
+            outputs={"output": "world"},
+            run_id="chain-1",
+            parent_run_id=None,
+        )
+        args = mock_run_session.client._request.call_args
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "chain_step"
+        assert obs["status"] == "succeeded"
+        assert obs["duration_ms"] >= 0
+
+    def test_on_chain_error_posts_failed(self, callback, mock_run_session) -> None:
+        callback._starts["chain-1"] = 0.0
+        callback.on_chain_error(
+            error=ValueError("boom"),
+            run_id="chain-1",
+            parent_run_id=None,
+        )
+        args = mock_run_session.client._request.call_args
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "chain_step"
+        assert obs["status"] == "failed"
+        assert obs["error_code"] == "ValueError"
+        assert "boom" in obs["error_summary"]
+
+
+class TestLLMCallbacks:
+    def test_on_llm_start_posts_observation(self, callback, mock_run_session) -> None:
+        callback.on_llm_start(
+            serialized={"name": "OpenAI"},
+            prompts=["hello"],
+            run_id="llm-1",
+            parent_run_id="chain-1",
+        )
+        args = mock_run_session.client._request.call_args
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "llm_call"
+        assert obs["status"] == "running"
+        assert obs["step_id"] == "llm-1"
+        assert obs["parent_step_id"] == "chain-1"
+
+    def test_on_llm_end_posts_observation(self, callback, mock_run_session) -> None:
+        callback._starts["llm-1"] = 0.0
+        # Mock LLMResult with OpenAI-style token usage
+        result = MagicMock()
+        result.llm_output = {"token_usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}
+        callback.on_llm_end(
+            response=result,
+            run_id="llm-1",
+            parent_run_id="chain-1",
+        )
+        args = mock_run_session.client._request.call_args
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "llm_call"
+        assert obs["status"] == "succeeded"
+        assert obs["metadata"]["prompt_tokens"] == 10
+        assert obs["metadata"]["completion_tokens"] == 5
+
+    def test_on_llm_error_posts_failed(self, callback, mock_run_session) -> None:
+        callback._starts["llm-1"] = 0.0
+        callback.on_llm_error(
+            error=RuntimeError("timeout"),
+            run_id="llm-1",
+            parent_run_id="chain-1",
+        )
+        args = mock_run_session.client._request.call_args
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "llm_call"
+        assert obs["status"] == "failed"
+        assert obs["error_code"] == "RuntimeError"
+
+    def test_extract_token_usage_openai(self, callback) -> None:
+        result = MagicMock()
+        result.llm_output = {"token_usage": {"prompt_tokens": 3, "completion_tokens": 2}}
+        usage = callback._extract_token_usage(result)
+        assert usage["prompt_tokens"] == 3
+        assert usage["completion_tokens"] == 2
+
+    def test_extract_token_usage_anthropic(self, callback) -> None:
+        result = MagicMock()
+        result.llm_output = {"token_usage": {"input_tokens": 4, "output_tokens": 1}}
+        usage = callback._extract_token_usage(result)
+        assert usage["prompt_tokens"] == 4
+        assert usage["completion_tokens"] == 1
+
+    def test_extract_token_usage_empty(self, callback) -> None:
+        assert callback._extract_token_usage(None) == {}
+
+
+class TestToolCallbacks:
+    def test_on_tool_start_posts_observation(self, callback, mock_run_session) -> None:
+        callback.on_tool_start(
+            serialized={"name": "search"},
+            input_str='{"q": "hello"}',
+            run_id="tool-1",
+            parent_run_id="llm-1",
+        )
+        args = mock_run_session.client._request.call_args
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "tool_call"
+        assert obs["status"] == "running"
+        assert obs["tool_name"] == "search"
+        assert obs["step_id"] == "tool-1"
+        assert obs["parent_step_id"] == "llm-1"
+
+    def test_on_tool_end_posts_result_and_observation(self, callback, mock_run_session) -> None:
+        callback._starts["tool-1"] = 0.0
+        callback.on_tool_end(
+            output="result here",
+            run_id="tool-1",
+            parent_run_id="llm-1",
+            metadata={"name": "search"},
+        )
+        assert mock_run_session.client._request.call_count == 2
+        # First call is tool-results
+        args0 = mock_run_session.client._request.call_args_list[0]
+        assert "/tool-results" in args0[0][1]
+        result = args0[0][2]["results"][0]
+        assert result["status"] == "succeeded"
+        assert result["name"] == "search"
+        # Second call is observation
+        args1 = mock_run_session.client._request.call_args_list[1]
+        obs = args1[0][2]["observations"][0]
+        assert obs["kind"] == "tool_call"
+        assert obs["status"] == "succeeded"
+
+    def test_on_tool_error_posts_failed(self, callback, mock_run_session) -> None:
+        callback._starts["tool-1"] = 0.0
+        callback.on_tool_error(
+            error=Exception("fail"),
+            run_id="tool-1",
+            parent_run_id="llm-1",
+            metadata={"name": "search"},
+        )
+        assert mock_run_session.client._request.call_count == 2
+        args0 = mock_run_session.client._request.call_args_list[0]
+        result = args0[0][2]["results"][0]
+        assert result["status"] == "failed"
+        assert result["error_code"] == "Exception"
+
+
+class TestRetrieverCallbacks:
+    def test_on_retriever_start(self, callback, mock_run_session) -> None:
+        callback.on_retriever_start(
+            serialized={"name": "VectorStore"},
+            query="AI agents",
+            run_id="ret-1",
+            parent_run_id="chain-1",
+        )
+        args = mock_run_session.client._request.call_args
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "retriever_query"
+        assert obs["status"] == "running"
+        assert "AI agents" in obs["input_summary"]
+
+    def test_on_retriever_end(self, callback, mock_run_session) -> None:
+        callback._starts["ret-1"] = 0.0
+        callback.on_retriever_end(
+            documents=[{"page_content": "doc1"}, {"page_content": "doc2"}],
+            run_id="ret-1",
+            parent_run_id="chain-1",
+        )
+        args = mock_run_session.client._request.call_args
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "retriever_query"
+        assert obs["status"] == "succeeded"
+        assert obs["metadata"]["document_count"] == 2
+
+    def test_retriever_events_disabled(self, mock_run_session) -> None:
+        cb = TraceAgentLangChainCallback(mock_run_session, record_retriever_events=False)
+        cb.on_retriever_start(
+            serialized={"name": "VectorStore"},
+            query="AI",
+            run_id="ret-1",
+        )
+        mock_run_session.client._request.assert_not_called()
+
+
+class TestAgentCallbacks:
+    def test_on_agent_action(self, callback, mock_run_session) -> None:
+        action = MagicMock()
+        action.tool = "search"
+        action.tool_input = {"q": "hello"}
+        callback.on_agent_action(
+            action=action,
+            run_id="agent-1",
+            parent_run_id="chain-1",
+        )
+        args = mock_run_session.client._request.call_args
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "agent_action"
+        assert obs["status"] == "running"
+        assert obs["tool_name"] == "search"
+        assert obs["metadata"]["tool_input"] == {"q": "hello"}
+
+    def test_on_agent_finish(self, callback, mock_run_session) -> None:
+        finish = MagicMock()
+        finish.return_values = {"output": "done"}
+        callback.on_agent_finish(
+            finish=finish,
+            run_id="agent-1",
+            parent_run_id="chain-1",
+        )
+        args = mock_run_session.client._request.call_args
+        obs = args[0][2]["observations"][0]
+        assert obs["kind"] == "agent_action"
+        assert obs["status"] == "succeeded"
+        assert obs["input_summary"] == "Agent finished"
+
+
+class TestResilience:
+    def test_backend_error_is_silenced(self, mock_run_session) -> None:
+        mock_run_session.client._request.side_effect = Exception("connection refused")
+        callback = TraceAgentLangChainCallback(mock_run_session)
+        # Should not raise
+        callback.on_chain_start(
+            serialized={"name": "Chain"},
+            inputs={},
+            run_id="chain-1",
+        )
+
+
+class TestSDKNotInstalled:
+    def test_raises_when_langchain_missing(self) -> None:
+        with patch("trace_agent_langchain.callback.BaseCallbackHandler", None):
+            mock_session = MagicMock()
+            with pytest.raises(RuntimeError, match="langchain-core is not installed"):
+                TraceAgentLangChainCallback(mock_session)

trace_agent_langchain-0.0.1/tests/test_langchain_integration.py

@@ -0,0 +1,167 @@
+"""Integration tests for LangChain callback with the TraceAgent backend.
+
+These tests require langchain to be installed (via `uv pip install -e ".[langchain]"`).
+They spin up the FastAPI app with an in-memory DB, create a run, invoke a LangChain
+chain with the TraceAgent callback, and verify that events are persisted.
+"""
+
+from __future__ import annotations
+
+import pytest
+from fastapi.testclient import TestClient
+
+from trace_agent_server.config import Settings
+from trace_agent_server.main import create_app
+from trace_agent_server.models import Event, Run
+
+# Skip entire module if langchain is not installed
+langchain = pytest.importorskip("langchain")
+langchain_core = pytest.importorskip("langchain_core")
+
+
+@pytest.fixture
+def test_client(tmp_path) -> TestClient:
+    settings = Settings(
+        database_url=f"sqlite:///{tmp_path / 'test.db'}",
+        cors_origins=["http://localhost:5173"],
+    )
+    app = create_app(settings)
+    return TestClient(app)
+
+
+class TestLangChainCallbackIntegration:
+    def test_chain_events_are_persisted(self, test_client) -> None:
+        client = test_client
+
+        # 1. Create a run via the API
+        resp = client.post(
+            "/api/runs",
+            json={"agent_name": "test-agent", "goal": "test chain observability"},
+        )
+        assert resp.status_code == 200
+        run_id = resp.json()["run_id"]
+
+        # 2. Build a simple LangChain chain and attach the callback
+        from langchain_core.prompts import ChatPromptTemplate
+        from langchain_core.runnables import RunnableLambda
+
+        from trace_agent_langchain import TraceAgentLangChainCallback
+        from trace_agent_sdk import RunSession, TraceAgentClient
+
+        TRACE_AGENT_client = TraceAgentClient(base_url="", http_client=client)
+        # Manually construct a RunSession since we bypass start_run()
+        run_session = RunSession(
+            client=TRACE_AGENT_client,
+            run_id=run_id,
+            agent_name="test-agent",
+            goal="test",
+        )
+        callback = TraceAgentLangChainCallback(run_session)
+
+        # 3. Invoke a simple Runnable sequence
+        chain = RunnableLambda(lambda x: x.upper())
+        result = chain.invoke("hello", config={"callbacks": [callback]})
+        assert result == "HELLO"
+
+        # 4. Verify events were persisted
+        with client.app.state.session_factory() as session:
+            run = session.query(Run).filter_by(id=run_id).first()
+            assert run is not None
+            events = session.query(Event).filter_by(run_id=run_id).order_by(Event.seq).all()
+            # Should have at least chain.started + chain.completed
+            assert len(events) >= 2
+            types = [e.type for e in events]
+            assert "chain.started" in types
+            assert "chain.completed" in types
+
+    def test_llm_events_with_mocked_provider(self, test_client) -> None:
+        """Test that LLM calls produce the correct event types.
+
+        We mock the LLM so no real API key is needed.
+        """
+        client = test_client
+
+        resp = client.post(
+            "/api/runs",
+            json={"agent_name": "test-llm-agent", "goal": "test llm events"},
+        )
+        run_id = resp.json()["run_id"]
+
+        from unittest.mock import MagicMock, patch
+
+        from langchain_core.outputs import Generation, LLMResult
+
+        from trace_agent_langchain import TraceAgentLangChainCallback
+        from trace_agent_sdk import RunSession, TraceAgentClient
+
+        TRACE_AGENT_client = TraceAgentClient(base_url="", http_client=client)
+        run_session = RunSession(
+            client=TRACE_AGENT_client,
+            run_id=run_id,
+            agent_name="test-llm-agent",
+            goal="test",
+        )
+        callback = TraceAgentLangChainCallback(run_session)
+
+        # Simulate an LLM call with token usage
+        mock_result = LLMResult(
+            generations=[[Generation(text="hello")]],
+            llm_output={"token_usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}},
+        )
+
+        callback.on_llm_start(
+            serialized={"name": "FakeLLM"},
+            prompts=["hi"],
+            run_id="llm-1",
+        )
+        callback.on_llm_end(
+            response=mock_result,
+            run_id="llm-1",
+        )
+
+        with client.app.state.session_factory() as session:
+            events = session.query(Event).filter_by(run_id=run_id).order_by(Event.seq).all()
+            types = [e.type for e in events]
+            assert "llm.started" in types
+            assert "llm.completed" in types
+            llm_completed = next(e for e in events if e.type == "llm.completed")
+            assert llm_completed.event_metadata.get("prompt_tokens") == 5
+            assert llm_completed.event_metadata.get("completion_tokens") == 2
+
+    def test_tool_events_are_persisted(self, test_client) -> None:
+        client = test_client
+
+        resp = client.post(
+            "/api/runs",
+            json={"agent_name": "test-tool-agent", "goal": "test tool events"},
+        )
+        run_id = resp.json()["run_id"]
+
+        from trace_agent_langchain import TraceAgentLangChainCallback
+        from trace_agent_sdk import RunSession, TraceAgentClient
+
+        TRACE_AGENT_client = TraceAgentClient(base_url="", http_client=client)
+        run_session = RunSession(
+            client=TRACE_AGENT_client,
+            run_id=run_id,
+            agent_name="test-tool-agent",
+            goal="test",
+        )
+        callback = TraceAgentLangChainCallback(run_session)
+
+        callback.on_tool_start(
+            serialized={"name": "search"},
+            input_str='{"q": "hello"}',
+            run_id="tool-1",
+        )
+        callback.on_tool_end(
+            output="found it",
+            run_id="tool-1",
+            metadata={"name": "search"},
+        )
+
+        with client.app.state.session_factory() as session:
+            events = session.query(Event).filter_by(run_id=run_id).order_by(Event.seq).all()
+            types = [e.type for e in events]
+            assert "tool.started" in types
+            assert "tool.succeeded" in types