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

rem/api/routers/chat/models.py

@@ -1,14 +1,38 @@
 """
 OpenAI-compatible API models for chat completions.
 
-Design Pattern 
+Design Pattern:
 - Full OpenAI compatibility for drop-in replacement
 - Support for streaming (SSE) and non-streaming modes
 - Response format control (text vs json_object)
-- Headers map to AgentContext (X-User-Id, X-Tenant-Id, X-Agent-Schema, etc.)
+- Headers map to AgentContext for session/context control
+- Body fields for OpenAI-compatible parameters + metadata
+
+Headers (context control):
+    X-User-Id        → context.user_id (user identifier)
+    X-Tenant-Id      → context.tenant_id (multi-tenancy, default: "default")
+    X-Session-Id     → context.session_id (conversation continuity)
+    X-Agent-Schema   → context.agent_schema_uri (which agent to use, default: "rem")
+    X-Model-Name     → context.default_model (model override)
+    X-Chat-Is-Audio  → triggers audio transcription ("true"/"false")
+    X-Is-Eval        → context.is_eval (marks session as evaluation, sets mode=EVALUATION)
+
+Body Fields (OpenAI-compatible + extensions):
+    model            → LLM model (e.g., "openai:gpt-4.1", "anthropic:claude-sonnet-4-5-20250929")
+    messages         → Chat conversation history
+    temperature      → Sampling temperature (0-2)
+    max_tokens       → Max tokens (deprecated, use max_completion_tokens)
+    max_completion_tokens → Max tokens to generate
+    stream           → Enable SSE streaming
+    metadata         → Key-value pairs merged with session metadata (for evals/experiments)
+    store            → Whether to store for distillation/evaluation
+    seed             → Deterministic sampling seed
+    top_p            → Nucleus sampling probability
+    reasoning_effort → low/medium/high for o-series models
+    service_tier     → auto/flex/priority/default
 """
 
-from typing import Literal
+from typing import Any, Literal
 
 from pydantic import BaseModel, Field
 
@@ -46,10 +70,17 @@ class ChatCompletionRequest(BaseModel):
     Compatible with OpenAI's /v1/chat/completions endpoint.
 
     Headers Map to AgentContext:
-    - X-User-Id → context.user_id
-    - X-Tenant-Id → context.tenant_id
-    - X-Session-Id → context.session_id
-    - X-Agent-Schema → context.agent_schema_uri
+        X-User-Id        → context.user_id
+        X-Tenant-Id      → context.tenant_id (default: "default")
+        X-Session-Id     → context.session_id
+        X-Agent-Schema   → context.agent_schema_uri (default: "rem")
+        X-Model-Name     → context.default_model
+        X-Chat-Is-Audio  → triggers audio transcription
+        X-Is-Eval        → context.is_eval (sets session mode=EVALUATION)
+
+    Body Fields for Metadata/Evals:
+        metadata         → Key-value pairs merged with session metadata
+        store            → Whether to store for distillation/evaluation
 
     Note: Model is specified in body.model (standard OpenAI field), not headers.
     """
@@ -73,6 +104,49 @@ class ChatCompletionRequest(BaseModel):
         default=None,
         description="Response format. Set type='json_object' to enable JSON mode.",
     )
+    # Additional OpenAI-compatible fields
+    metadata: dict[str, str] | None = Field(
+        default=None,
+        description="Key-value pairs attached to the request (max 16 keys, 64/512 char limits). "
+        "Merged with session metadata for persistence.",
+    )
+    store: bool | None = Field(
+        default=None,
+        description="Whether to store for distillation/evaluation purposes.",
+    )
+    max_completion_tokens: int | None = Field(
+        default=None,
+        ge=1,
+        description="Max tokens to generate (replaces deprecated max_tokens).",
+    )
+    seed: int | None = Field(
+        default=None,
+        description="Seed for deterministic sampling (best effort).",
+    )
+    top_p: float | None = Field(
+        default=None,
+        ge=0,
+        le=1,
+        description="Nucleus sampling probability. Use temperature OR top_p, not both.",
+    )
+    logprobs: bool | None = Field(
+        default=None,
+        description="Whether to return log probabilities for output tokens.",
+    )
+    top_logprobs: int | None = Field(
+        default=None,
+        ge=0,
+        le=20,
+        description="Number of most likely tokens to return at each position (requires logprobs=true).",
+    )
+    reasoning_effort: Literal["low", "medium", "high"] | None = Field(
+        default=None,
+        description="Reasoning effort for o-series models (low/medium/high).",
+    )
+    service_tier: Literal["auto", "flex", "priority", "default"] | None = Field(
+        default=None,
+        description="Service tier for processing (flex is 50% cheaper but slower).",
+    )
 
 
 # Response models

rem/api/routers/chat/otel_utils.py

@@ -0,0 +1,33 @@
+"""OTEL utilities for chat routers."""
+
+from loguru import logger
+
+
+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
+
+
+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()
+        ctx = span.get_span_context()
+        if ctx.is_valid:
+            trace_id = format(ctx.trace_id, '032x')
+            span_id = format(ctx.span_id, '016x')
+            return trace_id, span_id
+    except Exception as e:
+        logger.debug(f"Could not get trace context: {e}")
+
+    return None, None

rem/api/routers/chat/sse_events.py

@@ -321,7 +321,13 @@ class MetadataEvent(BaseModel):
     # Agent info
     agent_schema: str | None = Field(
         default=None,
-        description="Name of the agent schema used for this response (e.g., 'rem', 'Siggy')"
+        description="Name of the agent schema used for this response (e.g., 'rem', 'query-assistant')"
+    )
+
+    # Session info
+    session_name: str | None = Field(
+        default=None,
+        description="Short 1-3 phrase name for the session topic (e.g., 'Prescription Drug Questions', 'AWS Setup Help')"
     )
 
     # Quality indicators
@@ -350,6 +356,16 @@ class MetadataEvent(BaseModel):
         description="Token count for this response"
     )
 
+    # Trace context for observability (deterministic, captured from OTEL)
+    trace_id: str | None = Field(
+        default=None,
+        description="OTEL trace ID for correlating with Phoenix/observability systems"
+    )
+    span_id: str | None = Field(
+        default=None,
+        description="OTEL span ID for correlating with Phoenix/observability systems"
+    )
+
     # System flags
     flags: list[str] | None = Field(
         default=None,

rem/api/routers/chat/streaming.py

@@ -47,6 +47,7 @@ from pydantic_ai.messages import (
     ToolCallPart,
 )
 
+from .otel_utils import get_current_trace_context, get_tracer
 from .models import (
     ChatCompletionMessageDelta,
     ChatCompletionStreamChoice,
@@ -71,6 +72,10 @@ 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,
+    # Mutable container to capture trace context (deterministic, not AI-dependent)
+    trace_context_out: dict | None = None,
 ) -> AsyncGenerator[str, None]:
     """
     Stream Pydantic AI agent responses with rich SSE events.
@@ -154,6 +159,14 @@ async def stream_openai_response(
 
         # Use agent.iter() to get complete execution with tool calls
         async with agent.iter(prompt) as agent_run:
+            # Capture trace context IMMEDIATELY inside agent execution
+            # This is deterministic - it's the OTEL context from Pydantic AI instrumentation
+            # NOT dependent on any AI-generated content
+            captured_trace_id, captured_span_id = get_current_trace_context()
+            if trace_context_out is not None:
+                trace_context_out["trace_id"] = captured_trace_id
+                trace_context_out["span_id"] = captured_span_id
+
             async for node in agent_run:
                 # Check if this is a model request node (includes tool calls)
                 if Agent.is_model_request_node(node):
@@ -258,8 +271,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 +280,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 +377,47 @@ async def stream_openai_response(
                                     registered_sources = result_content.get("sources")
                                     registered_references = result_content.get("references")
                                     registered_flags = result_content.get("flags")
+                                    # Session naming
+                                    registered_session_name = result_content.get("session_name")
+                                    # 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"session_name={registered_session_name}, "
+                                        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,
+                                        session_name=registered_session_name,
                                         confidence=registered_confidence,
                                         sources=registered_sources,
                                         model_version=model,
                                         flags=registered_flags,
+                                        extra=extra_data if extra_data else None,
                                         hidden=False,
                                     ))
 
@@ -377,6 +426,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,10 +538,14 @@ 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,
                 token_count=token_count,
+                # Include deterministic trace context captured from OTEL
+                trace_id=captured_trace_id,
+                span_id=captured_span_id,
             ))
 
         # Mark all progress complete
@@ -606,3 +684,99 @@ 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
+
+    # Pre-generate message_id so it can be sent in metadata event
+    # This allows frontend to use it for feedback before DB persistence
+    message_id = str(uuid.uuid4())
+
+    # Mutable container for capturing trace context from inside agent execution
+    # This is deterministic - captured from OTEL instrumentation, not AI-generated
+    trace_context: dict = {}
+
+    # 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,
+        message_id=message_id,
+        trace_context_out=trace_context,  # Pass container to capture trace IDs
+    ):
+        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)
+        # Get captured trace context from container (deterministically captured inside agent execution)
+        captured_trace_id = trace_context.get("trace_id")
+        captured_span_id = trace_context.get("span_id")
+        assistant_message = {
+            "id": message_id,  # Use pre-generated ID for consistency with metadata event
+            "role": "assistant",
+            "content": full_content,
+            "timestamp": to_iso(utc_now()),
+            "trace_id": captured_trace_id,
+            "span_id": captured_span_id,
+        }
+        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 {message_id} to session {session_id} ({len(full_content)} chars)")
+        except Exception as e:
+            logger.error(f"Failed to save assistant response: {e}", exc_info=True)