Spaces:
Sleeping
Sleeping
File size: 2,490 Bytes
39555a8 | 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 | from flask import Flask, request, jsonify
import chromadb
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
import os
from openai import OpenAI
app = Flask(__name__)
# Configurar la API Key de OpenAI
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# Inicializar el cliente de OpenAI
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
# Cargar la base de datos de Chroma como un vector store
vectorstore = Chroma(
client=chroma_client,
collection_name="docs",
embedding_function=OpenAIEmbeddings(model="text-embedding-3-small", openai_api_key=OPENAI_API_KEY)
)
# Crear un retriever
retriever = vectorstore.as_retriever()
def obtener_extractos(pregunta):
"""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
# Obtener documentos relevantes
contexto = obtener_extractos(message)
# Construir el mensaje del sistema con el contexto
system_message_final = f"""{system_message}
Información relevante extraída de los documentos:
{contexto}
"""
messages = [
{"role": "system", "content": system_message_final},
{"role": "user", "content": message}
]
try:
# Llamar a la API de OpenAI
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=messages,
max_tokens=max_tokens,
temperature=temperature,
top_p=top_p
)
completion = response.choices[0].message.content
return jsonify({"response": completion, "context": contexto})
except Exception as e:
return jsonify({"error": str(e)}), 500
if __name__ == '__main__':
app.run(host="0.0.0.0", port=7860)
|