Spaces:
Sleeping
Sleeping
File size: 2,288 Bytes
6a040f1 |
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 |
import os
import chromadb
from groq import Groq
from dotenv import load_dotenv
from chroma_setup import initialize_client
from embedding import query_chroma_db
# Carga las variables de entorno desde el archivo .env
load_dotenv()
# Inicializa el cliente de Groq con la clave de API
groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
# Inicializa el cliente de ChromaDB
chroma_client = initialize_client()
# Verifica o crea la colecci贸n donde se guardan las incrustaciones (embeddings)
collection_name = "text_embeddings"
try:
collection = chroma_client.get_collection(name=collection_name)
except chromadb.api.segment.InvalidCollectionException:
collection = chroma_client.create_collection(name=collection_name)
def query_llm(query: str, user_id: str, document_id: str):
"""
Realiza una consulta al LLM (modelo de lenguaje) usando la informaci贸n
m谩s relevante obtenida de la base de datos ChromaDB.
Par谩metros:
-----------
query : str
Pregunta o consulta que se le har谩 al LLM.
user_id : str
Identificador 煤nico de usuario.
document_id : str
Identificador 煤nico de documento.
Retorna:
--------
str
Respuesta generada por el LLM en base al contexto encontrado.
"""
if not user_id or not document_id:
raise ValueError("Se requieren tanto user_id como document_id para realizar la consulta.")
# Consulta a la base de datos para obtener el contexto
context = query_chroma_db(user_id, document_id, query)
# Registra el contexto antes de enviarlo al LLM
print(f"Contexto para el LLM: {context}")
if not context.strip():
return "No se encontr贸 contexto para esta consulta."
prompt = f"""Eres un asistente inteligente. Basado en el siguiente contexto,
proporciona una respuesta resumida y bien interpretada a la pregunta.
Contexto: {context}
Pregunta: {query}
Por favor, ofrece una respuesta detallada y clara basada en el documento."""
chat_completion = groq_client.chat.completions.create(
messages=[
{
"role": "user",
"content": prompt,
}
],
model="llama3-8b-8192",
max_tokens=500,
)
return chat_completion.choices[0].message.content
|