[FEATURE] Unified LLM Observability & Interfaces — Cost Tracking, Streaming, Evaluation, and Natural Language Querying

api/app.py

@@ -60,6 +60,7 @@
     voice_router,
     ws_router,
     streaming_router,
+    cost_router,
     llm_usage_router,
     llm_cache_metrics_router,
     llm_router,
@@ -74,6 +75,7 @@
 
 
 from api.routers.ws import poll_and_broadcast_transactions
+from api.websocket.llm import router as ws_llm_router
 from astroml.llm import metrics as _llm_metrics
 from api.routers import health
 from api.routers import admin
@@ -198,6 +200,9 @@ async def _latency_middleware(request: Request, call_next):
 app.include_router(chat_router)
 app.include_router(ws_router)
 app.include_router(streaming_router)
+app.include_router(cost_router)
+app.include_router(ws_llm_router)
+app.include_router(query_router)
 app.include_router(llm_usage_router)
 app.include_router(llm_cache_metrics_router)
 app.include_router(voice_router)

api/routers/__init__.py

@@ -24,6 +24,7 @@
 from api.routers.voice import router as voice_router
 from api.routers.ws import router as ws_router
 from api.routers.streaming import router as streaming_router
+from api.routers.cost import router as cost_router
 from api.routers.llm_usage import router as llm_usage_router
 from api.routers.llm_cache_metrics import router as llm_cache_metrics_router
 from api.routers.llm import router as llm_router
@@ -59,6 +60,7 @@
     "voice_router",
     "ws_router",
     "streaming_router",
+    "cost_router",
     "llm_usage_router",
     "llm_router",
     "llm_health_router",

api/routers/cost.py

@@ -0,0 +1,105 @@
+"""Cost API Endpoints."""
+from __future__ import annotations
+
+from typing import Optional, Dict, Any, List
+from fastapi import APIRouter, Depends, HTTPException, status
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import select
+
+from api.database import get_db
+from api.auth.dependencies import AuthContext, get_current_auth, require_scopes
+from astroml.llm.cost import (
+    get_cost_summary,
+    forecast_cost,
+    check_budget,
+    set_emergency_override,
+)
+from astroml.db.models.cost import LLMBudget
+
+router = APIRouter(prefix="/api/v1/cost", tags=["cost"])
+
+
+@router.get("/summary")
+async def get_summary_endpoint(
+    days: int = 30,
+    auth: AuthContext = Depends(get_current_auth),
+    db: AsyncSession = Depends(get_db),
+):
+    """Retrieve cost summary for the authenticated user."""
+    user_id = str(auth.user_id or auth.subject)
+    summary = await get_cost_summary(db, user_id, days)
+    return summary
+
+
+@router.get("/forecast")
+async def get_forecast_endpoint(
+    days: int = 30,
+    auth: AuthContext = Depends(get_current_auth),
+    db: AsyncSession = Depends(get_db),
+):
+    """Get spend forecast for the next N days."""
+    user_id = str(auth.user_id or auth.subject)
+    forecast = await forecast_cost(db, user_id, days)
+    return forecast
+
+
+@router.post("/budget")
+async def configure_budget_endpoint(
+    limit_amount: float,
+    tier: str = "free",
+    period: str = "monthly",
+    auth: AuthContext = Depends(get_current_auth),
+    db: AsyncSession = Depends(get_db),
+):
+    """Configure or update budget limit and tier for the authenticated user."""
+    user_id = str(auth.user_id or auth.subject)
+
+    # Fetch or create budget
+    result = await db.execute(
+        select(LLMBudget).where(LLMBudget.entity_id == user_id)
+    )
+    budget = result.scalar_one_or_none()
+
+    if not budget:
+        budget = LLMBudget(
+            entity_id=user_id,
+            scope="user",
+            tier=tier,
+            limit_amount=limit_amount,
+            current_spend=0.0,
+            period=period,
+            is_blocked=False
+        )
+        db.add(budget)
+    else:
+        budget.limit_amount = limit_amount
+        budget.tier = tier
+        budget.period = period
+        if budget.current_spend < limit_amount:
+            budget.is_blocked = False
+
+    await db.commit()
+    return {
+        "status": "success",
+        "entity_id": user_id,
+        "limit_amount": budget.limit_amount,
+        "tier": budget.tier,
+        "is_blocked": budget.is_blocked,
+    }
+
+
+@router.post("/override")
+async def admin_override_endpoint(
+    entity_id: str,
+    override: bool = True,
+    auth: AuthContext = Depends(require_scopes("admin")),
+    db: AsyncSession = Depends(get_db),
+):
+    """Enable/disable emergency override for a budget (Admin only)."""
+    success = await set_emergency_override(db, entity_id, override)
+    if not success:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"Budget for entity '{entity_id}' not found."
+        )
+    return {"status": "success", "entity_id": entity_id, "emergency_override": override}

api/routers/query.py

@@ -1,3 +1,112 @@
+"""Natural Language Query API Router."""
+from __future__ import annotations
+
+from typing import Dict, Any, List, Optional
+from fastapi import APIRouter, Depends, HTTPException, Query, status
+from pydantic import BaseModel
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from api.database import get_db
+from api.auth.dependencies import AuthContext, get_current_auth
+from astroml.llm.query import (
+    generate_sql,
+    execute_safe_query,
+    generate_pipeline_config,
+    format_query_results,
+    get_query_suggestions,
+)
+from astroml.llm.cost import check_budget, track_request
+
+router = APIRouter(prefix="/api/v1/query", tags=["query"])
+
+
+class NLQueryIn(BaseModel):
+    query: str
+    model: str = "gpt-3.5-turbo"
+    mode: str = "sql"  # 'sql' or 'pipeline'
+    feature: str = "nlp_query"
+
+
+class NLQueryOut(BaseModel):
+    query: str
+    mode: str
+    sql: Optional[str] = None
+    pipeline_yaml: Optional[str] = None
+    results: Optional[Dict[str, Any]] = None
+    suggestions: List[str]
+
+
+@router.post("", response_model=NLQueryOut)
+async def post_natural_query(
+    body: NLQueryIn,
+    db: AsyncSession = Depends(get_db),
+    auth: AuthContext = Depends(get_current_auth),
+):
+    """
+    Query database or generate pipeline using natural language.
+    Includes validation, safety checks, audit logs and budgeting.
+    """
+    user_id = str(auth.user_id or auth.subject)
+
+    # 1. Check budget first
+    try:
+        await check_budget(db, user_id, body.model)
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=f"Budget exceeded or model access denied: {str(e)}"
+        )
+
+    start_time = 0.0
+    sql = None
+    pipeline_yaml = None
+    formatted_results = None
+
+    # 2. Process query
+    if body.mode == "sql":
+        # Translate to SQL
+        sql = generate_sql(body.query)
+        try:
+            # Execute safely
+            raw_rows = await execute_safe_query(db, sql)
+            formatted_results = format_query_results(raw_rows)
+        except Exception as e:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail=f"Database query execution failed: {str(e)}"
+            )
+    elif body.mode == "pipeline":
+        # Translate to ML Pipeline YAML configuration
+        pipeline_yaml = generate_pipeline_config(body.query)
+    else:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Invalid query mode '{body.mode}'. Supported: 'sql', 'pipeline'"
+        )
+
+    # 3. Track spending (mock usage metrics)
+    input_tokens = len(body.query) // 4 + 1
+    output_tokens = (len(sql or "") + len(pipeline_yaml or "")) // 4 + 1
+    await track_request(
+        db=db,
+        user_id=user_id,
+        feature=body.feature,
+        model_name=body.model,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        latency_ms=150.0,  # mock latency
+    )
+
+    suggestions = get_query_suggestions()
+
+    return NLQueryOut(
+        query=body.query,
+        mode=body.mode,
+        sql=sql,
+        pipeline_yaml=pipeline_yaml,
+        results=formatted_results,
+        suggestions=suggestions
+    )
 """LLM-powered SQL query optimization engine."""
 import re
 import json

api/routers/streaming.py

@@ -105,3 +105,78 @@ def get_streaming_health():
         last_updated=_stream_state.last_updated,
         streams=streams
     )
+
+
+# ---------------------------------------------------------------------------
+# LLM SSE Streaming Endpoint
+# ---------------------------------------------------------------------------
+import asyncio
+from fastapi.responses import StreamingResponse
+from fastapi import Depends, Query
+from api.auth.dependencies import AuthContext, get_current_auth
+from api.database import get_db
+from sqlalchemy.ext.asyncio import AsyncSession
+from astroml.llm.streaming import StreamHandler, format_sse
+from astroml.llm.cost import check_budget, track_request, route_request
+
+@router.get("/llm")
+async def stream_llm_response(
+    prompt: str = Query(..., description="Prompt for the LLM"),
+    model: str = Query("gpt-3.5-turbo", description="Model name"),
+    feature: str = Query("chatbot", description="Feature category"),
+    db: AsyncSession = Depends(get_db),
+    auth: AuthContext = Depends(get_current_auth),
+):
+    """
+    Stream LLM response token-by-token using Server-Sent Events (SSE).
+    Enforces cost tracking, dynamic model routing, and budgets.
+    """
+    user_id = str(auth.user_id or auth.subject)
+
+    # 1. Dynamic routing based on budget and complexity
+    routed_model = await route_request(db, user_id, model, prompt)
+
+    # 2. Check budget limits and model permissions
+    try:
+        await check_budget(db, user_id, routed_model)
+    except Exception as e:
+        raise HTTPException(status_code=403, detail=str(e))
+
+    async def sse_generator():
+        handler = StreamHandler(session_id=f"sse_{user_id}_{int(time.time())}")
+
+        # Simulate some generation tokens
+        mock_response = f"This is a progressive response for your query '{prompt}' using {routed_model}."
+        words = mock_response.split(" ")
+
+        total_tokens = 0
+        latency_ms = 100.0  # mock first token latency
+        start_time = time.perf_counter()
+
+        async def mock_word_gen():
+            for i, word in enumerate(words):
+                await asyncio.sleep(0.05)  # 50ms streaming latency
+                yield word + " " if i < len(words) - 1 else word
+
+        async for token in handler.process_stream(mock_word_gen()):
+            total_tokens += 1
+            yield format_sse(token=token, finished=False)
+
+        duration = (time.perf_counter() - start_time) * 1000
+
+        # Track cost and update budget
+        usage = {"prompt_tokens": len(prompt) // 4 + 1, "completion_tokens": total_tokens}
+        await track_request(
+            db=db,
+            user_id=user_id,
+            feature=feature,
+            model_name=routed_model,
+            input_tokens=usage["prompt_tokens"],
+            output_tokens=usage["completion_tokens"],
+            latency_ms=duration,
+        )
+
+        yield format_sse(token=None, finished=True, usage={"total_tokens": usage["prompt_tokens"] + usage["completion_tokens"]})
+
+    return StreamingResponse(sse_generator(), media_type="text/event-stream")
+