atareao / demo-sqlite-vec.py

#!/usr/bin/env python3 """ demo-sqlite-vec.py — Demo de sqlite-vec + Ollama embeddings Uso: pip install sqlite-vec ollama python demo-sqlite-vec.py Descripción: Demuestra cómo usar sqlite-vec para almacenar y buscar vectores de embeddings generados por Ollama. Este es el fundamento del RAG local que exploraremos en episodios posteriores. Flujo: 1. Carga la extensión sqlite-vec 2. Crea una tabla virtual para vectores de 768 dimensiones 3. Genera embeddings con Ollama (nomic-embed-text) 4. Inserta documentos con sus vectores 5. Busca por similitud semántica Requiere: - Ollama instalado y corriendo - Modelo nomic-embed-text descargado (ollama pull nomic-embed-text) - Python 3.10+ - pip install sqlite-vec ollama Referencia: https://github.com/asg017/sqlite-vec """ import sqlite3 import sys def main(): # --- Comprobar que tenemos los modelos necesarios -------------------- print("=" * 50) print("Demo: sqlite-vec + Ollama embeddings") print("=" * 50) print() # --- 1. Conectar a SQLite en memoria -------------------------------- print("[1/5] Conectando a SQLite en memoria...") conn = sqlite3.connect(":memory:") conn.enable_load_extension(True) try: conn.load_extension("vec0") print(" ✓ Extensión sqlite-vec cargada") except sqlite3.OperationalError as e: print(f" ✗ No se pudo cargar vec0: {e}") print(" Asegúrate de tener sqlite-vec instalado:") print(" pip install sqlite-vec") sys.exit(1) # --- 2. Crear tabla vectorial --------------------------------------- print() print("[2/5] Creando tabla vectorial vec_docs (768 dimensiones)...") conn.execute(""" CREATE VIRTUAL TABLE vec_docs USING vec0( embedding float[768] ) """) print(" ✓ Tabla vec_docs creada") # --- 3. Preparar documentos ----------------------------------------- print() print("[3/5] Preparando documentos de ejemplo...") documentos = [ { "id": 1, "titulo": "Linux", "texto": "Linux es un sistema operativo libre y de código abierto " "basado en el kernel de Linux. Fue creado por Linus Torvalds " "en 1991 y se distribuye bajo la GPL.", }, { "id": 2, "titulo": "Python", "texto": "Python es un lenguaje de programación interpretado, " "multiparadigma y de alto nivel. Creado por Guido van Rossum " "en 1991, es conocido por su sintaxis clara y legible.", }, { "id": 3, "titulo": "IA Local", "texto": "La inteligencia artificial local permite ejecutar modelos " "de lenguaje en tu propio hardware sin conexión a internet. " "Herramientas como Ollama facilitan este proceso.", }, { "id": 4, "titulo": "Docker", "texto": "Docker es una plataforma de contenedores que permite " "empaquetar aplicaciones y sus dependencias en unidades " "portables llamadas contenedores.", }, ] for doc in documentos: print(f" - [{doc['id']}] {doc['titulo']}") # --- 4. Generar embeddings e insertar -------------------------------- print() print("[4/5] Generando embeddings con Ollama (nomic-embed-text)...") print(" (Esto puede tardar unos segundos por documento)") print() try: import ollama except ImportError: print(" ✗ No se pudo importar ollama. Instálalo con:") print(" pip install ollama") sys.exit(1) for doc in documentos: # Generar embedding response = ollama.embeddings(model="nomic-embed-text", prompt=doc["texto"]) vector = response["embedding"] # Insertar en sqlite-vec conn.execute( "INSERT INTO vec_docs(rowid, embedding) VALUES (?, ?)", [doc["id"], str(vector)], ) print(f" ✓ Embedding generado para: {doc['titulo']}") # --- 5. Búsqueda semántica ------------------------------------------- print() print("[5/5] Búsqueda semántica...") print() # Consulta de ejemplo consulta = "¿Qué sistema operativo libre puedo usar?" print(f' Consulta: "{consulta}"') print() # Generar embedding de la consulta response = ollama.embeddings(model="nomic-embed-text", prompt=consulta) query_vector = response["embedding"] # Buscar en sqlite-vec cursor = conn.execute( """ SELECT rowid, distance FROM vec_docs WHERE embedding match ? ORDER BY distance LIMIT 3 """, [str(query_vector)], ) resultados = cursor.fetchall() print(" Resultados (más cercano primero):") print() for rowid, distancia in resultados: doc = documentos[rowid - 1] # rowid empieza en 1 print(f" [{rowid}] {doc['titulo']} (distancia: {distancia:.4f})") print(f" {doc['texto'][:80]}...") print() # --- Cerrar conexión ------------------------------------------------ conn.close() print("=" * 50) print("Demo completada") print() print(" Este es el fundamento del RAG local.") print(" En episodios posteriores veremos cómo:") print(" - Indexar documentos completos") print(" - Combinar búsqueda semántica + keywords") print(" - Integrar con Ollama para respuestas contextuales") print("=" * 50) if __name__ == "__main__": main()