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

rem/api/routers/auth.py

@@ -541,6 +541,9 @@ async def refresh_token(body: TokenRefreshRequest):
     """
     Refresh access token using refresh token.
 
+    Fetches the user's current role/tier from the database to ensure
+    the new access token reflects their actual permissions.
+
     Args:
         body: TokenRefreshRequest with refresh_token
 
@@ -548,7 +551,46 @@ async def refresh_token(body: TokenRefreshRequest):
         New access token or 401 if refresh token is invalid
     """
     jwt_service = get_jwt_service()
-    result = jwt_service.refresh_access_token(body.refresh_token)
+
+    # First decode the refresh token to get user_id (without full verification yet)
+    payload = jwt_service.decode_without_verification(body.refresh_token)
+    if not payload:
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid refresh token format"
+        )
+
+    user_id = payload.get("sub")
+    if not user_id:
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid refresh token: missing user ID"
+        )
+
+    # Fetch user from database to get current role/tier
+    user_override = None
+    if settings.postgres.enabled:
+        db = PostgresService()
+        try:
+            await db.connect()
+            user_service = UserService(db)
+            user_entity = await user_service.get_user_by_id(user_id)
+            if user_entity:
+                user_override = {
+                    "role": user_entity.role or "user",
+                    "roles": [user_entity.role] if user_entity.role else ["user"],
+                    "tier": user_entity.tier.value if user_entity.tier else "free",
+                    "name": user_entity.name,
+                }
+                logger.debug(f"Refresh token: fetched user {user_id} with role={user_override['role']}, tier={user_override['tier']}")
+        except Exception as e:
+            logger.warning(f"Could not fetch user for token refresh: {e}")
+            # Continue without override - will use defaults
+        finally:
+            await db.disconnect()
+
+    # Now do the actual refresh with proper verification
+    result = jwt_service.refresh_access_token(body.refresh_token, user_override=user_override)
 
     if not result:
         raise HTTPException(

rem/api/routers/chat/completions.py

@@ -503,16 +503,42 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
             logger.error(f"Failed to transcribe audio: {e}")
             # Fall through with original content (will likely fail at agent)
 
-    # Use ContextBuilder to construct complete message list with:
-    # 1. System context hint (date + user profile)
-    # 2. Session history (if session_id provided)
-    # 3. New messages from request body (transcribed if audio)
+    # Use ContextBuilder to construct context and basic messages
+    # Note: We load session history separately for proper pydantic-ai message_history
     context, messages = await ContextBuilder.build_from_headers(
         headers=dict(request.headers),
         new_messages=new_messages,
         user_id=temp_context.user_id,  # From JWT token (source of truth)
     )
 
+    # Load raw session history for proper pydantic-ai message_history format
+    # This enables proper tool call/return pairing for LLM API compatibility
+    from ....services.session import SessionMessageStore, session_to_pydantic_messages
+    from ....agentic.schema import get_system_prompt
+
+    pydantic_message_history = None
+    if context.session_id and settings.postgres.enabled:
+        try:
+            store = SessionMessageStore(user_id=context.user_id or settings.test.effective_user_id)
+            raw_session_history = await store.load_session_messages(
+                session_id=context.session_id,
+                user_id=context.user_id,
+                compress_on_load=False,  # Don't compress - we need full data for reconstruction
+            )
+            if raw_session_history:
+                # CRITICAL: Extract and pass the agent's system prompt
+                # pydantic-ai only auto-adds system prompts when message_history is empty
+                # When we pass message_history, we must include the system prompt ourselves
+                agent_system_prompt = get_system_prompt(agent_schema) if agent_schema else None
+                pydantic_message_history = session_to_pydantic_messages(
+                    raw_session_history,
+                    system_prompt=agent_system_prompt,
+                )
+                logger.debug(f"Converted {len(raw_session_history)} session messages to {len(pydantic_message_history)} pydantic-ai messages (with system prompt)")
+        except Exception as e:
+            logger.warning(f"Failed to load session history for message_history: {e}")
+            # Fall back to old behavior (concatenated prompt)
+
     logger.info(f"Built context with {len(messages)} total messages (includes history + user context)")
 
     # Ensure session exists with metadata and eval mode if applicable
@@ -533,9 +559,17 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
         model_override=body.model,  # type: ignore[arg-type]
     )
 
-    # Combine all messages into single prompt for agent
-    # ContextBuilder already assembled: system context + history + new messages
-    prompt = "\n".join(msg.content for msg in messages)
+    # Build the prompt for the agent
+    # If we have proper message_history, use just the latest user message as prompt
+    # Otherwise, fall back to concatenating all messages (legacy behavior)
+    if pydantic_message_history:
+        # Use the latest user message as the prompt, with history passed separately
+        user_prompt = body.messages[-1].content if body.messages else ""
+        prompt = user_prompt
+        logger.debug(f"Using message_history with {len(pydantic_message_history)} messages")
+    else:
+        # Legacy: Combine all messages into single prompt for agent
+        prompt = "\n".join(msg.content for msg in messages)
 
     # Generate OpenAI-compatible request ID
     request_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
@@ -570,6 +604,8 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
                 agent_schema=schema_name,
                 session_id=context.session_id,
                 user_id=context.user_id,
+                agent_context=context,  # Pass context for multi-agent support
+                message_history=pydantic_message_history,  # Native pydantic-ai message history
             ),
             media_type="text/event-stream",
             headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
@@ -592,10 +628,16 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
         ) as span:
             # Capture trace context from the span we just created
             trace_id, span_id = get_current_trace_context()
-            result = await agent.run(prompt)
+            if pydantic_message_history:
+                result = await agent.run(prompt, message_history=pydantic_message_history)
+            else:
+                result = await agent.run(prompt)
     else:
         # No tracer available, run without tracing
-        result = await agent.run(prompt)
+        if pydantic_message_history:
+            result = await agent.run(prompt, message_history=pydantic_message_history)
+        else:
+            result = await agent.run(prompt)
 
     # Determine content format based on response_format request
     if body.response_format and body.response_format.type == "json_object":

rem/api/routers/chat/sse_events.py

@@ -409,9 +409,9 @@ class ToolCallEvent(BaseModel):
         default=None,
         description="Tool arguments (for 'started' status)"
     )
-    result: str | None = Field(
+    result: str | dict[str, Any] | None = Field(
         default=None,
-        description="Tool result summary (for 'completed' status)"
+        description="Tool result - full dict for finalize_intake, summary string for others"
     )
     error: str | None = Field(
         default=None,

rem/api/routers/chat/streaming.py

@@ -15,6 +15,11 @@ Key Insight
 - Use PartEndEvent to detect tool completion
 - Use FunctionToolResultEvent to get tool results
 
+Multi-Agent Context Propagation:
+- AgentContext is set via agent_context_scope() before agent.iter()
+- Child agents (via ask_agent tool) can access parent context via get_current_context()
+- Context includes user_id, tenant_id, session_id, is_eval for proper scoping
+
 SSE Format (OpenAI-compatible):
     data: {"id": "chatcmpl-...", "choices": [{"delta": {"content": "..."}}]}\\n\\n
     data: [DONE]\\n\\n
@@ -28,10 +33,12 @@ Extended SSE Format (Custom Events):
 See sse_events.py for the full event type definitions.
 """
 
+from __future__ import annotations
+
 import json
 import time
 import uuid
-from typing import AsyncGenerator
+from typing import TYPE_CHECKING, AsyncGenerator
 
 from loguru import logger
 from pydantic_ai.agent import Agent
@@ -55,6 +62,7 @@ from .models import (
 )
 from .sse_events import (
     DoneEvent,
+    ErrorEvent,
     MetadataEvent,
     ProgressEvent,
     ReasoningEvent,
@@ -62,6 +70,9 @@ from .sse_events import (
     format_sse_event,
 )
 
+if TYPE_CHECKING:
+    from ....agentic.context import AgentContext
+
 
 async def stream_openai_response(
     agent: Agent,
@@ -79,6 +90,11 @@ async def stream_openai_response(
     # Mutable container to capture tool calls for persistence
     # Format: list of {"tool_name": str, "tool_id": str, "arguments": dict, "result": any}
     tool_calls_out: list | None = None,
+    # Agent context for multi-agent propagation
+    # When set, enables child agents to access parent context via get_current_context()
+    agent_context: "AgentContext | None" = None,
+    # Pydantic-ai native message history for proper tool call/return pairing
+    message_history: list | None = None,
 ) -> AsyncGenerator[str, None]:
     """
     Stream Pydantic AI agent responses with rich SSE events.
@@ -153,6 +169,17 @@ async def stream_openai_response(
     # Maps tool_id -> {"tool_name": str, "tool_id": str, "arguments": dict}
     pending_tool_data: dict[str, dict] = {}
 
+    # Import context functions for multi-agent support
+    from ....agentic.context import set_current_context
+
+    # Set up context for multi-agent propagation
+    # This allows child agents (via ask_agent tool) to access parent context
+    previous_context = None
+    if agent_context is not None:
+        from ....agentic.context import get_current_context
+        previous_context = get_current_context()
+        set_current_context(agent_context)
+
     try:
         # Emit initial progress event
         current_step = 1
@@ -164,7 +191,9 @@ async def stream_openai_response(
         ))
 
         # Use agent.iter() to get complete execution with tool calls
-        async with agent.iter(prompt) as agent_run:
+        # Pass message_history if available for proper tool call/return pairing
+        iter_kwargs = {"message_history": message_history} if message_history else {}
+        async with agent.iter(prompt, **iter_kwargs) 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
@@ -285,6 +314,12 @@ async def stream_openai_response(
                                         args_dict = event.part.args.args_dict
                                     elif isinstance(event.part.args, dict):
                                         args_dict = event.part.args
+                                    elif isinstance(event.part.args, str):
+                                        # Parse JSON string args (common with pydantic-ai)
+                                        try:
+                                            args_dict = json.loads(event.part.args)
+                                        except json.JSONDecodeError:
+                                            logger.warning(f"Failed to parse tool args as JSON: {event.part.args[:100]}")
 
                                 # Log tool call with key parameters
                                 if args_dict and tool_name == "search_rem":
@@ -330,8 +365,25 @@ async def stream_openai_response(
                             ):
                                 if event.index in active_tool_calls:
                                     tool_name, tool_id = active_tool_calls[event.index]
-                                    # Note: result comes from FunctionToolResultEvent below
-                                    # For now, mark as completed without result
+
+                                    # Extract full args from completed ToolCallPart
+                                    # (PartStartEvent only has empty/partial args during streaming)
+                                    args_dict = None
+                                    if event.part.args is not None:
+                                        if hasattr(event.part.args, 'args_dict'):
+                                            args_dict = event.part.args.args_dict
+                                        elif isinstance(event.part.args, dict):
+                                            args_dict = event.part.args
+                                        elif isinstance(event.part.args, str) and event.part.args:
+                                            try:
+                                                args_dict = json.loads(event.part.args)
+                                            except json.JSONDecodeError:
+                                                logger.warning(f"Failed to parse tool args: {event.part.args[:100]}")
+
+                                    # Update pending_tool_data with complete args
+                                    if tool_id in pending_tool_data:
+                                        pending_tool_data[tool_id]["arguments"] = args_dict
+
                                     del active_tool_calls[event.index]
 
                             # ============================================
@@ -434,6 +486,12 @@ async def stream_openai_response(
                                         hidden=False,
                                     ))
 
+                                # Get complete args from pending_tool_data BEFORE deleting
+                                # (captured at PartEndEvent with full args)
+                                completed_args = None
+                                if tool_id in pending_tool_data:
+                                    completed_args = pending_tool_data[tool_id].get("arguments")
+
                                 # Capture tool call with result for persistence
                                 # Special handling for register_metadata - always capture full data
                                 if tool_calls_out is not None and tool_id in pending_tool_data:
@@ -445,8 +503,12 @@ async def stream_openai_response(
 
                                 if not is_metadata_event:
                                     # Normal tool completion - emit ToolCallEvent
-                                    result_str = str(result_content)
-                                    result_summary = result_str[:200] + "..." if len(result_str) > 200 else result_str
+                                    # For finalize_intake, send full result dict for frontend
+                                    if tool_name == "finalize_intake" and isinstance(result_content, dict):
+                                        result_for_sse = result_content
+                                    else:
+                                        result_str = str(result_content)
+                                        result_for_sse = 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):
@@ -477,7 +539,8 @@ async def stream_openai_response(
                                         tool_name=tool_name,
                                         tool_id=tool_id,
                                         status="completed",
-                                        result=result_summary
+                                        arguments=completed_args,
+                                        result=result_for_sse
                                     ))
 
                                 # Update progress after tool completion
@@ -587,25 +650,77 @@ async def stream_openai_response(
 
     except Exception as e:
         import traceback
+        import re
 
         error_msg = str(e)
-        logger.error(f"Streaming error: {error_msg}")
-        logger.error(traceback.format_exc())
-
-        # Send error as final chunk
-        error_data = {
-            "error": {
-                "message": error_msg,
-                "type": "internal_error",
-                "code": "stream_error",
-            }
-        }
-        yield f"data: {json.dumps(error_data)}\n\n"
+
+        # Parse error details for better client handling
+        error_code = "stream_error"
+        error_details: dict = {}
+        recoverable = True
+
+        # Check for rate limit errors (OpenAI 429)
+        if "429" in error_msg or "rate_limit" in error_msg.lower() or "RateLimitError" in type(e).__name__:
+            error_code = "rate_limit_exceeded"
+            recoverable = True
+
+            # Extract retry-after time from error message
+            # Pattern: "Please try again in X.XXs" or "Please try again in Xs"
+            retry_match = re.search(r"try again in (\d+(?:\.\d+)?)\s*s", error_msg)
+            if retry_match:
+                retry_seconds = float(retry_match.group(1))
+                error_details["retry_after_seconds"] = retry_seconds
+                error_details["retry_after_ms"] = int(retry_seconds * 1000)
+
+            # Extract token usage info if available
+            used_match = re.search(r"Used (\d+)", error_msg)
+            limit_match = re.search(r"Limit (\d+)", error_msg)
+            requested_match = re.search(r"Requested (\d+)", error_msg)
+            if used_match:
+                error_details["tokens_used"] = int(used_match.group(1))
+            if limit_match:
+                error_details["tokens_limit"] = int(limit_match.group(1))
+            if requested_match:
+                error_details["tokens_requested"] = int(requested_match.group(1))
+
+            logger.error(f"🔴 Streaming error: status_code: 429, model_name: {model}, body: {error_msg[:200]}")
+
+        # Check for authentication errors
+        elif "401" in error_msg or "AuthenticationError" in type(e).__name__:
+            error_code = "authentication_error"
+            recoverable = False
+            logger.error(f"🔴 Streaming error: Authentication failed")
+
+        # Check for model not found / invalid model
+        elif "404" in error_msg or "model" in error_msg.lower() and "not found" in error_msg.lower():
+            error_code = "model_not_found"
+            recoverable = False
+            logger.error(f"🔴 Streaming error: Model not found")
+
+        # Generic error
+        else:
+            logger.error(f"🔴 Streaming error: {error_msg}")
+
+        logger.error(f"🔴 {traceback.format_exc()}")
+
+        # Emit proper ErrorEvent via SSE (with event: prefix for client parsing)
+        yield format_sse_event(ErrorEvent(
+            code=error_code,
+            message=error_msg,
+            details=error_details if error_details else None,
+            recoverable=recoverable,
+        ))
 
         # Emit done event with error reason
         yield format_sse_event(DoneEvent(reason="error"))
         yield "data: [DONE]\n\n"
 
+    finally:
+        # Restore previous context for multi-agent support
+        # This ensures nested agent calls don't pollute the parent's context
+        if agent_context is not None:
+            set_current_context(previous_context)
+
 
 async def stream_simulator_response(
     prompt: str,
@@ -716,6 +831,10 @@ async def stream_openai_response_with_save(
     agent_schema: str | None = None,
     session_id: str | None = None,
     user_id: str | None = None,
+    # Agent context for multi-agent propagation
+    agent_context: "AgentContext | None" = None,
+    # Pydantic-ai native message history for proper tool call/return pairing
+    message_history: list | None = None,
 ) -> AsyncGenerator[str, None]:
     """
     Wrapper around stream_openai_response that saves the assistant response after streaming.
@@ -731,6 +850,7 @@ async def stream_openai_response_with_save(
         agent_schema: Agent schema name
         session_id: Session ID for message storage
         user_id: User ID for message storage
+        agent_context: Agent context for multi-agent propagation (enables child agents)
 
     Yields:
         SSE-formatted strings
@@ -763,6 +883,8 @@ async def stream_openai_response_with_save(
         message_id=message_id,
         trace_context_out=trace_context,  # Pass container to capture trace IDs
         tool_calls_out=tool_calls,  # Capture tool calls for persistence
+        agent_context=agent_context,  # Pass context for multi-agent support
+        message_history=message_history,  # Native pydantic-ai message history
     ):
         yield chunk
 
@@ -793,6 +915,8 @@ async def stream_openai_response_with_save(
 
         # First, store tool call messages (message_type: "tool")
         for tool_call in tool_calls:
+            if not tool_call:
+                continue
             tool_message = {
                 "role": "tool",
                 "content": json.dumps(tool_call.get("result", {}), default=str),
@@ -838,8 +962,9 @@ async def stream_openai_response_with_save(
 
         # Update session description with session_name (non-blocking, after all yields)
         for tool_call in tool_calls:
-            if tool_call.get("tool_name") == "register_metadata" and tool_call.get("is_metadata"):
-                session_name = tool_call.get("arguments", {}).get("session_name")
+            if tool_call and tool_call.get("tool_name") == "register_metadata" and tool_call.get("is_metadata"):
+                arguments = tool_call.get("arguments") or {}
+                session_name = arguments.get("session_name")
                 if session_name:
                     try:
                         from ....models.entities import Session

rem/api/routers/messages.py

@@ -93,6 +93,23 @@ class SessionListResponse(BaseModel):
     has_more: bool
 
 
+class SessionWithUser(BaseModel):
+    """Session with user info for admin views."""
+
+    id: str
+    name: str
+    mode: str | None = None
+    description: str | None = None
+    user_id: str | None = None
+    user_name: str | None = None
+    user_email: str | None = None
+    message_count: int = 0
+    total_tokens: int | None = None
+    created_at: datetime | None = None
+    updated_at: datetime | None = None
+    metadata: dict | None = None
+
+
 class PaginationMetadata(BaseModel):
     """Pagination metadata for paginated responses."""
 
@@ -108,7 +125,7 @@ class SessionsQueryResponse(BaseModel):
     """Response for paginated sessions query."""
 
     object: Literal["list"] = "list"
-    data: list[Session] = Field(description="List of sessions for the current page")
+    data: list[SessionWithUser] = Field(description="List of sessions for the current page")
     metadata: PaginationMetadata = Field(description="Pagination metadata")
 
 
@@ -274,6 +291,8 @@ async def get_message(
 async def list_sessions(
     request: Request,
     user_id: str | None = Query(default=None, description="Filter by user ID (admin only for cross-user)"),
+    user_name: str | None = Query(default=None, description="Filter by user name (partial match, admin only)"),
+    user_email: str | None = Query(default=None, description="Filter by user email (partial match, admin only)"),
     mode: SessionMode | None = Query(default=None, description="Filter by session mode"),
     page: int = Query(default=1, ge=1, description="Page number (1-indexed)"),
     page_size: int = Query(default=50, ge=1, le=100, description="Number of results per page"),
@@ -283,46 +302,100 @@ async def list_sessions(
 
     Access Control:
     - Regular users: Only see their own sessions
-    - Admin users: Can filter by any user_id or see all sessions
+    - Admin users: Can filter by any user_id, user_name, user_email, or see all sessions
 
     Filters:
     - user_id: Filter by session owner (admin only for cross-user)
+    - user_name: Filter by user name partial match (admin only)
+    - user_email: Filter by user email partial match (admin only)
     - mode: Filter by session mode (normal or evaluation)
 
     Pagination:
     - page: Page number (1-indexed, default: 1)
     - page_size: Number of sessions per page (default: 50, max: 100)
 
-    Returns paginated results ordered by created_at descending with pagination metadata.
+    Returns paginated results with user info ordered by created_at descending.
     """
     if not settings.postgres.enabled:
         raise HTTPException(status_code=503, detail="Database not enabled")
 
-    repo = Repository(Session, table_name="sessions")
+    current_user = get_current_user(request)
+    admin = is_admin(current_user)
 
-    # Build user-scoped filters (admin can see all, regular users see only their own)
-    filters = await get_user_filter(request, x_user_id=user_id)
-    if mode:
-        filters["mode"] = mode.value
+    # Get postgres service for raw SQL query
+    db = get_postgres_service()
+    if not db:
+        raise HTTPException(status_code=503, detail="Database connection failed")
+    if not db.pool:
+        await db.connect()
 
-    # Use CTE-based pagination with ROW_NUMBER() OVER (PARTITION BY user_id ORDER BY created_at DESC)
-    result = await repo.find_paginated(
-        filters,
-        page=page,
-        page_size=page_size,
-        order_by="created_at DESC",
-        partition_by="user_id",
-    )
+    # Build effective filters based on user role
+    effective_user_id = user_id
+    effective_user_name = user_name if admin else None  # Only admin can search by name
+    effective_user_email = user_email if admin else None  # Only admin can search by email
+
+    if not admin:
+        # Non-admin users can only see their own sessions
+        effective_user_id = current_user.get("id") if current_user else None
+        if not effective_user_id:
+            # Anonymous user - return empty
+            return SessionsQueryResponse(
+                data=[],
+                metadata=PaginationMetadata(
+                    total=0, page=page, page_size=page_size,
+                    total_pages=0, has_next=False, has_previous=False,
+                ),
+            )
+
+    # Call the SQL function for sessions with user info
+    async with db.pool.acquire() as conn:
+        rows = await conn.fetch(
+            """
+            SELECT * FROM fn_list_sessions_with_user(
+                $1, $2, $3, $4, $5, $6
+            )
+            """,
+            effective_user_id,
+            effective_user_name,
+            effective_user_email,
+            mode.value if mode else None,
+            page,
+            page_size,
+        )
+
+    # Extract total from first row
+    total = rows[0]["total_count"] if rows else 0
+
+    # Convert rows to SessionWithUser
+    data = [
+        SessionWithUser(
+            id=str(row["id"]),
+            name=row["name"],
+            mode=row["mode"],
+            description=row["description"],
+            user_id=row["user_id"],
+            user_name=row["user_name"],
+            user_email=row["user_email"],
+            message_count=row["message_count"] or 0,
+            total_tokens=row["total_tokens"],
+            created_at=row["created_at"],
+            updated_at=row["updated_at"],
+            metadata=row["metadata"],
+        )
+        for row in rows
+    ]
+
+    total_pages = (total + page_size - 1) // page_size if total > 0 else 0
 
     return SessionsQueryResponse(
-        data=result["data"],
+        data=data,
         metadata=PaginationMetadata(
-            total=result["total"],
-            page=result["page"],
-            page_size=result["page_size"],
-            total_pages=result["total_pages"],
-            has_next=result["has_next"],
-            has_previous=result["has_previous"],
+            total=total,
+            page=page,
+            page_size=page_size,
+            total_pages=total_pages,
+            has_next=page < total_pages,
+            has_previous=page > 1,
         ),
     )