Create llm_query.py

llm_query.py

@@ -0,0 +1,74 @@
+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