tracebrain 1.0.0__py3-none-any.whl

tracebrain/__init__.py

@@ -0,0 +1,63 @@
+"""
+TraceBrain - Observability Platform for Agentic AI
+
+This package provides a complete observability solution for AI agents,
+allowing users to collect, store, and visualize execution traces.
+
+Philosophy: "Pip install and run"
+- Single package containing both backend (FastAPI) and frontend (React)
+- Support for SQLite (development) and PostgreSQL (production)
+- Custom TraceBrain Standard OTLP Trace Schema
+- Robust SDK client with automatic retries and fail-safe design
+
+Quick Start:
+    # Install
+    pip install tracebrain
+    
+    # Start infrastructure with Docker (recommended)
+    tracebrain up
+    
+    # Or use Python server directly for development
+    tracebrain init-db
+    tracebrain start
+    
+    # Use the SDK client in your code
+    from tracebrain import TraceClient
+    
+    client = TraceClient()
+    success = client.log_trace({
+        "trace_id": "abc123",
+        "attributes": {"system_prompt": "You are helpful"},
+        "spans": [...]
+    })
+
+Usage:
+    # Import the FastAPI app
+    from tracebrain import app
+    
+    # Import configuration
+    from tracebrain import settings
+    
+    # Import SDK client (recommended)
+    from tracebrain import TraceClient
+    
+    # Import TraceStore for programmatic access
+    from tracebrain.core.store import TraceStore
+"""
+
+__version__ = "1.0.0"
+__author__ = "TraceBrain Team"
+
+# Expose main components for easy import
+from .main import app
+from .config import settings
+from .sdk import TraceClient
+from .sdk.client import TraceScope
+
+__all__ = [
+    "app",
+    "settings",
+    "TraceClient",
+    "TraceScope",
+    "__version__",
+]

tracebrain/api/v1/__init__.py

@@ -0,0 +1,9 @@
+"""
+TraceBrain API v1
+
+This module exports the API v1 router for use in the main application.
+"""
+
+from .api_router import router
+
+__all__ = ["router"]

tracebrain/api/v1/ai_features.py

@@ -0,0 +1,168 @@
+"""AI-related endpoints for v1."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+import uuid
+
+from fastapi import APIRouter, HTTPException
+
+from ...evaluators.judge_agent import AIJudge
+from ...core.llm_providers import ProviderError
+from .common import build_ai_evaluation, get_librarian_agent, store
+from .schemas.api_models import (
+    AIEvaluationIn,
+    AIEvaluationOut,
+    ChatHistoryOut,
+    NaturalLanguageQuery,
+    NaturalLanguageResponse,
+)
+
+router = APIRouter()
+
+
+def _normalize_suggestions(value: Any) -> List[Dict[str, str]]:
+    if not isinstance(value, list):
+        return []
+
+    normalized: List[Dict[str, str]] = []
+    for item in value:
+        if not isinstance(item, dict):
+            continue
+        label = str(item.get("label", "")).strip()
+        suggestion_value = str(item.get("value", "")).strip()
+        if label and suggestion_value:
+            normalized.append({"label": label, "value": suggestion_value})
+    return normalized
+
+
+def _normalize_sources(value: Any) -> List[str]:
+    normalized: List[str] = []
+
+    def _append_candidate(candidate: Any) -> None:
+        if candidate is None:
+            return
+        if isinstance(candidate, str):
+            text = candidate.strip()
+            if text:
+                normalized.append(text)
+            return
+        if isinstance(candidate, dict):
+            source_id = candidate.get("trace_id") or candidate.get("id")
+            if source_id is not None:
+                text = str(source_id).strip()
+                if text:
+                    normalized.append(text)
+            return
+
+        text = str(candidate).strip()
+        if text:
+            normalized.append(text)
+
+    if isinstance(value, (list, tuple, set)):
+        for item in value:
+            _append_candidate(item)
+    elif value is not None:
+        _append_candidate(value)
+
+    return list(dict.fromkeys(normalized))
+
+
+def _normalize_filters(value: Any) -> Dict[str, Any]:
+    if not isinstance(value, dict):
+        return {}
+    return {str(k).strip(): v for k, v in value.items() if k and v is not None}
+
+
+def _build_nlq_response(result: Dict[str, Any], session_id: str) -> NaturalLanguageResponse:
+    return NaturalLanguageResponse(
+        answer=str(result.get("answer", "")),
+        session_id=session_id,
+        suggestions=_normalize_suggestions(result.get("suggestions")),
+        sources=_normalize_sources(result.get("sources")),
+        filters=_normalize_filters(result.get("filters")),
+        is_error=bool(result.get("is_error", False)),
+    )
+
+
+@router.get("/librarian_sessions/{session_id}", response_model=ChatHistoryOut, tags=["AI"])
+def get_librarian_session(session_id: str):
+    """Fetch the stored chat history for a Librarian session."""
+    try:
+        messages = store.get_chat_history(session_id)
+        if not messages:
+            raise HTTPException(status_code=404, detail="Session not found")
+
+        return ChatHistoryOut(session_id=session_id, messages=messages)
+    except HTTPException:
+        raise
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"Failed to load session: {str(exc)}")
+
+
+@router.post("/natural_language_query", response_model=NaturalLanguageResponse, tags=["AI"])
+def natural_language_query(query: NaturalLanguageQuery):
+    """
+    Process a natural language query about traces using the configured LLM provider.
+
+    The provider is selected via settings (LIBRARIAN_MODE/LLM_PROVIDER) and can route to
+    API-hosted or open-source backends. The agent uses function calling (when supported)
+    to query the TraceStore.
+    """
+    session_id = query.session_id or str(uuid.uuid4())
+
+    try:
+        agent = get_librarian_agent()
+        result = agent.query(query.query, session_id=session_id)
+        if not isinstance(result, dict):
+            result = {
+                "answer": str(result),
+                "suggestions": [],
+                "sources": [],
+                "filters": {},
+            }
+        return _build_nlq_response(result, session_id=session_id)
+
+    except Exception as exc:
+        if isinstance(exc, ProviderError):
+            answer = str(exc)
+        else:
+            answer = (
+                "Sorry, I encountered an error processing your query. "
+                "Please try rephrasing your question or check the server logs."
+            )
+        return _build_nlq_response(
+            {
+                "answer": answer,
+                "suggestions": [],
+                "sources": [],
+                "filters": {},
+                "is_error": True,
+            },
+            session_id=session_id,
+        )
+
+
+@router.post("/ai_evaluate/{trace_id}", response_model=AIEvaluationOut, tags=["AI Evaluation"])
+def evaluate_trace_with_ai(trace_id: str, payload: AIEvaluationIn):
+    """
+    Evaluate a trace with a judge model.
+
+    This endpoint is designed as a hook for more complex AI evaluation logic.
+    """
+    try:
+        judge = AIJudge(store)
+        result = judge.evaluate(trace_id, payload.judge_model_id)
+
+        ai_eval = build_ai_evaluation(result)
+        store.update_ai_evaluation(trace_id, ai_eval)
+
+        return AIEvaluationOut(**ai_eval)
+
+    except ValueError as exc:
+        message = str(exc)
+        if "Trace not found" in message:
+            raise HTTPException(status_code=404, detail=message)
+        raise HTTPException(status_code=400, detail=message)
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"AI evaluation failed: {exc}")

tracebrain/api/v1/api_router.py

@@ -0,0 +1,19 @@
+"""Main API router for v1."""
+
+from fastapi import APIRouter
+
+from .ai_features import router as ai_router
+from .curriculum import router as curriculum_router
+from .episodes import router as episodes_router
+from .operations import router as operations_router
+from .system import router as system_router
+from .traces import router as traces_router
+
+router = APIRouter()
+
+router.include_router(system_router)
+router.include_router(traces_router)
+router.include_router(episodes_router)
+router.include_router(curriculum_router)
+router.include_router(operations_router)
+router.include_router(ai_router)

tracebrain/api/v1/common.py

@@ -0,0 +1,53 @@
+"""Shared dependencies and helpers for v1 routers."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, Dict
+import logging
+
+from ...core.store import TraceStore
+from ...evaluators.judge_agent import AIJudge
+from ...core.librarian import LibrarianAgent
+from ...config import settings
+
+logger = logging.getLogger(__name__)
+
+store = TraceStore(
+    backend=settings.get_backend_type(),
+    db_url=settings.DATABASE_URL,
+)
+
+_librarian_agent = None
+
+
+def get_librarian_agent() -> LibrarianAgent:
+    """Lazy initialization of Librarian agent."""
+    global _librarian_agent
+    if _librarian_agent is None:
+        _librarian_agent = LibrarianAgent(store)
+    return _librarian_agent
+
+
+def build_ai_evaluation(result: Dict[str, Any]) -> Dict[str, Any]:
+    confidence = float(result.get("confidence", 0.0))
+    status_value = "auto_verified" if confidence > 0.8 else "pending_review"
+    return {
+        "rating": result.get("rating"),
+        "feedback": result.get("feedback"),
+        "confidence": confidence,
+        "error_type": result.get("error_type", "none"),
+        "status": status_value,
+        "priority": result.get("priority"),
+        "timestamp": datetime.utcnow().isoformat(),
+    }
+
+
+def run_bg_evaluation(trace_id: str) -> None:
+    try:
+        judge = AIJudge(store)
+        result = judge.evaluate(trace_id)
+        ai_eval = build_ai_evaluation(result)
+        store.update_ai_evaluation(trace_id, ai_eval)
+    except Exception as exc:
+        logger.error("Background evaluation failed for %s: %s", trace_id, exc)

tracebrain/api/v1/curriculum.py

@@ -0,0 +1,138 @@
+"""Curriculum endpoints for v1."""
+
+from __future__ import annotations
+
+from typing import List
+import json
+
+from fastapi import APIRouter, HTTPException, Query
+from fastapi.responses import Response
+
+from ...core.curator import CurriculumCurator
+from ...db.base import CurriculumTask
+from .common import store
+from .schemas.api_models import CurriculumTaskOut, GenerateCurriculumRequest
+
+router = APIRouter(prefix="/curriculum", tags=["Curriculum"])
+
+
+@router.post("/generate")
+def generate_curriculum(request: GenerateCurriculumRequest):
+    """Generate curriculum tasks from failed traces."""
+    try:
+        curator = CurriculumCurator(store)
+        provided_error_types = request.error_types or []
+        valid_error_types = [
+            value for value in provided_error_types if value in curator.VALID_ERROR_TYPES
+        ]
+        invalid_error_types = [
+            value for value in provided_error_types if value not in curator.VALID_ERROR_TYPES
+        ]
+        created = curator.generate_curriculum(
+            error_types=valid_error_types or None,
+            limit=request.limit,
+        )
+        response = {"status": "success", "tasks_generated": created}
+        if invalid_error_types:
+            response["warning"] = {
+                "message": "Some error_types were not recognized and were ignored.",
+                "invalid_error_types": invalid_error_types,
+            }
+        return response
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail=str(exc))
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"Failed to generate curriculum: {str(exc)}")
+
+
+@router.get("", response_model=List[CurriculumTaskOut])
+def list_curriculum_tasks():
+    """List all curriculum tasks ordered by creation time."""
+    session = store.get_session()
+    try:
+        return (
+            session.query(CurriculumTask)
+            .order_by(CurriculumTask.created_at.desc())
+            .all()
+        )
+    finally:
+        session.close()
+
+
+@router.get("/export")
+def export_curriculum(
+    format: str = Query("json", description="Export format: 'json' or 'jsonl'"),
+):
+    """Export pending curriculum tasks for training ingestion."""
+    format_value = format.lower().strip()
+    if format_value not in {"json", "jsonl"}:
+        raise HTTPException(status_code=400, detail="Invalid format. Use 'json' or 'jsonl'.")
+    try:
+        tasks = store.get_pending_curriculum(limit=100)
+
+        export_data = []
+        for task in tasks:
+            export_data.append(
+                {
+                    "id": task["id"],
+                    "role": "user",
+                    "content": task["instruction"],
+                    "metadata": {
+                        "difficulty": task["priority"],
+                        "focus": "auto_curriculum",
+                        "reasoning": task["context"],
+                    },
+                }
+            )
+
+        if format_value == "jsonl":
+            jsonl_content = "\n".join(json.dumps(item) for item in export_data)
+            return Response(content=jsonl_content, media_type="application/x-jsonlines")
+
+        return export_data
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=str(exc))
+
+@router.delete("/{task_id}", tags=["Curriculum"])
+def delete_curriculum_task(task_id: int):
+    """Delete a single curriculum task by ID."""
+    try:
+        deleted = store.delete_curriculum_task(task_id)
+        if not deleted:
+            raise HTTPException(status_code=404, detail="Task not found")
+        return Response(status_code=204)
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.delete("", tags=["Curriculum"])
+def delete_all_curriculum_tasks():
+    """Delete all curriculum tasks."""
+    try:
+        store.delete_all_curriculum_tasks()
+        return Response(status_code=204)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.patch("/{task_id}/complete", response_model=CurriculumTaskOut, tags=["Curriculum"])
+def mark_curriculum_task_complete(task_id: int):
+    """Mark a single curriculum task as complete."""
+    try:
+        task = store.mark_curriculum_task_complete(task_id)
+        if not task:
+            raise HTTPException(status_code=404, detail="Task not found")
+        return task
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.patch("/complete", tags=["Curriculum"])
+def mark_all_curriculum_tasks_complete():
+    """Mark all curriculum tasks as complete."""
+    try:
+        store.mark_all_curriculum_tasks_complete()
+        return Response(status_code=204)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

tracebrain/api/v1/episodes.py

@@ -0,0 +1,154 @@
+"""Episode endpoints for v1."""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from fastapi import APIRouter, HTTPException, Query
+
+from .common import store
+from .schemas.api_models import (
+    EpisodeAggregateOut,
+    EpisodeListOut,
+    EpisodeOut,
+    EpisodeSummaryListOut,
+    EpisodeTracesOut,
+    TraceSummaryOut,
+    trace_to_out,
+)
+
+router = APIRouter(prefix="/episodes", tags=["Episodes"])
+
+
+@router.get("", response_model=EpisodeListOut)
+def list_episodes(
+    skip: int = Query(0, ge=0, description="Number of episodes to skip"),
+    limit: int = Query(10, ge=1, le=100, description="Maximum number of episodes to return"),
+    query: Optional[str] = Query(None, description="Filter episodes by ID"),
+    min_confidence_lt: Optional[float] = Query(
+        None,
+        ge=0.0,
+        le=1.0,
+        description="Filter episodes where minimum confidence is below this value",
+    ),
+):
+    """List all episodes ordered by creation time, each with their traces."""
+    try:
+        episodes, total = store.list_episodes(
+            skip=skip,
+            limit=limit,
+            query=query,
+            include_spans=True,
+            min_confidence_lt=min_confidence_lt,
+        )
+
+        episode_outs = []
+        for episode_id, traces in episodes:
+            trace_outs = [trace_to_out(trace) for trace in traces]
+            episode_outs.append(EpisodeTracesOut(episode_id=episode_id, traces=trace_outs))
+
+        return EpisodeListOut(total=total, skip=skip, limit=limit, episodes=episode_outs)
+
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"Failed to list episodes: {str(exc)}")
+
+@router.delete("/{episode_id}")
+def delete_episode(episode_id: str):
+    """Delete an episode and all its traces."""
+    try:
+        store.delete_episode(episode_id)
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"Failed to delete episode: {str(exc)}")
+
+@router.get("/summary", response_model=EpisodeSummaryListOut)
+def list_episode_summaries(
+    skip: int = Query(0, ge=0, description="Number of episodes to skip"),
+    limit: int = Query(10, ge=1, le=100, description="Maximum number of episodes to return"),
+    query: Optional[str] = Query(None, description="Filter episodes by ID"),
+    min_confidence_lt: Optional[float] = Query(
+        None,
+        ge=0.0,
+        le=1.0,
+        description="Filter episodes where minimum confidence is below this value",
+    ),
+):
+    """List episodes with aggregated metrics."""
+    try:
+        episodes, total = store.list_episode_summaries(
+            skip=skip,
+            limit=limit,
+            query=query,
+            min_confidence_lt=min_confidence_lt,
+        )
+
+        episode_outs = [EpisodeAggregateOut(**episode) for episode in episodes]
+
+        return EpisodeSummaryListOut(total=total, skip=skip, limit=limit, episodes=episode_outs)
+
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"Failed to list episode summaries: {str(exc)}")
+
+
+@router.get("/{episode_id}", response_model=EpisodeOut)
+def get_episode_details(episode_id: str):
+    """Get episode details including the list of traces in that episode."""
+    try:
+        traces_in_episode = store.get_traces_by_episode_id(episode_id)
+
+        if not traces_in_episode:
+            raise HTTPException(status_code=404, detail="Episode not found")
+
+        trace_summaries: List[TraceSummaryOut] = []
+        for trace in traces_in_episode:
+            spans = trace.spans or []
+            span_count = len(spans)
+
+            start_times = [span.start_time for span in spans if span.start_time]
+            end_times = [span.end_time for span in spans if span.end_time]
+            duration_ms = 0.0
+            if start_times and end_times:
+                duration_ms = (max(end_times) - min(start_times)).total_seconds() * 1000
+
+            status = "OK"
+            for span in spans:
+                name = (span.name or "").lower()
+                span_type = (span.attributes or {}).get("tracebrain.span.type")
+                if "error" in name or span_type == "tool_error":
+                    status = "ERROR"
+                    break
+
+            trace_summaries.append(
+                TraceSummaryOut(
+                    trace_id=trace.id,
+                    status=status,
+                    duration_ms=round(duration_ms, 2),
+                    span_count=span_count,
+                    created_at=trace.created_at,
+                )
+            )
+
+        return EpisodeOut(episode_id=episode_id, traces=trace_summaries)
+
+    except HTTPException:
+        raise
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=str(exc))
+
+
+@router.get("/{episode_id}/traces", response_model=EpisodeTracesOut)
+def get_episode_traces(episode_id: str):
+    """Get all traces related to an episode."""
+    try:
+        traces_in_episode = store.get_traces_by_episode_id(episode_id)
+        if not traces_in_episode:
+            raise HTTPException(status_code=404, detail="Episode not found")
+
+        trace_ids = [trace.id for trace in traces_in_episode]
+        traces = store.get_traces_by_ids(trace_ids, include_spans=True)
+        trace_outs = [trace_to_out(trace) for trace in traces]
+        return EpisodeTracesOut(episode_id=episode_id, traces=trace_outs)
+
+    except HTTPException:
+        raise
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=str(exc))