remdb 0.3.103__py3-none-any.whl → 0.3.118__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

@@ -79,7 +79,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"])
 
@@ -256,7 +256,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"
@@ -317,8 +317,35 @@ 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"},
         )

rem/api/routers/chat/streaming.py

@@ -71,6 +71,8 @@ async def stream_openai_response(
     message_id: str | None = None,
     in_reply_to: str | None = None,
     session_id: str | None = None,
+    # Agent info for metadata
+    agent_schema: str | None = None,
 ) -> AsyncGenerator[str, None]:
     """
     Stream Pydantic AI agent responses with rich SSE events.
@@ -258,8 +260,6 @@ async def stream_openai_response(
                                 # Queue for completion matching (FIFO)
                                 pending_tool_completions.append((tool_name, tool_id))
 
-                                logger.info(f"🔧 {tool_name}")
-
                                 # Emit tool_call SSE event (started)
                                 # Try to get arguments as dict
                                 args_dict = None
@@ -269,6 +269,18 @@ async def stream_openai_response(
                                     elif isinstance(event.part.args, dict):
                                         args_dict = event.part.args
 
+                                # Log tool call with key parameters
+                                if args_dict and tool_name == "search_rem":
+                                    query_type = args_dict.get("query_type", "?")
+                                    limit = args_dict.get("limit", 20)
+                                    table = args_dict.get("table", "")
+                                    query_text = args_dict.get("query_text", args_dict.get("entity_key", ""))
+                                    if query_text and len(query_text) > 50:
+                                        query_text = query_text[:50] + "..."
+                                    logger.info(f"🔧 {tool_name} {query_type.upper()} '{query_text}' table={table} limit={limit}")
+                                else:
+                                    logger.info(f"🔧 {tool_name}")
+
                                 yield format_sse_event(ToolCallEvent(
                                     tool_name=tool_name,
                                     tool_id=tool_id,
@@ -354,21 +366,43 @@ async def stream_openai_response(
                                     registered_sources = result_content.get("sources")
                                     registered_references = result_content.get("references")
                                     registered_flags = result_content.get("flags")
+                                    # Risk assessment fields
+                                    registered_risk_level = result_content.get("risk_level")
+                                    registered_risk_score = result_content.get("risk_score")
+                                    registered_risk_reasoning = result_content.get("risk_reasoning")
+                                    registered_recommended_action = result_content.get("recommended_action")
+                                    # Extra fields
+                                    registered_extra = result_content.get("extra")
 
                                     logger.info(
                                         f"📊 Metadata registered: confidence={registered_confidence}, "
-                                        f"sources={registered_sources}"
+                                        f"risk_level={registered_risk_level}, sources={registered_sources}"
                                     )
 
+                                    # Build extra dict with risk fields and any custom extras
+                                    extra_data = {}
+                                    if registered_risk_level is not None:
+                                        extra_data["risk_level"] = registered_risk_level
+                                    if registered_risk_score is not None:
+                                        extra_data["risk_score"] = registered_risk_score
+                                    if registered_risk_reasoning is not None:
+                                        extra_data["risk_reasoning"] = registered_risk_reasoning
+                                    if registered_recommended_action is not None:
+                                        extra_data["recommended_action"] = registered_recommended_action
+                                    if registered_extra:
+                                        extra_data.update(registered_extra)
+
                                     # Emit metadata event immediately
                                     yield format_sse_event(MetadataEvent(
                                         message_id=message_id,
                                         in_reply_to=in_reply_to,
                                         session_id=session_id,
+                                        agent_schema=agent_schema,
                                         confidence=registered_confidence,
                                         sources=registered_sources,
                                         model_version=model,
                                         flags=registered_flags,
+                                        extra=extra_data if extra_data else None,
                                         hidden=False,
                                     ))
 
@@ -377,6 +411,31 @@ async def stream_openai_response(
                                     result_str = str(result_content)
                                     result_summary = result_str[:200] + "..." if len(result_str) > 200 else result_str
 
+                                    # Log result count for search_rem
+                                    if tool_name == "search_rem" and isinstance(result_content, dict):
+                                        results = result_content.get("results", {})
+                                        # Handle nested result structure: results may be a dict with 'results' list and 'count'
+                                        if isinstance(results, dict):
+                                            count = results.get("count", len(results.get("results", [])))
+                                            query_type = results.get("query_type", "?")
+                                            query_text = results.get("query_text", results.get("key", ""))
+                                            table = results.get("table_name", "")
+                                        elif isinstance(results, list):
+                                            count = len(results)
+                                            query_type = "?"
+                                            query_text = ""
+                                            table = ""
+                                        else:
+                                            count = "?"
+                                            query_type = "?"
+                                            query_text = ""
+                                            table = ""
+                                        status = result_content.get("status", "unknown")
+                                        # Truncate query text for logging
+                                        if query_text and len(str(query_text)) > 40:
+                                            query_text = str(query_text)[:40] + "..."
+                                        logger.info(f"  ↳ {tool_name} {query_type} '{query_text}' table={table} → {count} results")
+
                                     yield format_sse_event(ToolCallEvent(
                                         tool_name=tool_name,
                                         tool_id=tool_id,
@@ -464,6 +523,7 @@ async def stream_openai_response(
                 message_id=message_id,
                 in_reply_to=in_reply_to,
                 session_id=session_id,
+                agent_schema=agent_schema,
                 confidence=1.0,  # Default to 100% confidence
                 model_version=model,
                 latency_ms=latency_ms,
@@ -606,3 +666,83 @@ async def stream_minimal_simulator(
     # Simulator now yields SSE-formatted strings directly (OpenAI-compatible)
     async for sse_string in stream_minimal_demo(content=content, delay_ms=delay_ms):
         yield sse_string
+
+
+async def stream_openai_response_with_save(
+    agent: Agent,
+    prompt: str,
+    model: str,
+    request_id: str | None = None,
+    agent_schema: str | None = None,
+    session_id: str | None = None,
+    user_id: str | None = None,
+) -> AsyncGenerator[str, None]:
+    """
+    Wrapper around stream_openai_response that saves the assistant response after streaming.
+
+    This accumulates all text content during streaming and saves it to the database
+    after the stream completes.
+
+    Args:
+        agent: Pydantic AI agent instance
+        prompt: User prompt
+        model: Model name
+        request_id: Optional request ID
+        agent_schema: Agent schema name
+        session_id: Session ID for message storage
+        user_id: User ID for message storage
+
+    Yields:
+        SSE-formatted strings
+    """
+    from ....utils.date_utils import utc_now, to_iso
+    from ....services.session import SessionMessageStore
+    from ....settings import settings
+
+    # Accumulate content during streaming
+    accumulated_content = []
+
+    async for chunk in stream_openai_response(
+        agent=agent,
+        prompt=prompt,
+        model=model,
+        request_id=request_id,
+        agent_schema=agent_schema,
+        session_id=session_id,
+    ):
+        yield chunk
+
+        # Extract text content from OpenAI-format chunks
+        # Format: data: {"choices": [{"delta": {"content": "..."}}]}
+        if chunk.startswith("data: ") and not chunk.startswith("data: [DONE]"):
+            try:
+                data_str = chunk[6:].strip()  # Remove "data: " prefix
+                if data_str:
+                    data = json.loads(data_str)
+                    if "choices" in data and data["choices"]:
+                        delta = data["choices"][0].get("delta", {})
+                        content = delta.get("content")
+                        if content:
+                            accumulated_content.append(content)
+            except (json.JSONDecodeError, KeyError, IndexError):
+                pass  # Skip non-JSON or malformed chunks
+
+    # After streaming completes, save the assistant response
+    if settings.postgres.enabled and session_id and accumulated_content:
+        full_content = "".join(accumulated_content)
+        assistant_message = {
+            "role": "assistant",
+            "content": full_content,
+            "timestamp": to_iso(utc_now()),
+        }
+        try:
+            store = SessionMessageStore(user_id=user_id or settings.test.effective_user_id)
+            await store.store_session_messages(
+                session_id=session_id,
+                messages=[assistant_message],
+                user_id=user_id,
+                compress=True,  # Compress long assistant responses
+            )
+            logger.debug(f"Saved assistant response to session {session_id} ({len(full_content)} chars)")
+        except Exception as e:
+            logger.error(f"Failed to save assistant response: {e}", exc_info=True)