quraite 0.0.1__py3-none-any.whl

quraite/tracing/__init__.py

@@ -0,0 +1,24 @@
+"""Tracing infrastructure for OpenTelemetry span collection and processing."""
+
+from quraite.tracing.span_exporter import QuraiteInMemorySpanExporter
+from quraite.tracing.span_processor import QuraiteSimpleSpanProcessor
+from quraite.tracing.tool_extractors import Framework, ToolCallInfo, get_tool_extractor
+from quraite.tracing.trace import AgentSpan, AgentTrace, CostInfo, TokenInfo
+from quraite.tracing.types import Event, Link, Resource, SpanContext, Status
+
+__all__ = [
+    "AgentSpan",
+    "AgentTrace",
+    "CostInfo",
+    "TokenInfo",
+    "Framework",
+    "ToolCallInfo",
+    "get_tool_extractor",
+    "QuraiteInMemorySpanExporter",
+    "QuraiteSimpleSpanProcessor",
+    "Event",
+    "Link",
+    "Resource",
+    "SpanContext",
+    "Status",
+]

quraite/tracing/constants.py

@@ -0,0 +1,16 @@
+from enum import Enum
+
+QURAITE_ADAPTER_TRACE_PREFIX = "quraite-adapter"
+
+QURAITE_TRACER_NAME = "quraite.instrumentation"
+
+
+class Framework(str, Enum):
+    """Supported agent frameworks."""
+
+    PYDANTIC = "pydantic"
+    LANGGRAPH = "langgraph"
+    GOOGLE_ADK = "google_adk"
+    OPENAI_AGENTS = "openai_agents"
+    AGNO = "agno"
+    SMOLAGENTS = "smolagents"

quraite/tracing/span_exporter.py

@@ -0,0 +1,115 @@
+import json
+import os
+import threading
+import typing
+from collections import defaultdict
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+
+class QuraiteInMemorySpanExporter(SpanExporter):
+    def __init__(self) -> None:
+        # self.spans: typing.List[ReadableSpan] = []
+        self.traces: typing.Dict[int, typing.List[ReadableSpan]] = defaultdict(list)
+        self.testcase_to_trace: typing.Dict[str, int] = {}
+
+        self._stopped = False
+        self._lock = threading.Lock()
+
+    def handle_testcase_trace(self, span: ReadableSpan) -> None:
+        """Handle a testcase trace."""
+        # print(f"🟢 testcase trace received: {span.name}")
+        # print(f"🟢 Span: {span.context}")
+        formatted_trace_id = format(span.context.trace_id, "032x")[:8]
+        # print(f"🟢 testcase formatted trace id: {formatted_trace_id}")
+        self.traces[formatted_trace_id] = []
+        self.testcase_to_trace[span.name] = formatted_trace_id
+
+    def export(self, spans: typing.Sequence[ReadableSpan]) -> SpanExportResult:
+        """Stores a list of spans in memory."""
+        if self._stopped:
+            return SpanExportResult.FAILURE
+
+        with self._lock:
+            for span in spans:
+                formatted_trace_id = format(span.context.trace_id, "032x")[:8]
+                # print(f"🟢 span formatted context trace id: {formatted_trace_id}")
+                self.traces[formatted_trace_id].append(span)
+                # self.spans.append(span)
+
+        return SpanExportResult.SUCCESS
+
+    def shutdown(self) -> None:
+        """Shut downs the exporter.
+
+        Calls to export after the exporter has been shut down will fail.
+        """
+        print("Shutting down exporter")
+        self._stopped = True
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+    def get_traces(self) -> typing.Dict[int, typing.List[ReadableSpan]]:
+        """Get all spans grouped by trace ID"""
+        return dict(self.traces)
+
+    def get_trace_by_testcase(self, testcase_name: str) -> typing.List[ReadableSpan]:
+        """Get all spans for a specific testcase"""
+        with self._lock:
+            return self.traces.get(self.testcase_to_trace.get(testcase_name, None), [])
+
+    def get_trace(self, trace_id: int) -> typing.List[ReadableSpan]:
+        """Get all spans for a specific trace"""
+        return self.traces.get(trace_id, [])
+
+    def get_trace_count(self) -> int:
+        """Get the number of unique traces"""
+        return len(self.traces)
+
+    def print_trace_summary(self):
+        """Print a summary of all traces"""
+        print(f"\n{'='*60}")
+        print(f"Total Traces: {self.get_trace_count()}")
+        for trace_id, spans in self.traces.items():
+            print(f"Trace ID: {trace_id}...")
+            print(f"Spans in trace: {len(spans)}")
+        print(f"Total Testcases: {self.testcase_to_trace}")
+        print(f"TraceIDs: {self.traces.keys()}")
+        print(f"{'='*60}\n")
+
+        # for trace_id, spans in self.traces.items():
+        #     trace_id_hex = format(trace_id, '032x')
+        #     print(f"📊 Trace ID: {trace_id_hex}")
+        #     print(f"   Spans in trace: {len(spans)}")
+
+        #     # Sort spans by start time to show execution order
+        #     sorted_spans = sorted(spans, key=lambda s: s.start_time)
+
+        #     for span in sorted_spans:
+        #         duration_ms = (span.end_time - span.start_time) / 1e6
+        #         parent_id = format(span.parent.span_id, '016x') if span.parent else "None"
+        #         indent = "   " if span.parent else ""
+        #         print(f"   {indent}├─ {span.name}")
+        #         print(f"   {indent}   ├─ Span ID: {format(span.context.span_id, '016x')}")
+        #         print(f"   {indent}   ├─ Parent ID: {parent_id}")
+        #         print(f"   {indent}   ├─ Duration: {duration_ms:.2f}ms")
+        #         if span.attributes:
+        #             print(f"   {indent}   └─ Attributes: {dict(span.attributes)}")
+        #     print()
+
+    def save_traces_to_file(self, filename: typing.Optional[str] = "traces.json"):
+        """Save a trace to a file"""
+        traces = []
+        for trace_id, spans in self.traces.items():
+            traces.append(
+                {
+                    "trace_id": trace_id,
+                    "spans": [json.loads(span.to_json()) for span in spans],
+                }
+            )
+
+        os.makedirs(os.path.dirname(f"traces/{filename}"), exist_ok=True)
+        with open(f"traces/{filename}", "w") as f:
+            json.dump(traces, f, indent=2)

quraite/tracing/span_processor.py

@@ -0,0 +1,49 @@
+import typing
+
+from opentelemetry.context import (
+    _SUPPRESS_INSTRUMENTATION_KEY,
+    Context,
+    attach,
+    detach,
+    set_value,
+)
+from opentelemetry.sdk.trace.export import ReadableSpan, SpanExporter, SpanProcessor
+from opentelemetry.trace import Span, logger
+
+from quraite.tracing.constants import QURAITE_ADAPTER_TRACE_PREFIX
+from quraite.tracing.span_exporter import QuraiteInMemorySpanExporter
+
+
+class QuraiteSimpleSpanProcessor(SpanProcessor):
+    """Simple SpanProcessor implementation.
+
+    SimpleSpanProcessor is an implementation of `SpanProcessor` that
+    passes ended spans directly to the configured `SpanExporter`.
+    """
+
+    def __init__(self, span_exporter: SpanExporter):
+        self.span_exporter: QuraiteInMemorySpanExporter = span_exporter
+
+    def on_start(
+        self, span: Span, parent_context: typing.Optional[Context] = None
+    ) -> None:
+        if QURAITE_ADAPTER_TRACE_PREFIX in span.name:
+            self.span_exporter.handle_testcase_trace(span)
+
+    def on_end(self, span: ReadableSpan) -> None:
+        if not span.context.trace_flags.sampled:
+            return
+        token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True))
+        try:
+            self.span_exporter.export((span,))
+        # pylint: disable=broad-exception-caught
+        except Exception:
+            logger.exception("Exception while exporting Span.")
+        detach(token)
+
+    def shutdown(self) -> None:
+        self.span_exporter.shutdown()
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        # pylint: disable=unused-argument
+        return True

quraite/tracing/tool_extractors.py

@@ -0,0 +1,290 @@
+"""
+Framework-specific tool extractors for converting span attributes to standardized tool call information.
+
+These extractors handle the varying attribute structures across different agent frameworks
+(pydantic, langgraph, adk, openai_agents, agno, smolagents, etc.)
+"""
+
+import json
+from typing import Any, Protocol
+
+from quraite.tracing.constants import Framework
+
+
+class ToolCallInfo:
+    """Standardized tool call information extracted from a TOOL span."""
+
+    def __init__(
+        self,
+        tool_name: str,
+        tool_call_id: str | None,
+        arguments: str | dict,
+        response: Any,
+    ):
+        self.tool_name = tool_name
+        self.tool_call_id = tool_call_id
+        self.arguments = arguments
+        self.response = response
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "role": "tool",
+            "tool_name": self.tool_name,
+            "tool_call_id": self.tool_call_id,
+            "arguments": self.arguments,
+            "response": self.response,
+        }
+
+
+class ToolExtractor(Protocol):
+    """Protocol for framework-specific tool extractors."""
+
+    def __call__(self, span: dict[str, Any]) -> ToolCallInfo | None: ...
+
+
+# =============================================================================
+# Framework-specific tool extractors
+# =============================================================================
+
+
+def extract_tool_pydantic(span: dict[str, Any]) -> ToolCallInfo | None:
+    """
+    Extract tool info from Pydantic AI tool spans.
+
+    Attributes:
+        - tool.name: "customer_balance"
+        - tool_call.id: "call_xxx"
+        - tool_arguments: "{\"include_pending\":true}"
+        - tool_response: "$123.45"
+    """
+    attrs = span.get("attributes", {})
+
+    tool_name = attrs.get("tool.name") or attrs.get("gen_ai.tool.name")
+    if not tool_name:
+        return None
+
+    tool_call_id = attrs.get("tool_call.id") or attrs.get("gen_ai.tool.call.id")
+    arguments = attrs.get("tool_arguments", "{}")
+    response = attrs.get("tool_response", "")
+
+    return ToolCallInfo(
+        tool_name=tool_name,
+        tool_call_id=tool_call_id,
+        arguments=arguments,
+        response=response,
+    )
+
+
+def extract_tool_langgraph(span: dict[str, Any]) -> ToolCallInfo | None:
+    """
+    Extract tool info from LangGraph tool spans.
+
+    Attributes:
+        - tool.name: "add"
+        - tool.description: "Add two numbers."
+        - input.value: "{'b': 1, 'a': 1}"
+        - output.value: JSON with content
+    """
+    attrs = span.get("attributes", {})
+
+    tool_name = attrs.get("tool.name")
+    if not tool_name:
+        return None
+
+    arguments = attrs.get("input.value", "{}")
+    output_value = attrs.get("output.value", "")
+
+    # Also check for response attribute (some LangGraph spans store response here)
+    response_value = attrs.get("response", output_value)
+
+    # Try to parse output to extract content
+    response = response_value
+    if isinstance(response_value, str):
+        try:
+            parsed = json.loads(response_value)
+            if isinstance(parsed, dict):
+                # Check if response field contains JSON string (nested JSON)
+                if "response" in parsed and isinstance(parsed["response"], str):
+                    try:
+                        inner_parsed = json.loads(parsed["response"])
+                        if isinstance(inner_parsed, dict) and "update" in inner_parsed:
+                            parsed = inner_parsed
+                    except (json.JSONDecodeError, TypeError):
+                        pass
+
+                # First check for direct content field
+                if "content" in parsed:
+                    response = parsed.get("content", response_value)
+                # Check for update.messages structure (LangGraph graph updates)
+                # this comes when you use supervisor agent with multiple agents
+                elif "update" in parsed:
+                    update = parsed.get("update", {})
+                    messages = update.get("messages", [])
+                    # Find the last tool message
+                    for msg in reversed(messages):
+                        if isinstance(msg, dict) and msg.get("type") == "tool":
+                            content = msg.get("content", "")
+                            if content:
+                                response = content
+                                break
+                    else:
+                        # No tool message found, keep original response
+                        response = response_value
+                else:
+                    response = response_value
+        except json.JSONDecodeError:
+            pass
+
+    return ToolCallInfo(
+        tool_name=tool_name,
+        tool_call_id=None,  # LangGraph doesn't always have call IDs in tool spans
+        arguments=arguments,
+        response=response,
+    )
+
+
+def extract_tool_adk(span: dict[str, Any]) -> ToolCallInfo | None:
+    """
+    Extract tool info from Google ADK tool spans.
+
+    Attributes:
+        - tool.name: "get_weather"
+        - tool.parameters: "{\"city\": \"New York\"}"
+        - gcp.vertex.agent.tool_call_args: "{\"city\": \"New York\"}"
+        - gcp.vertex.agent.tool_response: JSON response
+        - output.value: JSON with id, name, response
+    """
+    attrs = span.get("attributes", {})
+
+    tool_name = attrs.get("tool.name") or attrs.get("gen_ai.tool.name")
+    if not tool_name:
+        return None
+
+    # Skip merged tool spans
+    if tool_name == "(merged tools)":
+        return None
+
+    tool_call_id = attrs.get("gen_ai.tool.call.id")
+    arguments = (
+        attrs.get("tool.parameters")
+        or attrs.get("gcp.vertex.agent.tool_call_args")
+        or attrs.get("input.value", "{}")
+    )
+
+    # Get response from various possible locations
+    response = attrs.get("gcp.vertex.agent.tool_response", "")
+    if not response or response == "<not serializable>":
+        output_value = attrs.get("output.value", "")
+        if isinstance(output_value, str):
+            try:
+                parsed = json.loads(output_value)
+                if isinstance(parsed, dict) and "response" in parsed:
+                    response = parsed.get("response", output_value)
+                else:
+                    response = output_value
+            except json.JSONDecodeError:
+                response = output_value
+        else:
+            response = output_value
+
+    return ToolCallInfo(
+        tool_name=tool_name,
+        tool_call_id=tool_call_id,
+        arguments=arguments,
+        response=response,
+    )
+
+
+def extract_tool_openai_agents(span: dict[str, Any]) -> ToolCallInfo | None:
+    """
+    Extract tool info from OpenAI Agents tool spans.
+
+    Attributes:
+        - tool.name: "multiply"
+        - input.value: "{\"a\":10,\"b\":2}"
+        - output.value: 20.0
+    """
+    attrs = span.get("attributes", {})
+
+    tool_name = attrs.get("tool.name")
+    if not tool_name:
+        return None
+
+    arguments = attrs.get("input.value", "{}")
+    response = attrs.get("output.value", "")
+
+    return ToolCallInfo(
+        tool_name=tool_name,
+        tool_call_id=None,  # OpenAI Agents SDK doesn't put call ID in tool span
+        arguments=arguments,
+        response=response,
+    )
+
+
+def extract_tool_agno(span: dict[str, Any]) -> ToolCallInfo | None:
+    """
+    Extract tool info from Agno tool spans.
+
+    Attributes:
+        - tool.name: "duckduckgo_search"
+        - tool.description: "..."
+        - tool.parameters: "{\"query\": \"...\", \"max_results\": 5}"
+        - input.value: same as parameters
+        - output.value: JSON response
+    """
+    attrs = span.get("attributes", {})
+
+    tool_name = attrs.get("tool.name")
+    if not tool_name:
+        return None
+
+    arguments = attrs.get("tool.parameters") or attrs.get("input.value", "{}")
+    response = attrs.get("output.value", "")
+
+    return ToolCallInfo(
+        tool_name=tool_name,
+        tool_call_id=None,
+        arguments=arguments,
+        response=response,
+    )
+
+
+def extract_tool_smolagents(span: dict[str, Any]) -> ToolCallInfo | None:
+    """
+    Extract tool info from SmolAgents tool spans.
+    """
+    attrs = span.get("attributes", {})
+
+    tool_name = attrs.get("tool.name")
+    if not tool_name:
+        return None
+
+    arguments = attrs.get("input.value", "{}")
+    response = attrs.get("output.value", "")
+
+    return ToolCallInfo(
+        tool_name=tool_name,
+        tool_call_id=None,
+        arguments=arguments,
+        response=response,
+    )
+
+
+# Registry of framework extractors
+TOOL_EXTRACTORS: dict[Framework, ToolExtractor] = {
+    Framework.PYDANTIC: extract_tool_pydantic,
+    Framework.LANGGRAPH: extract_tool_langgraph,
+    Framework.GOOGLE_ADK: extract_tool_adk,
+    Framework.OPENAI_AGENTS: extract_tool_openai_agents,
+    Framework.AGNO: extract_tool_agno,
+    Framework.SMOLAGENTS: extract_tool_smolagents,
+}
+
+
+def get_tool_extractor(framework: Framework | str) -> ToolExtractor:
+    """Get the appropriate tool extractor for the given framework."""
+    if isinstance(framework, str):
+        framework = Framework(framework.lower())
+    return TOOL_EXTRACTORS.get(
+        framework, extract_tool_langgraph
+    )  # Default to langgraph