agentsnap 0.1.0__py3-none-any.whl

agentsnap/__init__.py

@@ -0,0 +1,23 @@
+from agentsnap import config
+from agentsnap.core.asserter import AgentAsserter
+from agentsnap.core.diff import LLMJudge
+from agentsnap.core.recorder import AgentRecorder
+from agentsnap.exceptions import (
+    AdapterNotWrappedError,
+    AgentRegressionError,
+    SnapshotNotFoundError,
+)
+from agentsnap.patches import PatchSet
+from agentsnap.wrap import wrap
+
+__all__ = [
+    "wrap",
+    "AgentRecorder",
+    "AgentAsserter",
+    "LLMJudge",
+    "PatchSet",
+    "AgentRegressionError",
+    "SnapshotNotFoundError",
+    "AdapterNotWrappedError",
+    "config",
+]

agentsnap/adapters/anthropic.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from agentsnap.core.recorder import TraceAccumulator
+
+
+class _MessagesProxy:
+    def __init__(self, original) -> None:
+        self._original = original
+
+    def create(self, **kwargs):
+        acc = TraceAccumulator.current()
+        if acc is None:
+            return self._original.create(**kwargs)
+
+        messages = kwargs.get("messages", [])
+        response = self._original.create(**kwargs)
+
+        response_text = ""
+        tokens = 0
+        if hasattr(response, "content"):
+            for block in response.content:
+                if hasattr(block, "text"):
+                    response_text += block.text
+        if hasattr(response, "usage"):
+            tokens = getattr(response.usage, "input_tokens", 0) + getattr(
+                response.usage, "output_tokens", 0
+            )
+
+        acc.push(
+            {
+                "type": "llm_call",
+                "messages": messages,
+                "response": response_text,
+                "tokens": tokens,
+            }
+        )
+        return response
+
+    def __getattr__(self, name: str):
+        return getattr(self._original, name)
+
+
+class AnthropicAdapter:
+    """Wraps an anthropic.Anthropic() client to intercept .messages.create()."""
+
+    def __init__(self, client) -> None:
+        self._client = client
+        self.messages = _MessagesProxy(client.messages)
+
+    def __getattr__(self, name: str):
+        return getattr(self._client, name)

agentsnap/adapters/cohere.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from agentsnap.core.recorder import TraceAccumulator
+
+
+class _CohereV2ChatProxy:
+    """Intercepts cohere.ClientV2().chat()."""
+
+    def __init__(self, original_chat_fn) -> None:
+        self._original = original_chat_fn
+
+    def __call__(self, **kwargs):
+        acc = TraceAccumulator.current()
+        if acc is None:
+            return self._original(**kwargs)
+
+        messages = kwargs.get("messages", [])
+        response = self._original(**kwargs)
+
+        response_text = ""
+        tokens = 0
+        if hasattr(response, "message") and hasattr(response.message, "content"):
+            for block in response.message.content:
+                if hasattr(block, "text"):
+                    response_text += block.text
+        if hasattr(response, "usage"):
+            tokens = getattr(response.usage, "tokens", None)
+            if tokens is None:
+                inp = getattr(response.usage, "input_tokens", 0) or 0
+                out = getattr(response.usage, "output_tokens", 0) or 0
+                tokens = inp + out
+
+        acc.push(
+            {
+                "type": "llm_call",
+                "messages": messages,
+                "response": response_text,
+                "tokens": tokens or 0,
+            }
+        )
+        return response
+
+
+class CohereAdapter:
+    """Wraps a cohere.ClientV2() to intercept .chat().
+
+    Usage:
+        import cohere
+        from agentsnap.adapters.cohere import CohereAdapter
+        client = CohereAdapter(cohere.ClientV2(api_key="..."))
+        response = client.chat(model="command-r-plus", messages=[...])
+    """
+
+    def __init__(self, client) -> None:
+        self._client = client
+        self.chat = _CohereV2ChatProxy(client.chat)
+
+    def __getattr__(self, name: str):
+        return getattr(self._client, name)

agentsnap/adapters/google.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from agentsnap.core.recorder import TraceAccumulator
+
+
+class _ModelsProxy:
+    def __init__(self, original) -> None:
+        self._original = original
+
+    def generate_content(self, model: str, contents, **kwargs):
+        acc = TraceAccumulator.current()
+        if acc is None:
+            return self._original.generate_content(model=model, contents=contents, **kwargs)
+
+        if isinstance(contents, str):
+            messages = [{"role": "user", "content": contents}]
+        elif isinstance(contents, list):
+            messages = contents
+        else:
+            messages = [{"role": "user", "content": str(contents)}]
+
+        response = self._original.generate_content(model=model, contents=contents, **kwargs)
+
+        response_text = ""
+        tokens = 0
+        if hasattr(response, "text"):
+            response_text = response.text or ""
+        elif hasattr(response, "candidates") and response.candidates:
+            for part in response.candidates[0].content.parts:
+                if hasattr(part, "text"):
+                    response_text += part.text
+        if hasattr(response, "usage_metadata"):
+            tokens = getattr(response.usage_metadata, "total_token_count", 0)
+
+        acc.push(
+            {
+                "type": "llm_call",
+                "messages": messages,
+                "response": response_text,
+                "tokens": tokens,
+            }
+        )
+        return response
+
+    def __getattr__(self, name: str):
+        return getattr(self._original, name)
+
+
+class GeminiAdapter:
+    """Wraps a google.genai.Client() to intercept .models.generate_content().
+
+    Usage:
+        from google import genai
+        from agentsnap.adapters.google import GeminiAdapter
+        client = GeminiAdapter(genai.Client(api_key="..."))
+        response = client.models.generate_content(model="gemini-2.0-flash", contents="Hello")
+    """
+
+    def __init__(self, client) -> None:
+        self._client = client
+        self.models = _ModelsProxy(client.models)
+
+    def __getattr__(self, name: str):
+        return getattr(self._client, name)

agentsnap/adapters/groq.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from agentsnap.adapters.openai import OpenAIAdapter
+
+
+class GroqAdapter(OpenAIAdapter):
+    """Wraps a groq.Groq() client (OpenAI-compatible interface).
+
+    Groq uses the same .chat.completions.create() interface as OpenAI,
+    so this is a thin alias over OpenAIAdapter.
+
+    Usage:
+        from groq import Groq
+        from agentsnap.adapters.groq import GroqAdapter
+        client = GroqAdapter(Groq(api_key="..."))
+        response = client.chat.completions.create(
+            model="llama-3.3-70b-versatile", messages=[...]
+        )
+    """

agentsnap/adapters/langgraph.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from agentsnap.core.recorder import TraceAccumulator
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler as _Base
+except ImportError:
+    _Base = object  # type: ignore[assignment,misc]
+
+
+class AgentSnapCallback(_Base):
+    """LangChain callback handler that records LLM and tool events into TraceAccumulator.
+
+    Injected automatically by LangGraphAdapter. Also works standalone with any
+    LangChain model or chain that supports the callbacks API.
+
+    When langchain_core is not installed, _Base is object — the class is still
+    importable and its methods callable; it just won't be registered as a real
+    LangChain callback (duck-typing makes it work with our mock graphs in tests).
+    """
+
+    def on_llm_end(self, response, **kwargs) -> None:
+        acc = TraceAccumulator.current()
+        if acc is None:
+            return
+        text = ""
+        if hasattr(response, "generations") and response.generations:
+            gen = response.generations[0][0]
+            if hasattr(gen, "text") and gen.text:
+                text = gen.text
+            elif hasattr(gen, "message") and hasattr(gen.message, "content"):
+                text = gen.message.content or ""
+            else:
+                text = str(gen)
+        acc.push({"type": "llm_call", "messages": [], "response": text, "tokens": 0})
+
+    def on_tool_end(self, output, *, name: str = "", **kwargs) -> None:
+        acc = TraceAccumulator.current()
+        if acc is None:
+            return
+        acc.push({"type": "tool_call", "name": name, "args": {}, "result": str(output)})
+
+
+class LangGraphAdapter:
+    """Wraps a CompiledGraph to capture traces via the LangChain callbacks system.
+
+    Injects AgentSnapCallback into every invoke() call so that LLM and tool
+    events from individual graph nodes are recorded into the active
+    TraceAccumulator.
+
+    Falls back to recording the top-level result as a single llm_call event
+    when langchain_core is not installed (no callbacks support available).
+    """
+
+    def __init__(self, graph) -> None:
+        self._graph = graph
+
+    def invoke(self, input_data, **kwargs):
+        acc = TraceAccumulator.current()
+        if acc is None:
+            return self._graph.invoke(input_data, **kwargs)
+
+        # Inject AgentSnapCallback via config so node-level events are captured.
+        # AgentSnapCallback is always usable via duck-typing even without langchain_core.
+        # When langchain_core is absent (_Base is object), we still inject the callback
+        # so that graphs that accept a callbacks list (e.g. LangGraph) fire it correctly.
+        # The top-level fallback push only runs when langchain_core is absent, because
+        # in that case a real CompiledGraph won't know to call on_llm_end/on_tool_end.
+        config = dict(kwargs.pop("config", None) or {})
+        callbacks = list(config.get("callbacks") or [])
+        callbacks.append(AgentSnapCallback())
+        config["callbacks"] = callbacks
+
+        if _Base is not object:
+            # langchain_core available: callback will be invoked by the real runtime
+            return self._graph.invoke(input_data, config=config, **kwargs)
+
+        # Fallback: langchain_core absent — real CompiledGraph won't fire callbacks,
+        # so record the top-level invocation as a single llm_call event instead.
+        result = self._graph.invoke(input_data, config=config, **kwargs)
+        if not any(e["type"] == "llm_call" for e in acc.trace):
+            acc.push({
+                "type": "llm_call",
+                "messages": [{"role": "user", "content": str(input_data)}],
+                "response": str(result),
+                "tokens": 0,
+            })
+        return result
+
+    def stream(self, input_data, **kwargs):
+        return self._graph.stream(input_data, **kwargs)
+
+    def __getattr__(self, name: str):
+        return getattr(self._graph, name)

agentsnap/adapters/mistral.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from agentsnap.core.recorder import TraceAccumulator
+
+
+class _MistralCompletionsProxy:
+    def __init__(self, original) -> None:
+        self._original = original
+
+    def complete(self, **kwargs):
+        acc = TraceAccumulator.current()
+        if acc is None:
+            return self._original.complete(**kwargs)
+
+        messages = kwargs.get("messages", [])
+        kwargs["stream"] = False
+        response = self._original.complete(**kwargs)
+
+        response_text = ""
+        tokens = 0
+        if hasattr(response, "choices") and response.choices:
+            msg = response.choices[0].message
+            response_text = msg.content or ""
+        if hasattr(response, "usage"):
+            tokens = getattr(response.usage, "total_tokens", 0) or 0
+
+        acc.push(
+            {
+                "type": "llm_call",
+                "messages": messages,
+                "response": response_text,
+                "tokens": tokens,
+            }
+        )
+        return response
+
+    def __getattr__(self, name: str):
+        return getattr(self._original, name)
+
+
+class _MistralChatProxy:
+    def __init__(self, chat) -> None:
+        self._chat = chat
+        self.complete = _MistralCompletionsProxy(chat).complete
+
+    def __getattr__(self, name: str):
+        return getattr(self._chat, name)
+
+
+class MistralAdapter:
+    """Wraps a mistralai.Mistral() client to intercept .chat.complete().
+
+    Usage:
+        from mistralai import Mistral
+        from agentsnap.adapters.mistral import MistralAdapter
+        client = MistralAdapter(Mistral(api_key="..."))
+        response = client.chat.complete(model="mistral-large-latest", messages=[...])
+    """
+
+    def __init__(self, client) -> None:
+        self._client = client
+        self.chat = _MistralChatProxy(client.chat)
+
+    def __getattr__(self, name: str):
+        return getattr(self._client, name)

agentsnap/adapters/openai.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from agentsnap.core.recorder import TraceAccumulator
+
+
+class _CompletionsProxy:
+    def __init__(self, original) -> None:
+        self._original = original
+
+    def create(self, **kwargs):
+        acc = TraceAccumulator.current()
+        if acc is None:
+            return self._original.create(**kwargs)
+
+        messages = kwargs.get("messages", [])
+        # Force non-streaming so we always get a complete ChatCompletion object.
+        # Streaming responses expose deltas, not the full message content.
+        kwargs["stream"] = False
+        response = self._original.create(**kwargs)
+
+        response_text = ""
+        tokens = 0
+        if hasattr(response, "choices") and response.choices:
+            msg = response.choices[0].message
+            response_text = msg.content or ""
+        if hasattr(response, "usage"):
+            tokens = getattr(response.usage, "total_tokens", 0)
+
+        acc.push(
+            {
+                "type": "llm_call",
+                "messages": messages,
+                "response": response_text,
+                "tokens": tokens,
+            }
+        )
+        return response
+
+    def __getattr__(self, name: str):
+        return getattr(self._original, name)
+
+
+class _ChatProxy:
+    def __init__(self, chat) -> None:
+        self._chat = chat
+        self.completions = _CompletionsProxy(chat.completions)
+
+    def __getattr__(self, name: str):
+        return getattr(self._chat, name)
+
+
+class OpenAIAdapter:
+    """Wraps an openai.OpenAI() client to intercept .chat.completions.create()."""
+
+    def __init__(self, client) -> None:
+        self._client = client
+        self.chat = _ChatProxy(client.chat)
+
+    def __getattr__(self, name: str):
+        return getattr(self._client, name)

agentsnap/adapters/openrouter.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from agentsnap.adapters.openai import OpenAIAdapter
+
+OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
+
+
+class OpenRouterAdapter(OpenAIAdapter):
+    """Wraps an openai.OpenAI() client pointed at OpenRouter.
+
+    OpenRouter exposes an OpenAI-compatible API that proxies 300+ models
+    (Claude, GPT, Gemini, Llama, Mistral, etc.) under one key.
+
+    Usage:
+        import openai
+        from agentsnap.adapters.openrouter import OpenRouterAdapter, OPENROUTER_BASE_URL
+
+        client = OpenRouterAdapter(
+            openai.OpenAI(
+                api_key="sk-or-...",
+                base_url=OPENROUTER_BASE_URL,
+            )
+        )
+        response = client.chat.completions.create(
+            model="anthropic/claude-haiku-4-5",   # or any OpenRouter model slug
+            messages=[{"role": "user", "content": "Hello"}],
+        )
+    """

agentsnap/adapters/tool.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from agentsnap.core.recorder import TraceAccumulator
+
+
+class ToolAdapter:
+    """Wraps any callable to intercept and record tool calls."""
+
+    def __init__(self, func: Callable, name: str | None = None) -> None:
+        self._func = func
+        self._name = name or getattr(func, "__name__", "unknown_tool")
+
+    def __call__(self, **kwargs: Any) -> Any:
+        acc = TraceAccumulator.current()
+        if acc is None:
+            return self._func(**kwargs)
+
+        result = self._func(**kwargs)
+        acc.push(
+            {
+                "type": "tool_call",
+                "name": self._name,
+                "args": kwargs,
+                "result": str(result),
+            }
+        )
+        return result
+
+    @property
+    def __name__(self) -> str:
+        return self._name
+
+    def __repr__(self) -> str:
+        return f"ToolAdapter({self._name!r})"