zespan-fastapi 1.0.0__py3-none-any.whl

zespan_fastapi/__init__.py

@@ -0,0 +1,263 @@
+"""
+Zespan Trace - FastAPI Middleware
+Automatic Agent observability for FastAPI applications
+"""
+
+__version__ = "1.0.0"
+
+import time
+import json
+import uuid
+from typing import Any, Dict, List, Optional, Callable
+from dataclasses import dataclass, field, asdict
+from contextvars import ContextVar
+
+# Context for current trace
+current_trace: ContextVar[Optional["TraceContext"]] = ContextVar("current_trace", default=None)
+
+
+@dataclass
+class TraceContext:
+    """Context for a single trace/span"""
+    trace_id: str
+    span_id: str
+    parent_id: Optional[str] = None
+    name: str = ""
+    start_time: float = field(default_factory=time.time)
+    attributes: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ZespanConfig:
+    """Configuration for Zespan middleware"""
+    api_key: str
+    base_url: str = "https://api.zespan.com"
+    enabled: bool = True
+    sample_rate: float = 1.0
+    capture_request_body: bool = False
+    capture_response_body: bool = False
+    redact_fields: List[str] = field(default_factory=lambda: ["password", "token", "api_key"])
+    debug: bool = False
+
+
+class ZespanFastAPIMiddleware:
+    """
+    FastAPI middleware for automatic LLM observability.
+    
+    Usage:
+        from fastapi import FastAPI
+        from zespan_fastapi import ZespanFastAPIMiddleware, ZespanConfig
+        
+        app = FastAPI()
+        config = ZespanConfig(api_key="your-api-key")
+        app.add_middleware(ZespanFastAPIMiddleware, config=config)
+    """
+    
+    def __init__(self, app, config: ZespanConfig):
+        self.app = app
+        self.config = config
+        self._queue: List[Dict[str, Any]] = []
+        self._flush_interval = 5.0
+        self._last_flush = time.time()
+        
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http" or not self.config.enabled:
+            return await self.app(scope, receive, send)
+            
+        trace_id = f"trace_{uuid.uuid4().hex}"
+        span_id = f"span_{uuid.uuid4().hex[:16]}"
+        
+        trace_ctx = TraceContext(
+            trace_id=trace_id,
+            span_id=span_id,
+            name=f"{scope['method']} {scope['path']}",
+        )
+        
+        token = current_trace.set(trace_ctx)
+        
+        try:
+            start_time = time.time()
+            
+            # Capture request info
+            request_attrs = {
+                "http.method": scope.get("method"),
+                "http.url": scope.get("path"),
+                "http.route": scope.get("path"),
+            }
+            
+            # Wrap send to capture response
+            response_status = None
+            async def wrapped_send(message):
+                nonlocal response_status
+                if message["type"] == "http.response.start":
+                    response_status = message.get("status")
+                await send(message)
+            
+            await self.app(scope, receive, wrapped_send)
+            
+            duration_ms = (time.time() - start_time) * 1000
+            
+            # Build trace data
+            trace_data = {
+                "trace_id": trace_id,
+                "span_id": span_id,
+                "name": trace_ctx.name,
+                "type": "http_request",
+                "duration_ms": duration_ms,
+                "status": "success" if response_status and response_status < 400 else "error",
+                "http_status": response_status,
+                "attributes": {**request_attrs, **trace_ctx.attributes},
+                "timestamp": time.time(),
+            }
+            
+            self._queue.append(trace_data)
+            
+            # Flush if needed
+            if len(self._queue) >= 10 or time.time() - self._last_flush > self._flush_interval:
+                await self._flush()
+                
+        finally:
+            current_trace.reset(token)
+    
+    async def _flush(self):
+        """Send queued traces to Zespan"""
+        if not self._queue:
+            return
+            
+        traces = self._queue[:]
+        self._queue = []
+        self._last_flush = time.time()
+        
+        try:
+            import httpx
+            async with httpx.AsyncClient() as client:
+                await client.post(
+                    f"{self.config.base_url}/v1/ingest/batch",
+                    json={"traces": traces},
+                    headers={"X-API-Key": self.config.api_key},
+                    timeout=5.0
+                )
+        except Exception as e:
+            if self.config.debug:
+                print(f"[Zespan] Failed to send traces: {e}")
+
+
+def observe_span(name: str, span_type: str = "custom"):
+    """
+    Decorator to observe a function as a span.
+    
+    Usage:
+        @observe_span("my_llm_call", span_type="llm")
+        async def my_llm_function(prompt: str):
+            return await call_llm(prompt)
+    """
+    def decorator(func: Callable) -> Callable:
+        async def async_wrapper(*args, **kwargs):
+            trace_ctx = current_trace.get()
+            if not trace_ctx:
+                return await func(*args, **kwargs)
+                
+            span_id = f"span_{uuid.uuid4().hex[:16]}"
+            start = time.time()
+            
+            try:
+                result = await func(*args, **kwargs)
+                duration = (time.time() - start) * 1000
+                
+                # Add to parent trace
+                trace_ctx.attributes[f"span.{name}.duration_ms"] = duration
+                trace_ctx.attributes[f"span.{name}.status"] = "success"
+                
+                return result
+            except Exception as e:
+                duration = (time.time() - start) * 1000
+                trace_ctx.attributes[f"span.{name}.duration_ms"] = duration
+                trace_ctx.attributes[f"span.{name}.status"] = "error"
+                trace_ctx.attributes[f"span.{name}.error"] = str(e)
+                raise
+                
+        def sync_wrapper(*args, **kwargs):
+            import asyncio
+            return asyncio.run(async_wrapper(*args, **kwargs))
+            
+        return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
+    return decorator
+
+
+def observe_llm(model: Optional[str] = None, provider: Optional[str] = None):
+    """
+    Decorator specifically for LLM calls with automatic token tracking.
+    
+    Usage:
+        @observe_llm(model="gpt-4", provider="openai")
+        async def generate_text(prompt: str):
+            response = await openai.chat.completions.create(...)
+            return response
+    """
+    def decorator(func: Callable) -> Callable:
+        async def wrapper(*args, **kwargs):
+            trace_ctx = current_trace.get()
+            
+            span_id = f"span_{uuid.uuid4().hex[:16]}"
+            start = time.time()
+            
+            # Extract prompt from args/kwargs if possible
+            prompt = kwargs.get("prompt") or (args[0] if args else None)
+            
+            try:
+                result = await func(*args, **kwargs)
+                duration = (time.time() - start) * 1000
+                
+                if trace_ctx:
+                    trace_ctx.attributes.update({
+                        "llm.model": model or "unknown",
+                        "llm.provider": provider or "unknown",
+                        "llm.duration_ms": duration,
+                        "llm.prompt": str(prompt)[:1000] if prompt else None,
+                    })
+                    
+                    # Try to extract token usage from common response formats
+                    if hasattr(result, "usage"):
+                        usage = result.usage
+                        trace_ctx.attributes["llm.tokens.input"] = getattr(usage, "prompt_tokens", None)
+                        trace_ctx.attributes["llm.tokens.output"] = getattr(usage, "completion_tokens", None)
+                        trace_ctx.attributes["llm.tokens.total"] = getattr(usage, "total_tokens", None)
+                    elif isinstance(result, dict):
+                        usage = result.get("usage", {})
+                        trace_ctx.attributes["llm.tokens.input"] = usage.get("prompt_tokens")
+                        trace_ctx.attributes["llm.tokens.output"] = usage.get("completion_tokens")
+                        trace_ctx.attributes["llm.tokens.total"] = usage.get("total_tokens")
+                
+                return result
+            except Exception as e:
+                if trace_ctx:
+                    trace_ctx.attributes["llm.error"] = str(e)
+                raise
+                
+        return wrapper
+    return decorator
+
+
+def set_attribute(key: str, value: Any):
+    """Set an attribute on the current trace context"""
+    trace_ctx = current_trace.get()
+    if trace_ctx:
+        trace_ctx.attributes[key] = value
+
+
+def get_current_trace_id() -> Optional[str]:
+    """Get the current trace ID if available"""
+    trace_ctx = current_trace.get()
+    return trace_ctx.trace_id if trace_ctx else None
+
+
+# Re-exports
+__all__ = [
+    "ZespanFastAPIMiddleware",
+    "ZespanConfig",
+    "observe_span",
+    "observe_llm",
+    "set_attribute",
+    "get_current_trace_id",
+    "TraceContext",
+]

zespan_fastapi-1.0.0.dist-info/METADATA

@@ -0,0 +1,86 @@
+Metadata-Version: 2.4
+Name: zespan-fastapi
+Version: 1.0.0
+Summary: Zespan observability middleware for FastAPI
+Requires-Python: >=3.9
+Requires-Dist: fastapi>=0.100
+Requires-Dist: requests>=2.28
+Provides-Extra: dev
+Requires-Dist: httpx>=0.24; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# zespan-fastapi
+
+Automatic LLM observability middleware for FastAPI. Traces every request, LLM call, and custom span with zero code changes for basic usage.
+
+## Installation
+
+```bash
+pip install zespan-fastapi
+```
+
+## Quick Start
+
+```python
+from fastapi import FastAPI
+from zespan_fastapi import ZespanFastAPIMiddleware, ZespanConfig
+
+app = FastAPI()
+
+config = ZespanConfig(api_key="zsp_your_api_key")
+app.add_middleware(ZespanFastAPIMiddleware, config=config)
+```
+
+Every HTTP request is now automatically traced.
+
+## Trace LLM Calls
+
+```python
+from zespan_fastapi import observe_llm
+
+@observe_llm(model="gpt-4o", provider="openai")
+async def generate_response(prompt: str):
+    response = await openai.chat.completions.create(
+        model="gpt-4o",
+        messages=[{"role": "user", "content": prompt}]
+    )
+    return response
+```
+
+## Trace Custom Spans
+
+```python
+from zespan_fastapi import observe_span
+
+@observe_span("fetch_context", span_type="retrieval")
+async def fetch_context(query: str):
+    return await vector_db.search(query)
+```
+
+## Set Custom Attributes
+
+```python
+from zespan_fastapi import set_attribute
+
+@app.post("/chat")
+async def chat(request: ChatRequest):
+    set_attribute("user.id", request.user_id)
+    set_attribute("model", request.model)
+```
+
+## Configuration
+
+| Option | Default | Description |
+|---|---|---|
+| `api_key` | required | Zespan API key (`zsp_...`) |
+| `base_url` | `https://api.zespan.com` | API endpoint |
+| `enabled` | `True` | Toggle tracing on/off |
+| `sample_rate` | `1.0` | Fraction of requests to trace (0.0–1.0) |
+| `capture_request_body` | `False` | Include request body in traces |
+| `capture_response_body` | `False` | Include response body in traces |
+| `redact_fields` | `["password","token","api_key"]` | Fields redacted automatically |
+| `debug` | `False` | Print flush errors to stdout |
+
+Get your API key at [zespan.com](https://zespan.com).

zespan_fastapi-1.0.0.dist-info/RECORD

@@ -0,0 +1,4 @@
+zespan_fastapi/__init__.py,sha256=72wOuaVVOoPFli1NOSF2R_c_s_rG6g5_gLG2AB7fs4c,9410
+zespan_fastapi-1.0.0.dist-info/METADATA,sha256=GD_m539iqhR-dyNMiidcuJLYRZcNw_HvM--kIonoSWI,2303
+zespan_fastapi-1.0.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+zespan_fastapi-1.0.0.dist-info/RECORD,,

zespan_fastapi-1.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any