agent-observability-studio 0.1.0__tar.gz

agent_observability_studio-0.1.0/OVERVIEW.md

@@ -0,0 +1,3 @@
+# agent-observability-studio
+
+A real-time monitoring and debugging platform for multi-agent AI systems that provides live interaction tracing, token cost analysis, and performance bottleneck detection. Complements your existing mesh-health-dashboard by adding deep inspection tools for agent conversations, decision trees, and failure analysis — think "Chrome DevTools for AI agents."

agent_observability_studio-0.1.0/PKG-INFO

@@ -0,0 +1,159 @@
+Metadata-Version: 2.4
+Name: agent-observability-studio
+Version: 0.1.0
+Summary: Real-time monitoring and debugging platform for multi-agent AI systems
+Project-URL: Homepage, https://github.com/yourusername/agent-observability-studio
+Project-URL: Documentation, https://github.com/yourusername/agent-observability-studio#readme
+Project-URL: Repository, https://github.com/yourusername/agent-observability-studio
+Author-email: Zach <zach@example.com>
+License: MIT
+Keywords: agents,ai,debugging,llm,monitoring,observability
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.10
+Requires-Dist: alembic>=1.13.1
+Requires-Dist: anthropic>=0.18.0
+Requires-Dist: fastapi>=0.109.0
+Requires-Dist: httpx>=0.26.0
+Requires-Dist: openai>=1.12.0
+Requires-Dist: pydantic-settings>=2.1.0
+Requires-Dist: pydantic>=2.5.3
+Requires-Dist: python-multipart>=0.0.6
+Requires-Dist: redis>=5.0.1
+Requires-Dist: sqlalchemy>=2.0.25
+Requires-Dist: tiktoken>=0.5.2
+Requires-Dist: uvicorn[standard]>=0.27.0
+Requires-Dist: websockets>=12.0
+Provides-Extra: dev
+Requires-Dist: black>=24.1.1; extra == 'dev'
+Requires-Dist: mypy>=1.8.0; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.23.3; extra == 'dev'
+Requires-Dist: pytest>=7.4.4; extra == 'dev'
+Requires-Dist: ruff>=0.1.14; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# Agent Observability Studio
+
+Real-time monitoring and debugging platform for multi-agent AI systems—Chrome DevTools for AI agents.
+
+## What is this?
+
+Agent Observability Studio provides deep inspection tools for multi-agent AI systems, complementing existing infrastructure dashboards. It captures live interaction traces, analyzes token costs, detects performance bottlenecks, and helps debug agent decision trees and failures in production environments.
+
+## Features
+
+- **Live Interaction Tracing** — Real-time visibility into agent conversations and inter-agent communication
+- **Token Cost Analysis** — Granular per-agent, per-interaction token accounting with cost attribution
+- **Performance Bottleneck Detection** — Automatic identification of slow agents, network delays, and processing inefficiencies
+- **Decision Tree Inspection** — Visualize agent reasoning paths and fallback logic execution
+- **Failure Analysis** — Capture and replay failed interactions with full context
+- **WebSocket Streaming** — Push-based updates for low-latency monitoring
+- **RESTful API** — Programmatic access to traces, metrics, and historical data
+- **CLI Tools** — Command-line utilities for local debugging and integration
+
+## Quick Start
+
+### Installation
+
+```bash
+pip install agent-observability-studio
+```
+
+### Basic Setup
+
+```python
+from agent_observability_studio import ObservabilityClient
+
+# Initialize the client
+client = ObservabilityClient(
+    api_url="http://localhost:8000",
+    api_key="your-api-key"
+)
+
+# Start monitoring an agent interaction
+trace = client.start_trace(
+    agent_id="my-agent",
+    session_id="session-123"
+)
+
+# Log a step
+trace.log_step(
+    name="retrieve_documents",
+    duration_ms=245,
+    tokens_used=1024,
+    status="success"
+)
+
+# End trace
+trace.end()
+```
+
+### CLI Usage
+
+```bash
+# Start the monitoring server
+aos-server --port 8000 --db-url postgresql://localhost/observability
+
+# Stream live traces
+aos-trace watch
+
+# Export session data
+aos-export --session session-123 --format json --output trace.json
+
+# Analyze costs
+aos-costs --agent-id my-agent --date-range "2025-03-01:2025-03-18"
+```
+
+## Usage Examples
+
+**Monitor agent costs in real-time:**
+```python
+trace = client.start_trace(agent_id="researcher")
+# ... agent work ...
+cost_report = trace.get_cost_summary()
+print(f"Total tokens: {cost_report.total_tokens}")
+print(f"Estimated cost: ${cost_report.estimated_cost:.4f}")
+```
+
+**Capture and replay failures:**
+```python
+failed_traces = client.query_traces(status="error", limit=10)
+for trace in failed_traces:
+    print(f"Error: {trace.error_message}")
+    print(f"Decision tree: {trace.decision_path}")
+```
+
+**Subscribe to live events:**
+```python
+async def handle_trace(event):
+    print(f"Trace {event.trace_id}: {event.status}")
+
+client.subscribe("trace.completed", handle_trace)
+```
+
+## Tech Stack
+
+- **Runtime**: Python 3.12+
+- **API**: FastAPI with async support
+- **Database**: PostgreSQL with SQLAlchemy ORM
+- **Real-time**: WebSocket streams via websockets library
+- **Configuration**: Pydantic for settings management
+- **CLI**: Click for command-line interface
+- **Packaging**: Poetry for dependency management
+
+## Architecture
+
+The studio consists of four main components:
+
+- **API Server** (`api.py`) — RESTful endpoints for trace queries and configuration
+- **WebSocket Manager** (`websocket_manager.py`) — Real-time event streaming
+- **Database Layer** (`database.py`) — Persistent storage of traces and metrics
+- **Client SDK** (`client.py`) — Python library for agent instrumentation
+
+## License
+
+MIT

agent_observability_studio-0.1.0/README.md

@@ -0,0 +1,121 @@
+# Agent Observability Studio
+
+Real-time monitoring and debugging platform for multi-agent AI systems—Chrome DevTools for AI agents.
+
+## What is this?
+
+Agent Observability Studio provides deep inspection tools for multi-agent AI systems, complementing existing infrastructure dashboards. It captures live interaction traces, analyzes token costs, detects performance bottlenecks, and helps debug agent decision trees and failures in production environments.
+
+## Features
+
+- **Live Interaction Tracing** — Real-time visibility into agent conversations and inter-agent communication
+- **Token Cost Analysis** — Granular per-agent, per-interaction token accounting with cost attribution
+- **Performance Bottleneck Detection** — Automatic identification of slow agents, network delays, and processing inefficiencies
+- **Decision Tree Inspection** — Visualize agent reasoning paths and fallback logic execution
+- **Failure Analysis** — Capture and replay failed interactions with full context
+- **WebSocket Streaming** — Push-based updates for low-latency monitoring
+- **RESTful API** — Programmatic access to traces, metrics, and historical data
+- **CLI Tools** — Command-line utilities for local debugging and integration
+
+## Quick Start
+
+### Installation
+
+```bash
+pip install agent-observability-studio
+```
+
+### Basic Setup
+
+```python
+from agent_observability_studio import ObservabilityClient
+
+# Initialize the client
+client = ObservabilityClient(
+    api_url="http://localhost:8000",
+    api_key="your-api-key"
+)
+
+# Start monitoring an agent interaction
+trace = client.start_trace(
+    agent_id="my-agent",
+    session_id="session-123"
+)
+
+# Log a step
+trace.log_step(
+    name="retrieve_documents",
+    duration_ms=245,
+    tokens_used=1024,
+    status="success"
+)
+
+# End trace
+trace.end()
+```
+
+### CLI Usage
+
+```bash
+# Start the monitoring server
+aos-server --port 8000 --db-url postgresql://localhost/observability
+
+# Stream live traces
+aos-trace watch
+
+# Export session data
+aos-export --session session-123 --format json --output trace.json
+
+# Analyze costs
+aos-costs --agent-id my-agent --date-range "2025-03-01:2025-03-18"
+```
+
+## Usage Examples
+
+**Monitor agent costs in real-time:**
+```python
+trace = client.start_trace(agent_id="researcher")
+# ... agent work ...
+cost_report = trace.get_cost_summary()
+print(f"Total tokens: {cost_report.total_tokens}")
+print(f"Estimated cost: ${cost_report.estimated_cost:.4f}")
+```
+
+**Capture and replay failures:**
+```python
+failed_traces = client.query_traces(status="error", limit=10)
+for trace in failed_traces:
+    print(f"Error: {trace.error_message}")
+    print(f"Decision tree: {trace.decision_path}")
+```
+
+**Subscribe to live events:**
+```python
+async def handle_trace(event):
+    print(f"Trace {event.trace_id}: {event.status}")
+
+client.subscribe("trace.completed", handle_trace)
+```
+
+## Tech Stack
+
+- **Runtime**: Python 3.12+
+- **API**: FastAPI with async support
+- **Database**: PostgreSQL with SQLAlchemy ORM
+- **Real-time**: WebSocket streams via websockets library
+- **Configuration**: Pydantic for settings management
+- **CLI**: Click for command-line interface
+- **Packaging**: Poetry for dependency management
+
+## Architecture
+
+The studio consists of four main components:
+
+- **API Server** (`api.py`) — RESTful endpoints for trace queries and configuration
+- **WebSocket Manager** (`websocket_manager.py`) — Real-time event streaming
+- **Database Layer** (`database.py`) — Persistent storage of traces and metrics
+- **Client SDK** (`client.py`) — Python library for agent instrumentation
+
+## License
+
+MIT

agent_observability_studio-0.1.0/agent_observability_studio/__init__.py

@@ -0,0 +1,18 @@
+"""Agent Observability Studio - Deep inspection for multi-agent AI systems."""
+
+from agent_observability_studio.client import ObservabilityClient
+from agent_observability_studio.models import (
+    Session,
+    Interaction,
+    SessionStatus,
+    InteractionType,
+)
+
+__version__ = "0.1.0"
+__all__ = [
+    "ObservabilityClient",
+    "Session",
+    "Interaction",
+    "SessionStatus",
+    "InteractionType",
+]

agent_observability_studio-0.1.0/agent_observability_studio/api.py

@@ -0,0 +1,407 @@
+"""FastAPI application for observability platform."""
+
+from datetime import datetime
+from typing import List, Optional
+from uuid import UUID
+from fastapi import FastAPI, HTTPException, Depends, WebSocket, WebSocketDisconnect, Query
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from sqlalchemy.orm import Session
+from sqlalchemy import func, and_
+
+from agent_observability_studio.database import Database, SessionDB, InteractionDB
+from agent_observability_studio.models import (
+    Session as SessionModel,
+    Interaction as InteractionModel,
+    SessionCreate,
+    SessionUpdate,
+    InteractionCreate,
+    SessionQuery,
+    SessionStatus,
+    CostBreakdown,
+    DecisionNode,
+)
+from agent_observability_studio.websocket_manager import ConnectionManager
+from agent_observability_studio.config import get_settings
+
+
+settings = get_settings()
+db_manager = Database(settings.database_url)
+app = FastAPI(title="Agent Observability Studio", version="0.1.0")
+ws_manager = ConnectionManager()
+
+# CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+def get_db():
+    """Dependency to get database session."""
+    return next(db_manager.get_session())
+
+
+@app.get("/")
+async def root():
+    """Root endpoint."""
+    return {
+        "service": "Agent Observability Studio",
+        "version": "0.1.0",
+        "endpoints": {
+            "api": "/api/v1",
+            "docs": "/docs",
+            "websocket": "/ws",
+            "ui": "/ui",
+        },
+    }
+
+
+@app.post("/api/v1/sessions", response_model=SessionModel)
+async def create_session(
+    session_data: SessionCreate, db: Session = Depends(get_db)
+) -> SessionModel:
+    """Create a new agent session."""
+    db_session = SessionDB(
+        agent_name=session_data.agent_name,
+        task_id=session_data.task_id,
+        metadata=session_data.metadata,
+        parent_session_id=session_data.parent_session_id,
+    )
+    db.add(db_session)
+    db.commit()
+    db.refresh(db_session)
+
+    session_model = SessionModel(
+        id=db_session.id,
+        agent_name=db_session.agent_name,
+        task_id=db_session.task_id,
+        status=db_session.status,
+        start_time=db_session.start_time,
+        end_time=db_session.end_time,
+        total_tokens=db_session.total_tokens,
+        total_cost_usd=db_session.total_cost_usd,
+        metadata=db_session.metadata,
+        parent_session_id=db_session.parent_session_id,
+    )
+
+    await ws_manager.broadcast(
+        {"type": "session_created", "session": session_model.model_dump(mode="json")}
+    )
+
+    return session_model
+
+
+@app.patch("/api/v1/sessions/{session_id}", response_model=SessionModel)
+async def update_session(
+    session_id: UUID, update_data: SessionUpdate, db: Session = Depends(get_db)
+) -> SessionModel:
+    """Update an existing session."""
+    db_session = db.query(SessionDB).filter(SessionDB.id == session_id).first()
+    if not db_session:
+        raise HTTPException(status_code=404, detail="Session not found")
+
+    if update_data.status:
+        db_session.status = update_data.status
+    if update_data.end_time:
+        db_session.end_time = update_data.end_time
+    if update_data.metadata:
+        db_session.metadata.update(update_data.metadata)
+
+    db.commit()
+    db.refresh(db_session)
+
+    session_model = SessionModel(
+        id=db_session.id,
+        agent_name=db_session.agent_name,
+        task_id=db_session.task_id,
+        status=db_session.status,
+        start_time=db_session.start_time,
+        end_time=db_session.end_time,
+        total_tokens=db_session.total_tokens,
+        total_cost_usd=db_session.total_cost_usd,
+        metadata=db_session.metadata,
+        parent_session_id=db_session.parent_session_id,
+    )
+
+    await ws_manager.broadcast(
+        {"type": "session_updated", "session": session_model.model_dump(mode="json")}
+    )
+
+    return session_model
+
+
+@app.get("/api/v1/sessions", response_model=List[SessionModel])
+async def query_sessions(
+    agent_name: Optional[str] = None,
+    task_id: Optional[str] = None,
+    status: Optional[SessionStatus] = None,
+    start_date: Optional[datetime] = None,
+    end_date: Optional[datetime] = None,
+    limit: int = Query(100, le=1000),
+    offset: int = 0,
+    db: Session = Depends(get_db),
+) -> List[SessionModel]:
+    """Query sessions with filters."""
+    query = db.query(SessionDB)
+
+    if agent_name:
+        query = query.filter(SessionDB.agent_name == agent_name)
+    if task_id:
+        query = query.filter(SessionDB.task_id == task_id)
+    if status:
+        query = query.filter(SessionDB.status == status)
+    if start_date:
+        query = query.filter(SessionDB.start_time >= start_date)
+    if end_date:
+        query = query.filter(SessionDB.start_time <= end_date)
+
+    sessions = query.order_by(SessionDB.start_time.desc()).limit(limit).offset(offset).all()
+
+    return [
+        SessionModel(
+            id=s.id,
+            agent_name=s.agent_name,
+            task_id=s.task_id,
+            status=s.status,
+            start_time=s.start_time,
+            end_time=s.end_time,
+            total_tokens=s.total_tokens,
+            total_cost_usd=s.total_cost_usd,
+            metadata=s.metadata,
+            parent_session_id=s.parent_session_id,
+        )
+        for s in sessions
+    ]
+
+
+@app.get("/api/v1/sessions/{session_id}", response_model=SessionModel)
+async def get_session(session_id: UUID, db: Session = Depends(get_db)) -> SessionModel:
+    """Get a specific session by ID."""
+    db_session = db.query(SessionDB).filter(SessionDB.id == session_id).first()
+    if not db_session:
+        raise HTTPException(status_code=404, detail="Session not found")
+
+    return SessionModel(
+        id=db_session.id,
+        agent_name=db_session.agent_name,
+        task_id=db_session.task_id,
+        status=db_session.status,
+        start_time=db_session.start_time,
+        end_time=db_session.end_time,
+        total_tokens=db_session.total_tokens,
+        total_cost_usd=db_session.total_cost_usd,
+        metadata=db_session.metadata,
+        parent_session_id=db_session.parent_session_id,
+    )
+
+
+@app.post("/api/v1/interactions", response_model=InteractionModel)
+async def log_interaction(
+    interaction_data: InteractionCreate, db: Session = Depends(get_db)
+) -> InteractionModel:
+    """Log a new interaction for a session."""
+    # Verify session exists
+    db_session = (
+        db.query(SessionDB).filter(SessionDB.id == interaction_data.session_id).first()
+    )
+    if not db_session:
+        raise HTTPException(status_code=404, detail="Session not found")
+
+    # Create interaction
+    db_interaction = InteractionDB(
+        session_id=interaction_data.session_id,
+        type=interaction_data.type,
+        request=interaction_data.request,
+        response=interaction_data.response,
+        tokens_used=interaction_data.tokens_used,
+        cost_usd=interaction_data.cost_usd,
+        latency_ms=interaction_data.latency_ms,
+        model=interaction_data.model,
+        error=interaction_data.error,
+        metadata=interaction_data.metadata,
+    )
+    db.add(db_interaction)
+
+    # Update session totals
+    db_session.total_tokens += interaction_data.tokens_used
+    db_session.total_cost_usd += interaction_data.cost_usd
+
+    db.commit()
+    db.refresh(db_interaction)
+
+    interaction_model = InteractionModel(
+        id=db_interaction.id,
+        session_id=db_interaction.session_id,
+        type=db_interaction.type,
+        timestamp=db_interaction.timestamp,
+        request=db_interaction.request,
+        response=db_interaction.response,
+        tokens_used=db_interaction.tokens_used,
+        cost_usd=db_interaction.cost_usd,
+        latency_ms=db_interaction.latency_ms,
+        model=db_interaction.model,
+        error=db_interaction.error,
+        metadata=db_interaction.metadata,
+    )
+
+    await ws_manager.broadcast(
+        {"type": "interaction_logged", "interaction": interaction_model.model_dump(mode="json")}
+    )
+
+    return interaction_model
+
+
+@app.get("/api/v1/sessions/{session_id}/interactions", response_model=List[InteractionModel])
+async def get_session_interactions(
+    session_id: UUID, db: Session = Depends(get_db)
+) -> List[InteractionModel]:
+    """Get all interactions for a session."""
+    interactions = (
+        db.query(InteractionDB)
+        .filter(InteractionDB.session_id == session_id)
+        .order_by(InteractionDB.timestamp)
+        .all()
+    )
+
+    return [
+        InteractionModel(
+            id=i.id,
+            session_id=i.session_id,
+            type=i.type,
+            timestamp=i.timestamp,
+            request=i.request,
+            response=i.response,
+            tokens_used=i.tokens_used,
+            cost_usd=i.cost_usd,
+            latency_ms=i.latency_ms,
+            model=i.model,
+            error=i.error,
+            metadata=i.metadata,
+        )
+        for i in interactions
+    ]
+
+
+@app.get("/api/v1/analytics/cost-breakdown", response_model=CostBreakdown)
+async def get_cost_breakdown(
+    start_date: Optional[datetime] = None,
+    end_date: Optional[datetime] = None,
+    db: Session = Depends(get_db),
+) -> CostBreakdown:
+    """Get cost breakdown analytics."""
+    query = db.query(SessionDB)
+
+    if start_date:
+        query = query.filter(SessionDB.start_time >= start_date)
+    if end_date:
+        query = query.filter(SessionDB.start_time <= end_date)
+
+    sessions = query.all()
+
+    total_cost = sum(s.total_cost_usd for s in sessions)
+    total_tokens = sum(s.total_tokens for s in sessions)
+
+    by_agent = {}
+    by_task = {}
+    for s in sessions:
+        by_agent[s.agent_name] = by_agent.get(s.agent_name, 0) + s.total_cost_usd
+        if s.task_id:
+            by_task[s.task_id] = by_task.get(s.task_id, 0) + s.total_cost_usd
+
+    # Get model breakdown from interactions
+    interactions = (
+        db.query(InteractionDB.model, func.sum(InteractionDB.cost_usd))
+        .filter(
+            and_(
+                InteractionDB.session_id.in_([s.id for s in sessions]),
+                InteractionDB.model.isnot(None),
+            )
+        )
+        .group_by(InteractionDB.model)
+        .all()
+    )
+    by_model = {model: float(cost) for model, cost in interactions if model}
+
+    time_range = (
+        min(s.start_time for s in sessions) if sessions else datetime.utcnow(),
+        max(s.start_time for s in sessions) if sessions else datetime.utcnow(),
+    )
+
+    return CostBreakdown(
+        total_cost_usd=total_cost,
+        total_tokens=total_tokens,
+        by_agent=by_agent,
+        by_model=by_model,
+        by_task=by_task,
+        time_range=time_range,
+    )
+
+
+@app.get("/api/v1/analytics/decision-tree", response_model=List[DecisionNode])
+async def get_decision_tree(
+    root_session_id: Optional[UUID] = None, db: Session = Depends(get_db)
+) -> List[DecisionNode]:
+    """Get decision tree for agent collaboration."""
+    if root_session_id:
+        root = db.query(SessionDB).filter(SessionDB.id == root_session_id).first()
+        if not root:
+            raise HTTPException(status_code=404, detail="Root session not found")
+        sessions = [root] + _get_all_children(db, root.id)
+    else:
+        # Get all root sessions (no parent)
+        sessions = db.query(SessionDB).filter(SessionDB.parent_session_id.is_(None)).all()
+
+    nodes = []
+    for s in sessions:
+        children = (
+            db.query(SessionDB.id)
+            .filter(SessionDB.parent_session_id == s.id)
+            .all()
+        )
+        nodes.append(
+            DecisionNode(
+                session_id=s.id,
+                agent_name=s.agent_name,
+                task_id=s.task_id,
+                children=[c[0] for c in children],
+                status=s.status,
+                total_cost_usd=s.total_cost_usd,
+                metadata=s.metadata,
+            )
+        )
+
+    return nodes
+
+
+def _get_all_children(db: Session, parent_id: UUID) -> List[SessionDB]:
+    """Recursively get all child sessions."""
+    children = db.query(SessionDB).filter(SessionDB.parent_session_id == parent_id).all()
+    all_children = children.copy()
+    for child in children:
+        all_children.extend(_get_all_children(db, child.id))
+    return all_children
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    """WebSocket endpoint for real-time updates."""
+    await ws_manager.connect(websocket)
+    try:
+        while True:
+            # Keep connection alive and receive client filters
+            data = await websocket.receive_json()
+            if data.get("type") == "filter":
+                # Client can send filters to only receive certain events
+                await websocket.send_json({"status": "filter_applied", "filters": data.get("filters", {})})
+    except WebSocketDisconnect:
+        ws_manager.disconnect(websocket)
+
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint."""
+    return {"status": "healthy", "timestamp": datetime.utcnow().isoformat()}

agent_observability_studio-0.1.0/agent_observability_studio/cli.py

@@ -0,0 +1,32 @@
+"""CLI for running the observability platform."""
+
+import argparse
+import uvicorn
+
+
+def main():
+    """Main CLI entry point."""
+    parser = argparse.ArgumentParser(description="Agent Observability Studio CLI")
+    subparsers = parser.add_subparsers(dest="command", help="Command to run")
+
+    # Serve command
+    serve_parser = subparsers.add_parser("serve", help="Start the API server")
+    serve_parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
+    serve_parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
+    serve_parser.add_argument("--reload", action="store_true", help="Enable auto-reload")
+
+    args = parser.parse_args()
+
+    if args.command == "serve":
+        uvicorn.run(
+            "agent_observability_studio.api:app",
+            host=args.host,
+            port=args.port,
+            reload=args.reload,
+        )
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

agent_observability_studio-0.1.0/agent_observability_studio/client.py

@@ -0,0 +1,146 @@
+"""Python SDK for instrumenting agents."""
+
+from datetime import datetime
+from typing import Any, Dict, Optional
+from uuid import UUID
+import httpx
+
+from agent_observability_studio.models import (
+    SessionCreate,
+    SessionUpdate,
+    InteractionCreate,
+    SessionStatus,
+    InteractionType,
+)
+
+
+class ObservabilityClient:
+    """Client for logging agent interactions to the observability platform."""
+
+    def __init__(
+        self, api_url: str = "http://localhost:8000", timeout: float = 5.0
+    ):
+        """Initialize client.
+
+        Args:
+            api_url: Base URL of the observability API
+            timeout: Request timeout in seconds
+        """
+        self.api_url = api_url.rstrip("/")
+        self.timeout = timeout
+        self.client = httpx.Client(timeout=timeout)
+
+    def start_session(
+        self,
+        agent_name: str,
+        task_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        parent_session_id: Optional[UUID] = None,
+    ) -> UUID:
+        """Start a new agent session.
+
+        Args:
+            agent_name: Name of the agent
+            task_id: Optional task identifier
+            metadata: Optional metadata dictionary
+            parent_session_id: Optional parent session for nested agents
+
+        Returns:
+            Session UUID
+        """
+        session_data = SessionCreate(
+            agent_name=agent_name,
+            task_id=task_id,
+            metadata=metadata or {},
+            parent_session_id=parent_session_id,
+        )
+
+        response = self.client.post(
+            f"{self.api_url}/api/v1/sessions",
+            json=session_data.model_dump(mode="json"),
+        )
+        response.raise_for_status()
+
+        return UUID(response.json()["id"])
+
+    def end_session(
+        self,
+        session_id: UUID,
+        status: SessionStatus = SessionStatus.SUCCESS,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """End an agent session.
+
+        Args:
+            session_id: Session UUID to end
+            status: Final status of the session
+            metadata: Optional additional metadata
+        """
+        update_data = SessionUpdate(
+            status=status, end_time=datetime.utcnow(), metadata=metadata
+        )
+
+        response = self.client.patch(
+            f"{self.api_url}/api/v1/sessions/{session_id}",
+            json=update_data.model_dump(mode="json", exclude_none=True),
+        )
+        response.raise_for_status()
+
+    def log_interaction(
+        self,
+        session_id: UUID,
+        request: Dict[str, Any],
+        response: Dict[str, Any],
+        tokens_used: int,
+        cost_usd: float = 0.0,
+        latency_ms: Optional[int] = None,
+        model: Optional[str] = None,
+        error: Optional[str] = None,
+        interaction_type: InteractionType = InteractionType.LLM_CALL,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> UUID:
+        """Log an interaction within a session.
+
+        Args:
+            session_id: Session UUID this interaction belongs to
+            request: Request data (e.g., prompt, messages)
+            response: Response data (e.g., completion, tool output)
+            tokens_used: Number of tokens consumed
+            cost_usd: Cost in USD
+            latency_ms: Latency in milliseconds
+            model: Model identifier
+            error: Error message if interaction failed
+            interaction_type: Type of interaction
+            metadata: Optional additional metadata
+
+        Returns:
+            Interaction UUID
+        """
+        interaction_data = InteractionCreate(
+            session_id=session_id,
+            type=interaction_type,
+            request=request,
+            response=response,
+            tokens_used=tokens_used,
+            cost_usd=cost_usd,
+            latency_ms=latency_ms,
+            model=model,
+            error=error,
+            metadata=metadata or {},
+        )
+
+        response = self.client.post(
+            f"{self.api_url}/api/v1/interactions",
+            json=interaction_data.model_dump(mode="json"),
+        )
+        response.raise_for_status()
+
+        return UUID(response.json()["id"])
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.client.close()

agent_observability_studio-0.1.0/agent_observability_studio/config.py

@@ -0,0 +1,25 @@
+"""Configuration management."""
+
+from functools import lru_cache
+from pydantic_settings import BaseSettings
+
+
+class Settings(BaseSettings):
+    """Application settings."""
+
+    database_url: str = "sqlite:///./agent_studio.db"
+    redis_url: str = "redis://localhost:6379/0"
+    host: str = "0.0.0.0"
+    port: int = 8000
+    reload: bool = False
+    log_level: str = "info"
+
+    class Config:
+        env_file = ".env"
+        env_file_encoding = "utf-8"
+
+
+@lru_cache()
+def get_settings() -> Settings:
+    """Get cached settings instance."""
+    return Settings()

agent_observability_studio-0.1.0/agent_observability_studio/database.py

@@ -0,0 +1,86 @@
+"""Database configuration and models."""
+
+from datetime import datetime
+from typing import Optional
+from sqlalchemy import (
+    create_engine,
+    Column,
+    String,
+    Integer,
+    Float,
+    DateTime,
+    Enum,
+    JSON,
+    ForeignKey,
+)
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker, relationship
+from sqlalchemy.dialects.postgresql import UUID
+import uuid
+
+from agent_observability_studio.models import SessionStatus, InteractionType
+
+
+Base = declarative_base()
+
+
+class SessionDB(Base):
+    """Database model for sessions."""
+
+    __tablename__ = "sessions"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    agent_name = Column(String, nullable=False, index=True)
+    task_id = Column(String, nullable=True, index=True)
+    status = Column(Enum(SessionStatus), nullable=False, default=SessionStatus.ACTIVE)
+    start_time = Column(DateTime, nullable=False, default=datetime.utcnow, index=True)
+    end_time = Column(DateTime, nullable=True)
+    total_tokens = Column(Integer, nullable=False, default=0)
+    total_cost_usd = Column(Float, nullable=False, default=0.0)
+    metadata = Column(JSON, nullable=False, default=dict)
+    parent_session_id = Column(UUID(as_uuid=True), ForeignKey("sessions.id"), nullable=True)
+
+    interactions = relationship("InteractionDB", back_populates="session", cascade="all, delete")
+    children = relationship("SessionDB", backref="parent", remote_side=[id])
+
+
+class InteractionDB(Base):
+    """Database model for interactions."""
+
+    __tablename__ = "interactions"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    session_id = Column(
+        UUID(as_uuid=True), ForeignKey("sessions.id"), nullable=False, index=True
+    )
+    type = Column(Enum(InteractionType), nullable=False, default=InteractionType.LLM_CALL)
+    timestamp = Column(DateTime, nullable=False, default=datetime.utcnow, index=True)
+    request = Column(JSON, nullable=False)
+    response = Column(JSON, nullable=False)
+    tokens_used = Column(Integer, nullable=False)
+    cost_usd = Column(Float, nullable=False, default=0.0)
+    latency_ms = Column(Integer, nullable=True)
+    model = Column(String, nullable=True, index=True)
+    error = Column(String, nullable=True)
+    metadata = Column(JSON, nullable=False, default=dict)
+
+    session = relationship("SessionDB", back_populates="interactions")
+
+
+class Database:
+    """Database connection manager."""
+
+    def __init__(self, database_url: str = "sqlite:///./agent_studio.db"):
+        self.engine = create_engine(
+            database_url, connect_args={"check_same_thread": False} if "sqlite" in database_url else {}
+        )
+        self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine)
+        Base.metadata.create_all(bind=self.engine)
+
+    def get_session(self):
+        """Get database session."""
+        db = self.SessionLocal()
+        try:
+            yield db
+        finally:
+            db.close()

agent_observability_studio-0.1.0/agent_observability_studio/models.py

@@ -0,0 +1,126 @@
+"""Data models for observability platform."""
+
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, Optional
+from pydantic import BaseModel, Field
+from uuid import UUID, uuid4
+
+
+class SessionStatus(str, Enum):
+    """Status of an agent session."""
+
+    ACTIVE = "active"
+    SUCCESS = "success"
+    FAILED = "failed"
+    TIMEOUT = "timeout"
+
+
+class InteractionType(str, Enum):
+    """Type of agent interaction."""
+
+    LLM_CALL = "llm_call"
+    TOOL_USE = "tool_use"
+    MEMORY_ACCESS = "memory_access"
+    AGENT_COMMUNICATION = "agent_communication"
+
+
+class Session(BaseModel):
+    """An agent execution session."""
+
+    id: UUID = Field(default_factory=uuid4)
+    agent_name: str
+    task_id: Optional[str] = None
+    status: SessionStatus = SessionStatus.ACTIVE
+    start_time: datetime = Field(default_factory=datetime.utcnow)
+    end_time: Optional[datetime] = None
+    total_tokens: int = 0
+    total_cost_usd: float = 0.0
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    parent_session_id: Optional[UUID] = None
+
+
+class Interaction(BaseModel):
+    """A single LLM or tool interaction within a session."""
+
+    id: UUID = Field(default_factory=uuid4)
+    session_id: UUID
+    type: InteractionType = InteractionType.LLM_CALL
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
+    request: Dict[str, Any]
+    response: Dict[str, Any]
+    tokens_used: int
+    cost_usd: float = 0.0
+    latency_ms: Optional[int] = None
+    model: Optional[str] = None
+    error: Optional[str] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class SessionCreate(BaseModel):
+    """Request to create a new session."""
+
+    agent_name: str
+    task_id: Optional[str] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    parent_session_id: Optional[UUID] = None
+
+
+class SessionUpdate(BaseModel):
+    """Request to update a session."""
+
+    status: Optional[SessionStatus] = None
+    end_time: Optional[datetime] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class InteractionCreate(BaseModel):
+    """Request to log an interaction."""
+
+    session_id: UUID
+    type: InteractionType = InteractionType.LLM_CALL
+    request: Dict[str, Any]
+    response: Dict[str, Any]
+    tokens_used: int
+    cost_usd: float = 0.0
+    latency_ms: Optional[int] = None
+    model: Optional[str] = None
+    error: Optional[str] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class SessionQuery(BaseModel):
+    """Query parameters for filtering sessions."""
+
+    agent_name: Optional[str] = None
+    task_id: Optional[str] = None
+    status: Optional[SessionStatus] = None
+    start_date: Optional[datetime] = None
+    end_date: Optional[datetime] = None
+    min_tokens: Optional[int] = None
+    max_tokens: Optional[int] = None
+    limit: int = 100
+    offset: int = 0
+
+
+class CostBreakdown(BaseModel):
+    """Token cost analysis results."""
+
+    total_cost_usd: float
+    total_tokens: int
+    by_agent: Dict[str, float]
+    by_model: Dict[str, float]
+    by_task: Dict[str, float]
+    time_range: tuple[datetime, datetime]
+
+
+class DecisionNode(BaseModel):
+    """Node in an agent decision tree."""
+
+    session_id: UUID
+    agent_name: str
+    task_id: Optional[str]
+    children: list[UUID] = Field(default_factory=list)
+    status: SessionStatus
+    total_cost_usd: float
+    metadata: Dict[str, Any] = Field(default_factory=dict)

agent_observability_studio-0.1.0/agent_observability_studio/websocket_manager.py

@@ -0,0 +1,30 @@
+"""WebSocket connection manager for real-time updates."""
+
+from typing import List, Dict, Any
+from fastapi import WebSocket
+
+
+class ConnectionManager:
+    """Manages WebSocket connections and broadcasting."""
+
+    def __init__(self):
+        self.active_connections: List[WebSocket] = []
+
+    async def connect(self, websocket: WebSocket):
+        """Accept and store a new WebSocket connection."""
+        await websocket.accept()
+        self.active_connections.append(websocket)
+
+    def disconnect(self, websocket: WebSocket):
+        """Remove a WebSocket connection."""
+        if websocket in self.active_connections:
+            self.active_connections.remove(websocket)
+
+    async def broadcast(self, message: Dict[str, Any]):
+        """Broadcast a message to all connected clients."""
+        for connection in self.active_connections:
+            try:
+                await connection.send_json(message)
+            except Exception:
+                # Connection might be closed, will be removed on next disconnect
+                pass

agent_observability_studio-0.1.0/pyproject.toml

@@ -0,0 +1,72 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "agent-observability-studio"
+version = "0.1.0"
+description = "Real-time monitoring and debugging platform for multi-agent AI systems"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "MIT"}
+authors = [
+    {name = "Zach", email = "zach@example.com"}
+]
+keywords = ["ai", "agents", "observability", "monitoring", "debugging", "llm"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+]
+dependencies = [
+    "fastapi>=0.109.0",
+    "uvicorn[standard]>=0.27.0",
+    "websockets>=12.0",
+    "sqlalchemy>=2.0.25",
+    "alembic>=1.13.1",
+    "pydantic>=2.5.3",
+    "pydantic-settings>=2.1.0",
+    "httpx>=0.26.0",
+    "redis>=5.0.1",
+    "python-multipart>=0.0.6",
+    "tiktoken>=0.5.2",
+    "anthropic>=0.18.0",
+    "openai>=1.12.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.4.4",
+    "pytest-asyncio>=0.23.3",
+    "black>=24.1.1",
+    "ruff>=0.1.14",
+    "mypy>=1.8.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/yourusername/agent-observability-studio"
+Documentation = "https://github.com/yourusername/agent-observability-studio#readme"
+Repository = "https://github.com/yourusername/agent-observability-studio"
+
+[project.scripts]
+agent-studio = "agent_observability_studio.cli:main"
+
+[tool.hatch.build.targets.wheel]
+packages = ["agent_observability_studio"]
+
+[tool.black]
+line-length = 100
+target-version = ['py310']
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.mypy]
+python_version = "3.10"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true