Update app.py

app.py

@@ -1,11 +1,13 @@
-from flask import Flask, request, jsonify
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 import chromadb
 from langchain_community.vectorstores import Chroma
 from langchain_openai import OpenAIEmbeddings
 import os
 from openai import OpenAI
 
-app = Flask(__name__)
+# Inicializar FastAPI
+app = FastAPI()
 
 # Configurar la API Key de OpenAI
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
@@ -14,7 +16,7 @@ OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 client = OpenAI(api_key=OPENAI_API_KEY)
 
 # Inicializar el cliente de ChromaDB en Hugging Face Space
-chroma_client = chromadb.PersistentClient(path="/app/chroma_db")  # Usa la ruta dentro del contenedor
+chroma_client = chromadb.PersistentClient(path="/app/chroma_db")
 
 # Cargar la base de datos de Chroma como un vector store
 vectorstore = Chroma(
@@ -26,36 +28,35 @@ vectorstore = Chroma(
 # Crear un retriever
 retriever = vectorstore.as_retriever()
 
-def obtener_extractos(pregunta):
+def obtener_extractos(pregunta: str):
     """Obtiene documentos relevantes desde ChromaDB"""
     docs_relevantes = retriever.invoke(pregunta)
     return [(doc.page_content, doc.metadata.get("url", "URL no disponible")) for doc in docs_relevantes]
 
-@app.route('/chat', methods=['POST'])
-def chat():
-    """Endpoint para generar respuestas usando OpenAI y ChromaDB"""
-    data = request.json
-    message = data.get("message", "")
-    system_message = data.get("system_message", "Eres un asistente virtual.")
-    max_tokens = data.get("max_tokens", 512)
-    temperature = data.get("temperature", 0.7)
-    top_p = data.get("top_p", 0.95)
-
-    if not message:
-        return jsonify({"error": "El campo 'message' es obligatorio."}), 400
+# Modelo de datos para la solicitud
+class ChatRequest(BaseModel):
+    message: str
+    system_message: str = "Eres un asistente virtual."
+    max_tokens: int = 512
+    temperature: float = 0.7
+    top_p: float = 0.95
 
+@app.post("/chat")
+async def chat(request: ChatRequest):
+    """Endpoint para generar respuestas usando OpenAI y ChromaDB"""
+    
     # Obtener documentos relevantes
-    contexto = obtener_extractos(message)
+    contexto = obtener_extractos(request.message)
 
     # Construir el mensaje del sistema con el contexto
-    system_message_final = f"""{system_message}
+    system_message_final = f"""{request.system_message}
     Información relevante extraída de los documentos:
     {contexto}
     """
 
     messages = [
         {"role": "system", "content": system_message_final},
-        {"role": "user", "content": message}
+        {"role": "user", "content": request.message}
     ]
 
     try:
@@ -63,17 +64,18 @@ def chat():
         response = client.chat.completions.create(
             model="gpt-4o-mini",
             messages=messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p
+            max_tokens=request.max_tokens,
+            temperature=request.temperature,
+            top_p=request.top_p
         )
 
         completion = response.choices[0].message.content
-        return jsonify({"response": completion, "context": contexto})
+        return {"response": completion, "context": contexto}
 
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        raise HTTPException(status_code=500, detail=str(e))
 
+# Punto de entrada para ejecutar con Uvicorn en Hugging Face
 if __name__ == "__main__":
-    from waitress import serve  # Usa un servidor WSGI ligero
-    serve(app, host="0.0.0.0", port=7860)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)