65 lines
2.4 KiB
Python
65 lines
2.4 KiB
Python
"""Lettura del pgvector dietro un'unica interfaccia `.search(query_vec, top_n, kinds)`, così
|
|
`search.combined_search` resta agnostico al transport. Due implementazioni:
|
|
|
|
- `RestSearcher` → produzione: similarity search via REST (`search_similar`).
|
|
- `DirectSearcher` → dev/test: connessione diretta a Postgres/pgvector.
|
|
|
|
Entrambe mappano i `kind` sulle tabelle per-dominio dello schema `vectors`.
|
|
"""
|
|
|
|
from sqlalchemy import Engine
|
|
|
|
from tht.adapters.vector.legacy_direct import LegacyDirectVectorStore
|
|
from tht.adapters.vector.thoth_http import ThothHttpVectorStore
|
|
from tht.vectorstore.rest_client import VectorRestClient
|
|
from tht.vectorstore.store import VectorHit
|
|
|
|
# kind Thoth → tabella dello schema `vectors`.
|
|
KIND_TO_TABLE = {
|
|
"schema_table": "schema_records",
|
|
"schema_column": "schema_records",
|
|
"evidence": "evidence",
|
|
"memory": "memory",
|
|
"solved_question": "memory", # coppie domanda->SQL: stessa tabella, kind dedicato
|
|
}
|
|
ALL_TABLES = ["schema_records", "evidence", "memory"]
|
|
|
|
|
|
def tables_for_kinds(kinds: list[str] | None) -> list[str]:
|
|
"""Tabelle da interrogare per i kind richiesti (tutte se kinds è vuoto/None)."""
|
|
if not kinds:
|
|
return list(ALL_TABLES)
|
|
return sorted({KIND_TO_TABLE[k] for k in kinds if k in KIND_TO_TABLE})
|
|
|
|
|
|
class RestSearcher:
|
|
"""Similarity search via REST: una chiamata `search_similar` per tabella, poi fusione."""
|
|
|
|
def __init__(self, client: VectorRestClient):
|
|
self.client = client
|
|
self._store = ThothHttpVectorStore(reader=client, writer=None)
|
|
|
|
def search(
|
|
self, query_vec: list[float], top_n: int = 10, kinds: list[str] | None = None
|
|
) -> list[VectorHit]:
|
|
return self._store.search(
|
|
tables_for_kinds(kinds), query_vec, limit=top_n, kinds=kinds
|
|
)
|
|
|
|
|
|
class DirectSearcher:
|
|
"""Similarity search diretta su Postgres/pgvector, interrogando le tabelle per-dominio."""
|
|
|
|
def __init__(self, engine: Engine, schema: str = "vectors", dim: int = 768):
|
|
self.engine = engine
|
|
self.schema = schema
|
|
self.dim = dim
|
|
self._store = LegacyDirectVectorStore(engine, schema=schema, dim=dim)
|
|
|
|
def search(
|
|
self, query_vec: list[float], top_n: int = 10, kinds: list[str] | None = None
|
|
) -> list[VectorHit]:
|
|
return self._store.search(
|
|
tables_for_kinds(kinds), query_vec, limit=top_n, kinds=kinds
|
|
)
|