cortexdb-sdk 0.2.0b2__py3-none-any.whl

_cortexdb_client/grounding.py

@@ -0,0 +1,197 @@
+from __future__ import annotations
+
+from typing import Any
+
+from .models import (
+    AnswerGroundingReportResponse,
+    AnswerGroundingSpanResponse,
+    ContextPackResponse,
+    GroundedAnswerResponse,
+    VerificationReportResponse,
+)
+
+
+def _tokenize(text: str) -> tuple[str, ...]:
+    terms: list[str] = []
+    current: list[str] = []
+    for character in text.lower():
+        if character.isalnum():
+            current.append(character)
+        elif current:
+            term = "".join(current)
+            if term not in {"a", "an", "and", "the", "or", "of", "to", "in"}:
+                terms.append(term)
+            current = []
+    if current:
+        term = "".join(current)
+        if term not in {"a", "an", "and", "the", "or", "of", "to", "in"}:
+            terms.append(term)
+    return tuple(sorted(set(terms)))
+
+
+def _split_answer_spans(answer: str) -> tuple[tuple[str, int, int], ...]:
+    spans: list[tuple[str, int, int]] = []
+    start = 0
+    for index, character in enumerate(answer):
+        if character in {"!", "?", "\n"} or (
+            character == "."
+            and not (
+                index > 0
+                and index + 1 < len(answer)
+                and answer[index - 1].isdigit()
+                and answer[index + 1].isdigit()
+            )
+        ):
+            _push_answer_span(answer, start, index + 1, spans)
+            start = index + 1
+    _push_answer_span(answer, start, len(answer), spans)
+    return tuple(spans)
+
+
+def _push_answer_span(
+    answer: str,
+    start: int,
+    end: int,
+    spans: list[tuple[str, int, int]],
+) -> None:
+    raw = answer[start:end]
+    text = raw.strip()
+    if not text:
+        return
+    leading = len(raw) - len(raw.lstrip())
+    trailing = len(raw) - len(raw.rstrip())
+    spans.append((text, start + leading, end - trailing))
+
+
+def _q16_ratio(numerator: int, denominator: int) -> int:
+    if denominator == 0:
+        return 65535
+    return int(numerator * 65535 / denominator)
+
+
+def _unique(values: list[str] | list[int]) -> tuple[Any, ...]:
+    seen = set()
+    out = []
+    for value in values:
+        if value not in seen:
+            seen.add(value)
+            out.append(value)
+    return tuple(out)
+
+
+def ground_answer(
+    context: ContextPackResponse,
+    answer: str,
+    *,
+    min_span_support_q16: int,
+    require_citations: bool,
+    reject_unsupported: bool,
+) -> AnswerGroundingReportResponse:
+    spans: list[AnswerGroundingSpanResponse] = []
+    for text, start, end in _split_answer_spans(answer):
+        span_terms = _tokenize(text)
+        if not span_terms:
+            spans.append(
+                AnswerGroundingSpanResponse(
+                    text=text,
+                    start_byte=start,
+                    end_byte=end,
+                    support_q16=65535,
+                    supported=True,
+                    covered_terms=(),
+                    missing_terms=(),
+                    supported_by_cell_ids=(),
+                    citations=(),
+                )
+            )
+            continue
+        covered: set[str] = set()
+        cell_ids: list[int] = []
+        citations: list[str] = []
+        for cell in context.cells:
+            cell_terms = set(_tokenize(cell.payload_text))
+            matched = False
+            for term in span_terms:
+                if term in cell_terms:
+                    covered.add(term)
+                    matched = True
+            if matched:
+                cell_ids.append(cell.cell_id)
+                if cell.citation:
+                    citations.append(cell.citation)
+        support = _q16_ratio(len(covered), len(span_terms))
+        supported = support >= min_span_support_q16 and (
+            not require_citations or bool(citations)
+        )
+        spans.append(
+            AnswerGroundingSpanResponse(
+                text=text,
+                start_byte=start,
+                end_byte=end,
+                support_q16=support,
+                supported=supported,
+                covered_terms=tuple(sorted(covered)),
+                missing_terms=tuple(term for term in span_terms if term not in covered),
+                supported_by_cell_ids=_unique(cell_ids),
+                citations=_unique(citations),
+            )
+        )
+    supported_count = sum(1 for span in spans if span.supported)
+    unsupported_count = len(spans) - supported_count
+    average = int(sum(span.support_q16 for span in spans) / len(spans)) if spans else 65535
+    return AnswerGroundingReportResponse(
+        answer_supported=unsupported_count == 0,
+        rejected=reject_unsupported and unsupported_count > 0,
+        support_q16=average,
+        supported_span_count=supported_count,
+        unsupported_span_count=unsupported_count,
+        spans=tuple(spans),
+    )
+
+
+def _grounded_answer_response(
+    *,
+    question: str,
+    answer: str,
+    retrieve_statement: str,
+    verify_statement: str | None,
+    context: ContextPackResponse,
+    verification: "VerificationReportResponse | None",
+    require_citations: bool,
+    reject_unsupported: bool,
+) -> GroundedAnswerResponse:
+    grounding = context.ground_answer(
+        answer,
+        require_citations=require_citations,
+        reject_unsupported=reject_unsupported,
+    )
+    citations = _unique(
+        [
+            citation
+            for span in grounding.spans
+            for citation in span.citations
+        ]
+        + [cell.citation for cell in context.cells if cell.citation]
+    )
+    used_cell_ids = _unique(
+        [
+            cell_id
+            for span in grounding.spans
+            for cell_id in span.supported_by_cell_ids
+        ]
+        + [cell.cell_id for cell in context.cells]
+    )
+    return GroundedAnswerResponse(
+        question=question,
+        answer=answer,
+        retrieve_statement=retrieve_statement,
+        verify_statement=verify_statement,
+        context=context,
+        grounding=grounding,
+        verification=verification,
+        citations=citations,
+        used_context_cell_ids=used_cell_ids,
+        rejected=grounding.rejected,
+    )
+
+

_cortexdb_client/model_types/__init__.py

@@ -0,0 +1,66 @@
+from .context import (
+    AnswerGroundingReportResponse,
+    AnswerGroundingSpanResponse,
+    ContextPackAnomalyResponse,
+    ContextPackCellResponse,
+    ContextPackResponse,
+    ExplainResponse,
+    GroundedAnswerResponse,
+    SourceRefResponse,
+)
+from .core import (
+    AqlQueryCacheStatsResponse,
+    AqlCellResponse,
+    AqlResponse,
+    CellLookupResponse,
+    CellResponse,
+    HealthResponse,
+    PutCellResponse,
+    StatsResponse,
+    ValidationResponse,
+)
+from .ingestion import DeleteJobResponse, IngestResponse, IngestionJobResponse
+from .memory import RememberResponse
+from .search import (
+    AnnEvaluationResponse,
+    AnnNoFallbackDecision,
+    AnnSearchReport,
+    SearchResponse,
+    SearchResult,
+    SearchRoutingDecision,
+)
+from .verification import EvidenceResponse, GuardResponse, NumericConflictResponse, VerificationReportResponse
+
+__all__ = [
+    "AnnEvaluationResponse",
+    "AnnNoFallbackDecision",
+    "AnnSearchReport",
+    "AnswerGroundingReportResponse",
+    "AnswerGroundingSpanResponse",
+    "AqlCellResponse",
+    "AqlQueryCacheStatsResponse",
+    "AqlResponse",
+    "CellLookupResponse",
+    "CellResponse",
+    "ContextPackAnomalyResponse",
+    "ContextPackCellResponse",
+    "ContextPackResponse",
+    "DeleteJobResponse",
+    "EvidenceResponse",
+    "ExplainResponse",
+    "GroundedAnswerResponse",
+    "GuardResponse",
+    "HealthResponse",
+    "IngestResponse",
+    "IngestionJobResponse",
+    "NumericConflictResponse",
+    "PutCellResponse",
+    "RememberResponse",
+    "SearchResponse",
+    "SearchResult",
+    "SearchRoutingDecision",
+    "SourceRefResponse",
+    "StatsResponse",
+    "ValidationResponse",
+    "VerificationReportResponse",
+]

_cortexdb_client/model_types/context.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class ExplainResponse:
+    score: int
+    matched_terms: tuple[str, ...]
+    why_selected: str
+    base_bm25: int
+    source_trust_bonus: int
+    redundancy_penalty: int
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "ExplainResponse":
+        return cls(
+            score=int(value["score"]),
+            matched_terms=tuple(str(row) for row in value["matched_terms"]),
+            why_selected=str(value["why_selected"]),
+            base_bm25=int(value["base_bm25"]),
+            source_trust_bonus=int(value["source_trust_bonus"]),
+            redundancy_penalty=int(value["redundancy_penalty"]),
+        )
+
+
+@dataclass(frozen=True)
+class SourceRefResponse:
+    source_id: str
+    document_id: str | None
+    page: int | None
+    cell_range: str | None
+    json_path: str | None
+    confidence_q16: int
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "SourceRefResponse":
+        return cls(
+            source_id=str(value["source_id"]),
+            document_id=str(value["document_id"]) if value.get("document_id") is not None else None,
+            page=int(value["page"]) if value.get("page") is not None else None,
+            cell_range=str(value["cell_range"]) if value.get("cell_range") is not None else None,
+            json_path=str(value["json_path"]) if value.get("json_path") is not None else None,
+            confidence_q16=int(value["confidence_q16"]),
+        )
+
+
+@dataclass(frozen=True)
+class ContextPackCellResponse:
+    cell_id: int
+    estimated_tokens: int
+    citation: str | None
+    payload_text: str
+    explain: ExplainResponse | None
+    source_ref: SourceRefResponse | None
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "ContextPackCellResponse":
+        explain = value.get("explain")
+        source_ref = value.get("source_ref")
+        return cls(
+            cell_id=int(value["cell_id"]),
+            estimated_tokens=int(value["estimated_tokens"]),
+            citation=str(value["citation"]) if value.get("citation") is not None else None,
+            payload_text=str(value["payload_text"]),
+            explain=ExplainResponse.from_json(explain) if explain else None,
+            source_ref=SourceRefResponse.from_json(source_ref) if source_ref else None,
+        )
+
+
+@dataclass(frozen=True)
+class ContextPackAnomalyResponse:
+    cell_id: int | None
+    code: str
+    message: str
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "ContextPackAnomalyResponse":
+        return cls(
+            cell_id=int(value["cell_id"]) if value.get("cell_id") is not None else None,
+            code=str(value["code"]),
+            message=str(value["message"]),
+        )
+
+
+@dataclass(frozen=True)
+class ContextPackResponse:
+    schema_version: str
+    token_budget_tokens: int
+    estimated_tokens: int
+    truncated: bool
+    citations_required: bool
+    cells: tuple[ContextPackCellResponse, ...]
+    anomalies: tuple[ContextPackAnomalyResponse, ...]
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "ContextPackResponse":
+        return cls(
+            schema_version=str(value["schema_version"]),
+            token_budget_tokens=int(value["token_budget_tokens"]),
+            estimated_tokens=int(value["estimated_tokens"]),
+            truncated=bool(value["truncated"]),
+            citations_required=bool(value["citations_required"]),
+            cells=tuple(ContextPackCellResponse.from_json(row) for row in value["cells"]),
+            anomalies=tuple(ContextPackAnomalyResponse.from_json(row) for row in value.get("anomalies", [])),
+        )
+
+    def ground_answer(
+        self,
+        answer: str,
+        *,
+        min_span_support_q16: int = 65535,
+        require_citations: bool = False,
+        reject_unsupported: bool = False,
+    ) -> "AnswerGroundingReportResponse":
+        from ..grounding import ground_answer
+
+        return ground_answer(
+            self,
+            answer,
+            min_span_support_q16=min_span_support_q16,
+            require_citations=require_citations,
+            reject_unsupported=reject_unsupported,
+        )
+
+
+@dataclass(frozen=True)
+class AnswerGroundingSpanResponse:
+    text: str
+    start_byte: int
+    end_byte: int
+    support_q16: int
+    supported: bool
+    covered_terms: tuple[str, ...]
+    missing_terms: tuple[str, ...]
+    supported_by_cell_ids: tuple[int, ...]
+    citations: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class AnswerGroundingReportResponse:
+    answer_supported: bool
+    rejected: bool
+    support_q16: int
+    supported_span_count: int
+    unsupported_span_count: int
+    spans: tuple[AnswerGroundingSpanResponse, ...]
+
+
+@dataclass(frozen=True)
+class GroundedAnswerResponse:
+    question: str
+    answer: str
+    retrieve_statement: str
+    verify_statement: str | None
+    context: ContextPackResponse
+    grounding: AnswerGroundingReportResponse
+    verification: "VerificationReportResponse | None"
+    citations: tuple[str, ...]
+    used_context_cell_ids: tuple[int, ...]
+    rejected: bool
+

_cortexdb_client/model_types/core.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class HealthResponse:
+    status: str
+    version: str
+    server_version: str
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "HealthResponse":
+        return cls(
+            status=str(value["status"]),
+            version=str(value["version"]),
+            server_version=str(value["server_version"]),
+        )
+
+
+@dataclass(frozen=True)
+class AqlQueryCacheStatsResponse:
+    entries: int
+    max_entries: int
+    hits: int
+    misses: int
+    evictions: int
+    catalog_invalidations: int
+    hit_rate_q16: int
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "AqlQueryCacheStatsResponse":
+        return cls(
+            entries=int(value.get("entries", 0)),
+            max_entries=int(value.get("max_entries", 0)),
+            hits=int(value.get("hits", 0)),
+            misses=int(value.get("misses", 0)),
+            evictions=int(value.get("evictions", 0)),
+            catalog_invalidations=int(value.get("catalog_invalidations", 0)),
+            hit_rate_q16=int(value.get("hit_rate_q16", 0)),
+        )
+
+
+@dataclass(frozen=True)
+class StatsResponse:
+    current_seq: int
+    checkpoint_seq: int
+    live_segments: int
+    retired_segments: int
+    memtable_cells: int
+    memtable_versions: int
+    memtable_payload_bytes: int
+    estimated_memtable_bytes: int
+    estimated_index_bytes: int
+    estimated_context_pack_bytes: int
+    estimated_total_memory_bytes: int
+    live_segment_bytes: int
+    retired_segment_bytes: int
+    total_segment_bytes: int
+    durable_storage_bytes: int
+    live_segment_payload_bytes: int
+    logical_payload_bytes: int
+    space_amplification_q16: int
+    write_amplification_q16: int
+    compaction_pressure_q16: int
+    wal_size_bytes: int
+    wal_writer_records: int
+    wal_writer_bytes: int
+    wal_writer_fsyncs: int
+    wal_writer_batches: int
+    aql_query_cache: AqlQueryCacheStatsResponse
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "StatsResponse":
+        return cls(
+            current_seq=int(value["current_seq"]),
+            checkpoint_seq=int(value["checkpoint_seq"]),
+            live_segments=int(value["live_segments"]),
+            retired_segments=int(value["retired_segments"]),
+            memtable_cells=int(value["memtable_cells"]),
+            memtable_versions=int(value["memtable_versions"]),
+            memtable_payload_bytes=int(value.get("memtable_payload_bytes", 0)),
+            estimated_memtable_bytes=int(value.get("estimated_memtable_bytes", 0)),
+            estimated_index_bytes=int(value.get("estimated_index_bytes", 0)),
+            estimated_context_pack_bytes=int(value.get("estimated_context_pack_bytes", 0)),
+            estimated_total_memory_bytes=int(value.get("estimated_total_memory_bytes", 0)),
+            live_segment_bytes=int(value.get("live_segment_bytes", 0)),
+            retired_segment_bytes=int(value.get("retired_segment_bytes", 0)),
+            total_segment_bytes=int(value.get("total_segment_bytes", 0)),
+            durable_storage_bytes=int(value.get("durable_storage_bytes", 0)),
+            live_segment_payload_bytes=int(value.get("live_segment_payload_bytes", 0)),
+            logical_payload_bytes=int(value.get("logical_payload_bytes", 0)),
+            space_amplification_q16=int(value.get("space_amplification_q16", 0)),
+            write_amplification_q16=int(value.get("write_amplification_q16", 0)),
+            compaction_pressure_q16=int(value.get("compaction_pressure_q16", 0)),
+            wal_size_bytes=int(value["wal_size_bytes"]),
+            wal_writer_records=int(value["wal_writer_records"]),
+            wal_writer_bytes=int(value["wal_writer_bytes"]),
+            wal_writer_fsyncs=int(value["wal_writer_fsyncs"]),
+            wal_writer_batches=int(value["wal_writer_batches"]),
+            aql_query_cache=AqlQueryCacheStatsResponse.from_json(
+                value.get("aql_query_cache", {})
+            ),
+        )
+
+
+@dataclass(frozen=True)
+class ValidationResponse:
+    ok: bool
+    manifest_ok: bool
+    wal_ok: bool
+    live_segments_checked: int
+    bitmap_indexes_checked: int
+    lexical_indexes_checked: int
+    vector_indexes_checked: int
+    hnsw_graphs_checked: int
+    cells_checked: int
+    wal_records_checked: int
+    wal_safe_truncate_offset: int
+    errors: tuple[str, ...]
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "ValidationResponse":
+        return cls(
+            ok=bool(value["ok"]),
+            manifest_ok=bool(value["manifest_ok"]),
+            wal_ok=bool(value["wal_ok"]),
+            live_segments_checked=int(value["live_segments_checked"]),
+            bitmap_indexes_checked=int(value["bitmap_indexes_checked"]),
+            lexical_indexes_checked=int(value["lexical_indexes_checked"]),
+            vector_indexes_checked=int(value["vector_indexes_checked"]),
+            hnsw_graphs_checked=int(value["hnsw_graphs_checked"]),
+            cells_checked=int(value["cells_checked"]),
+            wal_records_checked=int(value["wal_records_checked"]),
+            wal_safe_truncate_offset=int(value["wal_safe_truncate_offset"]),
+            errors=tuple(str(row) for row in value.get("errors", [])),
+        )
+
+
+@dataclass(frozen=True)
+class PutCellResponse:
+    seq: int
+    cell_id: int
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "PutCellResponse":
+        return cls(seq=int(value["seq"]), cell_id=int(value["cell_id"]))
+
+
+@dataclass(frozen=True)
+class CellResponse:
+    cell_id: int
+    payload: str
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "CellResponse":
+        return cls(cell_id=int(value["cell_id"]), payload=str(value["payload"]))
+
+
+@dataclass(frozen=True)
+class CellLookupResponse:
+    cell: CellResponse | None
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "CellLookupResponse":
+        cell = value.get("cell")
+        return cls(cell=CellResponse.from_json(cell) if cell else None)
+
+
+@dataclass(frozen=True)
+class AqlCellResponse:
+    cell_id: int
+    payload: str
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "AqlCellResponse":
+        return cls(cell_id=int(value["cell_id"]), payload=str(value["payload"]))
+
+
+@dataclass(frozen=True)
+class AqlResponse:
+    cells: tuple[AqlCellResponse, ...]
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "AqlResponse":
+        return cls(cells=tuple(AqlCellResponse.from_json(row) for row in value["cells"]))

_cortexdb_client/model_types/ingestion.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class IngestResponse:
+    rows_ingested: int
+    chunks_ingested: int
+    facts_ingested: int
+    first_cell_id: int | None
+    job_id: int | None
+    validation_report: dict[str, Any]
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "IngestResponse":
+        return cls(
+            rows_ingested=int(value["rows_ingested"]),
+            chunks_ingested=int(value["chunks_ingested"]),
+            facts_ingested=int(value["facts_ingested"]),
+            first_cell_id=int(value["first_cell_id"]) if value.get("first_cell_id") is not None else None,
+            job_id=int(value["job_id"]) if value.get("job_id") is not None else None,
+            validation_report=dict(value.get("validation_report", {})),
+        )
+
+
+@dataclass(frozen=True)
+class IngestionJobResponse:
+    job_id: int
+    label: str
+    status: str
+    total_items: int | None
+    completed_items: int
+    failed_items: int
+    last_cell_id: int | None
+    message: str | None
+    retry_count: int = 0
+    max_retries: int = 3
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "IngestionJobResponse":
+        return cls(
+            job_id=int(value["job_id"]),
+            label=str(value["label"]),
+            status=str(value["status"]),
+            total_items=int(value["total_items"]) if value.get("total_items") is not None else None,
+            completed_items=int(value["completed_items"]),
+            failed_items=int(value["failed_items"]),
+            last_cell_id=int(value["last_cell_id"]) if value.get("last_cell_id") is not None else None,
+            message=str(value["message"]) if value.get("message") is not None else None,
+            retry_count=int(value.get("retry_count", 0)),
+            max_retries=int(value.get("max_retries", 3)),
+        )
+
+
+@dataclass(frozen=True)
+class DeleteJobResponse:
+    deleted: bool
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "DeleteJobResponse":
+        return cls(deleted=bool(value["deleted"]))

_cortexdb_client/model_types/memory.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class RememberResponse:
+    seq: int
+    cell_id: int
+    ttl_seconds: int | None
+
+    @classmethod
+    def from_json(cls, value: dict[str, Any]) -> "RememberResponse":
+        return cls(
+            seq=int(value["seq"]),
+            cell_id=int(value["cell_id"]),
+            ttl_seconds=int(value["ttl_seconds"]) if value.get("ttl_seconds") is not None else None,
+        )
+