File size: 3,150 Bytes
3949424 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import os
import shutil
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
# Configuraci贸n para ChromaDB
CHROMA_PATH = 'chroma'
def save_to_chroma_db(chunks: list[Document], embedding_model) -> Chroma:
"""
Guarda documentos en ChromaDB usando modo local con procesamiento por lotes
"""
print(f"Usando modo local de ChromaDB en {CHROMA_PATH}")
# Limpiar base de datos local existente
if os.path.exists(CHROMA_PATH):
try:
shutil.rmtree(CHROMA_PATH)
print(f"Base de datos local existente eliminada: {CHROMA_PATH}")
except Exception as e:
print(f"Error eliminando base de datos local: {e}")
try:
# Procesar en lotes para manejar gran volumen de datos
batch_size = 1000 # Procesar 1000 chunks por vez
total_chunks = len(chunks)
print(f"Procesando {total_chunks} chunks en lotes de {batch_size}...")
# Crear primera colecci贸n con el primer lote
first_batch = chunks[:batch_size]
print(f"Procesando primer lote: {len(first_batch)} chunks...")
db = Chroma.from_documents(
first_batch,
persist_directory=CHROMA_PATH,
embedding=embedding_model
)
print(f"Primer lote completado. Guardado en {CHROMA_PATH}")
# Procesar lotes restantes
for i in range(batch_size, total_chunks, batch_size):
end_idx = min(i + batch_size, total_chunks)
batch = chunks[i:end_idx]
batch_num = (i // batch_size) + 1
total_batches = (total_chunks + batch_size - 1) // batch_size
print(f"Procesando lote {batch_num}/{total_batches}: {len(batch)} chunks...")
try:
db.add_documents(batch)
print(f"Lote {batch_num}/{total_batches} completado")
except Exception as e:
print(f"Error procesando lote {batch_num}: {e}")
print("Continuando con siguiente lote...")
print(f"Procesamiento completado: {total_chunks} chunks guardados exitosamente")
return db
except Exception as e:
print(f"Error cr铆tico creando base de datos: {e}")
print("Verifica que Ollama est茅 funcionando y el modelo nomic-embed-text est茅 disponible")
return None
def get_chroma_client() -> Chroma:
"""
Obtiene un cliente ChromaDB para consultas
"""
try:
if os.path.exists(CHROMA_PATH):
# Crear funci贸n de embedding para consultas
from langchain_ollama import OllamaEmbeddings
embedding_model = OllamaEmbeddings(model="nomic-embed-text")
db = Chroma(
persist_directory=CHROMA_PATH,
embedding_function=embedding_model # Agregar funci贸n de embedding
)
print(f"Conectado a ChromaDB local en {CHROMA_PATH}")
return db
else:
print("Base de datos local no encontrada")
return None
except Exception as e:
print(f"Error conectando a ChromaDB local: {e}")
return None |