remdb 0.3.7__py3-none-any.whl

rem/api/routers/auth.py

@@ -0,0 +1,229 @@
+"""
+OAuth 2.1 Authentication Router.
+
+Leverages Authlib for standards-compliant OAuth/OIDC implementation.
+Minimal custom code - Authlib handles PKCE, token validation, JWKS.
+
+Endpoints:
+- GET  /api/auth/{provider}/login    - Initiate OAuth flow
+- GET  /api/auth/{provider}/callback - OAuth callback
+- POST /api/auth/logout              - Clear session
+- GET  /api/auth/me                  - Current user info
+
+Supported providers:
+- google: Google OAuth 2.0 / OIDC
+- microsoft: Microsoft Entra ID OIDC
+
+Design Pattern (OAuth 2.1 + PKCE):
+1. User clicks "Login with Google"
+2. /login generates state + PKCE code_verifier
+3. Store code_verifier in session
+4. Redirect to provider with code_challenge
+5. User authenticates and grants consent
+6. Provider redirects to /callback with code
+7. Exchange code + code_verifier for tokens
+8. Validate ID token signature with JWKS
+9. Store user info in session
+10. Redirect to application
+
+Dependencies:
+    pip install authlib httpx
+
+Environment variables:
+    AUTH__ENABLED=true
+    AUTH__SESSION_SECRET=<random-secret>
+    AUTH__GOOGLE__CLIENT_ID=<google-client-id>
+    AUTH__GOOGLE__CLIENT_SECRET=<google-client-secret>
+    AUTH__MICROSOFT__CLIENT_ID=<microsoft-client-id>
+    AUTH__MICROSOFT__CLIENT_SECRET=<microsoft-client-secret>
+    AUTH__MICROSOFT__TENANT=common
+
+References:
+- Authlib: https://docs.authlib.org/en/latest/
+- OAuth 2.1: https://datatracker.ietf.org/doc/html/draft-ietf-oauth-v2-1-11
+"""
+
+from fastapi import APIRouter, HTTPException, Request
+from fastapi.responses import RedirectResponse
+from authlib.integrations.starlette_client import OAuth
+from loguru import logger
+
+from ...settings import settings
+
+router = APIRouter(prefix="/api/auth", tags=["auth"])
+
+# Initialize Authlib OAuth client
+# Authlib handles PKCE, state, nonce, token validation automatically
+oauth = OAuth()
+
+# Register Google provider
+if settings.auth.google.client_id:
+    oauth.register(
+        name="google",
+        client_id=settings.auth.google.client_id,
+        client_secret=settings.auth.google.client_secret,
+        server_metadata_url="https://accounts.google.com/.well-known/openid-configuration",
+        client_kwargs={
+            "scope": "openid email profile",
+            # Authlib automatically adds PKCE to authorization request
+        },
+    )
+    logger.info("Google OAuth provider registered")
+
+# Register Microsoft provider
+if settings.auth.microsoft.client_id:
+    tenant = settings.auth.microsoft.tenant
+    oauth.register(
+        name="microsoft",
+        client_id=settings.auth.microsoft.client_id,
+        client_secret=settings.auth.microsoft.client_secret,
+        server_metadata_url=f"https://login.microsoftonline.com/{tenant}/v2.0/.well-known/openid-configuration",
+        client_kwargs={
+            "scope": "openid email profile User.Read",
+        },
+    )
+    logger.info(f"Microsoft OAuth provider registered (tenant: {tenant})")
+
+
+@router.get("/{provider}/login")
+async def login(provider: str, request: Request):
+    """
+    Initiate OAuth flow with provider.
+
+    Authlib automatically:
+    - Generates state for CSRF protection
+    - Generates PKCE code_verifier and code_challenge
+    - Stores state and code_verifier in session
+    - Redirects to provider's authorization endpoint
+
+    Args:
+        provider: OAuth provider (google, microsoft)
+        request: FastAPI request (for session access)
+
+    Returns:
+        Redirect to provider's authorization page
+    """
+    if not settings.auth.enabled:
+        raise HTTPException(status_code=501, detail="Authentication is disabled")
+
+    # Get OAuth client for provider
+    client = oauth.create_client(provider)
+    if not client:
+        raise HTTPException(status_code=400, detail=f"Unknown provider: {provider}")
+
+    # Get redirect URI from settings
+    if provider == "google":
+        redirect_uri = settings.auth.google.redirect_uri
+    elif provider == "microsoft":
+        redirect_uri = settings.auth.microsoft.redirect_uri
+    else:
+        raise HTTPException(status_code=400, detail=f"Unknown provider: {provider}")
+
+    # Authlib authorize_redirect() automatically:
+    # - Generates state parameter
+    # - Generates PKCE code_verifier and code_challenge
+    # - Stores state and code_verifier in session
+    # - Builds authorization URL with all required parameters
+    return await client.authorize_redirect(request, redirect_uri)
+
+
+@router.get("/{provider}/callback")
+async def callback(provider: str, request: Request):
+    """
+    OAuth callback endpoint.
+
+    Authlib automatically:
+    - Validates state parameter (CSRF protection)
+    - Exchanges code for tokens with PKCE code_verifier
+    - Validates ID token signature with JWKS
+    - Verifies ID token claims (iss, aud, exp, nonce)
+
+    Args:
+        provider: OAuth provider (google, microsoft)
+        request: FastAPI request (for session and query params)
+
+    Returns:
+        Redirect to application home page
+    """
+    if not settings.auth.enabled:
+        raise HTTPException(status_code=501, detail="Authentication is disabled")
+
+    # Get OAuth client for provider
+    client = oauth.create_client(provider)
+    if not client:
+        raise HTTPException(status_code=400, detail=f"Unknown provider: {provider}")
+
+    try:
+        # Authlib authorize_access_token() automatically:
+        # - Validates state from session (CSRF)
+        # - Retrieves code_verifier from session
+        # - Exchanges authorization code for tokens
+        # - Validates ID token signature with JWKS
+        # - Verifies ID token claims
+        token = await client.authorize_access_token(request)
+
+        # Parse user info from ID token or call userinfo endpoint
+        # Authlib parses ID token claims automatically
+        user_info = token.get("userinfo")
+        if not user_info:
+            # Fetch from userinfo endpoint if not in ID token
+            user_info = await client.userinfo(token=token)
+
+        # Store user info in session
+        request.session["user"] = {
+            "provider": provider,
+            "sub": user_info.get("sub"),
+            "email": user_info.get("email"),
+            "name": user_info.get("name"),
+            "picture": user_info.get("picture"),
+        }
+
+        # Store tokens in session for API access
+        request.session["tokens"] = {
+            "access_token": token.get("access_token"),
+            "refresh_token": token.get("refresh_token"),
+            "expires_at": token.get("expires_at"),
+        }
+
+        logger.info(f"User authenticated: {user_info.get('email')} via {provider}")
+
+        # Redirect to application
+        # TODO: Support custom redirect URL from state parameter
+        return RedirectResponse(url="/")
+
+    except Exception as e:
+        logger.error(f"OAuth callback error: {e}")
+        raise HTTPException(status_code=400, detail=f"Authentication failed: {str(e)}")
+
+
+@router.post("/logout")
+async def logout(request: Request):
+    """
+    Clear user session.
+
+    Args:
+        request: FastAPI request
+
+    Returns:
+        Success message
+    """
+    request.session.clear()
+    return {"message": "Logged out successfully"}
+
+
+@router.get("/me")
+async def me(request: Request):
+    """
+    Get current user information from session.
+
+    Args:
+        request: FastAPI request
+
+    Returns:
+        User information or 401 if not authenticated
+    """
+    user = request.session.get("user")
+    if not user:
+        raise HTTPException(status_code=401, detail="Not authenticated")
+
+    return user

rem/api/routers/chat/__init__.py

@@ -0,0 +1,5 @@
+"""Chat completions router with OpenAI-compatible API."""
+
+from .completions import router
+
+__all__ = ["router"]

rem/api/routers/chat/completions.py

@@ -0,0 +1,281 @@
+"""
+OpenAI-compatible chat completions router for REM.
+
+Design Pattern:
+- Headers map to AgentContext (X-User-Id, X-Tenant-Id, X-Session-Id, X-Agent-Schema)
+- ContextBuilder centralizes message construction with user profile + session history
+- Body.model is the LLM model for Pydantic AI
+- X-Agent-Schema header specifies which agent schema to use (defaults to 'rem')
+- Support for streaming (SSE) and non-streaming modes
+- Response format control (text vs json_object)
+
+Context Building Flow:
+1. ContextBuilder.build_from_headers() extracts user_id, session_id from headers
+2. Session history ALWAYS loaded with compression (if session_id provided)
+   - Uses SessionMessageStore with compression to keep context efficient
+   - Long messages include REM LOOKUP hints: "... [REM LOOKUP session-{id}-msg-{index}] ..."
+   - Agent can retrieve full content on-demand using REM LOOKUP
+3. User profile provided as REM LOOKUP hint (on-demand by default)
+   - Agent receives: "User ID: {user_id}. To load user profile: Use REM LOOKUP users/{user_id}"
+   - Agent decides whether to load profile based on query
+4. If CHAT__AUTO_INJECT_USER_CONTEXT=true: User profile auto-loaded and injected
+5. Combines: system context + compressed session history + new messages
+6. Agent receives complete message list ready for execution
+
+Headers Mapping
+    X-User-Id        → AgentContext.user_id
+    X-Tenant-Id      → AgentContext.tenant_id
+    X-Session-Id     → AgentContext.session_id
+    X-Model-Name     → AgentContext.default_model (overrides body.model)
+    X-Agent-Schema   → AgentContext.agent_schema_uri (defaults to 'rem')
+
+Default Agent:
+    If X-Agent-Schema header is not provided, the system loads 'rem' schema,
+    which is the REM expert assistant with comprehensive knowledge about:
+    - REM architecture and concepts
+    - Entity types and graph traversal
+    - REM queries (LOOKUP, FUZZY, TRAVERSE)
+    - Agent development with Pydantic AI
+    - Cloud infrastructure (EKS, Karpenter, CloudNativePG)
+
+Example Request:
+    POST /api/v1/chat/completions
+    X-Tenant-Id: acme-corp
+    X-User-Id: user123
+    X-Agent-Schema: rem  # Optional, this is the default
+
+    {
+      "model": "openai:gpt-4o-mini",
+      "messages": [
+        {"role": "user", "content": "How do I create a new REM entity?"}
+      ],
+      "stream": true
+    }
+"""
+
+import base64
+import tempfile
+import time
+import uuid
+from datetime import datetime
+from pathlib import Path
+
+from fastapi import APIRouter, Request
+from fastapi.responses import StreamingResponse
+from loguru import logger
+
+from ....agentic.context import AgentContext
+from ....agentic.context_builder import ContextBuilder
+from ....agentic.providers.pydantic_ai import create_agent
+from ....services.audio.transcriber import AudioTranscriber
+from ....services.session import SessionMessageStore, reload_session
+from ....settings import settings
+from ....utils.schema_loader import load_agent_schema
+from .json_utils import extract_json_resilient
+from .models import (
+    ChatCompletionChoice,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionUsage,
+    ChatMessage,
+)
+from .streaming import stream_openai_response
+
+router = APIRouter(prefix="/v1", tags=["chat"])
+
+# Default agent schema file
+DEFAULT_AGENT_SCHEMA = "rem"
+
+
+@router.post("/chat/completions", response_model=None)
+async def chat_completions(body: ChatCompletionRequest, request: Request):
+    """
+    OpenAI-compatible chat completions with REM agent support.
+
+    The 'model' field in the request body is the LLM model used by Pydantic AI.
+    The X-Agent-Schema header specifies which agent schema to use (defaults to 'rem').
+
+    Supported Headers:
+    | Header              | Description                          | Maps To                        | Default       |
+    |---------------------|--------------------------------------|--------------------------------|---------------|
+    | X-User-Id           | User identifier                      | AgentContext.user_id           | None          |
+    | X-Tenant-Id         | Tenant identifier (multi-tenancy)    | AgentContext.tenant_id         | "default"     |
+    | X-Session-Id        | Session/conversation identifier      | AgentContext.session_id        | None          |
+    | X-Agent-Schema      | Agent schema name                    | AgentContext.agent_schema_uri  | "rem"         |
+
+    Example Models:
+    - anthropic:claude-sonnet-4-5-20250929 (Claude 4.5 Sonnet)
+    - anthropic:claude-3-7-sonnet-20250219 (Claude 3.7 Sonnet)
+    - anthropic:claude-3-5-haiku-20241022 (Claude 3.5 Haiku)
+    - openai:gpt-4.1-turbo
+    - openai:gpt-4o
+    - openai:gpt-4o-mini
+
+    Response Formats:
+    - text (default): Plain text response
+    - json_object: Best-effort JSON extraction from agent output
+
+    Default Agent (rem):
+    - Expert assistant for REM system
+    - Comprehensive knowledge of REM architecture, concepts, and implementation
+    - Structured output with answer, confidence, and references
+
+    Session Management:
+    - Session history ALWAYS loaded with compression when X-Session-Id provided
+    - Uses SessionMessageStore with REM LOOKUP hints for long messages
+    - User profile provided as REM LOOKUP hint (on-demand by default)
+    - If CHAT__AUTO_INJECT_USER_CONTEXT=true: User profile auto-loaded and injected
+    - New messages saved to database with compression for session continuity
+    - When Postgres is disabled, session management is skipped
+    """
+    # Load agent schema: use header value from context or default
+    # Extract AgentContext first to get schema name
+    temp_context = AgentContext.from_headers(dict(request.headers))
+    schema_name = temp_context.agent_schema_uri or DEFAULT_AGENT_SCHEMA
+
+    # Load schema using centralized utility
+    try:
+        agent_schema = load_agent_schema(schema_name)
+    except FileNotFoundError:
+        # Fallback to default if specified schema not found
+        logger.warning(f"Schema '{schema_name}' not found, falling back to '{DEFAULT_AGENT_SCHEMA}'")
+        schema_name = DEFAULT_AGENT_SCHEMA
+        try:
+            agent_schema = load_agent_schema(schema_name)
+        except FileNotFoundError:
+            # No schema available at all
+            from fastapi import HTTPException
+
+            raise HTTPException(
+                status_code=500,
+                detail=f"Agent schema '{schema_name}' not found and default schema unavailable",
+            )
+
+    logger.info(f"Using agent schema: {schema_name}, model: {body.model}")
+
+    # Check for audio input
+    is_audio = request.headers.get("x-chat-is-audio", "").lower() == "true"
+
+    # Process messages (transcribe audio if needed)
+    new_messages = [msg.model_dump() for msg in body.messages]
+
+    if is_audio and new_messages and new_messages[0]["role"] == "user":
+        # First user message should be base64-encoded audio
+        try:
+            audio_b64 = new_messages[0]["content"]
+            audio_bytes = base64.b64decode(audio_b64)
+
+            # Write to temp file for transcription
+            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
+                tmp_file.write(audio_bytes)
+                tmp_path = tmp_file.name
+
+            # Transcribe audio
+            transcriber = AudioTranscriber()
+            result = transcriber.transcribe_file(tmp_path)
+
+            # Replace audio content with transcribed text
+            new_messages[0]["content"] = result.text
+            logger.info(f"Transcribed audio: {len(result.text)} characters")
+
+            # Clean up temp file
+            Path(tmp_path).unlink()
+
+        except Exception as e:
+            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)
+    context, messages = await ContextBuilder.build_from_headers(
+        headers=dict(request.headers),
+        new_messages=new_messages,
+    )
+
+    logger.info(f"Built context with {len(messages)} total messages (includes history + user context)")
+
+    # Create agent with schema and model override
+    agent = await create_agent(
+        context=context,
+        agent_schema_override=agent_schema,
+        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)
+
+    # Generate OpenAI-compatible request ID
+    request_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
+
+    # Streaming mode
+    if body.stream:
+        return StreamingResponse(
+            stream_openai_response(agent, prompt, body.model, request_id),
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
+        )
+
+    # Non-streaming mode
+    result = await agent.run(prompt)
+
+    # Determine content format based on response_format request
+    if body.response_format and body.response_format.type == "json_object":
+        # JSON mode: Best-effort extraction of JSON from agent output
+        content = extract_json_resilient(result.output)  # type: ignore[attr-defined]
+    else:
+        # Text mode: Return as string (handle structured output)
+        from rem.agentic.serialization import serialize_agent_result_json
+        content = serialize_agent_result_json(result.output)  # type: ignore[attr-defined]
+
+    # Get usage from result if available
+    usage = result.usage() if hasattr(result, "usage") else None
+    prompt_tokens = usage.input_tokens if usage else 0
+    completion_tokens = usage.output_tokens if usage else 0
+
+    # Save conversation messages to database (if session_id and postgres enabled)
+    if settings.postgres.enabled and context.session_id:
+        # Extract just the new user message (last message from body)
+        user_message = {
+            "role": "user",
+            "content": body.messages[-1].content if body.messages else "",
+            "timestamp": datetime.utcnow().isoformat(),
+        }
+
+        assistant_message = {
+            "role": "assistant",
+            "content": content,
+            "timestamp": datetime.utcnow().isoformat(),
+        }
+
+        # Store messages with compression
+        store = SessionMessageStore(user_id=context.user_id or "default")
+
+        await store.store_session_messages(
+            session_id=context.session_id,
+            messages=[user_message, assistant_message],
+            user_id=context.user_id,
+            compress=True,
+        )
+
+        logger.info(f"Saved conversation to session {context.session_id}")
+
+    return ChatCompletionResponse(
+        id=request_id,
+        created=int(time.time()),
+        model=body.model,  # Echo back the requested model
+        choices=[
+            ChatCompletionChoice(
+                index=0,
+                message=ChatMessage(role="assistant", content=content),
+                finish_reason="stop",
+            )
+        ],
+        usage=ChatCompletionUsage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_tokens + completion_tokens,
+        ),
+    )

rem/api/routers/chat/json_utils.py

@@ -0,0 +1,76 @@
+"""
+JSON extraction utilities for response_format='json_object' mode.
+
+Design Pattern:
+- Best-effort JSON extraction from agent output
+- Handles fenced code blocks (```json ... ```)
+- Handles raw JSON objects
+- Graceful fallback to string if extraction fails
+"""
+
+import json
+import re
+
+
+def extract_json_resilient(output: str | dict | list) -> str:
+    """
+    Extract JSON from agent output with multiple fallback strategies.
+
+    Strategies (in order):
+    1. If already dict/list, serialize directly
+    2. Extract from fenced JSON code blocks (```json ... ```)
+    3. Find JSON object/array in text ({...} or [...])
+    4. Return as-is if all strategies fail
+
+    Args:
+        output: Agent output (str, dict, or list)
+
+    Returns:
+        JSON string (best-effort)
+
+    Examples:
+        >>> extract_json_resilient({"answer": "test"})
+        '{"answer": "test"}'
+
+        >>> extract_json_resilient('Here is the result:\\n```json\\n{"answer": "test"}\\n```')
+        '{"answer": "test"}'
+
+        >>> extract_json_resilient('The answer is {"answer": "test"} as shown above.')
+        '{"answer": "test"}'
+    """
+    # Strategy 1: Already structured
+    if isinstance(output, (dict, list)):
+        return json.dumps(output)
+
+    text = str(output)
+
+    # Strategy 2: Extract from fenced code blocks
+    fenced_match = re.search(r"```json\s*\n(.*?)\n```", text, re.DOTALL)
+    if fenced_match:
+        try:
+            json_str = fenced_match.group(1).strip()
+            # Validate it's valid JSON
+            json.loads(json_str)
+            return json_str
+        except json.JSONDecodeError:
+            pass
+
+    # Strategy 3: Find JSON object or array
+    # Look for {...} or [...]
+    for pattern in [
+        r"\{[^{}]*\}",  # Simple object
+        r"\{.*\}",  # Nested object
+        r"\[.*\]",  # Array
+    ]:
+        match = re.search(pattern, text, re.DOTALL)
+        if match:
+            try:
+                json_str = match.group(0)
+                # Validate it's valid JSON
+                json.loads(json_str)
+                return json_str
+            except json.JSONDecodeError:
+                continue
+
+    # Strategy 4: Fallback to string
+    return text