api.py

"""
FastAPI endpoints for the Memory Unit.

Provides HTTP interface for:
- Hydrating memory from Google Drive
- Querying for context
- Context injection endpoints for Extension, Task Identifier, Workflow Builder

Auth token flow (from Chrome extension):
- Extension gets token via chrome.identity.getAuthToken()
- Sends in Authorization: Bearer <token> header
- Also sends X-User-Id and X-Thread-Id headers
"""

from typing import Optional, Dict, Any, List
from collections import OrderedDict
import os
import logging
import threading

from fastapi import FastAPI, HTTPException, Header, Depends, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from dotenv import load_dotenv

# Load OPENAI_API_KEY / CONFIDENT_API_KEY / PORT / HOST from .env (memory-unit had no
# config module that did this, so the .env file was previously inert).
load_dotenv()

from memory_unit import MemoryUnit, ContextQueryResult, DriveFolderConfig
from memory_unit.auth import verify_google_token, validation_enabled

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create FastAPI app
app = FastAPI(
    title="Agentic RAG Memory Unit API",
    description="Memory unit with Google Drive integration and agentic RAG",
    version="1.0.0"
)

# CORS middleware. Origins come from the ALLOWED_ORIGINS env var (comma-separated);
# we no longer ship `allow_origins=["*"]` — that is both a tenancy hole and an invalid
# combination with allow_credentials=True (browsers reject a credentialed wildcard).
_allowed_origins_env = os.getenv("ALLOWED_ORIGINS", "")
ALLOWED_ORIGINS = [o.strip() for o in _allowed_origins_env.split(",") if o.strip()] or [
    "http://localhost",
    "http://localhost:3000",
    "http://localhost:8080",
]
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Per-user memory units — one MemoryUnit per user_id, each with its own Chroma
# collection + BM25 + preferences (namespaced in MemoryUnit.__init__). This is the
# real multi-tenant isolation: a request routes only to its own user's unit, so
# there is no shared retrievable state to leak across users (the old single global
# `_memory_unit` + `_owner_user_id` lockout is replaced by this partitioning).
# LRU-ordered so we can evict the least-recently-used unit past MEMORY_MAX_USERS.
_memory_units: "OrderedDict[str, MemoryUnit]" = OrderedDict()
# Guards every read/mutation of `_memory_units` — sync endpoints run in a threadpool,
# so the registry is touched concurrently.
_units_lock = threading.Lock()


def _max_users() -> int:
    """Cap on resident per-user units (LRU-evicted past this). Bounds RAM growth."""
    try:
        return max(1, int(os.getenv("MEMORY_MAX_USERS", "100")))
    except ValueError:
        return 100


def _get_unit_for(
    user_id: str,
    *,
    create: bool = False,
    persist_dir: Optional[str] = None,
    model_name: str = "gpt-4o",
) -> Optional[MemoryUnit]:
    """Return this user's MemoryUnit (marking it most-recently-used).

    ``create=True`` lazily builds and registers one (for /hydrate, /learn, and now
    every read endpoint via ``require_user_unit``); otherwise returns None when the
    user has no unit yet. Evicts the LRU unit if creating pushes the registry over
    ``MEMORY_MAX_USERS``.

    When no ``persist_dir`` is given, fall back to ``MEMORY_PERSIST_DIR`` so a
    container can point storage at a writable path (Cloud Run: ``/tmp``) — the
    in-package default is not reliably writable there. Unset (local dev) keeps the
    in-package default. The base is namespaced per-user inside ``MemoryUnit``.
    """
    if persist_dir is None:
        persist_dir = os.getenv("MEMORY_PERSIST_DIR") or None
    with _units_lock:
        unit = _memory_units.get(user_id)
        if unit is not None:
            _memory_units.move_to_end(user_id)
            return unit
        if not create:
            return None
        unit = MemoryUnit(
            persist_dir=persist_dir, model_name=model_name, user_id=user_id
        )
        _memory_units[user_id] = unit

        # Re-ingest this user's durable "learned" blocks (write-back facts persisted
        # to the learned store, e.g. Postgres `planner.context_blocks`) right away,
        # so a cold instance that has never hydrated this user from Drive can still
        # serve resolve()/query() from what they already taught it. `_reload_learned`
        # only touches attributes MemoryUnit.__init__ has already set (documents,
        # vector_store, keyword_searcher, _learned_store), so it's safe immediately
        # after construction — no I/O happens inside __init__ itself.
        #
        # Double-index check: the only other caller of `_reload_learned` is
        # `hydrate_from_drive()` (core.py), and it always calls `self.clear()` first
        # (wiping vector_store + keyword index + self.documents) before its own
        # reload. So if this same unit later gets hydrated, that reload starts from a
        # clean slate and this constructor-time reload can never combine with it —
        # clear() is the reset point between the two call sites, so blocks reloaded
        # here are never re-added on top of themselves.
        #
        # Best-effort: a learned-store read failure (e.g. Postgres blip) must degrade
        # to an empty-but-usable unit, never break unit creation / turn a read into a
        # 500 further up the call chain.
        try:
            unit._reload_learned()
        except Exception as e:
            logger.error(f"Failed to reload learned context for user {user_id} at unit creation: {e}")

        while len(_memory_units) > _max_users():
            evicted_id, _ = _memory_units.popitem(last=False)  # LRU
            logger.info("Evicted LRU memory unit (cap=%d)", _max_users())
        return unit


# =============================================================================
# Pydantic Models
# =============================================================================

class HydrateRequest(BaseModel):
    root_folder_id: str = Field(..., description="Root Drive folder ID with 2 subfolders")
    persist_dir: Optional[str] = Field(None, description="Chroma persistence directory")
    model_name: str = Field("gpt-4o", description="OpenAI model for agent")


class HydrateResponse(BaseModel):
    status: str
    documents_indexed: int
    folder_structure: Dict[str, Any]
    stats: Dict[str, Any]


class QueryRequest(BaseModel):
    query: str = Field(..., description="Query text to retrieve context")
    n_results: int = Field(5, description="Number of results to retrieve")


class ContextResponse(BaseModel):
    answer: str
    sources: List[Dict[str, Any]]
    context_for_extension: str
    context_for_task_identifier: str
    context_for_workflow_builder: str

    # Machine-generated preference data
    user_preferences: List[str] = []
    task_patterns: List[str] = []
    workflow_trends: List[str] = []


class ResolveRequest(BaseModel):
    fields: List[str] = Field(..., description="Slot/parameter names to resolve to values")
    scope: Optional[List[str]] = Field(
        None, description="Ordered preferred scopes, most-specific first (e.g. [user, org, global])"
    )
    min_score: float = Field(0.0, description="Minimum BM25 score before a field counts as resolved")
    min_coverage: Optional[float] = Field(
        None,
        ge=0.0,
        le=1.0,
        description=(
            "Fraction of the slot name's tokens that must literally appear in the "
            "evidence before it counts as resolved. Omit to use the server default "
            "(MEMORY_RESOLVE_MIN_COVERAGE, itself 1.0); 0.0 accepts any hit."
        ),
    )


class ResolvedSlot(BaseModel):
    field: str
    value: Optional[str] = None
    evidence: Optional[str] = None  # the snippet `value` was extracted from
    source: Optional[str] = None
    confidence: float = 0.0
    scope: Optional[str] = None  # scope label of the winning evidence, if any
    status: str = "missing"


class ResolveResponse(BaseModel):
    slots: List[ResolvedSlot]


class LearnItem(BaseModel):
    text: str = Field(..., description="Distilled fact to remember (write-back)")
    category: Optional[str] = Field(
        None, description="user_preferences|task_patterns|workflow_trends"
    )
    task_id: Optional[str] = Field(None, description="Originating task, for provenance")
    scope: Optional[str] = Field(
        None, description="Scope label for hierarchical resolution, e.g. user|org|global"
    )


class LearnRequest(BaseModel):
    items: List[LearnItem]


class LearnResponse(BaseModel):
    learned: int


class ExtensionContextRequest(BaseModel):
    query: str


class TaskIdentifierContextRequest(BaseModel):
    task_description: str


class WorkflowBuilderContextRequest(BaseModel):
    task_description: str


class PreferencesResponse(BaseModel):
    user_preferences: List[str]
    task_patterns: List[str]
    workflow_trends: List[str]


class StatsResponse(BaseModel):
    is_hydrated: bool
    total_documents: int
    vector_store_count: int
    keyword_index_size: int


class HealthResponse(BaseModel):
    status: str
    memory_unit_initialized: bool
    documents_indexed: int


# =============================================================================
# Dependencies
# =============================================================================

def extract_bearer_token(authorization: Optional[str]) -> Optional[str]:
    """Extract token from Authorization: Bearer <token> header."""
    if not authorization:
        return None
    if authorization.startswith("Bearer "):
        return authorization[7:]
    return authorization


def authed_user(
    x_user_id: Optional[str] = Header(None),
    authorization: Optional[str] = Header(None),
) -> str:
    """Authenticate the caller and return their user_id.

    Requires ``X-User-Id`` always; when token validation is enabled
    (``MEMORY_VALIDATE_TOKEN``, default on) it also requires and verifies the Google
    bearer. This binds every data request to a real identity BEFORE it can reach that
    user's isolated store — essential now that routing is per-user (an unauthenticated
    ``X-User-Id`` would otherwise be a cross-user read vector). No-op bearer check when
    validation is off (offline/tests)."""
    if not x_user_id:
        raise HTTPException(status_code=400, detail="X-User-Id header required")
    if validation_enabled():
        token = extract_bearer_token(authorization)
        if not token:
            raise HTTPException(
                status_code=401,
                detail="Authorization header with Bearer token required",
            )
        verify_google_token(token, x_user_id)
    return x_user_id


def require_user_unit(user_id: str = Depends(authed_user)) -> MemoryUnit:
    """Route to the authenticated caller's own MemoryUnit, lazily creating one.

    A cold instance (fresh process, or this user simply hasn't hit any endpoint yet)
    used to 503 every read here until the caller ran /hydrate — but a user's durable
    "learned" (write-back) context lives in the learned store independent of any
    particular process, so there is no reason a read has to fail just because *this*
    instance hasn't served them before. `_get_unit_for(create=True)` builds the unit
    and re-ingests that learned context (see its docstring), so the unit this returns
    can already answer resolve()/query() from write-back facts with zero Drive I/O.
    A user who has never hydrated *and* never learned anything still gets a unit —
    just an empty one, so reads return normal empty/`missing` results instead of 503.
    """
    return _get_unit_for(user_id, create=True)


# =============================================================================
# Health & Status
# =============================================================================

@app.get("/health", response_model=HealthResponse)
def health_check():
    """Health check endpoint. Reports across all resident per-user units.

    Uses the in-RAM document lists (cheap) rather than hitting Chroma per unit, so a
    frequent liveness probe stays fast."""
    with _units_lock:
        units = list(_memory_units.values())
    total_docs = sum(len(u.documents) for u in units)
    return HealthResponse(
        status="healthy",
        memory_unit_initialized=len(units) > 0,
        documents_indexed=total_docs,
    )


@app.get("/stats", response_model=StatsResponse)
def get_stats(memory: MemoryUnit = Depends(require_user_unit)):
    """Get memory unit statistics for the calling user."""
    return StatsResponse(**memory.get_stats())


# =============================================================================
# Hydration
# =============================================================================

@app.post("/hydrate", response_model=HydrateResponse)
def hydrate_memory(
    request: HydrateRequest,
    authorization: Optional[str] = Header(None),
    x_user_id: Optional[str] = Header(None),
    x_thread_id: Optional[str] = Header(None)
):
    """
    Hydrate the calling user's memory unit from Google Drive.

    Auth token is passed in Authorization: Bearer <token> header (from extension).
    This fetches documents from the 2 subfolders and indexes them in *this user's*
    vector store + keyword index — it never touches another user's unit.
    """
    auth_token = extract_bearer_token(authorization)
    if not auth_token:
        raise HTTPException(status_code=401, detail="Authorization header with Bearer token required")
    if not x_user_id:
        raise HTTPException(status_code=400, detail="X-User-Id header required")

    # Verify the Google token (and that its `sub` matches X-User-Id) before we
    # trust it to read Drive. 401 on a bad token, 503 if Google is unreachable.
    verify_google_token(auth_token, x_user_id)

    try:
        logger.info(f"Hydrating user's memory unit from folder: {request.root_folder_id}")

        # Get-or-create THIS user's unit (namespaced Chroma collection + learned store).
        # The ephemeral auth token is NOT a constructor argument — it is supplied to
        # hydrate_from_drive() at call time. A re-hydrate reuses the existing unit and
        # rebuilds its store internally; other users' units are untouched.
        memory = _get_unit_for(
            x_user_id,
            create=True,
            persist_dir=request.persist_dir,
            model_name=request.model_name,
        )

        memory.folder_config = DriveFolderConfig(
            root_folder_id=request.root_folder_id,
            user_provided_folder_id="",
            machine_generated_folder_id=""
        )

        # Hydrate from Drive (root folder id + ephemeral token).
        result = memory.hydrate_from_drive(
            request.root_folder_id, auth_token, thread_id=x_thread_id
        )

        logger.info(f"Hydrated {result['documents_indexed']} documents")

        return HydrateResponse(**result)

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Hydration failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# =============================================================================
# Query Endpoints
# =============================================================================

@app.post("/query", response_model=ContextResponse)
def query_memory(
    request: QueryRequest,
    memory: MemoryUnit = Depends(require_user_unit),
    x_thread_id: Optional[str] = Header(None),
):
    """
    Query the memory unit using agentic RAG.

    Combines hybrid search over diverse user documents with targeted
    retrieval from machine-generated preference/trend files.
    """
    try:
        result = memory.query(request.query, thread_id=x_thread_id)
        return ContextResponse(
            answer=result.answer,
            sources=result.sources,
            context_for_extension=result.context_for_extension,
            context_for_task_identifier=result.context_for_task_identifier,
            context_for_workflow_builder=result.context_for_workflow_builder,
            user_preferences=result.user_preferences,
            task_patterns=result.task_patterns,
            workflow_trends=result.workflow_trends
        )
    except Exception as e:
        logger.error(f"Query failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/resolve", response_model=ResolveResponse)
def resolve_slots(
    request: ResolveRequest,
    memory: MemoryUnit = Depends(require_user_unit),
    x_thread_id: Optional[str] = Header(None),
):
    """Resolve task parameter slots to concrete values (structured field->value).

    This is the parameter-resolution surface the planner calls to pre-fill task
    slots from user context before falling back to HITL. Unlike /query it returns
    typed slots with source + confidence, not prose. Unresolved fields come back
    with status="missing" so the caller knows to ask the human. Retrieval is scoped
    to the calling user's own unit (`memory` is theirs).
    """
    try:
        results = memory.resolve(
            request.fields,
            user_id=memory.user_id,
            scope=request.scope,
            min_score=request.min_score,
            thread_id=x_thread_id,
            min_coverage=request.min_coverage,
        )
        return ResolveResponse(slots=[ResolvedSlot(**r) for r in results])
    except Exception as e:
        logger.error(f"Resolve failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/learn", response_model=LearnResponse)
def learn_context(
    request: LearnRequest,
    authorization: Optional[str] = Header(None),
    x_user_id: Optional[str] = Header(None),
    x_thread_id: Optional[str] = Header(None),
):
    """Write-back: ingest distilled context learned from completed tasks so future
    resolve()/query() calls benefit. In-repo self-learning; durable Drive
    persistence is a follow-up (extension-owned).

    /learn lazily creates the *calling user's* unit — so a unit can be seeded via
    write-back without a Drive hydrate. Because it writes, it authenticates the
    Google token just like /hydrate (verification is a no-op when MEMORY_VALIDATE_TOKEN
    is off, but a bearer token is still required)."""
    if not x_user_id:
        raise HTTPException(status_code=400, detail="X-User-Id header required")

    # Authenticate before creating the unit (401/503 here must not be swallowed by
    # the 500 handler below).
    auth_token = extract_bearer_token(authorization)
    if not auth_token:
        raise HTTPException(status_code=401, detail="Authorization header with Bearer token required")
    verify_google_token(auth_token, x_user_id)

    try:
        # Get-or-create this user's own unit; user_id is bound in the constructor and
        # scopes the write in the shared pg store (inert for JSONL).
        memory = _get_unit_for(x_user_id, create=True)
        count = memory.learn(
            [item.model_dump() for item in request.items], thread_id=x_thread_id
        )
        return LearnResponse(learned=count)
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Learn failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/vector-search")
def vector_search(
    query: str,
    n_results: int = 5,
    memory: MemoryUnit = Depends(require_user_unit)
):
    """Direct vector search (semantic similarity)."""
    try:
        results = memory.vector_store.query(query, n_results=n_results)
        return {
            "query": query,
            "results": [
                {
                    "content": doc,
                    "metadata": meta,
                    "distance": dist
                }
                for doc, meta, dist in zip(
                    results["documents"][0],
                    results["metadatas"][0],
                    results["distances"][0]
                )
            ]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/keyword-search")
def keyword_search(
    query: str,
    top_k: int = 5,
    memory: MemoryUnit = Depends(require_user_unit)
):
    """Direct keyword search (BM25)."""
    try:
        results = memory.keyword_searcher.search(query, top_k=top_k)
        return {
            "query": query,
            "results": results
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# =============================================================================
# Context Injection Endpoints (for System Diagram integration)
# =============================================================================

@app.post("/context/extension")
def get_extension_context(
    request: ExtensionContextRequest,
    memory: MemoryUnit = Depends(require_user_unit)
):
    """
    Get additional context for Extension component.

    From system diagram:
    [Memory Unit] --> [Additional ctxt for Task/Workflow] --> [Extension]
    """
    try:
        context = memory.get_context_for_extension(request.query)
        return {
            "context": context,
            "target": "extension"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/context/task-identifier")
def get_task_identifier_context(
    request: TaskIdentifierContextRequest,
    memory: MemoryUnit = Depends(require_user_unit)
):
    """
    Get additional context for Task Identifier component.

    From system diagram:
    [Memory Unit] --> [Additional ctxt for Task/Workflow] --> [Task Identifier]
    """
    try:
        context = memory.get_context_for_task_identifier(request.task_description)
        return {
            "context": context,
            "target": "task_identifier"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/context/workflow-builder")
def get_workflow_builder_context(
    request: WorkflowBuilderContextRequest,
    memory: MemoryUnit = Depends(require_user_unit)
):
    """
    Get additional context for Workflow Builder component.

    From system diagram:
    [Memory Unit] --> [Additional ctxt for Workflow] --> [Workflow Builder]
    """
    try:
        context = memory.get_context_for_workflow_builder(request.task_description)
        return {
            "context": context,
            "target": "workflow_builder"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/preferences", response_model=PreferencesResponse)
def get_preferences(
    category: Optional[str] = None,
    memory: MemoryUnit = Depends(require_user_unit)
):
    """
    Get raw machine-generated preferences by category.

    Categories:
    - user_preferences: User's style, habits, likes/dislikes
    - task_patterns: Common task types, frequencies, patterns
    - workflow_trends: Successful workflows, optimization opportunities
    """
    try:
        prefs = memory.get_direct_preferences(category)
        return PreferencesResponse(
            user_preferences=prefs.get("user_preferences", []),
            task_patterns=prefs.get("task_patterns", []),
            workflow_trends=prefs.get("workflow_trends", [])
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# =============================================================================
# Management
# =============================================================================

@app.post("/clear")
def clear_memory(memory: MemoryUnit = Depends(require_user_unit)):
    """Clear the calling user's indexed documents."""
    try:
        memory.clear()
        return {"status": "cleared", "message": "All documents removed from memory"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/refresh")
def refresh_memory(
    authorization: Optional[str] = Header(None),
    x_user_id: Optional[str] = Header(None),
    x_thread_id: Optional[str] = Header(None),
):
    """Refresh the calling user's memory by re-hydrating from Drive."""
    # Routed manually (not via require_user_unit): /refresh may fall back to the token
    # captured at hydrate, so it can't require a fresh bearer up front.
    if not x_user_id:
        raise HTTPException(status_code=400, detail="X-User-Id header required")
    memory = _get_unit_for(x_user_id, create=False)
    if memory is None:
        raise HTTPException(
            status_code=503,
            detail="Memory unit not initialized for this user. Call /hydrate first.",
        )
    try:
        # Use a freshly-supplied token if present, else fall back to the one captured
        # at hydrate. Verify whichever we're about to use: the stored token may have
        # expired since hydrate, so verifying it turns expiry into a clean 401
        # ("re-authenticate") instead of a 500 from a failed Drive call downstream.
        supplied_token = extract_bearer_token(authorization)
        auth_token = supplied_token or memory.auth_token
        if not auth_token:
            raise HTTPException(status_code=401, detail="Authorization header with Bearer token required")
        verify_google_token(auth_token, x_user_id)

        # Resolve the root folder id captured at the last hydrate. Without it there
        # is nothing to refresh — the caller must hydrate first.
        root_folder_id = getattr(memory, "root_folder_id", None)
        if not root_folder_id and memory.folder_config:
            root_folder_id = memory.folder_config.root_folder_id
        if not root_folder_id:
            raise HTTPException(
                status_code=400,
                detail="Nothing to refresh; call /hydrate first."
            )

        # Re-hydrate (clears old data internally) with the root folder + token.
        result = memory.hydrate_from_drive(root_folder_id, auth_token, thread_id=x_thread_id)

        # Spread result first so its "status": "success" does not clobber "refreshed".
        return {
            **result,
            "status": "refreshed"
        }
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# =============================================================================
# Main
# =============================================================================

if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("PORT", 8000))
    host = os.getenv("HOST", "0.0.0.0")

    uvicorn.run(app, host=host, port=port)