remdb 0.3.103__py3-none-any.whl → 0.3.141__py3-none-any.whl

rem/api/routers/admin.py

@@ -9,6 +9,9 @@ Endpoints:
     GET  /api/admin/messages       - List all messages across users (admin only)
     GET  /api/admin/stats          - System statistics (admin only)
 
+Internal Endpoints (hidden from Swagger, secret-protected):
+    POST /api/admin/internal/rebuild-kv  - Trigger kv_store rebuild (called by pg_net)
+
 All endpoints require:
 1. Authentication (valid session)
 2. Admin role in user's roles list
@@ -17,11 +20,14 @@ Design Pattern:
 - Uses require_admin dependency for role enforcement
 - Cross-tenant queries (no user_id filtering)
 - Audit logging for admin actions
+- Internal endpoints use X-Internal-Secret header for authentication
 """
 
+import asyncio
+import threading
 from typing import Literal
 
-from fastapi import APIRouter, Depends, HTTPException, Query
+from fastapi import APIRouter, Depends, Header, HTTPException, Query, BackgroundTasks
 from loguru import logger
 from pydantic import BaseModel
 
@@ -32,6 +38,12 @@ from ...settings import settings
 
 router = APIRouter(prefix="/api/admin", tags=["admin"])
 
+# =============================================================================
+# Internal Router (hidden from Swagger)
+# =============================================================================
+
+internal_router = APIRouter(prefix="/internal", include_in_schema=False)
+
 
 # =============================================================================
 # Response Models
@@ -275,3 +287,208 @@ async def get_system_stats(
         active_sessions_24h=0,  # TODO: implement
         messages_24h=0,  # TODO: implement
     )
+
+
+# =============================================================================
+# Internal Endpoints (hidden from Swagger, secret-protected)
+# =============================================================================
+
+
+class RebuildKVRequest(BaseModel):
+    """Request body for kv_store rebuild trigger."""
+
+    user_id: str | None = None
+    triggered_by: str = "api"
+    timestamp: str | None = None
+
+
+class RebuildKVResponse(BaseModel):
+    """Response from kv_store rebuild trigger."""
+
+    status: Literal["submitted", "started", "skipped"]
+    message: str
+    job_method: str | None = None  # "sqs" or "thread"
+
+
+async def _get_internal_secret() -> str | None:
+    """
+    Get the internal API secret from cache_system_state table.
+
+    Returns None if the table doesn't exist or secret not found.
+    """
+    from ...services.postgres import get_postgres_service
+
+    db = get_postgres_service()
+    if not db:
+        return None
+
+    try:
+        await db.connect()
+        secret = await db.fetchval("SELECT rem_get_cache_api_secret()")
+        return secret
+    except Exception as e:
+        logger.warning(f"Could not get internal API secret: {e}")
+        return None
+    finally:
+        await db.disconnect()
+
+
+async def _validate_internal_secret(x_internal_secret: str | None = Header(None)):
+    """
+    Dependency to validate the X-Internal-Secret header.
+
+    Raises 401 if secret is missing or invalid.
+    """
+    if not x_internal_secret:
+        logger.warning("Internal endpoint called without X-Internal-Secret header")
+        raise HTTPException(status_code=401, detail="Missing X-Internal-Secret header")
+
+    expected_secret = await _get_internal_secret()
+    if not expected_secret:
+        logger.error("Could not retrieve internal secret from database")
+        raise HTTPException(status_code=503, detail="Internal secret not configured")
+
+    if x_internal_secret != expected_secret:
+        logger.warning("Internal endpoint called with invalid secret")
+        raise HTTPException(status_code=401, detail="Invalid X-Internal-Secret")
+
+    return True
+
+
+def _run_rebuild_in_thread():
+    """
+    Run the kv_store rebuild in a background thread.
+
+    This is the fallback when SQS is not available.
+    """
+
+    def rebuild_task():
+        """Thread target function."""
+        import asyncio
+        from ...workers.unlogged_maintainer import UnloggedMaintainer
+
+        async def _run():
+            maintainer = UnloggedMaintainer()
+            if not maintainer.db:
+                logger.error("Database not configured, cannot rebuild")
+                return
+            try:
+                await maintainer.db.connect()
+                await maintainer.rebuild_with_lock()
+            except Exception as e:
+                logger.error(f"Background rebuild failed: {e}")
+            finally:
+                await maintainer.db.disconnect()
+
+        # Create new event loop for this thread
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            loop.run_until_complete(_run())
+        finally:
+            loop.close()
+
+    thread = threading.Thread(target=rebuild_task, name="kv-rebuild-worker")
+    thread.daemon = True
+    thread.start()
+    logger.info(f"Started background rebuild thread: {thread.name}")
+
+
+def _submit_sqs_rebuild_job_sync(request: RebuildKVRequest) -> bool:
+    """
+    Submit rebuild job to SQS queue (synchronous).
+
+    Returns True if job was submitted, False if SQS unavailable.
+    """
+    import json
+
+    import boto3
+    from botocore.exceptions import ClientError
+
+    if not settings.sqs.queue_url:
+        logger.debug("SQS queue URL not configured, cannot submit SQS job")
+        return False
+
+    try:
+        sqs = boto3.client("sqs", region_name=settings.sqs.region)
+
+        message_body = {
+            "action": "rebuild_kv_store",
+            "user_id": request.user_id,
+            "triggered_by": request.triggered_by,
+            "timestamp": request.timestamp,
+        }
+
+        response = sqs.send_message(
+            QueueUrl=settings.sqs.queue_url,
+            MessageBody=json.dumps(message_body),
+            MessageAttributes={
+                "action": {"DataType": "String", "StringValue": "rebuild_kv_store"},
+            },
+        )
+
+        message_id = response.get("MessageId")
+        logger.info(f"Submitted rebuild job to SQS: {message_id}")
+        return True
+
+    except ClientError as e:
+        logger.warning(f"Failed to submit SQS job: {e}")
+        return False
+    except Exception as e:
+        logger.warning(f"SQS submission error: {e}")
+        return False
+
+
+async def _submit_sqs_rebuild_job(request: RebuildKVRequest) -> bool:
+    """
+    Submit rebuild job to SQS queue (async wrapper).
+
+    Runs boto3 call in thread pool to avoid blocking event loop.
+    """
+    import asyncio
+
+    return await asyncio.to_thread(_submit_sqs_rebuild_job_sync, request)
+
+
+@internal_router.post("/rebuild-kv", response_model=RebuildKVResponse)
+async def trigger_kv_rebuild(
+    request: RebuildKVRequest,
+    _: bool = Depends(_validate_internal_secret),
+) -> RebuildKVResponse:
+    """
+    Trigger kv_store rebuild (internal endpoint, not shown in Swagger).
+
+    Called by pg_net from PostgreSQL when self-healing detects empty cache.
+    Authentication: X-Internal-Secret header must match secret in cache_system_state.
+
+    Priority:
+    1. Submit job to SQS (if configured) - scales with KEDA
+    2. Fallback to background thread - runs in same process
+
+    Note: This endpoint returns immediately. Rebuild happens asynchronously.
+    """
+    logger.info(
+        f"Rebuild kv_store requested by {request.triggered_by} "
+        f"(user_id={request.user_id})"
+    )
+
+    # Try SQS first
+    if await _submit_sqs_rebuild_job(request):
+        return RebuildKVResponse(
+            status="submitted",
+            message="Rebuild job submitted to SQS queue",
+            job_method="sqs",
+        )
+
+    # Fallback to background thread
+    _run_rebuild_in_thread()
+
+    return RebuildKVResponse(
+        status="started",
+        message="Rebuild started in background thread (SQS unavailable)",
+        job_method="thread",
+    )
+
+
+# Include internal router in main router
+router.include_router(internal_router)

rem/api/routers/chat/completions.py

@@ -1,13 +1,94 @@
 """
 OpenAI-compatible chat completions router for REM.
 
-Design Pattern:
-- Headers map to AgentContext (X-User-Id, X-Tenant-Id, X-Session-Id, X-Agent-Schema)
+Quick Start (Local Development)
+===============================
+
+NOTE: Local dev uses LOCAL databases (Postgres via Docker Compose on port 5050).
+      Do NOT port-forward databases. Only port-forward observability services.
+
+IMPORTANT: Session IDs MUST be UUIDs. Non-UUID session IDs will cause message
+           storage issues and feedback will not work correctly.
+
+1. Port Forwarding (REQUIRED for trace capture and Phoenix sync):
+
+    # Terminal 1: OTEL Collector (HTTP) - sends traces to Phoenix
+    kubectl port-forward -n observability svc/otel-collector-collector 4318:4318
+
+    # Terminal 2: Phoenix UI - view traces at http://localhost:6006
+    kubectl port-forward -n siggy svc/phoenix 6006:6006
+
+2. Get Phoenix API Key (REQUIRED for feedback->Phoenix sync):
+
+    export PHOENIX_API_KEY=$(kubectl get secret -n siggy rem-phoenix-api-key \\
+      -o jsonpath='{.data.PHOENIX_API_KEY}' | base64 -d)
+
+3. Start API with OTEL and Phoenix enabled:
+
+    cd /path/to/remstack/rem
+    source .venv/bin/activate
+    OTEL__ENABLED=true \\
+    PHOENIX__ENABLED=true \\
+    PHOENIX_API_KEY="$PHOENIX_API_KEY" \\
+    uvicorn rem.api.main:app --host 0.0.0.0 --port 8000 --app-dir src
+
+4. Test Chat Request (session_id MUST be a UUID):
+
+    SESSION_ID=$(python3 -c "import uuid; print(uuid.uuid4())")
+    curl -s -N -X POST http://localhost:8000/api/v1/chat/completions \\
+      -H 'Content-Type: application/json' \\
+      -H "X-Session-Id: $SESSION_ID" \\
+      -H 'X-Agent-Schema: rem' \\
+      -d '{"messages": [{"role": "user", "content": "Hello"}], "stream": true}'
+
+    # Note: Use 'rem' agent schema (default) for real LLM responses.
+    # The 'simulator' agent is for testing SSE events without LLM calls.
+
+5. Submit Feedback on Response:
+
+    The metadata SSE event contains message_id and trace_id for feedback:
+        event: metadata
+        data: {"message_id": "728882f8-...", "trace_id": "e53c701c...", ...}
+
+    Use session_id (UUID you generated) and message_id to submit feedback:
+
+    curl -X POST http://localhost:8000/api/v1/messages/feedback \\
+      -H 'Content-Type: application/json' \\
+      -H 'X-Tenant-Id: default' \\
+      -d '{
+        "session_id": "<your-uuid-session-id>",
+        "message_id": "<message-id-from-metadata>",
+        "rating": 1,
+        "categories": ["helpful"],
+        "comment": "Good response"
+      }'
+
+    Expected response (201 = synced to Phoenix):
+        {"phoenix_synced": true, "trace_id": "e53c701c...", "span_id": "6432d497..."}
+
+OTEL Architecture
+=================
+
+    REM API --[OTLP/HTTP]--> OTEL Collector --[relay]--> Phoenix
+             (port 4318)    (k8s: observability)         (k8s: siggy)
+
+Environment Variables:
+    OTEL__ENABLED=true              Enable OTEL tracing (required for trace capture)
+    PHOENIX__ENABLED=true           Enable Phoenix integration (required for feedback sync)
+    PHOENIX_API_KEY=<jwt>           Phoenix API key (required for feedback->Phoenix sync)
+    OTEL__COLLECTOR_ENDPOINT        Default: http://localhost:4318
+    OTEL__PROTOCOL                  Default: http (use port 4318, not gRPC 4317)
+
+Design Pattern
+==============
+
+- Headers map to AgentContext (X-User-Id, X-Tenant-Id, X-Session-Id, X-Agent-Schema, X-Is-Eval)
 - ContextBuilder centralizes message construction with user profile + session history
 - Body.model is the LLM model for Pydantic AI
 - X-Agent-Schema header specifies which agent schema to use (defaults to 'rem')
 - Support for streaming (SSE) and non-streaming modes
 - Response format control (text vs json_object)
+- OpenAI-compatible body fields: metadata, store, reasoning_effort, etc.
 
 Context Building Flow:
 1. ContextBuilder.build_from_headers() extracts user_id, session_id from headers
@@ -25,9 +106,10 @@ Context Building Flow:
 Headers Mapping
     X-User-Id        → AgentContext.user_id
     X-Tenant-Id      → AgentContext.tenant_id
-    X-Session-Id     → AgentContext.session_id
+    X-Session-Id     → AgentContext.session_id (use UUID for new sessions)
     X-Model-Name     → AgentContext.default_model (overrides body.model)
     X-Agent-Schema   → AgentContext.agent_schema_uri (defaults to 'rem')
+    X-Is-Eval        → AgentContext.is_eval (sets session mode to EVALUATION)
 
 Default Agent:
     If X-Agent-Schema header is not provided, the system loads 'rem' schema,
@@ -42,6 +124,7 @@ Example Request:
     POST /api/v1/chat/completions
     X-Tenant-Id: acme-corp
     X-User-Id: user123
+    X-Session-Id: a1b2c3d4-e5f6-7890-abcd-ef1234567890  # UUID
     X-Agent-Schema: rem  # Optional, this is the default
 
     {
@@ -67,7 +150,9 @@ from loguru import logger
 from ....agentic.context import AgentContext
 from ....agentic.context_builder import ContextBuilder
 from ....agentic.providers.pydantic_ai import create_agent
+from ....models.entities.session import Session, SessionMode
 from ....services.audio.transcriber import AudioTranscriber
+from ....services.postgres.repository import Repository
 from ....services.session import SessionMessageStore, reload_session
 from ....settings import settings
 from ....utils.schema_loader import load_agent_schema, load_agent_schema_async
@@ -79,7 +164,7 @@ from .models import (
     ChatCompletionUsage,
     ChatMessage,
 )
-from .streaming import stream_openai_response, stream_simulator_response
+from .streaming import stream_openai_response, stream_openai_response_with_save, stream_simulator_response
 
 router = APIRouter(prefix="/api/v1", tags=["chat"])
 
@@ -87,6 +172,105 @@ router = APIRouter(prefix="/api/v1", tags=["chat"])
 DEFAULT_AGENT_SCHEMA = "rem"
 
 
+def get_current_trace_context() -> tuple[str | None, str | None]:
+    """Get trace_id and span_id from current OTEL context.
+
+    Returns:
+        Tuple of (trace_id, span_id) as hex strings, or (None, None) if not available.
+    """
+    try:
+        from opentelemetry import trace
+        span = trace.get_current_span()
+        if span and span.get_span_context().is_valid:
+            ctx = span.get_span_context()
+            trace_id = format(ctx.trace_id, '032x')
+            span_id = format(ctx.span_id, '016x')
+            return trace_id, span_id
+    except Exception:
+        pass
+    return None, None
+
+
+def get_tracer():
+    """Get the OpenTelemetry tracer for chat completions."""
+    try:
+        from opentelemetry import trace
+        return trace.get_tracer("rem.chat.completions")
+    except Exception:
+        return None
+
+
+async def ensure_session_with_metadata(
+    session_id: str,
+    user_id: str | None,
+    tenant_id: str,
+    is_eval: bool,
+    request_metadata: dict[str, str] | None,
+    agent_schema: str | None = None,
+) -> None:
+    """
+    Ensure session exists and update with metadata/mode.
+
+    If X-Is-Eval header is true, sets session mode to EVALUATION.
+    Merges request metadata with existing session metadata.
+
+    Args:
+        session_id: Session identifier (maps to Session.name)
+        user_id: User identifier
+        tenant_id: Tenant identifier
+        is_eval: Whether this is an evaluation session
+        request_metadata: Metadata from request body to merge
+        agent_schema: Optional agent schema being used
+    """
+    if not settings.postgres.enabled:
+        return
+
+    try:
+        repo = Repository(Session, table_name="sessions")
+
+        # Try to load existing session by name (session_id is the name field)
+        existing_list = await repo.find(
+            filters={"name": session_id, "tenant_id": tenant_id},
+            limit=1,
+        )
+        existing = existing_list[0] if existing_list else None
+
+        if existing:
+            # Merge metadata if provided
+            merged_metadata = existing.metadata or {}
+            if request_metadata:
+                merged_metadata.update(request_metadata)
+
+            # Update session if eval flag or new metadata
+            needs_update = False
+            if is_eval and existing.mode != SessionMode.EVALUATION:
+                existing.mode = SessionMode.EVALUATION
+                needs_update = True
+            if request_metadata:
+                existing.metadata = merged_metadata
+                needs_update = True
+
+            if needs_update:
+                await repo.upsert(existing)
+                logger.debug(f"Updated session {session_id} (eval={is_eval}, metadata keys={list(merged_metadata.keys())})")
+        else:
+            # Create new session
+            session = Session(
+                name=session_id,
+                mode=SessionMode.EVALUATION if is_eval else SessionMode.NORMAL,
+                user_id=user_id,
+                tenant_id=tenant_id,
+                agent_schema_uri=agent_schema,
+                metadata=request_metadata or {},
+            )
+            await repo.upsert(session)
+            logger.info(f"Created session {session_id} (eval={is_eval})")
+
+    except Exception as e:
+        # Non-critical - log but don't fail the request
+        logger.error(f"Failed to ensure session metadata: {e}", exc_info=True)
+
+
 @router.post("/chat/completions", response_model=None)
 async def chat_completions(body: ChatCompletionRequest, request: Request):
     """
@@ -102,6 +286,17 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
     | X-Tenant-Id         | Tenant identifier (multi-tenancy)    | AgentContext.tenant_id         | "default"     |
     | X-Session-Id        | Session/conversation identifier      | AgentContext.session_id        | None          |
     | X-Agent-Schema      | Agent schema name                    | AgentContext.agent_schema_uri  | "rem"         |
+    | X-Is-Eval           | Mark as evaluation session           | AgentContext.is_eval           | false         |
+
+    Additional OpenAI-compatible Body Fields:
+    - metadata: Key-value pairs merged with session metadata (max 16 keys)
+    - store: Whether to store for distillation/evaluation
+    - max_completion_tokens: Max tokens to generate (replaces max_tokens)
+    - seed: Seed for deterministic sampling
+    - top_p: Nucleus sampling probability
+    - logprobs: Return log probabilities
+    - reasoning_effort: low/medium/high for o-series models
+    - service_tier: auto/flex/priority/default
 
     Example Models:
     - anthropic:claude-sonnet-4-5-20250929 (Claude 4.5 Sonnet)
@@ -127,6 +322,12 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
     - If CHAT__AUTO_INJECT_USER_CONTEXT=true: User profile auto-loaded and injected
     - New messages saved to database with compression for session continuity
     - When Postgres is disabled, session management is skipped
+
+    Evaluation Sessions:
+    - Set X-Is-Eval: true header to mark session as evaluation
+    - Session mode will be set to EVALUATION
+    - Request metadata is merged with session metadata
+    - Useful for A/B testing, model comparison, and feedback collection
     """
     # Load agent schema: use header value from context or default
     # Extract AgentContext first to get schema name
@@ -151,6 +352,17 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
             new_messages=new_messages,
         )
 
+        # Ensure session exists with metadata and eval mode if applicable
+        if context.session_id:
+            await ensure_session_with_metadata(
+                session_id=context.session_id,
+                user_id=context.user_id,
+                tenant_id=context.tenant_id,
+                is_eval=context.is_eval,
+                request_metadata=body.metadata,
+                agent_schema="simulator",
+            )
+
         # Get the last user message as prompt
         prompt = body.messages[-1].content if body.messages else "demo"
         request_id = f"sim-{uuid.uuid4().hex[:24]}"
@@ -256,7 +468,7 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
                 detail=f"Agent schema '{schema_name}' not found and default schema unavailable",
             )
 
-    logger.info(f"Using agent schema: {schema_name}, model: {body.model}")
+    logger.debug(f"Using agent schema: {schema_name}, model: {body.model}")
 
     # Check for audio input
     is_audio = request.headers.get("x-chat-is-audio", "").lower() == "true"
@@ -301,6 +513,17 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
 
     logger.info(f"Built context with {len(messages)} total messages (includes history + user context)")
 
+    # Ensure session exists with metadata and eval mode if applicable
+    if context.session_id:
+        await ensure_session_with_metadata(
+            session_id=context.session_id,
+            user_id=context.user_id,
+            tenant_id=context.tenant_id,
+            is_eval=context.is_eval,
+            request_metadata=body.metadata,
+            agent_schema=schema_name,
+        )
+
     # Create agent with schema and model override
     agent = await create_agent(
         context=context,
@@ -317,14 +540,60 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
 
     # Streaming mode
     if body.stream:
+        # Save user message before streaming starts
+        if settings.postgres.enabled and context.session_id:
+            user_message = {
+                "role": "user",
+                "content": body.messages[-1].content if body.messages else "",
+                "timestamp": datetime.utcnow().isoformat(),
+            }
+            try:
+                store = SessionMessageStore(user_id=context.user_id or settings.test.effective_user_id)
+                await store.store_session_messages(
+                    session_id=context.session_id,
+                    messages=[user_message],
+                    user_id=context.user_id,
+                    compress=False,  # User messages are typically short
+                )
+                logger.debug(f"Saved user message to session {context.session_id}")
+            except Exception as e:
+                logger.error(f"Failed to save user message: {e}", exc_info=True)
+
         return StreamingResponse(
-            stream_openai_response(agent, prompt, body.model, request_id),
+            stream_openai_response_with_save(
+                agent=agent,
+                prompt=prompt,
+                model=body.model,
+                request_id=request_id,
+                agent_schema=schema_name,
+                session_id=context.session_id,
+                user_id=context.user_id,
+            ),
             media_type="text/event-stream",
             headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
         )
 
     # Non-streaming mode
-    result = await agent.run(prompt)
+    # Create a parent span to capture trace context for message storage
+    trace_id, span_id = None, None
+    tracer = get_tracer()
+
+    if tracer:
+        with tracer.start_as_current_span(
+            "chat_completion",
+            attributes={
+                "session.id": context.session_id or "",
+                "user.id": context.user_id or "",
+                "model": body.model,
+                "agent.schema": context.agent_schema_uri or DEFAULT_AGENT_SCHEMA,
+            }
+        ) as span:
+            # Capture trace context from the span we just created
+            trace_id, span_id = get_current_trace_context()
+            result = await agent.run(prompt)
+    else:
+        # No tracer available, run without tracing
+        result = await agent.run(prompt)
 
     # Determine content format based on response_format request
     if body.response_format and body.response_format.type == "json_object":
@@ -347,12 +616,16 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
             "role": "user",
             "content": body.messages[-1].content if body.messages else "",
             "timestamp": datetime.utcnow().isoformat(),
+            "trace_id": trace_id,
+            "span_id": span_id,
         }
 
         assistant_message = {
             "role": "assistant",
             "content": content,
             "timestamp": datetime.utcnow().isoformat(),
+            "trace_id": trace_id,
+            "span_id": span_id,
         }
 
         try: