vikingrag 0.4.0__py3-none-any.whl

vikingrag/__init__.py

@@ -0,0 +1,7 @@
+"""VikingRAG: production hierarchical retrieval platform."""
+
+from vikingrag.client import VikingRAGClient
+
+__version__ = "0.4.0"
+
+__all__ = ["VikingRAGClient", "__version__"]

vikingrag/api/__init__.py

@@ -0,0 +1 @@
+"""HTTP API surface."""

vikingrag/api/app.py

@@ -0,0 +1,180 @@
+"""FastAPI application factory and lifespan."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any
+
+from fastapi import FastAPI, Request, Response
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+
+from vikingrag import __version__
+from vikingrag.api.errors import register_exception_handlers
+from vikingrag.api.routers import (
+    answers_router,
+    documents_router,
+    health_router,
+    retrieval_router,
+    search_router,
+)
+from vikingrag.application.experience.support_select import build_support_selector
+from vikingrag.domain.errors import NotImplementedCapabilityError
+from vikingrag.infrastructure.cache.redis import create_redis_client
+from vikingrag.infrastructure.database.engine import create_database
+from vikingrag.infrastructure.object_store.factory import build_object_store
+from vikingrag.observability.context import bind_request_context, clear_request_context
+from vikingrag.observability.logging import configure_logging, get_logger
+from vikingrag.providers.factory import build_embedding_provider, build_llm_provider
+from vikingrag.settings.config import Settings, get_settings
+from vikingrag.workers.edge_builder import EdgeBuilderWorker
+
+logger = get_logger(__name__)
+
+
+class CorrelationMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        request_id = request.headers.get("x-request-id")
+        trace_id = request.headers.get("x-trace-id")
+        rid, tid = bind_request_context(request_id=request_id, trace_id=trace_id)
+        try:
+            response = await call_next(request)
+            response.headers["x-request-id"] = rid
+            response.headers["x-trace-id"] = tid
+            return response
+        finally:
+            clear_request_context()
+
+
+def _provider_configured(name: str) -> bool:
+    return name.lower().strip() not in {"unimplemented", "", "none"}
+
+
+async def _aclose_maybe(obj: Any) -> None:
+    if obj is None:
+        return
+    close = getattr(obj, "aclose", None)
+    if close is not None:
+        await close()
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI) -> AsyncIterator[None]:
+    settings: Settings = app.state.settings
+    configure_logging(
+        level=settings.app.log_level,
+        json_logs=settings.app.env != "development",
+    )
+    logger.info("starting_application", version=__version__, env=settings.app.env)
+
+    database = create_database(settings.database)
+    redis = create_redis_client(settings.redis)
+    object_store = build_object_store(settings.object_store)
+
+    embedding_provider = None
+    llm_provider = None
+    if _provider_configured(settings.embedding.provider):
+        try:
+            embedding_provider = build_embedding_provider(settings)
+        except NotImplementedCapabilityError:
+            logger.warning(
+                "embedding_provider_unconfigured",
+                provider=settings.embedding.provider,
+            )
+    if _provider_configured(settings.llm.provider):
+        try:
+            llm_provider = build_llm_provider(settings)
+        except NotImplementedCapabilityError:
+            logger.warning("llm_provider_unconfigured", provider=settings.llm.provider)
+
+    app.state.database = database
+    app.state.redis = redis
+    app.state.object_store = object_store
+    app.state.embedding_provider = embedding_provider
+    app.state.llm_provider = llm_provider
+
+    edge_worker: EdgeBuilderWorker | None = None
+    edge_task: asyncio.Task[None] | None = None
+    if settings.workers.edge_builder_enabled:
+        selector = build_support_selector(
+            mode=settings.retrieval.support_selector,
+            llm=llm_provider,
+            model=settings.llm.model,
+            max_support=settings.retrieval.experience_max_support,
+        )
+        edge_worker = EdgeBuilderWorker(
+            database=database,
+            poll_interval_seconds=settings.workers.edge_builder_poll_seconds,
+            batch_size=settings.workers.edge_builder_batch_size,
+            max_support=settings.retrieval.experience_max_support,
+            support_selector=selector,
+        )
+        edge_task = asyncio.create_task(edge_worker.run_forever(), name="edge_builder")
+        app.state.edge_builder_worker = edge_worker
+        logger.info(
+            "edge_builder_worker_attached",
+            poll_seconds=settings.workers.edge_builder_poll_seconds,
+        )
+    else:
+        app.state.edge_builder_worker = None
+        logger.info("edge_builder_worker_disabled")
+
+    try:
+        yield
+    finally:
+        logger.info("shutting_down_application")
+        if edge_worker is not None:
+            edge_worker.request_stop()
+        if edge_task is not None:
+            try:
+                await asyncio.wait_for(edge_task, timeout=10.0)
+            except (TimeoutError, asyncio.CancelledError):
+                edge_task.cancel()
+        await _aclose_maybe(embedding_provider)
+        await _aclose_maybe(llm_provider)
+        await _aclose_maybe(object_store)
+        await redis.close()
+        await database.dispose()
+
+
+def create_app(settings: Settings | None = None) -> FastAPI:
+    cfg = settings or get_settings()
+    app = FastAPI(
+        title="VikingRAG",
+        description=(
+            "Production-oriented hierarchical retrieval platform inspired by VikingRAG. "
+            "Hierarchical ingestion, semantic Search, List/Grep/Read, and evidence sufficiency."
+        ),
+        version=__version__,
+        lifespan=lifespan,
+        docs_url="/docs",
+        redoc_url="/redoc",
+        openapi_url="/openapi.json",
+    )
+    app.state.settings = cfg
+    app.add_middleware(CorrelationMiddleware)
+    register_exception_handlers(app)
+    app.include_router(health_router)
+    app.include_router(documents_router)
+    app.include_router(search_router)
+    app.include_router(retrieval_router)
+    app.include_router(answers_router)
+    return app
+
+
+def run() -> None:
+    import uvicorn
+
+    settings = get_settings()
+    uvicorn.run(
+        "vikingrag.api.app:create_app",
+        factory=True,
+        host=settings.app.host,
+        port=settings.app.port,
+        reload=settings.app.debug,
+    )
+
+
+# ASGI entry used by Docker / uvicorn without factory flag when needed
+app: Any = None

vikingrag/api/auth.py

@@ -0,0 +1,82 @@
+"""Single-tenant API-key auth and server-derived document ACL."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from uuid import UUID
+
+from fastapi import Request
+from fastapi.security.utils import get_authorization_scheme_param
+
+from vikingrag.domain.errors import AuthenticationError, ScopeDeniedError
+from vikingrag.domain.models.document import DocumentId
+from vikingrag.settings.config import (
+    Settings,
+    derive_permitted_document_ids,
+    parse_allowed_document_ids,
+)
+
+__all__ = [
+    "AuthContext",
+    "derive_permitted_document_ids",
+    "parse_allowed_document_ids",
+    "require_auth",
+    "resolve_auth_context",
+]
+
+
+@dataclass(frozen=True, slots=True)
+class AuthContext:
+    """Resolved auth for a request.
+
+    ``permitted_document_ids``:
+      - ``None`` — unrestricted
+      - ``frozenset()`` — allow-nothing
+      - nonempty frozenset — allowlist
+    """
+
+    authenticated: bool
+    permitted_document_ids: frozenset[DocumentId] | None
+
+    def ensure_document_allowed(self, document_id: UUID | DocumentId) -> None:
+        did = DocumentId(document_id) if isinstance(document_id, UUID) else document_id
+        if self.permitted_document_ids is None:
+            return
+        if did not in self.permitted_document_ids:
+            raise ScopeDeniedError(f"Document {did} is outside permitted scope")
+
+
+def _extract_api_key(request: Request) -> str | None:
+    header = request.headers.get("x-api-key")
+    if header:
+        return header.strip() or None
+    authorization = request.headers.get("authorization")
+    if not authorization:
+        return None
+    scheme, param = get_authorization_scheme_param(authorization)
+    if scheme.lower() != "bearer":
+        return None
+    return param.strip() or None
+
+
+def resolve_auth_context(settings: Settings, request: Request) -> AuthContext:
+    auth = settings.auth
+    permitted = derive_permitted_document_ids(auth)
+    if not auth.enabled:
+        return AuthContext(authenticated=False, permitted_document_ids=None)
+
+    expected = auth.api_key
+    if not expected:
+        raise AuthenticationError("Auth is enabled but VIKINGRAG_AUTH_API_KEY is not configured")
+
+    provided = _extract_api_key(request)
+    if provided is None or provided != expected:
+        raise AuthenticationError("Invalid or missing API key")
+
+    return AuthContext(authenticated=True, permitted_document_ids=permitted)
+
+
+async def require_auth(request: Request) -> AuthContext:
+    """FastAPI dependency: enforce API key when auth is enabled."""
+    settings: Settings = request.app.state.settings
+    return resolve_auth_context(settings, request)

vikingrag/api/dependencies.py

@@ -0,0 +1,54 @@
+"""FastAPI dependency injection - resources live on app.state."""
+
+from __future__ import annotations
+
+from typing import Annotated, Any, cast
+
+from fastapi import Depends, Request
+
+from vikingrag.api.auth import AuthContext, require_auth
+from vikingrag.infrastructure.cache.redis import RedisClient
+from vikingrag.infrastructure.database.engine import Database
+from vikingrag.infrastructure.object_store.base import ObjectStore
+from vikingrag.providers.embeddings.base import EmbeddingProvider
+from vikingrag.providers.factory import build_embedding_provider, build_llm_provider
+from vikingrag.settings.config import Settings
+
+
+def get_settings_dep(request: Request) -> Settings:
+    return cast(Settings, request.app.state.settings)
+
+
+def get_database(request: Request) -> Database:
+    return cast(Database, request.app.state.database)
+
+
+def get_redis(request: Request) -> RedisClient:
+    return cast(RedisClient, request.app.state.redis)
+
+
+def get_object_store(request: Request) -> ObjectStore:
+    return cast(ObjectStore, request.app.state.object_store)
+
+
+def get_embedding_provider(request: Request) -> EmbeddingProvider:
+    owned = getattr(request.app.state, "embedding_provider", None)
+    if owned is not None:
+        return cast(EmbeddingProvider, owned)
+    return build_embedding_provider(cast(Settings, request.app.state.settings))
+
+
+def get_llm_provider(request: Request) -> Any:
+    owned = getattr(request.app.state, "llm_provider", None)
+    if owned is not None:
+        return owned
+    return build_llm_provider(cast(Settings, request.app.state.settings))
+
+
+SettingsDep = Annotated[Settings, Depends(get_settings_dep)]
+DatabaseDep = Annotated[Database, Depends(get_database)]
+RedisDep = Annotated[RedisClient, Depends(get_redis)]
+ObjectStoreDep = Annotated[ObjectStore, Depends(get_object_store)]
+AuthDep = Annotated[AuthContext, Depends(require_auth)]
+EmbeddingProviderDep = Annotated[EmbeddingProvider, Depends(get_embedding_provider)]
+LLMProviderDep = Annotated[Any, Depends(get_llm_provider)]

vikingrag/api/errors.py

@@ -0,0 +1,98 @@
+"""Map domain errors to HTTP responses."""
+
+from __future__ import annotations
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+
+from vikingrag.domain.errors import (
+    AuthenticationError,
+    BudgetExhaustedError,
+    ConflictError,
+    DocumentAlreadyExists,
+    DocumentNotFound,
+    DocumentParseError,
+    DomainError,
+    HierarchyConstructionError,
+    InvalidVikingURI,
+    NodeNotFound,
+    NotFoundError,
+    ScopeDeniedError,
+    StaleSourceError,
+    UnsupportedDocumentType,
+    ValidationDomainError,
+)
+from vikingrag.observability.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def _payload(exc: DomainError) -> dict[str, str]:
+    return {"error": exc.code, "message": exc.message}
+
+
+def register_exception_handlers(app: FastAPI) -> None:
+    @app.exception_handler(AuthenticationError)
+    async def authentication(_request: Request, exc: AuthenticationError) -> JSONResponse:
+        return JSONResponse(
+            status_code=401,
+            content=_payload(exc),
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    @app.exception_handler(UnsupportedDocumentType)
+    async def unsupported_type(_request: Request, exc: UnsupportedDocumentType) -> JSONResponse:
+        return JSONResponse(status_code=415, content=_payload(exc))
+
+    @app.exception_handler(DocumentParseError)
+    async def parse_error(_request: Request, exc: DocumentParseError) -> JSONResponse:
+        return JSONResponse(status_code=422, content=_payload(exc))
+
+    @app.exception_handler(HierarchyConstructionError)
+    async def hierarchy_error(_request: Request, exc: HierarchyConstructionError) -> JSONResponse:
+        return JSONResponse(status_code=422, content=_payload(exc))
+
+    @app.exception_handler(InvalidVikingURI)
+    async def invalid_uri(_request: Request, exc: InvalidVikingURI) -> JSONResponse:
+        return JSONResponse(status_code=400, content=_payload(exc))
+
+    @app.exception_handler(DocumentAlreadyExists)
+    async def already_exists(_request: Request, exc: DocumentAlreadyExists) -> JSONResponse:
+        return JSONResponse(status_code=409, content=_payload(exc))
+
+    @app.exception_handler(DocumentNotFound)
+    async def document_not_found(_request: Request, exc: DocumentNotFound) -> JSONResponse:
+        return JSONResponse(status_code=404, content=_payload(exc))
+
+    @app.exception_handler(NodeNotFound)
+    async def node_not_found(_request: Request, exc: NodeNotFound) -> JSONResponse:
+        return JSONResponse(status_code=404, content=_payload(exc))
+
+    @app.exception_handler(NotFoundError)
+    async def not_found(_request: Request, exc: NotFoundError) -> JSONResponse:
+        return JSONResponse(status_code=404, content=_payload(exc))
+
+    @app.exception_handler(ConflictError)
+    async def conflict(_request: Request, exc: ConflictError) -> JSONResponse:
+        return JSONResponse(status_code=409, content=_payload(exc))
+
+    @app.exception_handler(ValidationDomainError)
+    async def validation(_request: Request, exc: ValidationDomainError) -> JSONResponse:
+        return JSONResponse(status_code=422, content=_payload(exc))
+
+    @app.exception_handler(ScopeDeniedError)
+    async def scope_denied(_request: Request, exc: ScopeDeniedError) -> JSONResponse:
+        return JSONResponse(status_code=403, content=_payload(exc))
+
+    @app.exception_handler(StaleSourceError)
+    async def stale_source(_request: Request, exc: StaleSourceError) -> JSONResponse:
+        return JSONResponse(status_code=409, content=_payload(exc))
+
+    @app.exception_handler(BudgetExhaustedError)
+    async def budget_exhausted(_request: Request, exc: BudgetExhaustedError) -> JSONResponse:
+        return JSONResponse(status_code=429, content=_payload(exc))
+
+    @app.exception_handler(DomainError)
+    async def domain(_request: Request, exc: DomainError) -> JSONResponse:
+        logger.warning("domain_error", code=exc.code, message=exc.message)
+        return JSONResponse(status_code=400, content=_payload(exc))

vikingrag/api/routers/__init__.py

@@ -0,0 +1,13 @@
+from vikingrag.api.routers.answers import router as answers_router
+from vikingrag.api.routers.documents import router as documents_router
+from vikingrag.api.routers.health import router as health_router
+from vikingrag.api.routers.retrieval import router as retrieval_router
+from vikingrag.api.routers.search import router as search_router
+
+__all__ = [
+    "answers_router",
+    "documents_router",
+    "health_router",
+    "retrieval_router",
+    "search_router",
+]

vikingrag/api/routers/answers.py

@@ -0,0 +1,132 @@
+"""HTTP endpoints for answers and query execution modes."""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Annotated, Any
+from uuid import uuid4
+
+from fastapi import APIRouter, Depends, Request
+
+from vikingrag.api.auth import AuthContext, require_auth
+from vikingrag.api.schemas.answers import (
+    AnswerRequestBody,
+    AnswerResponseBody,
+    QueryRequestBody,
+)
+from vikingrag.application.answer.generate import AnswerGenerator
+from vikingrag.application.grep_primitive import GrepService
+from vikingrag.application.list_primitive import ListService
+from vikingrag.application.read_primitive import ReadService
+from vikingrag.application.search import SemanticSearchService
+from vikingrag.domain.models.answer import AnswerRequest, ExecutionMode
+from vikingrag.domain.models.document import DocumentId
+from vikingrag.providers.factory import (
+    build_embedding_provider,
+    build_llm_provider,
+    build_reranker,
+)
+
+router = APIRouter(
+    prefix="/v1",
+    tags=["answers"],
+    dependencies=[Depends(require_auth)],
+)
+
+
+def _embedding_provider(request: Request) -> Any:
+    owned = getattr(request.app.state, "embedding_provider", None)
+    if owned is not None:
+        return owned
+    return build_embedding_provider(request.app.state.settings)
+
+
+def _llm_provider(request: Request) -> Any:
+    owned = getattr(request.app.state, "llm_provider", None)
+    if owned is not None:
+        return owned
+    return build_llm_provider(request.app.state.settings)
+
+
+def _search_service(request: Request) -> SemanticSearchService:
+    settings = request.app.state.settings
+    reranker = build_reranker(settings) if settings.retrieval.rerank_enabled else None
+    return SemanticSearchService(
+        database=request.app.state.database,
+        embedding_provider=_embedding_provider(request),
+        retrieval_settings=settings.retrieval,
+        embedding_settings=settings.embedding,
+        reranker=reranker,
+    )
+
+
+def _maybe_search_plus(request: Request) -> Any | None:
+    """Prefer Search+ hook when application/search_plus exists."""
+    try:
+        mod = import_module("vikingrag.application.search_plus")
+    except ModuleNotFoundError:
+        return None
+    factory = getattr(mod, "build_search_plus", None) or getattr(mod, "SearchPlusService", None)
+    if factory is None:
+        return None
+    try:
+        if callable(factory) and getattr(factory, "__name__", "") == "build_search_plus":
+            return factory(request.app.state.settings, request.app.state.database)
+        return factory(
+            database=request.app.state.database,
+            settings=request.app.state.settings,
+        )
+    except (TypeError, Exception):
+        return None
+
+
+def _answer_generator(request: Request) -> AnswerGenerator:
+    settings = request.app.state.settings
+    return AnswerGenerator(
+        llm=_llm_provider(request),
+        search=_search_service(request),
+        list_service=ListService(database=request.app.state.database),
+        grep_service=GrepService(database=request.app.state.database),
+        read_service=ReadService(database=request.app.state.database),
+        settings=settings,
+        search_plus=_maybe_search_plus(request),
+        database=request.app.state.database,
+    )
+
+
+def _to_domain(body: AnswerRequestBody) -> AnswerRequest:
+    return AnswerRequest(
+        question=body.question,
+        document_ids=tuple(DocumentId(x) for x in body.document_ids),
+        execution_mode=ExecutionMode(body.execution_mode),
+        instructions=body.instructions,
+        max_rounds=body.max_rounds,
+        query_id=uuid4(),
+    )
+
+
+@router.post("/answers", response_model=AnswerResponseBody)
+async def create_answer(
+    request: Request,
+    body: AnswerRequestBody,
+    auth: Annotated[AuthContext, Depends(require_auth)],
+) -> AnswerResponseBody:
+    result = await _answer_generator(request).generate(
+        _to_domain(body),
+        permitted_document_ids=auth.permitted_document_ids,
+    )
+    return AnswerResponseBody.from_domain(result)
+
+
+@router.post("/query", response_model=AnswerResponseBody)
+async def query(
+    request: Request,
+    body: QueryRequestBody,
+    auth: Annotated[AuthContext, Depends(require_auth)],
+) -> AnswerResponseBody:
+    """Execute vikingrag | vikingrag_e | vikingrag_e_plus."""
+    result = await _answer_generator(request).generate(
+        _to_domain(body),
+        permitted_document_ids=auth.permitted_document_ids,
+    )
+    return AnswerResponseBody.from_domain(result)