"""Direct PostgreSQL/pgvector implementation of the vector port.""" import json import re from psycopg2 import sql from sqlalchemy import Engine from tht.config import DatabaseConfig from tht.db.connection import make_engine from tht.ports.vector import ( VectorCapabilities, VectorHealth, VectorReadUnavailable, VectorStoreError, VectorWriteRecord, VectorWriteUnavailable, require_positive_limit, ) from tht.vectorstore.store import VectorHit, hit_from_metadata COLLECTION_KINDS = { "schema_records": {"schema_table", "schema_column"}, "evidence": {"evidence"}, "memory": {"memory", "solved_question"}, } ALLOWED_COLLECTIONS = frozenset(COLLECTION_KINDS) _VECTOR_DIMENSION = re.compile(r"^vector\((\d+)\)$") def _collection(schema: str, name: str) -> sql.Identifier: if name not in ALLOWED_COLLECTIONS: raise VectorStoreError(f"Collection not allowed: {name}") return sql.Identifier(schema, name) def _vector_literal(values: list[float]) -> str: return "[" + ",".join(str(float(value)) for value in values) + "]" def _validate_kinds(collection: str, kinds: list[str]) -> None: invalid = set(kinds) - COLLECTION_KINDS[collection] if invalid: raise VectorStoreError(f"Kind not allowed for {collection}: {', '.join(sorted(invalid))}") class PgVectorStore: """Direct store with independent reader and writer database credentials.""" def __init__( self, read_config: DatabaseConfig | None, write_config: DatabaseConfig | None = None, *, expected_dimension: int | None = None, ): self._reader = make_engine(read_config) if read_config is not None else None self._writer = make_engine(write_config) if write_config is not None else None config = read_config or write_config self._schema = config.db_schema if config is not None else "vectors" if read_config and write_config and read_config.db_schema != write_config.db_schema: raise VectorStoreError("Reader and writer vector schemas must match") self._expected_dimension = expected_dimension @property def capabilities(self) -> VectorCapabilities: writable = self._writer is not None return VectorCapabilities( search=self._reader is not None, existing_hashes=writable, upsert=writable, ) def _probe(self, engine: Engine | None) -> tuple[bool | None, str | None, set[int]]: if engine is None: return None, None, set() try: raw = engine.raw_connection() try: with raw.cursor() as cursor: cursor.execute("SELECT 1") cursor.execute( """SELECT format_type(a.atttypid, a.atttypmod) FROM pg_attribute a JOIN pg_class c ON c.oid = a.attrelid JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = %s AND c.relname = ANY(%s) AND a.attname = 'embedding' AND NOT a.attisdropped""", (self._schema, list(ALLOWED_COLLECTIONS)), ) dimensions = { int(match.group(1)) for (type_name,) in cursor.fetchall() if (match := _VECTOR_DIMENSION.match(type_name)) } return True, None, dimensions finally: raw.close() except Exception as exc: return False, str(exc), set() def health(self) -> VectorHealth: read_ok, read_detail, read_dimensions = self._probe(self._reader) write_ok, write_detail, write_dimensions = self._probe(self._writer) dimensions = tuple(sorted(read_dimensions | write_dimensions)) compatible = ( None if self._expected_dimension is None or not dimensions else dimensions == (self._expected_dimension,) ) reachable = [value for value in (read_ok, write_ok) if value is not None] details = [value for value in (read_detail, write_detail) if value] return VectorHealth( ok=bool(reachable) and all(reachable) and compatible is not False, detail="; ".join(details) or None, read_configured=self._reader is not None, read_reachable=read_ok, read_detail=read_detail, write_configured=self._writer is not None, write_reachable=write_ok, write_detail=write_detail, expected_dimension=self._expected_dimension, observed_dimensions=dimensions, dimension_compatible=compatible, ) def search( self, collections: list[str], embedding: list[float], *, limit: int, kinds: list[str] | None = None, ) -> list[VectorHit]: require_positive_limit(limit) if self._reader is None: raise VectorReadUnavailable("Vector reader credential is not configured") if self._expected_dimension is not None and len(embedding) != self._expected_dimension: raise VectorStoreError("Query embedding dimension does not match configured dimension") hits: list[VectorHit] = [] raw = self._reader.raw_connection() try: with raw.cursor() as cursor: for collection in collections: table = _collection(self._schema, collection) if kinds: _validate_kinds(collection, kinds) where = sql.SQL(" WHERE kind = ANY(%s)") if kinds else sql.SQL("") query = sql.SQL( "SELECT metadata, 1 - (embedding <=> %s::vector) AS similarity " "FROM {}{} ORDER BY embedding <=> %s::vector LIMIT %s" ).format(table, where) params = [_vector_literal(embedding)] if kinds: params.append(kinds) params.extend((_vector_literal(embedding), limit)) cursor.execute(query, params) hits.extend(hit_from_metadata(row[1], row[0]) for row in cursor.fetchall()) finally: raw.close() return sorted(hits, key=lambda hit: hit.similarity, reverse=True)[:limit] def _require_writer(self) -> Engine: if self._writer is None: raise VectorWriteUnavailable("Vector writer credential is not configured") return self._writer def existing_hashes(self, collection: str, kinds: list[str]) -> dict[str, str]: engine = self._require_writer() table = _collection(self._schema, collection) _validate_kinds(collection, kinds) raw = engine.raw_connection() try: with raw.cursor() as cursor: cursor.execute( sql.SQL("SELECT record_key, content_hash FROM {} WHERE kind = ANY(%s)").format( table ), (kinds,), ) return dict(cursor.fetchall()) finally: raw.close() def upsert(self, collection: str, records: list[VectorWriteRecord]) -> int: engine = self._require_writer() table = _collection(self._schema, collection) for write_record in records: _validate_kinds(collection, [write_record.record.kind]) if ( self._expected_dimension is not None and len(write_record.embedding) != self._expected_dimension ): raise VectorStoreError("Embedding dimension does not match configured dimension") query = sql.SQL( "INSERT INTO {} (record_key, kind, content_hash, metadata, embedding) " "VALUES (%s, %s, %s, %s::jsonb, %s::vector) " "ON CONFLICT (record_key) DO UPDATE SET kind = EXCLUDED.kind, " "content_hash = EXCLUDED.content_hash, metadata = EXCLUDED.metadata, " "embedding = EXCLUDED.embedding, indexed_at = now()" ).format(table) raw = engine.raw_connection() try: with raw.cursor() as cursor: for write_record in records: record = write_record.record metadata = { "kind": record.kind, "ref": record.ref, "record_key": record.id, "title": record.title, "content": record.content, **record.metadata, } cursor.execute( query, ( record.id, record.kind, write_record.content_hash, json.dumps(metadata), _vector_literal(write_record.embedding), ), ) raw.commit() except Exception: raw.rollback() raise finally: raw.close() return len(records) __all__ = ["ALLOWED_COLLECTIONS", "PgVectorStore"]