remdb 0.3.202__py3-none-any.whl → 0.3.245__py3-none-any.whl

rem/api/routers/chat/streaming_utils.py

@@ -0,0 +1,327 @@
+"""
+Streaming Utilities.
+
+Pure functions and data structures for SSE streaming.
+No I/O, no database calls - just data transformation.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any
+
+from loguru import logger
+
+from .models import (
+    ChatCompletionMessageDelta,
+    ChatCompletionStreamChoice,
+    ChatCompletionStreamResponse,
+)
+from .sse_events import (
+    MetadataEvent,
+    ProgressEvent,
+    ReasoningEvent,
+    ToolCallEvent,
+    format_sse_event,
+)
+
+
+# =============================================================================
+# STREAMING STATE
+# =============================================================================
+
+@dataclass
+class StreamingState:
+    """
+    Tracks state during SSE streaming.
+
+    This is a pure data container - no methods that do I/O.
+    """
+    request_id: str
+    created_at: int
+    model: str
+    start_time: float = field(default_factory=time.time)
+
+    # Content tracking
+    is_first_chunk: bool = True
+    token_count: int = 0
+
+    # Child agent tracking - KEY FOR DUPLICATION FIX
+    child_content_streamed: bool = False
+    responding_agent: str | None = None
+
+    # Tool tracking
+    active_tool_calls: dict = field(default_factory=dict)  # index -> (name, id)
+    pending_tool_completions: list = field(default_factory=list)  # FIFO queue
+    pending_tool_data: dict = field(default_factory=dict)  # tool_id -> data
+
+    # Reasoning tracking
+    reasoning_step: int = 0
+
+    # Progress tracking
+    current_step: int = 0
+    total_steps: int = 3
+
+    # Metadata tracking
+    metadata_registered: bool = False
+
+    # Trace context (captured from OTEL)
+    trace_id: str | None = None
+    span_id: str | None = None
+
+    @classmethod
+    def create(cls, model: str, request_id: str | None = None) -> "StreamingState":
+        """Create a new streaming state."""
+        return cls(
+            request_id=request_id or f"chatcmpl-{uuid.uuid4().hex[:24]}",
+            created_at=int(time.time()),
+            model=model,
+        )
+
+    def latency_ms(self) -> int:
+        """Calculate latency since start."""
+        return int((time.time() - self.start_time) * 1000)
+
+
+# =============================================================================
+# SSE CHUNK BUILDERS
+# =============================================================================
+
+def build_content_chunk(state: StreamingState, content: str) -> str:
+    """
+    Build an SSE content chunk in OpenAI format.
+
+    Updates state.is_first_chunk and state.token_count.
+    """
+    state.token_count += len(content.split())
+
+    chunk = ChatCompletionStreamResponse(
+        id=state.request_id,
+        created=state.created_at,
+        model=state.model,
+        choices=[
+            ChatCompletionStreamChoice(
+                index=0,
+                delta=ChatCompletionMessageDelta(
+                    role="assistant" if state.is_first_chunk else None,
+                    content=content,
+                ),
+                finish_reason=None,
+            )
+        ],
+    )
+    state.is_first_chunk = False
+    return f"data: {chunk.model_dump_json()}\n\n"
+
+
+def build_final_chunk(state: StreamingState) -> str:
+    """Build the final SSE chunk with finish_reason=stop."""
+    chunk = ChatCompletionStreamResponse(
+        id=state.request_id,
+        created=state.created_at,
+        model=state.model,
+        choices=[
+            ChatCompletionStreamChoice(
+                index=0,
+                delta=ChatCompletionMessageDelta(),
+                finish_reason="stop",
+            )
+        ],
+    )
+    return f"data: {chunk.model_dump_json()}\n\n"
+
+
+def build_reasoning_event(state: StreamingState, content: str) -> str:
+    """Build a reasoning SSE event."""
+    return format_sse_event(ReasoningEvent(
+        content=content,
+        step=state.reasoning_step,
+    ))
+
+
+def build_progress_event(
+    step: int,
+    total_steps: int,
+    label: str,
+    status: str = "in_progress",
+) -> str:
+    """Build a progress SSE event."""
+    return format_sse_event(ProgressEvent(
+        step=step,
+        total_steps=total_steps,
+        label=label,
+        status=status,
+    ))
+
+
+def build_tool_start_event(
+    tool_name: str,
+    tool_id: str,
+    arguments: dict | None = None,
+) -> str:
+    """Build a tool call started SSE event."""
+    return format_sse_event(ToolCallEvent(
+        tool_name=tool_name,
+        tool_id=tool_id,
+        status="started",
+        arguments=arguments,
+    ))
+
+
+def build_tool_complete_event(
+    tool_name: str,
+    tool_id: str,
+    arguments: dict | None = None,
+    result: Any = None,
+) -> str:
+    """Build a tool call completed SSE event."""
+    result_str = None
+    if result is not None:
+        result_str = str(result)
+        if len(result_str) > 200:
+            result_str = result_str[:200] + "..."
+
+    return format_sse_event(ToolCallEvent(
+        tool_name=tool_name,
+        tool_id=tool_id,
+        status="completed",
+        arguments=arguments,
+        result=result_str,
+    ))
+
+
+def build_metadata_event(
+    message_id: str | None = None,
+    in_reply_to: str | None = None,
+    session_id: str | None = None,
+    agent_schema: str | None = None,
+    responding_agent: str | None = None,
+    confidence: float | None = None,
+    sources: list | None = None,
+    model_version: str | None = None,
+    latency_ms: int | None = None,
+    token_count: int | None = None,
+    trace_id: str | None = None,
+    span_id: str | None = None,
+    extra: dict | None = None,
+) -> str:
+    """Build a metadata SSE event."""
+    return format_sse_event(MetadataEvent(
+        message_id=message_id,
+        in_reply_to=in_reply_to,
+        session_id=session_id,
+        agent_schema=agent_schema,
+        responding_agent=responding_agent,
+        confidence=confidence,
+        sources=sources,
+        model_version=model_version,
+        latency_ms=latency_ms,
+        token_count=token_count,
+        trace_id=trace_id,
+        span_id=span_id,
+        extra=extra,
+    ))
+
+
+# =============================================================================
+# TOOL ARGUMENT EXTRACTION
+# =============================================================================
+
+def extract_tool_args(part) -> dict | None:
+    """
+    Extract arguments from a ToolCallPart.
+
+    Handles various formats:
+    - ArgsDict object with args_dict attribute
+    - Plain dict
+    - JSON string
+    """
+    if part.args is None:
+        return None
+
+    if hasattr(part.args, 'args_dict'):
+        return part.args.args_dict
+
+    if isinstance(part.args, dict):
+        return part.args
+
+    if isinstance(part.args, str) and part.args:
+        try:
+            return json.loads(part.args)
+        except json.JSONDecodeError:
+            logger.warning(f"Failed to parse tool args: {part.args[:100]}")
+
+    return None
+
+
+def log_tool_call(tool_name: str, args_dict: dict | None) -> None:
+    """Log a 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(str(query_text)) > 50:
+            query_text = str(query_text)[:50] + "..."
+        logger.info(f"🔧 {tool_name} {query_type.upper()} '{query_text}' table={table} limit={limit}")
+    else:
+        logger.info(f"🔧 {tool_name}")
+
+
+def log_tool_result(tool_name: str, result_content: Any) -> None:
+    """Log a tool result with key metrics."""
+    if tool_name == "search_rem" and isinstance(result_content, dict):
+        results = result_content.get("results", {})
+        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 = ""
+
+        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")
+
+
+# =============================================================================
+# METADATA EXTRACTION
+# =============================================================================
+
+def extract_metadata_from_result(result_content: Any) -> dict | None:
+    """
+    Extract metadata from a register_metadata tool result.
+
+    Returns dict with extracted fields or None if not a metadata event.
+    """
+    if not isinstance(result_content, dict):
+        return None
+
+    if not result_content.get("_metadata_event"):
+        return None
+
+    return {
+        "confidence": result_content.get("confidence"),
+        "sources": result_content.get("sources"),
+        "references": result_content.get("references"),
+        "flags": result_content.get("flags"),
+        "session_name": result_content.get("session_name"),
+        "risk_level": result_content.get("risk_level"),
+        "risk_score": result_content.get("risk_score"),
+        "risk_reasoning": result_content.get("risk_reasoning"),
+        "recommended_action": result_content.get("recommended_action"),
+        "agent_schema": result_content.get("agent_schema"),
+        "extra": result_content.get("extra"),
+    }

rem/api/routers/common.py

@@ -0,0 +1,18 @@
+"""
+Common models shared across API routers.
+"""
+
+from pydantic import BaseModel, Field
+
+
+class ErrorResponse(BaseModel):
+    """Standard error response format for HTTPException errors.
+
+    This is different from FastAPI's HTTPValidationError which is used
+    for Pydantic validation failures (422 errors with loc/msg/type array).
+
+    HTTPException errors return this simpler format:
+        {"detail": "Error message here"}
+    """
+
+    detail: str = Field(description="Error message describing what went wrong")

rem/api/routers/dev.py

@@ -11,6 +11,7 @@ Endpoints:
 from fastapi import APIRouter, HTTPException, Request
 from loguru import logger
 
+from .common import ErrorResponse
 from ...settings import settings
 
 router = APIRouter(prefix="/api/dev", tags=["dev"])
@@ -45,7 +46,12 @@ def verify_dev_token(token: str) -> bool:
     return token == expected
 
 
-@router.get("/token")
+@router.get(
+    "/token",
+    responses={
+        401: {"model": ErrorResponse, "description": "Dev tokens not available in production"},
+    },
+)
 async def get_dev_token(request: Request):
     """
     Get a development token for testing (non-production only).

rem/api/routers/feedback.py

@@ -63,6 +63,8 @@ from fastapi import APIRouter, Header, HTTPException, Request, Response
 from loguru import logger
 from pydantic import BaseModel, Field
 
+from .common import ErrorResponse
+
 from ..deps import get_user_id_from_request
 from ...models.entities import Feedback
 from ...services.postgres import Repository
@@ -121,7 +123,13 @@ class FeedbackResponse(BaseModel):
 # =============================================================================
 
 
-@router.post("/messages/feedback", response_model=FeedbackResponse)
+@router.post(
+    "/messages/feedback",
+    response_model=FeedbackResponse,
+    responses={
+        503: {"model": ErrorResponse, "description": "Database not enabled"},
+    },
+)
 async def submit_feedback(
     request: Request,
     response: Response,

rem/api/routers/messages.py

@@ -16,6 +16,7 @@ Endpoints:
 """
 
 from datetime import datetime
+from enum import Enum
 from typing import Literal
 from uuid import UUID
 
@@ -23,6 +24,8 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request
 from loguru import logger
 from pydantic import BaseModel, Field
 
+from .common import ErrorResponse
+
 from ..deps import (
     get_current_user,
     get_user_filter,
@@ -38,6 +41,18 @@ from ...utils.date_utils import parse_iso, utc_now
 router = APIRouter(prefix="/api/v1")
 
 
+# =============================================================================
+# Enums
+# =============================================================================
+
+
+class SortOrder(str, Enum):
+    """Sort order for list queries."""
+
+    ASC = "asc"
+    DESC = "desc"
+
+
 # =============================================================================
 # Request/Response Models
 # =============================================================================
@@ -134,7 +149,14 @@ class SessionsQueryResponse(BaseModel):
 # =============================================================================
 
 
-@router.get("/messages", response_model=MessageListResponse, tags=["messages"])
+@router.get(
+    "/messages",
+    response_model=MessageListResponse,
+    tags=["messages"],
+    responses={
+        503: {"model": ErrorResponse, "description": "Database not enabled"},
+    },
+)
 async def list_messages(
     request: Request,
     mine: bool = Query(default=False, description="Only show my messages (uses JWT identity)"),
@@ -151,6 +173,7 @@ async def list_messages(
     ),
     limit: int = Query(default=50, ge=1, le=100, description="Max results to return"),
     offset: int = Query(default=0, ge=0, description="Offset for pagination"),
+    sort: SortOrder = Query(default=SortOrder.DESC, description="Sort order by created_at (asc or desc)"),
 ) -> MessageListResponse:
     """
     List messages with optional filters.
@@ -166,8 +189,9 @@ async def list_messages(
     - session_id: Filter by conversation session
     - start_date/end_date: Filter by creation time range (ISO 8601 format)
     - message_type: Filter by role (user, assistant, system, tool)
+    - sort: Sort order by created_at (asc or desc, default: desc)
 
-    Returns paginated results ordered by created_at descending.
+    Returns paginated results ordered by created_at.
     """
     if not settings.postgres.enabled:
         raise HTTPException(status_code=503, detail="Database not enabled")
@@ -189,6 +213,7 @@ async def list_messages(
 
     # Apply optional filters
     if session_id:
+        # session_id is the session UUID - use directly
         filters["session_id"] = session_id
     if message_type:
         filters["message_type"] = message_type
@@ -200,12 +225,15 @@ async def list_messages(
         f"filters={filters}"
     )
 
+    # Build order_by clause based on sort parameter
+    order_by = f"created_at {sort.value.upper()}"
+
     # For date filtering, we need custom SQL (not supported by basic Repository)
     # For now, fetch all matching base filters and filter in Python
     # TODO: Extend Repository to support date range filters
     messages = await repo.find(
         filters,
-        order_by="created_at DESC",
+        order_by=order_by,
         limit=limit + 1,  # Fetch one extra to determine has_more
         offset=offset,
     )
@@ -241,7 +269,16 @@ async def list_messages(
     return MessageListResponse(data=messages, total=total, has_more=has_more)
 
 
-@router.get("/messages/{message_id}", response_model=Message, tags=["messages"])
+@router.get(
+    "/messages/{message_id}",
+    response_model=Message,
+    tags=["messages"],
+    responses={
+        403: {"model": ErrorResponse, "description": "Access denied: not owner"},
+        404: {"model": ErrorResponse, "description": "Message not found"},
+        503: {"model": ErrorResponse, "description": "Database not enabled"},
+    },
+)
 async def get_message(
     request: Request,
     message_id: str,
@@ -287,7 +324,14 @@ async def get_message(
 # =============================================================================
 
 
-@router.get("/sessions", response_model=SessionsQueryResponse, tags=["sessions"])
+@router.get(
+    "/sessions",
+    response_model=SessionsQueryResponse,
+    tags=["sessions"],
+    responses={
+        503: {"model": ErrorResponse, "description": "Database not enabled or connection failed"},
+    },
+)
 async def list_sessions(
     request: Request,
     user_id: str | None = Query(default=None, description="Filter by user ID (admin only for cross-user)"),
@@ -400,7 +444,15 @@ async def list_sessions(
     )
 
 
-@router.post("/sessions", response_model=Session, status_code=201, tags=["sessions"])
+@router.post(
+    "/sessions",
+    response_model=Session,
+    status_code=201,
+    tags=["sessions"],
+    responses={
+        503: {"model": ErrorResponse, "description": "Database not enabled"},
+    },
+)
 async def create_session(
     request_body: SessionCreateRequest,
     user: dict = Depends(require_admin),
@@ -452,7 +504,16 @@ async def create_session(
     return result  # type: ignore
 
 
-@router.get("/sessions/{session_id}", response_model=Session, tags=["sessions"])
+@router.get(
+    "/sessions/{session_id}",
+    response_model=Session,
+    tags=["sessions"],
+    responses={
+        403: {"model": ErrorResponse, "description": "Access denied: not owner"},
+        404: {"model": ErrorResponse, "description": "Session not found"},
+        503: {"model": ErrorResponse, "description": "Database not enabled"},
+    },
+)
 async def get_session(
     request: Request,
     session_id: str,
@@ -465,7 +526,7 @@ async def get_session(
     - Admin users: Can access any session
 
     Args:
-        session_id: UUID or name of the session
+        session_id: UUID of the session
 
     Returns:
         Session object if found
@@ -481,12 +542,7 @@ async def get_session(
     session = await repo.get_by_id(session_id)
 
     if not session:
-        # Try finding by name
-        sessions = await repo.find({"name": session_id}, limit=1)
-        if sessions:
-            session = sessions[0]
-        else:
-            raise HTTPException(status_code=404, detail=f"Session '{session_id}' not found")
+        raise HTTPException(status_code=404, detail=f"Session '{session_id}' not found")
 
     # Check access: admin or owner
     current_user = get_current_user(request)
@@ -498,7 +554,16 @@ async def get_session(
     return session
 
 
-@router.put("/sessions/{session_id}", response_model=Session, tags=["sessions"])
+@router.put(
+    "/sessions/{session_id}",
+    response_model=Session,
+    tags=["sessions"],
+    responses={
+        403: {"model": ErrorResponse, "description": "Access denied: not owner"},
+        404: {"model": ErrorResponse, "description": "Session not found"},
+        503: {"model": ErrorResponse, "description": "Database not enabled"},
+    },
+)
 async def update_session(
     request: Request,
     session_id: str,

rem/api/routers/models.py

@@ -15,6 +15,8 @@ from typing import Literal
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel, Field
 
+from .common import ErrorResponse
+
 from rem.agentic.llm_provider_models import (
     ModelInfo,
     AVAILABLE_MODELS,
@@ -57,7 +59,13 @@ async def list_models() -> ModelsResponse:
     return ModelsResponse(data=AVAILABLE_MODELS)
 
 
-@router.get("/models/{model_id:path}", response_model=ModelInfo)
+@router.get(
+    "/models/{model_id:path}",
+    response_model=ModelInfo,
+    responses={
+        404: {"model": ErrorResponse, "description": "Model not found"},
+    },
+)
 async def get_model(model_id: str) -> ModelInfo:
     """
     Get information about a specific model.

rem/api/routers/query.py

@@ -86,10 +86,10 @@ from fastapi import APIRouter, Header, HTTPException
 from loguru import logger
 from pydantic import BaseModel, Field
 
+from .common import ErrorResponse
+
 from ...services.postgres import get_postgres_service
 from ...services.rem.service import RemService
-from ...services.rem.parser import RemQueryParser
-from ...models.core import RemQuery
 from ...settings import settings
 
 router = APIRouter(prefix="/api/v1", tags=["query"])
@@ -213,7 +213,16 @@ class QueryResponse(BaseModel):
     )
 
 
-@router.post("/query", response_model=QueryResponse)
+@router.post(
+    "/query",
+    response_model=QueryResponse,
+    responses={
+        400: {"model": ErrorResponse, "description": "Invalid query or missing required fields"},
+        500: {"model": ErrorResponse, "description": "Query execution failed"},
+        501: {"model": ErrorResponse, "description": "Feature not yet implemented"},
+        503: {"model": ErrorResponse, "description": "Database not configured or unavailable"},
+    },
+)
 async def execute_query(
     request: QueryRequest,
     x_user_id: str | None = Header(default=None, description="User ID for query isolation (optional, uses default if not provided)"),
@@ -320,7 +329,7 @@ async def execute_query(
             return response
 
         else:
-            # REM dialect mode - parse and execute directly
+            # REM dialect mode - use unified execute_query_string
             if not request.query:
                 raise HTTPException(
                     status_code=400,
@@ -329,17 +338,10 @@ async def execute_query(
 
             logger.info(f"REM dialect query: {request.query[:100]}...")
 
-            parser = RemQueryParser()
-            query_type, parameters = parser.parse(request.query)
-
-            # Create and execute RemQuery
-            rem_query = RemQuery.model_validate({
-                "query_type": query_type,
-                "parameters": parameters,
-                "user_id": effective_user_id,
-            })
-
-            result = await rem_service.execute_query(rem_query)
+            # Use the unified execute_query_string method
+            result = await rem_service.execute_query_string(
+                request.query, user_id=effective_user_id
+            )
 
             return QueryResponse(
                 query_type=result["query_type"],