agentreplay 0.1.2__py3-none-any.whl

agentreplay/langchain_tracer.py

@@ -0,0 +1,385 @@
+# Copyright 2025 Sushanth (https://github.com/sushanthpy)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+LangChain callback handler that emits OpenTelemetry spans with proper hierarchy.
+
+This creates parent-child span relationships for:
+- Chains → LLM calls
+- Agents → Tool calls → LLM calls
+- RAG pipelines → Retrieval → LLM synthesis
+
+Integrates with Agentreplay's zero-code instrumentation.
+"""
+
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID
+import time
+
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain_core.outputs import LLMResult
+from langchain_core.agents import AgentAction, AgentFinish
+from langchain_core.documents import Document
+
+try:
+    from opentelemetry import trace
+    from opentelemetry.trace import Status, StatusCode
+    OTEL_AVAILABLE = True
+except ImportError:
+    OTEL_AVAILABLE = False
+
+
+class AgentreplayCallbackHandler(BaseCallbackHandler):
+    """LangChain callback handler that creates hierarchical OTEL spans."""
+    
+    def __init__(self):
+        super().__init__()
+        self.tracer = trace.get_tracer(__name__) if OTEL_AVAILABLE else None
+        self.spans: Dict[str, Any] = {}  # run_id -> span
+        self.parent_map: Dict[str, str] = {}  # run_id -> parent_run_id
+        
+    def _get_parent_span(self, parent_run_id: Optional[UUID]) -> Optional[Any]:
+        """Get parent span from run_id."""
+        if not parent_run_id or not self.tracer:
+            return None
+        parent_id = str(parent_run_id)
+        return self.spans.get(parent_id)
+    
+    def _start_span(self, name: str, run_id: UUID, parent_run_id: Optional[UUID] = None, **attributes) -> Any:
+        """Start a new OTEL span with optional parent."""
+        if not self.tracer:
+            return None
+            
+        run_id_str = str(run_id)
+        parent_span = self._get_parent_span(parent_run_id)
+        
+        # Create span with parent context
+        if parent_span:
+            ctx = trace.set_span_in_context(parent_span)
+            span = self.tracer.start_span(name, context=ctx)
+        else:
+            span = self.tracer.start_span(name)
+        
+        # Set attributes
+        for key, value in attributes.items():
+            if value is not None:
+                span.set_attribute(key, str(value))
+        
+        self.spans[run_id_str] = span
+        if parent_run_id:
+            self.parent_map[run_id_str] = str(parent_run_id)
+            
+        return span
+    
+    def _end_span(self, run_id: UUID, status: Optional[StatusCode] = None, error: Optional[str] = None):
+        """End a span and clean up."""
+        if not self.tracer:
+            return
+            
+        run_id_str = str(run_id)
+        span = self.spans.pop(run_id_str, None)
+        
+        if span:
+            if error:
+                span.set_status(Status(StatusCode.ERROR, error))
+                span.record_exception(Exception(error))
+            elif status:
+                span.set_status(Status(status))
+            else:
+                span.set_status(Status(StatusCode.OK))
+            span.end()
+        
+        self.parent_map.pop(run_id_str, None)
+    
+    # ===== Chain Callbacks =====
+    
+    def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a chain starts running."""
+        chain_name = serialized.get("name", serialized.get("id", ["unknown"])[-1])
+        
+        self._start_span(
+            f"chain.{chain_name}",
+            run_id,
+            parent_run_id,
+            **{
+                "chain.name": chain_name,
+                "chain.type": serialized.get("id", ["unknown"])[0],
+                "chain.inputs": str(inputs)[:1000],  # Truncate
+                "span.type": "chain",
+            }
+        )
+    
+    def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a chain finishes running."""
+        span = self.spans.get(str(run_id))
+        if span:
+            span.set_attribute("chain.outputs", str(outputs)[:1000])
+        self._end_span(run_id, StatusCode.OK)
+    
+    def on_chain_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a chain errors."""
+        self._end_span(run_id, StatusCode.ERROR, str(error))
+    
+    # ===== LLM Callbacks =====
+    
+    def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when LLM starts running."""
+        model_name = serialized.get("name", "unknown")
+        
+        self._start_span(
+            f"llm.{model_name}",
+            run_id,
+            parent_run_id,
+            **{
+                "llm.model": model_name,
+                "llm.prompts": str(prompts)[:2000],
+                "llm.prompt_count": len(prompts),
+                "span.type": "llm",
+            }
+        )
+    
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when LLM finishes running."""
+        span = self.spans.get(str(run_id))
+        if span:
+            # Extract token usage
+            if hasattr(response, "llm_output") and response.llm_output:
+                token_usage = response.llm_output.get("token_usage", {})
+                if token_usage:
+                    span.set_attribute("llm.tokens.prompt", token_usage.get("prompt_tokens", 0))
+                    span.set_attribute("llm.tokens.completion", token_usage.get("completion_tokens", 0))
+                    span.set_attribute("llm.tokens.total", token_usage.get("total_tokens", 0))
+            
+            # Extract generations
+            if response.generations:
+                first_gen = response.generations[0][0] if response.generations[0] else None
+                if first_gen:
+                    span.set_attribute("llm.response", str(first_gen.text)[:2000])
+        
+        self._end_span(run_id, StatusCode.OK)
+    
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when LLM errors."""
+        self._end_span(run_id, StatusCode.ERROR, str(error))
+    
+    # ===== Tool Callbacks =====
+    
+    def on_tool_start(
+        self,
+        serialized: Dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        inputs: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a tool starts running."""
+        tool_name = serialized.get("name", "unknown_tool")
+        
+        self._start_span(
+            f"tool.{tool_name}",
+            run_id,
+            parent_run_id,
+            **{
+                "tool.name": tool_name,
+                "tool.description": serialized.get("description", ""),
+                "tool.input": input_str[:1000],
+                "span.type": "tool",
+            }
+        )
+    
+    def on_tool_end(
+        self,
+        output: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a tool finishes running."""
+        span = self.spans.get(str(run_id))
+        if span:
+            span.set_attribute("tool.output", str(output)[:2000])
+        self._end_span(run_id, StatusCode.OK)
+    
+    def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a tool errors."""
+        self._end_span(run_id, StatusCode.ERROR, str(error))
+    
+    # ===== Agent Callbacks =====
+    
+    def on_agent_action(
+        self,
+        action: AgentAction,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when an agent takes an action."""
+        span = self.spans.get(str(parent_run_id) if parent_run_id else str(run_id))
+        if span:
+            span.add_event(
+                "agent.action",
+                attributes={
+                    "tool": action.tool,
+                    "tool_input": str(action.tool_input)[:1000],
+                    "log": action.log[:500] if action.log else "",
+                }
+            )
+    
+    def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when an agent finishes."""
+        span = self.spans.get(str(parent_run_id) if parent_run_id else str(run_id))
+        if span:
+            span.add_event(
+                "agent.finish",
+                attributes={
+                    "return_values": str(finish.return_values)[:1000],
+                    "log": finish.log[:500] if finish.log else "",
+                }
+            )
+    
+    # ===== Retriever Callbacks =====
+    
+    def on_retriever_start(
+        self,
+        serialized: Dict[str, Any],
+        query: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when retriever starts."""
+        self._start_span(
+            "retriever.search",
+            run_id,
+            parent_run_id,
+            **{
+                "retriever.query": query[:1000],
+                "retriever.type": serialized.get("name", "unknown"),
+                "span.type": "retrieval",
+            }
+        )
+    
+    def on_retriever_end(
+        self,
+        documents: List[Document],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when retriever finishes."""
+        span = self.spans.get(str(run_id))
+        if span:
+            span.set_attribute("retriever.document_count", len(documents))
+            # Store top 3 document snippets
+            for i, doc in enumerate(documents[:3]):
+                span.set_attribute(
+                    f"retriever.doc_{i+1}",
+                    doc.page_content[:500]
+                )
+                if doc.metadata:
+                    span.set_attribute(
+                        f"retriever.doc_{i+1}_metadata",
+                        str(doc.metadata)[:200]
+                    )
+        self._end_span(run_id, StatusCode.OK)
+    
+    def on_retriever_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when retriever errors."""
+        self._end_span(run_id, StatusCode.ERROR, str(error))
+
+
+def get_agentreplay_callback() -> Optional[AgentreplayCallbackHandler]:
+    """Get Agentreplay callback handler if OTEL is available."""
+    if not OTEL_AVAILABLE:
+        return None
+    return AgentreplayCallbackHandler()

agentreplay/models.py

@@ -0,0 +1,120 @@
+# Copyright 2025 Sushanth (https://github.com/sushanthpy)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Data models for Agentreplay SDK."""
+
+from enum import IntEnum
+from typing import Optional
+from pydantic import BaseModel, Field
+import time
+
+
+class SpanType(IntEnum):
+    """Agent execution span types."""
+
+    ROOT = 0
+    PLANNING = 1
+    REASONING = 2
+    TOOL_CALL = 3
+    TOOL_RESPONSE = 4
+    SYNTHESIS = 5
+    RESPONSE = 6
+    ERROR = 7
+    CUSTOM = 255
+    
+    # Backward compatibility aliases
+    AGENT = 0  # Alias for ROOT
+    TOOL = 3   # Alias for TOOL_CALL
+
+
+class SensitivityFlags(IntEnum):
+    """Sensitivity flags for PII and redaction control."""
+
+    NONE = 0
+    PII = 1 << 0  # Contains personally identifiable information
+    SECRET = 1 << 1  # Contains secrets/credentials
+    INTERNAL = 1 << 2  # Internal-only data
+    NO_EMBED = 1 << 3  # Never embed in vector index
+
+
+class AgentFlowEdge(BaseModel):
+    """AgentFlow Edge - represents one step in agent execution.
+    
+    This is the fundamental unit of data in Agentreplay.
+    Fixed 128-byte format when serialized.
+    """
+
+    # Identity & Causality
+    edge_id: int = Field(default=0, description="Unique edge identifier (u128)")
+    causal_parent: int = Field(default=0, description="Parent edge ID (0 for root)")
+
+    # Temporal
+    timestamp_us: int = Field(
+        default_factory=lambda: int(time.time() * 1_000_000),
+        description="Timestamp in microseconds since epoch",
+    )
+    logical_clock: int = Field(default=0, description="Lamport logical clock")
+
+    # Multi-tenancy
+    tenant_id: int = Field(description="Tenant identifier")
+    project_id: int = Field(default=0, description="Project identifier within tenant")
+    schema_version: int = Field(default=2, description="AFF schema version")
+    sensitivity_flags: int = Field(
+        default=SensitivityFlags.NONE, description="Sensitivity/privacy flags"
+    )
+
+    # Context
+    agent_id: int = Field(description="Agent identifier")
+    session_id: int = Field(description="Session/conversation identifier")
+    span_type: SpanType = Field(description="Type of agent execution span")
+    parent_count: int = Field(default=1, description="Number of parents (>1 for DAG fan-in)")
+
+    # Probabilistic / Cost
+    confidence: float = Field(default=1.0, ge=0.0, le=1.0, description="Confidence score")
+    token_count: int = Field(default=0, ge=0, description="Number of tokens used")
+    duration_us: int = Field(default=0, ge=0, description="Duration in microseconds")
+    sampling_rate: float = Field(default=1.0, ge=0.0, le=1.0, description="Sampling rate")
+
+    # Payload metadata
+    compression_type: int = Field(default=0, description="Compression type (0=None, 1=LZ4, 2=ZSTD)")
+    has_payload: bool = Field(default=False, description="Whether payload data exists")
+
+    # Metadata
+    flags: int = Field(default=0, description="General purpose flags")
+    checksum: int = Field(default=0, description="BLAKE3 checksum for integrity")
+
+    class Config:
+        """Pydantic config."""
+
+        use_enum_values = True
+
+
+class QueryFilter(BaseModel):
+    """Filter for querying edges."""
+
+    tenant_id: Optional[int] = None
+    project_id: Optional[int] = None
+    agent_id: Optional[int] = None
+    session_id: Optional[int] = None
+    span_type: Optional[SpanType] = None
+    min_confidence: Optional[float] = None
+    exclude_pii: bool = False
+
+
+class QueryResponse(BaseModel):
+    """Response from query operations."""
+
+    edges: list[AgentFlowEdge]
+    total_count: int
+    has_more: bool = False