Spaces:
Sleeping
Sleeping
File size: 4,487 Bytes
2e818da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """Repository boundaries for Chroma-backed evidence and project memory."""
from __future__ import annotations
import json
from dataclasses import dataclass
from typing import Any, Iterable, Sequence
from app.observability.operation import observe_operation
from app.rag.chromadb_client import ChromaDBClient
from app.rag.models import EvidenceType, EvidenceUnit, SourceKind
PAPER_EVIDENCE_COLLECTION = "paper_evidence"
@dataclass(frozen=True)
class VectorCandidate:
item_id: str
text: str
metadata: dict[str, Any]
distance: float | None
score: float
class PaperEvidenceIndex:
def __init__(self, db: ChromaDBClient | None = None) -> None:
self.db = db or ChromaDBClient()
def replace_document(
self,
project_id: str,
document_id: str,
units: Sequence[EvidenceUnit],
) -> int:
where = {"$and": [{"project_id": project_id}, {"document_id": document_id}]}
self.db.delete_where(PAPER_EVIDENCE_COLLECTION, where)
indexable = [
unit
for unit in units
if unit.index_text.strip() and unit.source_kind is not SourceKind.STUDENT_ANNOTATION
]
for batch in _batches(indexable, 128):
self.db.upsert(
PAPER_EVIDENCE_COLLECTION,
documents=[unit.index_text for unit in batch],
metadatas=[_evidence_metadata(unit) for unit in batch],
ids=[unit.evidence_id for unit in batch],
)
return len(indexable)
def search(
self,
project_id: str,
query: str,
limit: int = 30,
*,
document_ids: Sequence[str] | None = None,
evidence_types: set[EvidenceType] | None = None,
consumer: str | None = None,
) -> list[VectorCandidate]:
if not project_id or not query.strip():
return []
clauses: list[dict[str, Any]] = [{"project_id": project_id}]
if document_ids:
clauses.append({"document_id": {"$in": list(document_ids)}})
if evidence_types:
clauses.append({"element_type": {"$in": [item.value for item in evidence_types]}})
where = clauses[0] if len(clauses) == 1 else {"$and": clauses}
with observe_operation(
"embedding.query",
subsystem="embedding",
consumer=consumer,
attributes={"collection_role": "paper_evidence"},
) as op:
embedding = self.db.embedder.embed([query])[0]
op.add_count("query_chars", len(query))
rows = self.db.query_raw(
PAPER_EVIDENCE_COLLECTION,
embedding,
n_results=limit,
where=where,
consumer=consumer,
raise_on_failure=True,
)
return [_candidate(row) for row in rows]
def delete_document(self, project_id: str, document_id: str) -> None:
self.db.delete_where(
PAPER_EVIDENCE_COLLECTION,
{"$and": [{"project_id": project_id}, {"document_id": document_id}]},
)
def delete_project(self, project_id: str) -> None:
self.db.delete_where(PAPER_EVIDENCE_COLLECTION, {"project_id": project_id})
def _candidate(row: dict[str, Any]) -> VectorCandidate:
distance = row.get("distance")
return VectorCandidate(
item_id=str(row["id"]),
text=str(row.get("text") or ""),
metadata=dict(row.get("metadata") or {}),
distance=float(distance) if distance is not None else None,
score=_distance_score(distance),
)
def _distance_score(distance: Any) -> float:
if distance is None:
return 0.0
value = max(0.0, float(distance))
return 1.0 / (1.0 + value)
def _evidence_metadata(unit: EvidenceUnit) -> dict[str, Any]:
return {
"project_id": unit.project_id,
"document_id": unit.document_id,
"evidence_id": unit.evidence_id,
"element_type": unit.element_type.value,
"page_start": unit.page_start,
"page_end": unit.page_end,
"parent_id": unit.parent_id or "",
"section_path": " > ".join(unit.section_path),
"source_kind": unit.source_kind.value,
"bbox_json": json.dumps(unit.bbox_norm.rounded(), separators=(",", ":")) if unit.bbox_norm else "",
}
def _batches(items: Sequence[Any], size: int) -> Iterable[Sequence[Any]]:
for start in range(0, len(items), size):
yield items[start : start + size]
|