Rifqi Hafizuddin
Claude Opus 5.5
feat(documents): v2 unstructured ingestion, image-token, scoped retrieval
b848f73 | """DocumentRetriever — dense similarity over prose chunks. | |
| For unstructured sources only (PDF / DOCX / TXT). Backed by PGVector via | |
| raw SQL to avoid LangChain ORM / asyncpg type-mapping issues (id UUID vs | |
| String mismatch, jsonb_path_match asyncpg binding quirks). | |
| Collection `document_embeddings`. Methods: cosine | manhattan. | |
| Scoped and v2-only since 2026-09-23 (UNSTRUCTURED_V2_PLAN U5, U-D3 / U-D6): | |
| only rows written by the v2 ingestion (`parser_version = 'v2'`) of documents | |
| bound to the caller's analysis are searched. Both filters run IN SQL, before the | |
| top-k — filtering an already-top-k'd result set would return too few rows, or | |
| none, whenever unbound documents rank higher. | |
| """ | |
| import functools | |
| import math | |
| from langchain_openai import AzureOpenAIEmbeddings | |
| from sqlalchemy import text | |
| from src.config.settings import settings | |
| from src.db.postgres.connection import _pgvector_engine | |
| from src.middlewares.logging import get_logger | |
| from src.retrieval.base import BaseRetriever, RetrievalResult | |
| logger = get_logger("document_retriever") | |
| # Change this one line to switch retrieval method | |
| # Options: "cosine" | "manhattan" | |
| _RETRIEVAL_METHOD = "cosine" | |
| _TABULAR_TYPES = {"csv", "xlsx"} | |
| _COLLECTION_NAME = "documents" | |
| _COSINE_SQL = text(""" | |
| SELECT | |
| lpe.document, | |
| lpe.cmetadata, | |
| lpe.embedding <=> CAST(:embedding AS vector) AS distance | |
| FROM langchain_pg_embedding lpe | |
| JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid | |
| WHERE lpc.name = :collection | |
| AND lpe.cmetadata->>'user_id' = :user_id | |
| AND lpe.cmetadata->>'source_type' = 'document' | |
| AND lpe.cmetadata->>'parser_version' = 'v2' | |
| AND lpe.cmetadata->'data'->>'document_id' = ANY(:document_ids) | |
| ORDER BY distance ASC | |
| LIMIT :k | |
| """) | |
| _MANHATTAN_SQL = text(""" | |
| SELECT | |
| lpe.document, | |
| lpe.cmetadata, | |
| lpe.embedding <+> CAST(:embedding AS vector) AS distance | |
| FROM langchain_pg_embedding lpe | |
| JOIN langchain_pg_collection lpc ON lpe.collection_id = lpc.uuid | |
| WHERE lpc.name = :collection | |
| AND lpe.cmetadata->>'user_id' = :user_id | |
| AND lpe.cmetadata->>'source_type' = 'document' | |
| AND lpe.cmetadata->>'parser_version' = 'v2' | |
| AND lpe.cmetadata->'data'->>'document_id' = ANY(:document_ids) | |
| ORDER BY distance ASC | |
| LIMIT :k | |
| """) | |
| def _get_embeddings() -> AzureOpenAIEmbeddings: | |
| return AzureOpenAIEmbeddings( | |
| azure_deployment=settings.azureai_deployment_name_embedding, | |
| openai_api_version=settings.azureai_api_version_embedding, | |
| azure_endpoint=settings.azureai_endpoint_url_embedding, | |
| api_key=settings.azureai_api_key_embedding, | |
| ) | |
| class DocumentRetriever(BaseRetriever): | |
| async def retrieve( | |
| self, | |
| query: str, | |
| user_id: str, | |
| k: int = 5, | |
| document_ids: list[str] | None = None, | |
| ) -> list[RetrievalResult]: | |
| # No scope → nothing to search (U-D3). `None` is refused rather than | |
| # read as "everything": an unscoped caller is a bug, not a request. | |
| if not document_ids: | |
| logger.info("retrieve skipped — no documents in scope", user_id=user_id) | |
| return [] | |
| query_vector = await _get_embeddings().aembed_query(query) | |
| if not all(math.isfinite(v) for v in query_vector): | |
| raise ValueError("Embedding vector contains NaN or Infinity values.") | |
| vector_str = "[" + ",".join(str(v) for v in query_vector) + "]" | |
| fetch_k = k + len(_TABULAR_TYPES) | |
| sql = _COSINE_SQL if _RETRIEVAL_METHOD == "cosine" else _MANHATTAN_SQL | |
| logger.info("retrieve called", user_id=user_id, collection=_COLLECTION_NAME, fetch_k=fetch_k) | |
| async with _pgvector_engine.connect() as conn: | |
| result = await conn.execute(sql, { | |
| "embedding": vector_str, | |
| "collection": _COLLECTION_NAME, | |
| "user_id": user_id, | |
| "document_ids": list(document_ids), | |
| "k": fetch_k, | |
| }) | |
| rows = result.fetchall() | |
| logger.info("raw rows from db", row_count=len(rows)) | |
| results = [] | |
| for row in rows: | |
| file_type = row.cmetadata.get("data", {}).get("file_type", "") | |
| if file_type not in _TABULAR_TYPES: | |
| results.append(RetrievalResult( | |
| content=row.document, | |
| metadata=row.cmetadata, | |
| score=float(row.distance), | |
| source_type="document", | |
| )) | |
| if len(results) == k: | |
| break | |
| logger.info("retrieved chunks", method=_RETRIEVAL_METHOD, count=len(results)) | |
| return results | |
| document_retriever = DocumentRetriever() | |