feat: harden runtime readiness and session workflow

This commit is contained in:
User
2026-07-14 10:27:25 +02:00
parent 6d7738d538
commit 6dbf93fff9
52 changed files with 1464 additions and 169 deletions
+50 -13
View File
@@ -48,6 +48,39 @@ def _cosine_operator(schema: str) -> sql.Composed:
return sql.SQL("OPERATOR({}.<=>)").format(sql.Identifier(schema))
def _vector_sql_names(cursor, table_schema: str, collection: str) -> tuple[str, str]:
"""Discover pgvector type and operator namespaces from the embedding column."""
cursor.execute(
"""SELECT type_ns.nspname, operator_ns.nspname
FROM pg_catalog.pg_attribute attribute
JOIN pg_catalog.pg_class table_class
ON table_class.oid = attribute.attrelid
JOIN pg_catalog.pg_namespace table_ns
ON table_ns.oid = table_class.relnamespace
JOIN pg_catalog.pg_type vector_type
ON vector_type.oid = attribute.atttypid
JOIN pg_catalog.pg_namespace type_ns
ON type_ns.oid = vector_type.typnamespace
JOIN pg_catalog.pg_operator cosine
ON cosine.oprname = %s
AND cosine.oprleft = vector_type.oid
AND cosine.oprright = vector_type.oid
JOIN pg_catalog.pg_namespace operator_ns
ON operator_ns.oid = cosine.oprnamespace
WHERE table_ns.nspname = %s
AND table_class.relname = %s
AND attribute.attname = %s
AND NOT attribute.attisdropped
ORDER BY cosine.oid
LIMIT 1""",
("<=>", table_schema, collection, "embedding"),
)
row = cursor.fetchone()
if row is None:
raise VectorStoreError(f"Collection {collection} has no usable pgvector embedding")
return row[0], row[1]
def _validate_collection_kinds(collection: str, kinds: list[str]) -> None:
invalid = set(kinds) - COLLECTION_KINDS[collection]
if invalid:
@@ -261,6 +294,9 @@ class PgVectorStore:
with raw.cursor() as cursor:
for collection in collections:
table = _collection(self._schema, collection)
type_schema, operator_schema = _vector_sql_names(
cursor, self._schema, collection
)
collection_kinds = (
sorted(set(kinds) & COLLECTION_KINDS[collection]) if kinds else None
)
@@ -293,12 +329,12 @@ class PgVectorStore:
"SELECT metadata, 1 - (embedding {} %s::{}) AS similarity "
"FROM {}{} ORDER BY embedding {} %s::{}, record_key LIMIT %s"
).format(
_cosine_operator(self._schema),
_vector_type(self._schema),
_cosine_operator(operator_schema),
_vector_type(type_schema),
table,
where,
_cosine_operator(self._schema),
_vector_type(self._schema),
_cosine_operator(operator_schema),
_vector_type(type_schema),
)
params = [_vector_literal(embedding)]
params.extend(filter_params)
@@ -352,19 +388,20 @@ class PgVectorStore:
and len(write_record.embedding) != self._expected_dimension
):
raise VectorStoreError("Embedding dimension does not match configured dimension")
insert = sql.SQL(
"INSERT INTO {} (record_key, kind, content_hash, metadata, embedding) "
"VALUES (%s, %s, %s, %s::jsonb, %s::{}) "
"ON CONFLICT (record_key) DO NOTHING"
).format(table, _vector_type(self._schema))
update = sql.SQL(
"UPDATE {} SET kind = %s, content_hash = %s, metadata = %s::jsonb, "
"embedding = %s::{}, indexed_at = pg_catalog.now() WHERE record_key = %s"
).format(table, _vector_type(self._schema))
raw = None
try:
raw = engine.raw_connection()
with raw.cursor() as cursor:
type_schema, _ = _vector_sql_names(cursor, self._schema, collection)
insert = sql.SQL(
"INSERT INTO {} (record_key, kind, content_hash, metadata, embedding) "
"VALUES (%s, %s, %s, %s::jsonb, %s::{}) "
"ON CONFLICT (record_key) DO NOTHING"
).format(table, _vector_type(type_schema))
update = sql.SQL(
"UPDATE {} SET kind = %s, content_hash = %s, metadata = %s::jsonb, "
"embedding = %s::{}, indexed_at = pg_catalog.now() WHERE record_key = %s"
).format(table, _vector_type(type_schema))
for write_record in records:
record = write_record.record
metadata = {