feat: auto-detect embedding dimensions per table in consult endpoint
- Detect VECTOR_DIMENSION_COUNT per table before querying - Map dimensions to correct embedding model (1536→small, 3072→large) - Cache query embeddings by dimension to avoid redundant API calls - All tables now searchable regardless of which model generated them
This commit is contained in:
@@ -2458,6 +2458,22 @@ def _embed_text(text: str, genai_cfg: dict, embedding_model_id: str) -> list:
|
||||
response = client.embed_text(embed_detail)
|
||||
return response.data.embeddings[0]
|
||||
|
||||
_DIM_TO_MODEL = {1536: "openai.text-embedding-3-small", 3072: "openai.text-embedding-3-large"}
|
||||
|
||||
def _get_table_embedding_dim(cfg: dict, table_name: str) -> int:
|
||||
"""Detect the embedding dimension of a table by sampling one row."""
|
||||
conn = _get_adb_connection(cfg)
|
||||
try:
|
||||
cur = conn.cursor()
|
||||
cur.execute(f'SELECT VECTOR_DIMENSION_COUNT(EMBEDDING) FROM "{table_name}" FETCH FIRST 1 ROWS ONLY')
|
||||
row = cur.fetchone()
|
||||
cur.close()
|
||||
return int(row[0]) if row and row[0] else 0
|
||||
except Exception:
|
||||
return 0
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _vector_search(cfg: dict, query_embedding: list, top_k: int = 5, table_name: str = None, tenancy: str = None) -> list:
|
||||
"""Search ADB vector store using cosine similarity. Returns top-K documents.
|
||||
If tenancy is provided, filters METADATA JSON to match that tenancy only."""
|
||||
@@ -3335,35 +3351,37 @@ async def consult_embeddings(req: ConsultQuery, u=Depends(current_user)):
|
||||
# Collect results from all active ADB configs + tables
|
||||
all_docs = []
|
||||
for adb_cfg in adb_configs:
|
||||
log.info(f"Consult: processing ADB {adb_cfg['id']} ({adb_cfg.get('config_name','')})")
|
||||
try:
|
||||
emb_genai = _resolve_embed_config(oci_config_id=adb_cfg.get("oci_config_id"))
|
||||
log.info(f"Consult: resolved embed config, compartment={emb_genai.get('compartment_id','?')[:30]}, region={emb_genai.get('genai_region','?')}")
|
||||
except Exception as e:
|
||||
log.warning(f"Consult: resolve config failed for {adb_cfg['id']}: {e}")
|
||||
continue
|
||||
emb_model = adb_cfg.get("embedding_model_id", "")
|
||||
if not emb_model:
|
||||
log.warning(f"Consult: no embedding_model_id for {adb_cfg['id']}")
|
||||
continue
|
||||
try:
|
||||
query_embedding = _embed_text(req.query, emb_genai, emb_model)
|
||||
log.info(f"Consult: embedded query, dims={len(query_embedding)}")
|
||||
except Exception as e:
|
||||
log.warning(f"Consult: embed failed for {adb_cfg['id']}: {e}")
|
||||
continue
|
||||
tables = _get_tables_for_config(adb_cfg["id"], active_only=True)
|
||||
log.info(f"Consult: {len(tables)} active tables, filter='{req.table_name}'")
|
||||
if req.table_name:
|
||||
tables = [t for t in tables if t["table_name"] == req.table_name]
|
||||
# Group tables by embedding dimension and generate one query embedding per dimension
|
||||
embeddings_cache = {} # dim -> query_embedding
|
||||
for tbl in tables:
|
||||
tbl_name = tbl["table_name"]
|
||||
try:
|
||||
docs = _vector_search(adb_cfg, query_embedding, top_k=req.top_k, table_name=tbl["table_name"])
|
||||
dim = _get_table_embedding_dim(adb_cfg, tbl_name)
|
||||
if not dim:
|
||||
log.warning(f"Consult: table {tbl_name} is empty or has no embeddings")
|
||||
continue
|
||||
if dim not in embeddings_cache:
|
||||
model = _DIM_TO_MODEL.get(dim, adb_cfg.get("embedding_model_id", ""))
|
||||
if not model:
|
||||
log.warning(f"Consult: no model for dim={dim} on {tbl_name}")
|
||||
continue
|
||||
embeddings_cache[dim] = _embed_text(req.query, emb_genai, model)
|
||||
log.info(f"Consult: embedded query for dim={dim} using {model}")
|
||||
query_embedding = embeddings_cache[dim]
|
||||
docs = _vector_search(adb_cfg, query_embedding, top_k=req.top_k, table_name=tbl_name)
|
||||
for d in docs:
|
||||
d["source"] = f"{tbl['table_name']}"
|
||||
d["source"] = tbl_name
|
||||
all_docs.extend(docs)
|
||||
except Exception as e:
|
||||
log.warning(f"Consult: search failed on {tbl['table_name']}: {type(e).__name__}: {e}")
|
||||
log.warning(f"Consult: failed on {tbl_name}: {type(e).__name__}: {str(e)[:200]}")
|
||||
if not all_docs:
|
||||
return {"answer": "Nenhum resultado encontrado nas bases vetoriais.", "documents": [], "total": 0}
|
||||
# Sort by distance and take top results
|
||||
|
||||
Reference in New Issue
Block a user