feat: harden runtime readiness and session workflow
This commit is contained in:
@@ -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 = {
|
||||
|
||||
Reference in New Issue
Block a user