cortexdb-sdk 0.2.0b2__py3-none-any.whl

_cortexdb_client/__init__.py

@@ -0,0 +1,73 @@
+from .aql import build_remember_aql, build_retrieve_context_aql, build_verify_fact_aql
+from .client import CortexDBClient
+from .errors import CortexDBError
+from .generated import openapi_types
+from .models import (
+    AnnEvaluationResponse,
+    AnnSearchReport,
+    AnswerGroundingReportResponse,
+    AnswerGroundingSpanResponse,
+    AqlCellResponse,
+    AqlQueryCacheStatsResponse,
+    AqlResponse,
+    CellLookupResponse,
+    CellResponse,
+    ContextPackAnomalyResponse,
+    ContextPackCellResponse,
+    ContextPackResponse,
+    DeleteJobResponse,
+    EvidenceResponse,
+    ExplainResponse,
+    GroundedAnswerResponse,
+    GuardResponse,
+    HealthResponse,
+    IngestResponse,
+    IngestionJobResponse,
+    NumericConflictResponse,
+    PutCellResponse,
+    RememberResponse,
+    SearchResponse,
+    SearchResult,
+    SourceRefResponse,
+    StatsResponse,
+    ValidationResponse,
+    VerificationReportResponse,
+)
+
+__all__ = [
+    "AnnEvaluationResponse",
+    "AnnSearchReport",
+    "AnswerGroundingReportResponse",
+    "AnswerGroundingSpanResponse",
+    "AqlCellResponse",
+    "AqlQueryCacheStatsResponse",
+    "AqlResponse",
+    "CellLookupResponse",
+    "CellResponse",
+    "ContextPackAnomalyResponse",
+    "ContextPackCellResponse",
+    "ContextPackResponse",
+    "CortexDBClient",
+    "CortexDBError",
+    "DeleteJobResponse",
+    "EvidenceResponse",
+    "ExplainResponse",
+    "GroundedAnswerResponse",
+    "GuardResponse",
+    "HealthResponse",
+    "IngestResponse",
+    "IngestionJobResponse",
+    "NumericConflictResponse",
+    "PutCellResponse",
+    "RememberResponse",
+    "SearchResponse",
+    "SearchResult",
+    "SourceRefResponse",
+    "StatsResponse",
+    "ValidationResponse",
+    "VerificationReportResponse",
+    "build_remember_aql",
+    "build_retrieve_context_aql",
+    "build_verify_fact_aql",
+    "openapi_types",
+]

_cortexdb_client/answering.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import Callable, Protocol
+
+from .aql import build_retrieve_context_aql, build_verify_fact_aql
+from .grounding import _grounded_answer_response
+from .models import ContextPackResponse, GroundedAnswerResponse, VerificationReportResponse
+
+
+class GroundedAnswerClient(Protocol):
+    def context_response(self, scope: str, statement: str) -> ContextPackResponse: ...
+
+    def verify_response(self, scope: str, statement: str) -> VerificationReportResponse: ...
+
+
+def answer_with_grounded_context(
+    client: GroundedAnswerClient,
+    scope: str,
+    brain: str,
+    question: str,
+    answerer: Callable[[ContextPackResponse], str],
+    *,
+    mode: str | None = "balanced",
+    budget_tokens: int | None = None,
+    limit_candidates: int | None = None,
+    where_clause: str | None = None,
+    require_citations: bool = True,
+    reject_unsupported: bool = False,
+    verify_answer: bool = True,
+) -> GroundedAnswerResponse:
+    retrieve_statement = build_retrieve_context_aql(
+        question,
+        brain,
+        mode=mode,
+        budget_tokens=budget_tokens,
+        limit_candidates=limit_candidates,
+        where_clause=where_clause,
+        require_citations=require_citations,
+    )
+    context = client.context_response(scope, retrieve_statement)
+    answer = answerer(context)
+    verify_statement = (
+        build_verify_fact_aql(answer, brain) if verify_answer and answer.strip() else None
+    )
+    verification = (
+        client.verify_response(scope, verify_statement)
+        if verify_statement is not None
+        else None
+    )
+    return _grounded_answer_response(
+        question=question,
+        answer=answer,
+        retrieve_statement=retrieve_statement,
+        verify_statement=verify_statement,
+        context=context,
+        verification=verification,
+        require_citations=require_citations,
+        reject_unsupported=reject_unsupported,
+    )

_cortexdb_client/aql.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+
+def _quote_aql_string(value: str) -> str:
+    escaped = (
+        value.replace("\\", "\\\\")
+        .replace('"', '\\"')
+        .replace("\n", "\\n")
+        .replace("\r", "\\r")
+        .replace("\t", "\\t")
+    )
+    return f'"{escaped}"'
+
+
+def _validate_aql_identifier(field: str, value: str) -> None:
+    if not value:
+        raise ValueError(f"{field} must be a non-empty AQL identifier")
+    first = value[0]
+    if not (first == "_" or first.isascii() and first.isalpha()):
+        raise ValueError(f"{field} must start with '_' or an ASCII letter")
+    for character in value[1:]:
+        if not (
+            character == "_"
+            or character == "-"
+            or character == ":"
+            or character.isascii()
+            and character.isalnum()
+        ):
+            raise ValueError(f"{field} contains an invalid AQL identifier character")
+
+
+def _validate_decimal(field: str, value: str | None) -> None:
+    if value is None:
+        return
+    left, separator, right = value.partition(".")
+    if separator != "." or not left.isdigit() or not right.isdigit():
+        raise ValueError(f"{field} must be a decimal literal")
+
+
+def build_retrieve_context_aql(
+    task: str,
+    brain: str,
+    *,
+    mode: str | None = None,
+    budget_tokens: int | None = None,
+    limit_candidates: int | None = None,
+    where_clause: str | None = None,
+    require_citations: bool = False,
+    min_confidence: str | None = None,
+    source_trust: str | None = None,
+    freshness_seconds: int | None = None,
+    explain: bool = False,
+) -> str:
+    _validate_aql_identifier("brain", brain)
+    if mode is not None and mode not in {"fast", "balanced", "semantic", "audit"}:
+        raise ValueError("mode must be fast, balanced, semantic, or audit")
+    if where_clause is not None and not where_clause.strip():
+        raise ValueError("where_clause must not be empty")
+    _validate_decimal("min_confidence", min_confidence)
+    _validate_decimal("source_trust", source_trust)
+
+    parts = []
+    if explain:
+        parts.append("EXPLAIN")
+    parts.extend(["RETRIEVE CONTEXT FOR TASK", _quote_aql_string(task), "IN BRAIN", brain])
+    if mode is not None:
+        parts.extend(["USING MODE", mode])
+    if budget_tokens is not None:
+        parts.extend(["BUDGET", str(budget_tokens), "TOKENS"])
+    if limit_candidates is not None:
+        parts.extend(["LIMIT", str(limit_candidates), "CANDIDATES"])
+    if where_clause is not None:
+        parts.extend(["WHERE", where_clause.strip()])
+    if require_citations:
+        parts.extend(["REQUIRE", "citations"])
+    if min_confidence is not None:
+        parts.extend(["REQUIRE", "confidence", ">=", min_confidence])
+    if source_trust is not None:
+        parts.extend(["REQUIRE", "source_trust", ">=", source_trust])
+    if freshness_seconds is not None:
+        parts.extend(["REQUIRE", "freshness", "<=", str(freshness_seconds), "SECONDS"])
+    return " ".join(parts) + ";"
+
+
+def build_verify_fact_aql(fact: str, brain: str) -> str:
+    _validate_aql_identifier("brain", brain)
+    return f"VERIFY FACT {_quote_aql_string(fact)} IN BRAIN {brain};"
+
+
+def build_remember_aql(
+    content: str,
+    scope: str,
+    memory_type: str,
+    *,
+    ttl_seconds: int | None = None,
+) -> str:
+    _validate_aql_identifier("scope", scope)
+    _validate_aql_identifier("memory_type", memory_type)
+    statement = (
+        f"REMEMBER {_quote_aql_string(content)} IN SCOPE {scope} AS TYPE {memory_type}"
+    )
+    if ttl_seconds is not None:
+        statement += f" TTL {ttl_seconds} SECONDS"
+    return statement + ";"
+
+

_cortexdb_client/client.py

@@ -0,0 +1,291 @@
+from __future__ import annotations
+
+import urllib.parse
+from dataclasses import dataclass, field, replace
+from typing import Any, Callable
+
+from .aql import build_remember_aql, build_retrieve_context_aql, build_verify_fact_aql
+from .answering import answer_with_grounded_context as _answer_with_grounded_context
+from .models import (
+    AnnEvaluationResponse,
+    AqlResponse,
+    CellLookupResponse,
+    ContextPackResponse,
+    DeleteJobResponse,
+    GroundedAnswerResponse,
+    HealthResponse,
+    IngestResponse,
+    IngestionJobResponse,
+    PutCellResponse,
+    RememberResponse,
+    SearchResponse,
+    StatsResponse,
+    ValidationResponse,
+    VerificationReportResponse,
+)
+from .transport import build_opener, close_opener, request_json, scoped_path
+
+
+@dataclass(frozen=True)
+class CortexDBClient:
+    base_url: str = "http://127.0.0.1:8181"
+    token: str | None = None
+    tenant: str | None = None
+    max_retries: int = 0
+    retry_delay_seconds: float = 0.5
+    timeout_seconds: float = 10.0
+    _opener: Any | None = field(default=None, repr=False, compare=False)
+
+    def with_tenant(self, tenant: str) -> "CortexDBClient":
+        return replace(self, tenant=tenant)
+
+    def with_retries(self, max_retries: int, retry_delay_seconds: float = 0.5) -> "CortexDBClient":
+        return replace(self, max_retries=max_retries, retry_delay_seconds=retry_delay_seconds)
+
+    def with_timeout(self, timeout_seconds: float) -> "CortexDBClient":
+        return replace(self, timeout_seconds=timeout_seconds)
+
+    def with_session(self) -> "CortexDBClient":
+        if self._opener is not None:
+            return self
+        return replace(self, _opener=build_opener())
+
+    def __enter__(self) -> "CortexDBClient":
+        if self._opener is None:
+            object.__setattr__(self, "_opener", build_opener())
+        return self
+
+    def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
+        self.close()
+
+    def close(self) -> None:
+        close_opener(self._opener)
+        object.__setattr__(self, "_opener", None)
+
+    build_retrieve_context_aql = staticmethod(build_retrieve_context_aql)
+    build_verify_fact_aql = staticmethod(build_verify_fact_aql)
+    build_remember_aql = staticmethod(build_remember_aql)
+
+    def health(self) -> dict[str, Any]:
+        return self._request("GET", "/v1/health", b"")
+
+    def health_response(self) -> HealthResponse:
+        return HealthResponse.from_json(self.health())
+
+    def put_cell(self, cell_id: int, payload: str) -> dict[str, Any]:
+        return self._request("POST", self._path("/v1/cell", cell_id=cell_id), payload.encode())
+
+    def put_cell_response(self, cell_id: int, payload: str) -> PutCellResponse:
+        return PutCellResponse.from_json(self.put_cell(cell_id, payload))
+
+    def get_cell(self, cell_id: int) -> dict[str, Any]:
+        return self._request("GET", self._path("/v1/cell", cell_id=cell_id), b"")
+
+    def get_cell_response(self, cell_id: int) -> CellLookupResponse:
+        return CellLookupResponse.from_json(self.get_cell(cell_id))
+
+    def tombstone_cell(self, cell_id: int) -> dict[str, Any]:
+        return self._request("DELETE", self._path("/v1/cell", cell_id=cell_id), b"")
+
+    def flush(self) -> dict[str, Any]:
+        return self._request("POST", "/v1/flush", b"")
+
+    def compact(self) -> dict[str, Any]:
+        return self._request("POST", "/v1/compact", b"")
+
+    def search(self, scope: str, query: str, limit: int = 20) -> dict[str, Any]:
+        path = self._path("/v1/search", scope=scope, mode="keyword", q=query, limit=limit)
+        return self._request("POST", path, b"")
+
+    def search_response(self, scope: str, query: str, limit: int = 20) -> SearchResponse:
+        return SearchResponse.from_json(self.search(scope, query, limit))
+
+    def search_vector(
+        self,
+        scope: str,
+        vector: list[int] | tuple[int, ...],
+        limit: int = 20,
+        algorithm: str = "ann",
+    ) -> dict[str, Any]:
+        literal = ",".join(str(value) for value in vector)
+        path = self._path(
+            "/v1/search",
+            scope=scope,
+            mode="vector",
+            algorithm=algorithm,
+            vector=literal,
+            limit=limit,
+        )
+        return self._request("POST", path, b"")
+
+    def search_vector_response(
+        self,
+        scope: str,
+        vector: list[int] | tuple[int, ...],
+        limit: int = 20,
+        algorithm: str = "ann",
+    ) -> SearchResponse:
+        return SearchResponse.from_json(self.search_vector(scope, vector, limit, algorithm))
+
+    def evaluate_ann(
+        self,
+        scope: str,
+        vector: list[int] | tuple[int, ...],
+        limit: int = 20,
+    ) -> dict[str, Any]:
+        literal = ",".join(str(value) for value in vector)
+        path = self._path(
+            "/v1/search/ann-evaluate",
+            scope=scope,
+            vector=literal,
+            limit=limit,
+        )
+        return self._request("POST", path, b"")
+
+    def evaluate_ann_response(
+        self,
+        scope: str,
+        vector: list[int] | tuple[int, ...],
+        limit: int = 20,
+    ) -> AnnEvaluationResponse:
+        return AnnEvaluationResponse.from_json(self.evaluate_ann(scope, vector, limit))
+
+    def aql(self, scope: str, statement: str) -> dict[str, Any]:
+        return self._request("POST", self._path("/v1/aql", scope=scope), statement.encode())
+
+    def aql_response(self, scope: str, statement: str) -> AqlResponse:
+        return AqlResponse.from_json(self.aql(scope, statement))
+
+    def context(self, scope: str, statement: str) -> dict[str, Any]:
+        return self._request("POST", self._path("/v1/context", scope=scope), statement.encode())
+
+    def context_response(self, scope: str, statement: str) -> ContextPackResponse:
+        return ContextPackResponse.from_json(self.context(scope, statement))
+
+    def answer_with_grounded_context(
+        self,
+        scope: str,
+        brain: str,
+        question: str,
+        answerer: Callable[[ContextPackResponse], str],
+        *,
+        mode: str | None = "balanced",
+        budget_tokens: int | None = None,
+        limit_candidates: int | None = None,
+        where_clause: str | None = None,
+        require_citations: bool = True,
+        reject_unsupported: bool = False,
+        verify_answer: bool = True,
+    ) -> GroundedAnswerResponse:
+        return _answer_with_grounded_context(
+            self,
+            scope,
+            brain,
+            question,
+            answerer,
+            mode=mode,
+            budget_tokens=budget_tokens,
+            limit_candidates=limit_candidates,
+            where_clause=where_clause,
+            require_citations=require_citations,
+            reject_unsupported=reject_unsupported,
+            verify_answer=verify_answer,
+        )
+
+    def verify(self, scope: str, statement: str) -> dict[str, Any]:
+        return self._request("POST", self._path("/v1/verify", scope=scope), statement.encode())
+
+    def verify_response(self, scope: str, statement: str) -> VerificationReportResponse:
+        return VerificationReportResponse.from_json(self.verify(scope, statement))
+
+    def remember(self, scope: str, statement: str) -> dict[str, Any]:
+        return self._request("POST", self._path("/v1/remember", scope=scope), statement.encode())
+
+    def remember_response(self, scope: str, statement: str) -> RememberResponse:
+        return RememberResponse.from_json(self.remember(scope, statement))
+
+    def ingest_text(
+        self,
+        scope: str,
+        text: str,
+        source: str = "python_sdk",
+    ) -> dict[str, Any]:
+        path = self._path("/v1/ingest/text", scope=scope, source=source)
+        return self._request("POST", path, text.encode())
+
+    def ingest_text_response(self, scope: str, text: str, source: str = "python_sdk") -> IngestResponse:
+        return IngestResponse.from_json(self.ingest_text(scope, text, source))
+
+    def ingest_json(
+        self,
+        scope: str,
+        document: str,
+        source: str = "python_sdk",
+    ) -> dict[str, Any]:
+        path = self._path("/v1/ingest/json", scope=scope, source=source)
+        return self._request("POST", path, document.encode())
+
+    def ingest_json_response(self, scope: str, document: str, source: str = "python_sdk") -> IngestResponse:
+        return IngestResponse.from_json(self.ingest_json(scope, document, source))
+
+    def ingest_csv(
+        self,
+        scope: str,
+        document: str,
+        source: str = "python_sdk",
+    ) -> dict[str, Any]:
+        path = self._path("/v1/ingest/csv", scope=scope, source=source)
+        return self._request("POST", path, document.encode())
+
+    def ingest_csv_response(self, scope: str, document: str, source: str = "python_sdk") -> IngestResponse:
+        return IngestResponse.from_json(self.ingest_csv(scope, document, source))
+
+    def ingestion_job(self, job_id: int) -> dict[str, Any]:
+        return self._request("GET", f"/v1/ingest/jobs/{job_id}", b"")
+
+    def ingestion_job_response(self, job_id: int) -> IngestionJobResponse:
+        return IngestionJobResponse.from_json(self.ingestion_job(job_id))
+
+    def delete_ingestion_job(self, job_id: int) -> DeleteJobResponse:
+        return DeleteJobResponse.from_json(
+            self._request("DELETE", f"/v1/ingest/jobs/{job_id}", b"")
+        )
+
+    def retry_ingestion_job(self, job_id: int) -> IngestionJobResponse:
+        return IngestionJobResponse.from_json(
+            self._request("POST", f"/v1/ingest/jobs/{job_id}/retry", b"")
+        )
+
+    def validate(self) -> dict[str, Any]:
+        return self._request("GET", "/v1/validate", b"")
+
+    def validate_response(self) -> ValidationResponse:
+        return ValidationResponse.from_json(self.validate())
+
+    def stats(self) -> dict[str, Any]:
+        return self._request("GET", "/v1/stats", b"")
+
+    def stats_response(self) -> StatsResponse:
+        return StatsResponse.from_json(self.stats())
+
+    def _request(self, method: str, path: str, body: bytes) -> dict[str, Any]:
+        return request_json(
+            base_url=self.base_url,
+            tenant=self.tenant,
+            token=self.token,
+            timeout_seconds=self.timeout_seconds,
+            max_retries=self.max_retries,
+            retry_delay_seconds=self.retry_delay_seconds,
+            opener=self._opener,
+            method=method,
+            path=path,
+            body=body,
+        )
+
+    def _scoped(self, path: str) -> str:
+        return scoped_path(path, self.tenant)
+
+    @staticmethod
+    def _path(path: str, **query: object) -> str:
+        encoded = urllib.parse.urlencode({key: str(value) for key, value in query.items()})
+        return f"{path}?{encoded}" if encoded else path

_cortexdb_client/errors.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import json
+
+
+class CortexDBError(Exception):
+    """Typed exception raised for CortexDB HTTP errors."""
+
+    def __init__(
+        self,
+        message: str,
+        code: str | None = None,
+        status: int | None = None,
+        body: str | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.code = code
+        self.status = status
+        self.body = body
+
+    @classmethod
+    def from_response(cls, status: int, body: str) -> "CortexDBError":
+        try:
+            data = json.loads(body)
+            return cls(
+                message=str(data.get("message", body)),
+                code=str(data.get("code", "unknown")),
+                status=status,
+                body=body,
+            )
+        except (json.JSONDecodeError, KeyError):
+            return cls(message=body, code=None, status=status, body=body)
+
+

_cortexdb_client/generated/__init__.py

@@ -0,0 +1,3 @@
+"""Generated CortexDB SDK type artifacts."""
+
+from .openapi_types import *  # noqa: F401,F403