albert-memory 0.4.0__py3-none-any.whl

albert/api.py

@@ -0,0 +1,698 @@
+from __future__ import annotations
+
+import logging
+import math
+from contextlib import asynccontextmanager
+from datetime import UTC, datetime
+from uuid import UUID
+
+import uvicorn
+from fastapi import Depends, FastAPI, Header, HTTPException, Response, status
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+from albert import __version__
+from albert.capture import router as capture_router
+from albert.config import get_settings
+from albert.console import router as console_router
+from albert.console_backends import router as console_backends_router
+from albert.console_map import router as console_map_router
+from albert.console_ops import router as console_ops_router
+from albert.credentials import onboard_router, settings_router
+from albert.db import get_session
+from albert.graph_store import get_graph_store
+from albert.models import Episode
+from albert.portability import export_bundle, import_bundle
+from albert.provenance import normalize_provenance
+from albert.recorder import CANDIDATE_LIMIT, cap_candidates, current_recorder, span
+from albert.schemas import (
+    ContextRequest,
+    ContextResponse,
+    EpisodeCreate,
+    EpisodeRead,
+    ExportBundle,
+    GraphQuery,
+    HealthResponse,
+    ImportRequest,
+    ImportResult,
+    LockAcquire,
+    LockRead,
+    LockRelease,
+    LockRenew,
+    MemoryCreate,
+    MemoryRead,
+    MemorySafetyReview,
+    MemoryUpdate,
+    RelationshipCreate,
+    RelationshipRead,
+    SearchRequest,
+    SearchResponse,
+    SessionContextRequest,
+    SessionContextResponse,
+    SubgraphQuery,
+    WorkingMemoryCreate,
+    WorkingMemoryCreateRead,
+    WorkingMemoryFinish,
+    WorkingMemoryRead,
+    WorkingMemoryUpdate,
+)
+from albert.scope import covers
+from albert.search import hybrid_search
+from albert.security import AuthContext, authenticate
+from albert.services import (
+    acquire_lock,
+    add_explicit_relationship,
+    audit,
+    create_episode,
+    create_memory,
+    create_working_memory,
+    delete_episode,
+    delete_memory,
+    episode_read,
+    finish_working_memory,
+    get_episode,
+    get_memory,
+    get_working_memory,
+    release_lock,
+    renew_lock,
+    resolve_workspace,
+    review_memory_safety,
+    update_memory,
+    update_working_memory,
+    wait_for_indexing,
+)
+from albert.session_context import assemble_session_context
+from albert.trace_writer import get_trace_writer
+from albert.tracing import RecordingMiddleware
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("albert.api")
+
+
+@asynccontextmanager
+async def lifespan(_app: FastAPI):  # type: ignore[no-untyped-def]
+    logger.info("Albert API starting", extra={"version": __version__})
+    writer = get_trace_writer()
+    writer.start()
+    yield
+    writer.stop()
+
+
+app = FastAPI(
+    title="Albert Memory API",
+    version=__version__,
+    description="Independent hybrid memory service for AI agents",
+    lifespan=lifespan,
+)
+app.add_middleware(RecordingMiddleware)
+app.include_router(capture_router)
+app.include_router(console_router)
+app.include_router(console_backends_router)
+app.include_router(console_map_router)
+app.include_router(console_ops_router)
+app.include_router(settings_router)
+app.include_router(onboard_router)
+
+
+def _memory_ids(result: SearchResponse) -> list[str]:
+    """Final memory hits in rank order, ids only, for the trace summary."""
+    ids: list[str] = []
+    for hit in result.hits:
+        if hit.memory_id is not None and str(hit.memory_id) not in ids:
+            ids.append(str(hit.memory_id))
+    return ids[:CANDIDATE_LIMIT]
+
+
+def _note(**fields: object) -> None:
+    recorder = current_recorder()
+    if recorder is not None:
+        recorder.note(**fields)
+
+
+@app.get("/v1/health/live", response_model=HealthResponse, tags=["health"])
+def live() -> HealthResponse:
+    from albert.graph_builders import builder_status
+
+    settings = get_settings()
+    builder = builder_status()
+    return HealthResponse(
+        status="ok",
+        database="unchecked",
+        version=__version__,
+        embedding_provider=settings.embedding_provider,
+        embedding_model=settings.embedding_model,
+        embedding_dimensions=settings.embedding_dimensions,
+        classifier_provider=settings.llm_provider,
+        graph_store=get_graph_store().name,
+        graph_builder=settings.graph_builder,
+        graph_builder_version=builder["version"] or "unavailable",
+        graph_builder_error=builder["error"],
+    )
+
+
+@app.get("/v1/health/ready", response_model=HealthResponse, tags=["health"])
+def ready(session: Session = Depends(get_session)) -> HealthResponse:
+    try:
+        session.execute(text("SELECT 1"))
+    except Exception as exc:
+        raise HTTPException(status_code=503, detail="Database unavailable") from exc
+    from albert.graph_builders import builder_status
+
+    settings = get_settings()
+    builder = builder_status()
+    try:
+        get_graph_store().ping(session)
+    except Exception as exc:
+        raise HTTPException(
+            status_code=503, detail=f"Graph backend unavailable ({type(exc).__name__})"
+        ) from exc
+    return HealthResponse(
+        status="ok",
+        database="ok",
+        version=__version__,
+        embedding_provider=settings.embedding_provider,
+        embedding_model=settings.embedding_model,
+        embedding_dimensions=settings.embedding_dimensions,
+        classifier_provider=settings.llm_provider,
+        graph_store=get_graph_store().name,
+        graph_builder=settings.graph_builder,
+        graph_builder_version=builder["version"] or "unavailable",
+        graph_builder_error=builder["error"],
+    )
+
+
+def _preferred_wait(prefer: str | None) -> float | None:
+    """Seconds from an RFC 7240 `Prefer: wait=N` header, capped by configuration."""
+    for preference in (prefer or "").split(","):
+        name, _, value = preference.strip().partition("=")
+        if name.strip().lower() == "wait":
+            try:
+                seconds = float(value.strip().strip('"'))
+            except ValueError:
+                return None
+            if seconds != seconds or seconds < 0:  # NaN or negative
+                return None
+            return min(seconds, float(get_settings().max_index_wait_seconds))
+    return None
+
+
+def _episode_response(
+    session: Session, episode: Episode, response: Response, prefer: str | None
+) -> EpisodeRead:
+    """The episode with its indexing status, after waiting if the caller asked to.
+
+    Location always points at the episode; Retry-After is set while it is not
+    yet searchable, so a client that cannot wait knows when to look again.
+    """
+    wait = _preferred_wait(prefer)
+    if wait is not None:
+        view = wait_for_indexing(session, episode, wait)
+        response.headers["Preference-Applied"] = f"wait={wait:g}"
+    else:
+        view = episode_read(session, episode)
+    response.headers["Location"] = f"/v1/episodes/{view.id}"
+    indexing = view.indexing
+    if indexing is not None and indexing.state not in ("complete", "failed"):
+        retry_after = 1
+        if indexing.next_attempt_at is not None:
+            delta = indexing.next_attempt_at.replace(
+                tzinfo=indexing.next_attempt_at.tzinfo or UTC
+            ) - datetime.now(UTC)
+            retry_after = max(1, math.ceil(delta.total_seconds()))
+        response.headers["Retry-After"] = str(retry_after)
+    return view
+
+
+@app.post("/v1/episodes", response_model=EpisodeRead, status_code=201, tags=["episodes"])
+def post_episode(
+    data: EpisodeCreate,
+    response: Response,
+    prefer: str | None = Header(default=None),
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> EpisodeRead:
+    episode, created = create_episode(session, auth, data)
+    if not created:
+        response.status_code = status.HTTP_200_OK
+    return _episode_response(session, episode, response, prefer)
+
+
+@app.get("/v1/episodes/{episode_id}", response_model=EpisodeRead, tags=["episodes"])
+def read_episode(
+    episode_id: UUID,
+    response: Response,
+    prefer: str | None = Header(default=None),
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> EpisodeRead:
+    episode = get_episode(session, auth, episode_id)
+    audit(session, auth, "episode.read", resource_type="episode", resource_id=str(episode.id))
+    session.commit()
+    return _episode_response(session, episode, response, prefer)
+
+
+@app.delete("/v1/episodes/{episode_id}", status_code=204, tags=["episodes"])
+def remove_episode(
+    episode_id: UUID,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> Response:
+    delete_episode(session, auth, episode_id)
+    return Response(status_code=status.HTTP_204_NO_CONTENT)
+
+
+@app.post("/v1/memories", response_model=MemoryRead, status_code=201, tags=["memories"])
+def post_memory(
+    data: MemoryCreate,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+):  # type: ignore[no-untyped-def]
+    memory = create_memory(session, auth, data)
+    _note(stored_ids=[str(memory.id)])
+    return memory
+
+
+@app.get("/v1/memories/{memory_id}", response_model=MemoryRead, tags=["memories"])
+def read_memory(
+    memory_id: UUID,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+):  # type: ignore[no-untyped-def]
+    memory = get_memory(session, auth, memory_id)
+    audit(session, auth, "memory.read", resource_type="memory", resource_id=str(memory.id))
+    session.commit()
+    return memory
+
+
+@app.patch("/v1/memories/{memory_id}", response_model=MemoryRead, tags=["memories"])
+def patch_memory(
+    memory_id: UUID,
+    data: MemoryUpdate,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+):  # type: ignore[no-untyped-def]
+    return update_memory(session, auth, memory_id, data)
+
+
+@app.post(
+    "/v1/memories/{memory_id}/quarantine/review",
+    response_model=MemoryRead,
+    tags=["memories"],
+)
+def review_memory_quarantine(
+    memory_id: UUID,
+    data: MemorySafetyReview,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+):  # type: ignore[no-untyped-def]
+    return review_memory_safety(session, auth, memory_id, data)
+
+
+@app.delete("/v1/memories/{memory_id}", status_code=204, tags=["memories"])
+def remove_memory(
+    memory_id: UUID,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> Response:
+    delete_memory(session, auth, memory_id)
+    return Response(status_code=status.HTTP_204_NO_CONTENT)
+
+
+@app.post("/v1/search", response_model=SearchResponse, tags=["retrieval"])
+def search_memories(
+    data: SearchRequest,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> SearchResponse:
+    auth.require("memory.read")
+    workspace_id = resolve_workspace(session, auth, data.workspace_id)
+    allowed_sensitivities = auth.allowed_sensitivities()
+    if data.sensitivity and not set(data.sensitivity).issubset(allowed_sensitivities):
+        raise HTTPException(status_code=403, detail="Requested sensitivity is not authorized")
+    effective = data.model_copy(update={"sensitivity": data.sensitivity or allowed_sensitivities})
+    _note(query=data.query)
+    result = hybrid_search(
+        session,
+        organization_id=auth.organization_id,
+        workspace_id=workspace_id,
+        request=effective,
+        allowed_sensitivities=allowed_sensitivities,
+    )
+    _note(hits=len(result.hits), degraded=result.degraded, memory_ids=_memory_ids(result))
+    audit(
+        session,
+        auth,
+        "memory.searched",
+        detail={"result_count": len(result.hits), "degraded": result.degraded},
+    )
+    session.commit()
+    return result
+
+
+@app.post("/v1/context/assemble", response_model=ContextResponse, tags=["retrieval"])
+def assemble_context(
+    data: ContextRequest,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> ContextResponse:
+    auth.require("memory.read")
+    workspace_id = resolve_workspace(session, auth, data.workspace_id)
+    allowed_sensitivities = auth.allowed_sensitivities()
+    if data.sensitivity and not set(data.sensitivity).issubset(allowed_sensitivities):
+        raise HTTPException(status_code=403, detail="Requested sensitivity is not authorized")
+    effective = data.model_copy(update={"sensitivity": data.sensitivity or allowed_sensitivities})
+    _note(query=data.query)
+    result = hybrid_search(
+        session,
+        organization_id=auth.organization_id,
+        workspace_id=workspace_id,
+        request=effective,
+        allowed_sensitivities=allowed_sensitivities,
+    )
+    _note(hits=len(result.hits), degraded=result.degraded, memory_ids=_memory_ids(result))
+    sections: list[str] = []
+    citations: list[dict] = []
+    used = 0
+    for index, hit in enumerate(result.hits, start=1):
+        section = f"[{index}] {hit.subject}\n{hit.content}".strip()
+        if used + len(section) > data.max_characters:
+            break
+        sections.append(section)
+        used += len(section) + 2
+        citations.append(
+            {
+                "index": index,
+                "kind": hit.kind,
+                "memory_id": str(hit.memory_id) if hit.memory_id else None,
+                "episode_id": str(hit.source_episode_id) if hit.source_episode_id else None,
+                "backends": hit.backends,
+            }
+        )
+    audit(
+        session,
+        auth,
+        "context.assembled",
+        detail={"citations": len(citations), "characters": used, "degraded": result.degraded},
+    )
+    session.commit()
+    return ContextResponse(
+        context="\n\n".join(sections), citations=citations, degraded=result.degraded
+    )
+
+
+@app.post("/v1/context/session-start", response_model=SessionContextResponse, tags=["retrieval"])
+def session_start_context(
+    data: SessionContextRequest,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> SessionContextResponse:
+    auth.require("memory.read")
+    workspace_id = resolve_workspace(session, auth, data.workspace_id)
+    allowed_sensitivities = auth.allowed_sensitivities()
+    if data.sensitivity and not set(data.sensitivity).issubset(allowed_sensitivities):
+        raise HTTPException(status_code=403, detail="Requested sensitivity is not authorized")
+    if data.provenance is not None:
+        # Same normalization as capture, so a hook's recall scope matches the
+        # project identity its captured memories were stored under.
+        observed = normalize_provenance(
+            data.provenance,
+            organization_id=auth.organization_id,
+            allow_absolute_path=False,
+        )
+        data = data.model_copy(update={
+            "project_ref": data.project_ref or observed.get("project_id"),
+            "branch_ref": data.branch_ref or observed.get("branch") or observed.get("ref"),
+        })
+    # Task and run are ranking signals here, not hard filters: durable project truth
+    # remains eligible alongside the most specific current-task evidence.
+    request = SearchRequest(
+        query=data.query,
+        limit=min(100, max(data.max_items * 4, 20)),
+        workspace_id=data.workspace_id,
+        project_ref=data.project_ref,
+        task_ref=None,
+        run_ref=None,
+        branch_ref=data.branch_ref,
+        memory_types=data.memory_types,
+        sensitivity=data.sensitivity or allowed_sensitivities,
+        include_graph=False,
+    )
+    result = hybrid_search(
+        session,
+        organization_id=auth.organization_id,
+        workspace_id=workspace_id,
+        request=request,
+        allowed_sensitivities=allowed_sensitivities,
+    )
+    response = assemble_session_context(result, data)
+    _note(
+        capture_stage="injection",
+        injection_characters=response.characters,
+        injection_tokens=response.tokens,
+        injection_citations=len(response.citations),
+        injection_truncated=response.truncated,
+    )
+    audit(
+        session,
+        auth,
+        "context.session_started",
+        detail={
+            "citations": len(response.citations),
+            "characters": response.characters,
+            "tokens": response.tokens,
+            "truncated": response.truncated,
+            "degraded": response.degraded,
+        },
+    )
+    session.commit()
+    return response
+
+
+@app.post("/v1/relationships", response_model=RelationshipRead, status_code=201, tags=["graph"])
+def post_relationship(
+    data: RelationshipCreate,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> RelationshipRead:
+    return add_explicit_relationship(session, auth, data).read()
+
+
+@app.post("/v1/graph/query", response_model=list[RelationshipRead], tags=["graph"])
+def query_graph(
+    data: GraphQuery,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> list[RelationshipRead]:
+    auth.require("graph.query")
+    workspace_id = resolve_workspace(session, auth, data.workspace_id)
+    _note(query=data.query)
+    with span("graph") as handle:
+        result = get_graph_store().query(
+            session,
+            organization_id=auth.organization_id,
+            workspace_id=workspace_id,
+            query=data.query,
+            temporal_as_of=data.temporal_as_of,
+            limit=data.limit,
+            sensitivities=auth.allowed_sensitivities(),
+        )
+        handle.set(
+            candidates_total=len(result),
+            candidates=cap_candidates(
+                [
+                    {"id": str(r.id), "kind": "relationship", "score": r.confidence, "rank": i}
+                    for i, r in enumerate(result, start=1)
+                ]
+            ),
+        )
+    _note(hits=len(result))
+    audit(session, auth, "graph.queried", detail={"result_count": len(result)})
+    session.commit()
+    return result
+
+
+@app.get(
+    "/v1/graph/entities/{entity_id}/subgraph",
+    response_model=list[RelationshipRead],
+    tags=["graph"],
+)
+def get_subgraph(
+    entity_id: UUID,
+    query: SubgraphQuery = Depends(),
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> list[RelationshipRead]:
+    auth.require("graph.query")
+    graph = get_graph_store()
+    entity = graph.get_entity(session, entity_id=entity_id, organization_id=auth.organization_id)
+    if entity is None or not covers(auth.workspace_id, entity.workspace_id):
+        raise HTTPException(status_code=404, detail="Entity not found")
+    return graph.subgraph(
+        session,
+        organization_id=auth.organization_id,
+        workspace_id=auth.workspace_id,
+        entity_id=entity.id,
+        hops=query.hops,
+        temporal_as_of=query.temporal_as_of,
+        limit=query.limit,
+        sensitivities=auth.allowed_sensitivities(),
+    )
+
+
+@app.post(
+    "/v1/working-memory",
+    response_model=WorkingMemoryCreateRead,
+    status_code=201,
+    tags=["working-memory"],
+)
+def post_working_memory(
+    data: WorkingMemoryCreate,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+):  # type: ignore[no-untyped-def]
+    working, lock_result = create_working_memory(session, auth, data)
+    lock_read = None
+    if lock_result is not None:
+        lock, token = lock_result
+        lock_read = LockRead(
+            id=lock.id,
+            resource_type=lock.resource_type,
+            resource_ref=lock.resource_ref,
+            fence=lock.fence,
+            expires_at=lock.expires_at,
+            token=token,
+        )
+    return WorkingMemoryCreateRead(
+        **WorkingMemoryRead.model_validate(working).model_dump(), lock=lock_read
+    )
+
+
+@app.patch(
+    "/v1/working-memory/{working_id}",
+    response_model=WorkingMemoryRead,
+    tags=["working-memory"],
+)
+def patch_working_memory(
+    working_id: UUID,
+    data: WorkingMemoryUpdate,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+):  # type: ignore[no-untyped-def]
+    return update_working_memory(session, auth, working_id, data)
+
+
+@app.get(
+    "/v1/working-memory/{working_id}",
+    response_model=WorkingMemoryRead,
+    tags=["working-memory"],
+)
+def read_working_memory(
+    working_id: UUID,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+):  # type: ignore[no-untyped-def]
+    if "working_memory.read" not in auth.capabilities:
+        auth.require("working_memory.write")
+    return get_working_memory(session, auth, working_id)
+
+
+@app.post(
+    "/v1/working-memory/{working_id}/complete",
+    response_model=WorkingMemoryRead,
+    tags=["working-memory"],
+)
+def complete_working_memory(
+    working_id: UUID,
+    data: WorkingMemoryFinish,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+):  # type: ignore[no-untyped-def]
+    return finish_working_memory(session, auth, working_id, data, failed=False)
+
+
+@app.post(
+    "/v1/working-memory/{working_id}/fail",
+    response_model=WorkingMemoryRead,
+    tags=["working-memory"],
+)
+def fail_working_memory(
+    working_id: UUID,
+    data: WorkingMemoryFinish,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+):  # type: ignore[no-untyped-def]
+    return finish_working_memory(session, auth, working_id, data, failed=True)
+
+
+@app.post("/v1/locks/acquire", response_model=LockRead, status_code=201, tags=["locks"])
+def post_lock(
+    data: LockAcquire,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> LockRead:
+    lock, token = acquire_lock(session, auth, data)
+    return LockRead(
+        id=lock.id,
+        resource_type=lock.resource_type,
+        resource_ref=lock.resource_ref,
+        fence=lock.fence,
+        expires_at=lock.expires_at,
+        token=token,
+    )
+
+
+@app.get("/v1/export", response_model=ExportBundle, tags=["portability"])
+def get_export(
+    workspace_id: UUID | None = None,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> ExportBundle:
+    return export_bundle(session, auth, workspace_id)
+
+
+@app.post("/v1/import", response_model=ImportResult, tags=["portability"])
+def post_import(
+    data: ImportRequest,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> ImportResult:
+    return import_bundle(session, auth, data)
+
+
+@app.post("/v1/locks/{lock_id}/renew", response_model=LockRead, tags=["locks"])
+def post_lock_renew(
+    lock_id: UUID,
+    data: LockRenew,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> LockRead:
+    lock = renew_lock(
+        session,
+        auth,
+        lock_id,
+        token=data.token,
+        fence=data.fence,
+        ttl_seconds=data.ttl_seconds,
+    )
+    return LockRead(
+        id=lock.id,
+        resource_type=lock.resource_type,
+        resource_ref=lock.resource_ref,
+        fence=lock.fence,
+        expires_at=lock.expires_at,
+    )
+
+
+@app.post("/v1/locks/{lock_id}/release", status_code=204, tags=["locks"])
+def post_lock_release(
+    lock_id: UUID,
+    data: LockRelease,
+    auth: AuthContext = Depends(authenticate),
+    session: Session = Depends(get_session),
+) -> Response:
+    release_lock(session, auth, lock_id, token=data.token, fence=data.fence)
+    return Response(status_code=204)
+
+
+def run() -> None:
+    uvicorn.run("albert.api:app", host="0.0.0.0", port=8080)