trace-agent-server 0.0.1__py3-none-any.whl

trace_agent_server/cli.py

@@ -0,0 +1,29 @@
+"""CLI entry point for trace-agent-server."""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="TraceAgent Server")
+    parser.add_argument("--host", default="0.0.0.0", help="Bind host (default: 0.0.0.0)")
+    parser.add_argument("--port", type=int, default=8000, help="Bind port (default: 8000)")
+    parser.add_argument("--reload", action="store_true", help="Enable auto-reload")
+    parser.add_argument("--workers", type=int, default=1, help="Number of worker processes")
+    args = parser.parse_args()
+
+    import uvicorn
+
+    uvicorn.run(
+        "trace_agent_server.main:app",
+        host=args.host,
+        port=args.port,
+        reload=args.reload,
+        workers=args.workers if not args.reload else 1,
+    )
+
+
+if __name__ == "__main__":
+    main()

trace_agent_server/config.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from functools import lru_cache
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class Settings(BaseSettings):
+    app_name: str = "TraceAgent"
+    database_url: str = "sqlite:///./trace_agent.db"
+    openai_api_key: str | None = None
+    openai_base_url: str | None = None
+    anthropic_api_key: str | None = None
+    anthropic_base_url: str | None = None
+    google_api_key: str | None = None
+    kimi_api_key: str | None = None
+    kimi_base_url: str | None = None
+    lmstudio_base_url: str | None = None
+    lmstudio_mode: str = "openai"
+    lmstudio_api_key: str | None = None
+    cors_origins: list[str] = Field(default_factory=lambda: ["http://localhost:5173"])
+    model_pricing: dict[str, dict[str, float]] = Field(default_factory=dict)
+    audit_enable_guardrails: bool = False
+    audit_fail_on_prompt_injection: bool = False
+    audit_enable_pii_anonymization: bool = False
+    audit_prompt_injection_threshold: float = 0.85
+    audit_pii_entities: list[str] = Field(
+        default_factory=lambda: [
+            "CREDIT_CARD",
+            "CRYPTO",
+            "EMAIL_ADDRESS",
+            "IBAN_CODE",
+            "IP_ADDRESS",
+            "PERSON",
+            "PHONE_NUMBER",
+            "US_SSN",
+        ]
+    )
+    audit_metrics_enabled: bool = True
+    audit_enable_otel: bool = False
+    audit_otel_service_name: str = "trace-agent-backend"
+    audit_otel_exporter_endpoint: str | None = None
+    audit_otel_exporter_protocol: str = "grpc"
+    capture_full_payloads: bool = True
+
+    model_config = SettingsConfigDict(
+        env_prefix="TRACE_AGENT_",
+        env_file=".env",
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+
+@lru_cache(maxsize=1)
+def get_settings() -> Settings:
+    return Settings()

trace_agent_server/db.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from collections.abc import Generator
+
+from fastapi import Request
+from sqlalchemy import create_engine
+from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+def create_session_factory(database_url: str) -> tuple[object, sessionmaker[Session]]:
+    connect_args = {"check_same_thread": False} if database_url.startswith("sqlite") else {}
+    engine = create_engine(database_url, connect_args=connect_args, future=True)
+    session_factory = sessionmaker(bind=engine, autoflush=False, autocommit=False, expire_on_commit=False)
+    return engine, session_factory
+
+
+def get_session(request: Request) -> Generator[Session, None, None]:
+    session_factory: sessionmaker[Session] = request.app.state.session_factory
+    with session_factory() as session:
+        yield session

trace_agent_server/main.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+from trace_agent_server.api.runs import router as runs_router
+from trace_agent_server.config import Settings, get_settings
+from trace_agent_server.db import Base, create_session_factory
+from trace_agent_server.services.audit_guardrails import AuditGuardrails
+from trace_agent_server.services.observability import AppObservability, configure_opentelemetry
+from trace_agent_server.services.provider import (
+    AnthropicProviderAdapter,
+    BaseProviderAdapter,
+    GoogleProviderAdapter,
+    KimiProviderAdapter,
+    LMStudioProviderAdapter,
+    OpenAIProviderAdapter,
+)
+
+
+def build_provider_registry(settings: Settings) -> dict[str, BaseProviderAdapter]:
+    """Dynamically instantiate provider adapters based on available configuration."""
+    registry: dict[str, BaseProviderAdapter] = {}
+
+    # OpenAI is always registered (it can work with just a base_url for local proxies)
+    registry["openai"] = OpenAIProviderAdapter(settings)
+
+    if settings.anthropic_api_key or settings.anthropic_base_url:
+        registry["anthropic"] = AnthropicProviderAdapter(settings)
+
+    if settings.google_api_key:
+        registry["google"] = GoogleProviderAdapter(settings)
+
+    if settings.kimi_api_key or settings.kimi_base_url:
+        registry["kimi"] = KimiProviderAdapter(settings)
+
+    if settings.lmstudio_base_url:
+        registry["lmstudio"] = LMStudioProviderAdapter(settings)
+
+    return registry
+
+
+def create_app(settings: Settings | None = None) -> FastAPI:
+    current_settings = settings or get_settings()
+    app = FastAPI(title=current_settings.app_name)
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=current_settings.cors_origins,
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    engine, session_factory = create_session_factory(current_settings.database_url)
+    Base.metadata.create_all(bind=engine)
+    app.state.engine = engine
+    app.state.session_factory = session_factory
+    app.state.settings = current_settings
+    app.state.provider_registry = build_provider_registry(current_settings)
+    app.state.observability = AppObservability(current_settings.audit_metrics_enabled)
+    app.state.observability.instrument_app(app)
+    app.state.guardrails = AuditGuardrails(current_settings, app.state.observability)
+    app.state.otel = configure_opentelemetry(app, current_settings)
+
+    @app.get("/health")
+    def healthcheck() -> dict[str, str]:
+        return {"status": "ok"}
+
+    app.include_router(runs_router)
+    return app
+
+
+app = create_app()

trace_agent_server/models.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from uuid import uuid4
+
+from sqlalchemy import DateTime, ForeignKey, Integer, JSON, String, Text, UniqueConstraint
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from trace_agent_server.db import Base
+
+
+def utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class Run(Base):
+    __tablename__ = "runs"
+
+    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid4()))
+    agent_name: Mapped[str] = mapped_column(String(255))
+    goal: Mapped[str] = mapped_column(Text)
+    status: Mapped[str] = mapped_column(String(32), default="running")
+    provider: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    model: Mapped[str | None] = mapped_column(String(128), nullable=True)
+    redaction_level: Mapped[str] = mapped_column(String(32), default="summary")
+    final_output_summary: Mapped[str | None] = mapped_column(Text, nullable=True)
+    run_metadata: Mapped[dict] = mapped_column("metadata", JSON, default=dict)
+    explanation_cache: Mapped[dict] = mapped_column(JSON, default=dict)
+    flag_cache: Mapped[list] = mapped_column(JSON, default=list)
+    tool_count: Mapped[int] = mapped_column(Integer, default=0)
+    error_count: Mapped[int] = mapped_column(Integer, default=0)
+    retry_count: Mapped[int] = mapped_column(Integer, default=0)
+    artifact_count: Mapped[int] = mapped_column(Integer, default=0)
+    total_prompt_tokens: Mapped[int] = mapped_column(Integer, default=0)
+    total_completion_tokens: Mapped[int] = mapped_column(Integer, default=0)
+    last_event_seq: Mapped[int] = mapped_column(Integer, default=0)
+    started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+    ended_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow, onupdate=utcnow)
+
+    events: Mapped[list["Event"]] = relationship(back_populates="run", cascade="all, delete-orphan")
+    decisions: Mapped[list["Decision"]] = relationship(back_populates="run", cascade="all, delete-orphan")
+    artifacts: Mapped[list["Artifact"]] = relationship(back_populates="run", cascade="all, delete-orphan")
+    turns: Mapped[list["Turn"]] = relationship(back_populates="run", cascade="all, delete-orphan")
+
+
+class Turn(Base):
+    __tablename__ = "turns"
+
+    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid4()))
+    run_id: Mapped[str] = mapped_column(ForeignKey("runs.id", ondelete="CASCADE"), index=True)
+    seq: Mapped[int] = mapped_column(Integer, index=True)
+    provider: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    model: Mapped[str | None] = mapped_column(String(128), nullable=True)
+    messages: Mapped[list] = mapped_column(JSON, default=list)
+    assistant_message: Mapped[dict | None] = mapped_column(JSON, nullable=True)
+    tool_calls: Mapped[list] = mapped_column(JSON, default=list)
+    prompt_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    completion_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    total_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    reasoning_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    finish_reason: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    response_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+    run: Mapped["Run"] = relationship(back_populates="turns")
+
+
+class Event(Base):
+    __tablename__ = "events"
+    __table_args__ = (UniqueConstraint("run_id", "seq", name="uq_event_run_seq"),)
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
+    run_id: Mapped[str] = mapped_column(ForeignKey("runs.id", ondelete="CASCADE"), index=True)
+    turn_id: Mapped[str | None] = mapped_column(ForeignKey("turns.id", ondelete="SET NULL"), nullable=True, index=True)
+    seq: Mapped[int] = mapped_column(Integer, index=True)
+    timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+    step_id: Mapped[str | None] = mapped_column(String(128), nullable=True, index=True)
+    parent_step_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
+    actor: Mapped[str] = mapped_column(String(64))
+    type: Mapped[str] = mapped_column(String(64), index=True)
+    summary: Mapped[str] = mapped_column(Text)
+    status: Mapped[str | None] = mapped_column(String(32), nullable=True)
+    duration_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    provider: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    model: Mapped[str | None] = mapped_column(String(128), nullable=True)
+    tool_name: Mapped[str | None] = mapped_column(String(128), nullable=True)
+    artifact_refs: Mapped[list] = mapped_column(JSON, default=list)
+    error_code: Mapped[str | None] = mapped_column(String(128), nullable=True)
+    payload_hash: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    payload_full: Mapped[dict | None] = mapped_column(JSON, nullable=True)
+    event_metadata: Mapped[dict] = mapped_column("metadata", JSON, default=dict)
+
+    run: Mapped["Run"] = relationship(back_populates="events")
+
+
+class Decision(Base):
+    __tablename__ = "decisions"
+
+    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid4()))
+    run_id: Mapped[str] = mapped_column(ForeignKey("runs.id", ondelete="CASCADE"), index=True)
+    seq: Mapped[int] = mapped_column(Integer, index=True)
+    kind: Mapped[str] = mapped_column(String(64))
+    label: Mapped[str] = mapped_column(String(255))
+    chosen_action: Mapped[str] = mapped_column(String(255))
+    rationale_summary: Mapped[str] = mapped_column(Text)
+    trigger_event_ids: Mapped[list] = mapped_column(JSON, default=list)
+    evidence_event_ids: Mapped[list] = mapped_column(JSON, default=list)
+    alternatives: Mapped[list] = mapped_column(JSON, default=list)
+    outcome: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    decision_metadata: Mapped[dict] = mapped_column("metadata", JSON, default=dict)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+    run: Mapped["Run"] = relationship(back_populates="decisions")
+
+
+class Artifact(Base):
+    __tablename__ = "artifacts"
+
+    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid4()))
+    run_id: Mapped[str] = mapped_column(ForeignKey("runs.id", ondelete="CASCADE"), index=True)
+    source_event_id: Mapped[int | None] = mapped_column(ForeignKey("events.id", ondelete="SET NULL"), nullable=True)
+    kind: Mapped[str] = mapped_column(String(64))
+    label: Mapped[str] = mapped_column(String(255))
+    summary: Mapped[str] = mapped_column(Text)
+    artifact_metadata: Mapped[dict] = mapped_column("metadata", JSON, default=dict)
+    content_hash: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+    run: Mapped["Run"] = relationship(back_populates="artifacts")