prismtrace-sdk 0.3.2__py3-none-any.whl

prismtrace/__init__.py

@@ -0,0 +1,34 @@
+from .client import PRISMtrace
+from .langchain_handler import PRISMtraceCallbackHandler
+from .claude_tracer import ClaudeAgentTracer
+from .langgraph_helper import PRISMtraceLangGraphHandler, wrap_graph as wrap_langgraph
+from .google_adk import PRISMtraceADKAdapter
+from .litellm_callback import install as install_litellm
+from .openai_agents import PRISMtraceTracingProcessor, install as install_openai_agents
+from .elevenlabs_voice import (
+    PRISMtraceVoiceTracer,
+    install as install_elevenlabs_voice,
+)
+
+__version__ = "0.3.2"
+__all__ = [
+    # Core
+    "PRISMtrace",
+    # LangChain (Supported)
+    "PRISMtraceCallbackHandler",
+    # Anthropic direct (Supported)
+    "ClaudeAgentTracer",
+    # LangGraph (Supported)
+    "PRISMtraceLangGraphHandler",
+    "wrap_langgraph",
+    # Google ADK (Beta)
+    "PRISMtraceADKAdapter",
+    # LiteLLM (Universal Gateway, Supported)
+    "install_litellm",
+    # OpenAI Agents SDK (Preview)
+    "PRISMtraceTracingProcessor",
+    "install_openai_agents",
+    # ElevenLabs Agents — voice (Supported)
+    "PRISMtraceVoiceTracer",
+    "install_elevenlabs_voice",
+]

prismtrace/claude_tracer.py

@@ -0,0 +1,363 @@
+"""Claude agent tracer — wraps Anthropic client to auto-trace agentic tool use loops.
+
+Emits both spans (for the trace detail view) AND a trajectory (for PRISM
+evaluation + trajectory analytics) after each agentic run.
+"""
+
+import json
+import time
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Callable, Optional
+
+import httpx
+
+
+class ClaudeAgentTracer:
+    """Wraps an Anthropic client to trace agentic tool use loops automatically."""
+
+    def __init__(
+        self,
+        anthropic_client: Any,
+        api_key: str,
+        project_id: str,
+        endpoint: str = "https://prismtrace-production.up.railway.app",
+        agent_name: str = "claude-agent",
+        emit_trajectory: bool = True,
+    ):
+        self.client = anthropic_client
+        self.api_key = api_key
+        self.project_id = project_id
+        self.endpoint = endpoint.rstrip("/")
+        self.agent_name = agent_name
+        self.emit_trajectory = emit_trajectory
+        self._http = httpx.Client(
+            headers={
+                "x-prismtrace-key": api_key,
+                "Content-Type": "application/json",
+            },
+            timeout=15,
+        )
+
+    def _now_iso(self) -> str:
+        return datetime.now(timezone.utc).isoformat()
+
+    def run(
+        self,
+        messages: list[dict],
+        tools: list[dict],
+        system: str = "",
+        model: str = "claude-sonnet-4-5-20250514",
+        max_iterations: int = 10,
+        tool_executor: Optional[Callable[[str, dict], str]] = None,
+        session_id: Optional[str] = None,
+    ) -> dict:
+        """Run the full agentic loop, tracing every step.
+
+        Emits spans to /api/spans/ingest AND a trajectory to /api/trajectories.
+        """
+        trace_id = str(uuid.uuid4())
+        conversation_id = session_id or str(uuid.uuid4())
+        root_span_id = str(uuid.uuid4())
+        spans: list[dict] = []
+        trajectory_steps: list[dict] = []
+        current_messages = list(messages)
+        iteration = 0
+        total_duration_ms = 0
+
+        loop_start = self._now_iso()
+
+        while iteration < max_iterations:
+            iteration += 1
+            step_span_id = str(uuid.uuid4())
+
+            # --- LLM call ---
+            llm_span_id = str(uuid.uuid4())
+            llm_start = self._now_iso()
+            t0 = time.perf_counter()
+
+            kwargs: dict[str, Any] = {
+                "model": model,
+                "max_tokens": 4096,
+                "messages": current_messages,
+                "tools": tools,
+            }
+            if system:
+                kwargs["system"] = system
+
+            response = self.client.messages.create(**kwargs)
+            t1 = time.perf_counter()
+
+            llm_end = self._now_iso()
+            llm_duration = (t1 - t0) * 1000
+            total_duration_ms += int(llm_duration)
+
+            usage = getattr(response, "usage", None)
+            input_tokens = getattr(usage, "input_tokens", 0) if usage else 0
+            output_tokens = getattr(usage, "output_tokens", 0) if usage else 0
+
+            resp_text_parts = []
+            tool_uses = []
+            for block in response.content:
+                if block.type == "text":
+                    resp_text_parts.append(block.text)
+                elif block.type == "tool_use":
+                    tool_uses.append(block)
+
+            spans.append({
+                "span_id": llm_span_id,
+                "parent_span_id": root_span_id,
+                "name": f"llm:{model}",
+                "span_type": "llm",
+                "input_text": json.dumps(current_messages[-3:]),
+                "output_text": "\n".join(resp_text_parts) or None,
+                "metadata": {"iteration": iteration},
+                "start_time": llm_start,
+                "end_time": llm_end,
+                "duration_ms": llm_duration,
+                "status": "ok",
+                "token_count_input": input_tokens,
+                "token_count_output": output_tokens,
+                "model": model,
+            })
+
+            # Trajectory: record the LLM reasoning step
+            step_type = "reasoning" if tool_uses else "final_answer"
+            trajectory_steps.append({
+                "step_type": step_type,
+                "label": f"LLM call (iteration {iteration})",
+                "input_summary": (current_messages[-1].get("content", "") or "")[:200]
+                    if isinstance(current_messages[-1].get("content"), str)
+                    else f"[{len(current_messages)} messages]",
+                "output_summary": (" ".join(resp_text_parts))[:200] or None,
+                "duration_ms": int(llm_duration),
+                "token_count": input_tokens + output_tokens,
+                "status": "success",
+            })
+
+            if response.stop_reason != "tool_use" or not tool_uses:
+                break
+
+            # --- Tool calls ---
+            assistant_content = []
+            for block in response.content:
+                if block.type == "text":
+                    assistant_content.append({"type": "text", "text": block.text})
+                elif block.type == "tool_use":
+                    assistant_content.append({
+                        "type": "tool_use",
+                        "id": block.id,
+                        "name": block.name,
+                        "input": block.input,
+                    })
+
+            current_messages.append({"role": "assistant", "content": assistant_content})
+
+            tool_results = []
+            for tu in tool_uses:
+                tool_span_id = str(uuid.uuid4())
+                tool_start = self._now_iso()
+                t_tool_0 = time.perf_counter()
+
+                result_text = ""
+                status = "ok"
+                error_msg = None
+                if tool_executor:
+                    try:
+                        result_text = tool_executor(tu.name, tu.input)
+                    except Exception as e:
+                        result_text = str(e)
+                        status = "error"
+                        error_msg = str(e)
+                else:
+                    result_text = json.dumps({"info": "no tool_executor provided"})
+
+                t_tool_1 = time.perf_counter()
+                tool_end = self._now_iso()
+                tool_dur = (t_tool_1 - t_tool_0) * 1000
+                total_duration_ms += int(tool_dur)
+
+                spans.append({
+                    "span_id": tool_span_id,
+                    "parent_span_id": root_span_id,
+                    "name": tu.name,
+                    "span_type": "tool",
+                    "input_text": json.dumps(tu.input),
+                    "output_text": result_text[:10000] if result_text else None,
+                    "metadata": {"tool_id": tu.id},
+                    "start_time": tool_start,
+                    "end_time": tool_end,
+                    "duration_ms": tool_dur,
+                    "status": status,
+                    "error_message": error_msg,
+                })
+
+                # Trajectory: record each tool call
+                trajectory_steps.append({
+                    "step_type": "tool_call",
+                    "label": tu.name,
+                    "tool_name": tu.name,
+                    "input_summary": json.dumps(tu.input)[:200],
+                    "output_summary": (result_text or "")[:200],
+                    "duration_ms": int(tool_dur),
+                    "token_count": 0,
+                    "status": "success" if status == "ok" else "error",
+                })
+
+                tool_results.append({
+                    "type": "tool_result",
+                    "tool_use_id": tu.id,
+                    "content": result_text,
+                })
+
+            current_messages.append({"role": "user", "content": tool_results})
+
+            spans.append({
+                "span_id": step_span_id,
+                "parent_span_id": root_span_id,
+                "name": f"agent_step_{iteration}",
+                "span_type": "agent",
+                "input_text": json.dumps([tu.name for tu in tool_uses]),
+                "output_text": None,
+                "metadata": {"iteration": iteration, "tool_count": len(tool_uses)},
+                "start_time": llm_start,
+                "end_time": self._now_iso(),
+                "duration_ms": (time.perf_counter() - t0) * 1000,
+                "status": "ok",
+            })
+
+        # Root span
+        loop_end = self._now_iso()
+        spans.insert(0, {
+            "span_id": root_span_id,
+            "parent_span_id": None,
+            "name": "claude_agent_run",
+            "span_type": "chain",
+            "input_text": json.dumps(messages[:2]),
+            "output_text": "\n".join(resp_text_parts) if resp_text_parts else None,
+            "metadata": {"total_iterations": iteration, "model": model},
+            "start_time": loop_start,
+            "end_time": loop_end,
+            "duration_ms": None,
+            "status": "ok",
+        })
+
+        # Flush spans
+        try:
+            self._http.post(
+                f"{self.endpoint}/api/spans/ingest",
+                json={
+                    "trace_id": trace_id,
+                    "project_id": self.project_id,
+                    "spans": spans,
+                },
+            )
+        except Exception as e:
+            print(f"PRISMtrace flush error (spans): {e}")
+
+        # Flush trajectory for PRISM evaluation
+        trajectory_id = None
+        if self.emit_trajectory and trajectory_steps:
+            has_errors = any(s.get("status") == "error" for s in trajectory_steps)
+            try:
+                resp = self._http.post(
+                    f"{self.endpoint}/api/trajectories",
+                    json={
+                        "project_id": self.project_id,
+                        "conversation_id": conversation_id,
+                        "request_id": trace_id,
+                        "agent_id": self.agent_name,
+                        "agent_name": self.agent_name,
+                        "steps": trajectory_steps,
+                        "total_duration_ms": total_duration_ms,
+                        "final_status": "error" if has_errors else "success",
+                        "model": model,
+                    },
+                )
+                if resp.status_code in (200, 201):
+                    trajectory_id = resp.json().get("id")
+            except Exception as e:
+                print(f"PRISMtrace flush error (trajectory): {e}")
+
+        return {
+            "response": response,
+            "trace_id": trace_id,
+            "trajectory_id": trajectory_id,
+            "iterations": iteration,
+            "spans_count": len(spans),
+            "trajectory_steps": len(trajectory_steps),
+        }
+
+    def instrument_client(self) -> None:
+        """Wraps the Anthropic client's messages.create so every call is auto-traced."""
+        original_create = self.client.messages.create
+
+        tracer = self
+
+        def traced_create(*args: Any, **kwargs: Any) -> Any:
+            trace_id = str(uuid.uuid4())
+            span_id = str(uuid.uuid4())
+            start = tracer._now_iso()
+            t0 = time.perf_counter()
+
+            response = original_create(*args, **kwargs)
+
+            t1 = time.perf_counter()
+            end = tracer._now_iso()
+            dur_ms = (t1 - t0) * 1000
+
+            usage = getattr(response, "usage", None)
+            model = kwargs.get("model", "unknown")
+
+            tool_spans: list[dict] = []
+            if hasattr(response, "content"):
+                for block in response.content:
+                    if hasattr(block, "type") and block.type == "tool_use":
+                        tool_spans.append({
+                            "span_id": str(uuid.uuid4()),
+                            "parent_span_id": span_id,
+                            "name": block.name,
+                            "span_type": "tool",
+                            "input_text": json.dumps(block.input),
+                            "metadata": {"tool_id": block.id},
+                            "start_time": start,
+                            "end_time": end,
+                            "duration_ms": dur_ms,
+                            "status": "ok",
+                        })
+
+            spans = [
+                {
+                    "span_id": span_id,
+                    "parent_span_id": None,
+                    "name": f"llm:{model}",
+                    "span_type": "llm",
+                    "input_text": json.dumps(kwargs.get("messages", [])[-3:]),
+                    "output_text": None,
+                    "metadata": {},
+                    "start_time": start,
+                    "end_time": end,
+                    "duration_ms": dur_ms,
+                    "status": "ok",
+                    "token_count_input": getattr(usage, "input_tokens", None) if usage else None,
+                    "token_count_output": getattr(usage, "output_tokens", None) if usage else None,
+                    "model": model,
+                },
+                *tool_spans,
+            ]
+
+            try:
+                tracer._http.post(
+                    f"{tracer.endpoint}/api/spans/ingest",
+                    json={
+                        "trace_id": trace_id,
+                        "project_id": tracer.project_id,
+                        "spans": spans,
+                    },
+                )
+            except Exception as e:
+                print(f"PRISMtrace auto-trace error: {e}")
+
+            return response
+
+        self.client.messages.create = traced_create

prismtrace/client.py

@@ -0,0 +1,265 @@
+"""PRISMtrace client for AI observability."""
+
+import json
+import sys
+import threading
+import uuid
+from functools import wraps
+from typing import Any, Callable, Optional
+
+import httpx
+
+
+class PRISMtrace:
+    def __init__(
+        self,
+        api_key: str,
+        host: str,
+        project_id: str,
+        timeout: int = 10,
+    ):
+        self.api_key = api_key
+        self.host = host.rstrip("/")
+        self.project_id = project_id
+        self.timeout = timeout
+        self._client = httpx.Client(
+            base_url=self.host,
+            headers={
+                "x-prismtrace-key": api_key,
+                "Content-Type": "application/json",
+            },
+            timeout=self.timeout,
+        )
+
+    # ------------------------------------------------------------------
+    # Traces (existing)
+    # ------------------------------------------------------------------
+
+    def trace_llm(
+        self,
+        model: str,
+        input_messages: list,
+        output: str,
+        latency_ms: int,
+        token_count_input: int = 0,
+        token_count_output: int = 0,
+        trace_id: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> None:
+        payload = {
+            "project_id": self.project_id,
+            "model": model,
+            "input_messages": input_messages,
+            "output_message": output,
+            "latency_ms": latency_ms,
+            "token_count_input": token_count_input,
+            "token_count_output": token_count_output,
+        }
+        if trace_id is not None:
+            payload["trace_id"] = trace_id
+        if metadata is not None:
+            payload["metadata"] = metadata
+        threading.Thread(target=self._post, args=("/api/traces", payload), daemon=True).start()
+
+    def _post(self, path: str, payload: dict) -> Optional[dict]:
+        try:
+            resp = self._client.post(path, json=payload)
+            if resp.status_code in (200, 201):
+                return resp.json()
+            print(
+                f"PRISMtrace warning: {path} returned {resp.status_code}: {resp.text[:200]}",
+                file=sys.stderr,
+            )
+        except Exception as e:
+            print(f"PRISMtrace warning: {path} failed: {e}", file=sys.stderr)
+        return None
+
+    def _post_sync(self, path: str, payload: dict) -> Optional[dict]:
+        try:
+            resp = self._client.post(path, json=payload)
+            if resp.status_code in (200, 201):
+                return resp.json()
+            print(
+                f"PRISMtrace warning: {path} returned {resp.status_code}: {resp.text[:200]}",
+                file=sys.stderr,
+            )
+        except Exception as e:
+            print(f"PRISMtrace warning: {path} failed: {e}", file=sys.stderr)
+        return None
+
+    def _get_sync(self, path: str) -> Optional[dict]:
+        try:
+            resp = self._client.get(path)
+            if resp.status_code == 200:
+                return resp.json()
+            print(
+                f"PRISMtrace warning: GET {path} returned {resp.status_code}: {resp.text[:200]}",
+                file=sys.stderr,
+            )
+        except Exception as e:
+            print(f"PRISMtrace warning: GET {path} failed: {e}", file=sys.stderr)
+        return None
+
+    def trace(self, name: Optional[str] = None) -> Callable:
+        def decorator(fn: Callable) -> Callable:
+            @wraps(fn)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                import time
+                start = time.perf_counter()
+                result = fn(*args, **kwargs)
+                elapsed_ms = int((time.perf_counter() - start) * 1000)
+                self.trace_llm(
+                    model=name or fn.__name__,
+                    input_messages=[{"role": "user", "content": str(args)}],
+                    output=str(result),
+                    latency_ms=elapsed_ms,
+                )
+                return result
+            return wrapper
+        return decorator
+
+    # ------------------------------------------------------------------
+    # Trajectories
+    # ------------------------------------------------------------------
+
+    def submit_trajectory(
+        self,
+        steps: list[dict],
+        *,
+        agent_name: str = "default-agent",
+        agent_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        request_id: Optional[str] = None,
+        model: Optional[str] = None,
+        final_status: str = "success",
+        async_send: bool = False,
+    ) -> Optional[dict]:
+        """Submit a trajectory (ordered step list) for PRISM evaluation.
+
+        Each step dict should contain:
+            step_type:       "reasoning" | "tool_call" | "final_answer" | ...
+            label:           short description
+            output_summary:  what this step produced (optional but recommended)
+            tool_name:       tool identifier (required when step_type is tool_call)
+            input_summary:   what went in (optional)
+            duration_ms:     step latency (optional)
+            token_count:     tokens used (optional)
+            status:          "success" | "error" (default "success")
+
+        Returns the API response dict (trajectory id, step_count, created_at)
+        or None on failure.
+        """
+        payload = {
+            "project_id": self.project_id,
+            "conversation_id": conversation_id or str(uuid.uuid4()),
+            "request_id": request_id or str(uuid.uuid4()),
+            "agent_id": agent_id or agent_name,
+            "agent_name": agent_name,
+            "steps": steps,
+            "total_duration_ms": sum(s.get("duration_ms") or 0 for s in steps),
+            "final_status": final_status,
+            "model": model,
+        }
+        if async_send:
+            threading.Thread(
+                target=self._post, args=("/api/trajectories", payload), daemon=True,
+            ).start()
+            return None
+        return self._post_sync("/api/trajectories", payload)
+
+    def get_trajectory(self, trajectory_id: str) -> Optional[dict]:
+        """Fetch a trajectory and its steps."""
+        return self._get_sync(f"/api/trajectories/{trajectory_id}")
+
+    def get_trajectory_evaluation(self, trajectory_id: str) -> Optional[dict]:
+        """Fetch PRISM evaluation results for a trajectory."""
+        return self._get_sync(f"/api/trajectories/{trajectory_id}/evaluation")
+
+    def retrigger_evaluation(self, trajectory_id: str) -> Optional[dict]:
+        """Re-run PRISM evaluation for an existing trajectory."""
+        return self._post_sync(
+            f"/api/trajectories/{trajectory_id}/evaluate?project_id={self.project_id}",
+            {},
+        )
+
+    # ------------------------------------------------------------------
+    # Knowledge Base
+    # ------------------------------------------------------------------
+
+    def kb_upload(
+        self,
+        filename: str,
+        content: bytes | str,
+        *,
+        agent_id: Optional[str] = None,
+        description: Optional[str] = None,
+        content_type: str = "text/plain",
+    ) -> Optional[dict]:
+        """Upload a document to the project knowledge base.
+
+        Returns {"id", "name", "chunk_count", ...} or None on failure.
+        """
+        if isinstance(content, str):
+            content = content.encode("utf-8")
+
+        files = {"file": (filename, content, content_type)}
+        data: dict[str, str] = {"project_id": self.project_id}
+        if agent_id:
+            data["agent_id"] = agent_id
+        if description:
+            data["description"] = description
+
+        try:
+            resp = httpx.post(
+                f"{self.host}/api/knowledge-base/documents",
+                files=files,
+                data=data,
+                headers={"x-prismtrace-key": self.api_key},
+                timeout=self.timeout,
+            )
+            if resp.status_code in (200, 201):
+                return resp.json()
+            print(f"PRISMtrace warning: kb_upload returned {resp.status_code}")
+        except Exception as e:
+            print(f"PRISMtrace warning: kb_upload failed: {e}")
+        return None
+
+    def kb_search(
+        self,
+        query: str,
+        *,
+        agent_id: Optional[str] = None,
+        limit: int = 5,
+    ) -> list[dict]:
+        """Search the project knowledge base. Returns list of chunk results."""
+        payload: dict[str, Any] = {
+            "query": query,
+            "project_id": self.project_id,
+            "limit": limit,
+        }
+        if agent_id:
+            payload["agent_id"] = agent_id
+        result = self._post_sync("/api/knowledge-base/search", payload)
+        if result:
+            return result.get("results", [])
+        return []
+
+    def kb_list_documents(self, agent_id: Optional[str] = None) -> list[dict]:
+        """List documents in the project knowledge base."""
+        path = f"/api/knowledge-base/documents?project_id={self.project_id}"
+        if agent_id:
+            path += f"&agent_id={agent_id}"
+        result = self._get_sync(path)
+        if result:
+            return result.get("documents", [])
+        return []
+
+    def kb_delete_document(self, doc_id: str) -> bool:
+        """Delete a document from the knowledge base."""
+        try:
+            resp = self._client.delete(
+                f"/api/knowledge-base/documents/{doc_id}?project_id={self.project_id}",
+            )
+            return resp.status_code == 200
+        except Exception:
+            return False