atareao / rag_hyde.py — HyDE (Hypothetical Document Embeddings) para RAG

#!/usr/bin/env python3 """ rag_hyde.py - HyDE (Hypothetical Document Embeddings) para RAG Genera un documento hipotético con Ollama a partir de una consulta, lo embeddea con bge-m3, y busca por similitud coseno en la BD. Uso: python rag_hyde.py "¿qué es una variable de entorno?" python rag_hyde.py "explica la configuración de red" --top 10 """ import argparse import sqlite3 import sys import numpy as np from pathlib import Path DB_PATH = Path.home() / ".cerebro" / "rag_conocimiento.db" OLLAMA_MODEL = "llama3.2" EMBED_MODEL = "bge-m3" HYDE_PROMPT = """Eres un asistente experto en Linux y administración de sistemas. Dada una pregunta, genera un documento hipotético de 3-5 párrafos que RESPONDA a la pregunta como si fuera un artículo del blog o la documentación oficial. Sé técnico pero claro. Usa markdown si es necesario. No digas "basado en" ni "hipotético". Responde DIRECTAMENTE. Pregunta: {consulta} Documento:""" def _ollama(): try: import ollama except ImportError: print("Necesitas instalar ollama: pip install ollama") sys.exit(1) return ollama def generar_doc_hipotetico(consulta: str) -> str: ollama = _ollama() prompt = HYDE_PROMPT.format(consulta=consulta) respuesta = ollama.generate(model=OLLAMA_MODEL, prompt=prompt) return respuesta["response"] def embed_text(texto: str) -> list[float]: ollama = _ollama() respuesta = ollama.embed(model=EMBED_MODEL, input=texto) return respuesta["embeddings"][0] def cosine_sim(a: np.ndarray, b: np.ndarray) -> float: return float(np.dot(a, b)) def cargar_embeddings_db() -> list[tuple[int, str, np.ndarray]]: if not DB_PATH.exists(): print(f"BD no encontrada: {DB_PATH}") sys.exit(1) conn = sqlite3.connect(str(DB_PATH)) cursor = conn.cursor() cursor.execute("""SELECT c.id, c.content, e.vector FROM chunks c JOIN embeddings e ON e.chunk_id = c.id""") filas = cursor.fetchall() conn.close() if not filas: print("No hay chunks con embeddings en la BD.") sys.exit(1) return [(chunk_id, contenido, np.frombuffer(blob, dtype=np.float32).copy()) for chunk_id, contenido, blob in filas] def buscar_semantica(query_emb: np.ndarray, chunks: list, top_k: int = 5) -> list: query_norm = query_emb / (np.linalg.norm(query_emb) + 1e-10) scores = [] for chunk_id, contenido, vector in chunks: vec_norm = vector / (np.linalg.norm(vector) + 1e-10) sim = cosine_sim(query_norm, vec_norm) scores.append((chunk_id, contenido, sim)) scores.sort(key=lambda x: x[2], reverse=True) return scores[:top_k] def main(): parser = argparse.ArgumentParser(description="HyDE: genera un doc hipotético y busca en la BD") parser.add_argument("consulta", help="Consulta en lenguaje natural") parser.add_argument("--top", type=int, default=5, help="Resultados a mostrar") args = parser.parse_args() doc_hipotetico = generar_doc_hipotetico(args.consulta) print(f"Documento generado:\n{─ * 50}\n{doc_hipotetico[:500]}...\n{─ * 50}\n") query_emb = np.array(embed_text(doc_hipotetico), dtype=np.float32) chunks = cargar_embeddings_db() resultados = buscar_semantica(query_emb, chunks, args.top) for i, (chunk_id, contenido, score) in enumerate(resultados, 1): preview = contenido[:150].replace("\n", " ") print(f" {i}. [{score:.4f}] (chunk #{chunk_id})") print(f" {preview}...") if __name__ == "__main__": main()