274 lines
9.0 KiB
Python
274 lines
9.0 KiB
Python
"""Vectoriza chunks y los sube a Neo4j con embeddings para vector search."""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from db import run_write, run_query
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
PARSED_DIR = os.path.join(os.path.dirname(__file__), "parsed")
|
|
CATALOG_PATH = os.path.join(os.path.dirname(__file__), "catalog.json")
|
|
|
|
# Embedding config
|
|
EMBEDDING_MODEL = "gemini-embedding-2-preview"
|
|
EMBEDDING_DIMS = 3072
|
|
GCP_API_KEY = os.getenv("GCP_API_KEY", "")
|
|
BATCH_SIZE = 10 # Smaller batches for gemini-embedding-2 (token limits)
|
|
|
|
|
|
def load_catalog():
|
|
with open(CATALOG_PATH, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def save_catalog(catalog):
|
|
with open(CATALOG_PATH, "w", encoding="utf-8") as f:
|
|
json.dump(catalog, f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
def init_embeddings():
|
|
"""Inicializa cliente de embeddings (google-genai con API key)."""
|
|
from google import genai
|
|
client = genai.Client(api_key=GCP_API_KEY)
|
|
print(f" Embedding model: {EMBEDDING_MODEL} ({EMBEDDING_DIMS} dims)")
|
|
return client
|
|
|
|
|
|
def create_vector_index():
|
|
"""Crea el indice vectorial en Neo4j si no existe."""
|
|
try:
|
|
run_write(f"""
|
|
CREATE VECTOR INDEX chunk_embeddings IF NOT EXISTS
|
|
FOR (c:Chunk)
|
|
ON (c.embedding)
|
|
OPTIONS {{
|
|
indexConfig: {{
|
|
`vector.dimensions`: {EMBEDDING_DIMS},
|
|
`vector.similarity_function`: 'cosine'
|
|
}}
|
|
}}
|
|
""")
|
|
print(f" Indice vectorial creado/verificado ({EMBEDDING_DIMS} dims)")
|
|
except Exception as e:
|
|
if "already exists" in str(e).lower() or "equivalent" in str(e).lower():
|
|
print(" Indice vectorial ya existe")
|
|
else:
|
|
print(f" Advertencia indice: {e}")
|
|
|
|
|
|
def build_embedding_text(chunk: dict) -> str:
|
|
"""Prepend metadata al texto para embedding contextual (v2).
|
|
|
|
Esto mejora la calidad del embedding porque reduce ambigüedad semántica.
|
|
Ej: "Capítulo: Nefrología. Sección: Inmunodepresores. <texto>"
|
|
"""
|
|
parts = []
|
|
if chunk.get("titulo_capitulo"):
|
|
parts.append(f"Capítulo: {chunk['titulo_capitulo']}")
|
|
if chunk.get("titulo_seccion"):
|
|
parts.append(f"Sección: {chunk['titulo_seccion']}")
|
|
if chunk.get("tipo_contenido") and chunk["tipo_contenido"] != "body":
|
|
parts.append(f"Tipo: {chunk['tipo_contenido']}")
|
|
prefix = ". ".join(parts)
|
|
return f"{prefix}. {chunk['text']}" if prefix else chunk["text"]
|
|
|
|
|
|
def generate_embeddings(client, texts: list) -> list:
|
|
"""Genera embeddings con gemini-embedding-2-preview via google-genai."""
|
|
truncated = [t[:2000] if len(t) > 2000 else t for t in texts]
|
|
result = client.models.embed_content(
|
|
model=EMBEDDING_MODEL,
|
|
contents=truncated,
|
|
)
|
|
return [e.values for e in result.embeddings]
|
|
|
|
|
|
def vectorize_libro(libro_id: str, model=None, use_v2: bool = None,
|
|
on_progress: callable = None) -> dict:
|
|
"""Vectoriza chunks de un libro que ya están en Neo4j (les agrega embedding).
|
|
|
|
Soporta tanto chunks v1 (text plano) como v2 (con metadata contextual).
|
|
|
|
Args:
|
|
libro_id: ID del libro
|
|
model: Modelo de embeddings (se inicializa si no se pasa)
|
|
use_v2: Si True, usa build_embedding_text. Auto-detecta si None.
|
|
on_progress: Callback(step, pct, msg)
|
|
|
|
Returns:
|
|
dict con estadísticas: {total, skipped, errors}
|
|
"""
|
|
def report(pct, msg):
|
|
if on_progress:
|
|
on_progress("vectorize", pct, msg)
|
|
print(f" [{pct}%] {msg}")
|
|
|
|
# Obtener chunks sin embedding de Neo4j
|
|
chunks = run_query("""
|
|
MATCH (c:Chunk {libro_id: $lid})
|
|
WHERE c.embedding IS NULL
|
|
RETURN c.id AS id, c.text AS text,
|
|
c.titulo_capitulo AS titulo_capitulo,
|
|
c.titulo_seccion AS titulo_seccion,
|
|
c.tipo_contenido AS tipo_contenido
|
|
ORDER BY c.chunk_index
|
|
""", {"lid": libro_id})
|
|
|
|
if not chunks:
|
|
report(100, f"{libro_id}: todos los chunks ya tienen embedding")
|
|
return {"total": 0, "skipped": 0, "errors": 0}
|
|
|
|
# Auto-detectar v2: si tiene titulo_capitulo, es v2
|
|
if use_v2 is None:
|
|
use_v2 = any(c.get("titulo_capitulo") for c in chunks)
|
|
|
|
report(0, f"{libro_id}: vectorizando {len(chunks)} chunks {'(v2 contextual)' if use_v2 else '(v1)'}")
|
|
|
|
if model is None:
|
|
model = init_embeddings()
|
|
|
|
total = 0
|
|
errors = 0
|
|
|
|
for i in range(0, len(chunks), BATCH_SIZE):
|
|
batch = chunks[i:i + BATCH_SIZE]
|
|
|
|
# Preparar textos para embedding
|
|
if use_v2:
|
|
texts = [build_embedding_text(c) for c in batch]
|
|
else:
|
|
texts = [c["text"] for c in batch]
|
|
|
|
# Generar embeddings
|
|
try:
|
|
embeddings = generate_embeddings(model, texts)
|
|
except Exception as e:
|
|
print(f" Error generando embeddings batch {i}: {e}")
|
|
time.sleep(5)
|
|
try:
|
|
embeddings = generate_embeddings(model, texts)
|
|
except Exception as e2:
|
|
print(f" Error retry: {e2}. Saltando batch.")
|
|
errors += len(batch)
|
|
continue
|
|
|
|
# Actualizar embedding en Neo4j (batch update)
|
|
updates = [{"id": c["id"], "embedding": emb}
|
|
for c, emb in zip(batch, embeddings)]
|
|
try:
|
|
run_write("""
|
|
UNWIND $updates AS u
|
|
MATCH (c:Chunk {id: u.id})
|
|
SET c.embedding = u.embedding
|
|
""", {"updates": updates})
|
|
total += len(batch)
|
|
except Exception as e:
|
|
print(f" Error subiendo batch {i}: {e}")
|
|
errors += len(batch)
|
|
|
|
pct = min(99, int((i + len(batch)) / len(chunks) * 100))
|
|
if (i + BATCH_SIZE) % 100 < BATCH_SIZE:
|
|
report(pct, f" {i + len(batch)}/{len(chunks)} embeddings")
|
|
|
|
# Rate limiting
|
|
if i + BATCH_SIZE < len(chunks):
|
|
time.sleep(1)
|
|
|
|
report(100, f"{libro_id}: {total} embeddings generados, {errors} errores")
|
|
return {"total": total, "skipped": 0, "errors": errors}
|
|
|
|
|
|
def vectorize_all():
|
|
"""Vectoriza todos los libros parseados."""
|
|
catalog = load_catalog()
|
|
|
|
print("Inicializando Google GenAI...")
|
|
model = init_embeddings()
|
|
|
|
print("Creando indice vectorial en Neo4j...")
|
|
create_vector_index()
|
|
|
|
total = 0
|
|
for libro in catalog["libros"]:
|
|
if libro["estado"] == "parseado":
|
|
print(f"\n{'='*60}")
|
|
print(f" {libro['titulo']} ({libro['chunks_generados']} chunks)")
|
|
print(f"{'='*60}")
|
|
n = vectorize_libro(libro["id"], model)
|
|
total += n
|
|
|
|
# Stats finales
|
|
stats = run_query("MATCH (c:Chunk) RETURN c.libro_id AS libro, count(c) AS chunks ORDER BY chunks DESC")
|
|
print(f"\n{'='*60}")
|
|
print(f" VECTORIZACION COMPLETA: {total} chunks nuevos")
|
|
print(f"{'='*60}")
|
|
for s in stats:
|
|
print(f" {s['libro']}: {s['chunks']} chunks en Neo4j")
|
|
|
|
|
|
def search_vector(query: str, top_k: int = 5, libro_id: str = None):
|
|
"""Busqueda semantica sobre los chunks vectorizados."""
|
|
model = init_embeddings()
|
|
query_embedding = generate_embeddings(model, [query])[0]
|
|
|
|
filter_clause = ""
|
|
params = {"embedding": query_embedding, "top_k": top_k}
|
|
|
|
if libro_id:
|
|
filter_clause = "WHERE c.libro_id = $libro_id"
|
|
params["libro_id"] = libro_id
|
|
|
|
results = run_query(f"""
|
|
CALL db.index.vector.queryNodes('chunk_embeddings', $top_k, $embedding)
|
|
YIELD node AS c, score
|
|
{filter_clause}
|
|
RETURN c.id AS id, c.libro_id AS libro, c.page_start AS pag_inicio,
|
|
c.page_end AS pag_fin, c.text AS texto, score
|
|
ORDER BY score DESC
|
|
LIMIT $top_k
|
|
""", params)
|
|
|
|
return results
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = sys.argv[1:]
|
|
|
|
if not args:
|
|
print("\nUso:")
|
|
print(" python vectorize.py all - Vectorizar todos los libros")
|
|
print(" python vectorize.py <libro_id> - Vectorizar un libro")
|
|
print(' python vectorize.py search "query" - Busqueda semantica')
|
|
print(' python vectorize.py search "query" --libro farreras-2020')
|
|
sys.exit(0)
|
|
|
|
if args[0] == "all":
|
|
vectorize_all()
|
|
elif args[0] == "search":
|
|
query = args[1] if len(args) > 1 else ""
|
|
libro_id = None
|
|
if "--libro" in args:
|
|
idx = args.index("--libro")
|
|
libro_id = args[idx + 1] if idx + 1 < len(args) else None
|
|
|
|
if not query:
|
|
print("Error: query vacia")
|
|
sys.exit(1)
|
|
|
|
print(f"\n Busqueda semantica: \"{query}\"")
|
|
results = search_vector(query, top_k=5, libro_id=libro_id)
|
|
|
|
if not results:
|
|
print(" Sin resultados")
|
|
else:
|
|
for i, r in enumerate(results):
|
|
print(f"\n [{i+1}] {r['libro']} - pags {r['pag_inicio']}-{r['pag_fin']} (score: {r['score']:.4f})")
|
|
preview = r['texto'][:300].replace('\n', ' ')
|
|
try:
|
|
print(f" {preview}...")
|
|
except UnicodeEncodeError:
|
|
print(f" {preview.encode('ascii', errors='replace').decode()}...")
|