chatbot / rag_api.py
adonairlariva's picture
Upload rag_api.py
b521581 verified
Raw
History Blame Contribute Delete
7.24 kB
import os
import requests
import shutil
import re
from langchain_community.vectorstores import FAISS
from fastapi import FastAPI
from pydantic import BaseModel
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langchain_groq import ChatGroq
# --------------------------------------------------------
# CACHÉ EN /tmp
# --------------------------------------------------------
TEMP_CACHE_DIR = '/tmp/huggingface_cache'
os.environ['TRANSFORMERS_CACHE'] = TEMP_CACHE_DIR
os.environ['HF_HOME'] = TEMP_CACHE_DIR
os.environ['SENTENCE_TRANSFORMERS_HOME'] = TEMP_CACHE_DIR
os.makedirs(TEMP_CACHE_DIR, exist_ok=True)
def extraer_id(url):
# Buscamos el patrΓ³n que estΓ‘ entre 'd/' y '/view'
match = re.search(r'd/(.*?)/view', url)
if match:
return match.group(1)
else:
return "No se pudo encontrar el ID en el enlace proporcionado."
# Ejemplo de uso:
faiss_id = extraer_id("https://drive.google.com/file/d/1nCvXViXaRn_LyRxzJ4OR8Q8RPVPP5HER/view?usp=drive_link")
pkl_id = extraer_id("https://drive.google.com/file/d/1P8M8mYf-SXNAaZNV4ANYJh8oL6xOEI-z/view?usp=drive_link")
# --------------------------------------------------------
# 1. CONFIGURACIΓ“N
# --------------------------------------------------------
URL_FAISS = (f"https://drive.google.com/uc?export=download&id={faiss_id}")
URL_PKL = (f"https://drive.google.com/uc?export=download&id={pkl_id}")
DOWNLOAD_DIR = "/tmp/db_faiss"
DB_FAISS_PATH = DOWNLOAD_DIR
# --------------------------------------------------------
# 2. CLASIFICADOR DE INTENCIΓ“N ← NUEVO
# --------------------------------------------------------
INTENT_PROMPT = PromptTemplate(
template="""Eres un clasificador de intenciones para un asistente de nutriciΓ³n llamado NutriActive.
Analiza el mensaje del usuario y clasifΓ­calo en UNA de estas categorΓ­as:
- SALUDO: saludos, despedidas, conversaciΓ³n casual ("hola", "gracias", "adiΓ³s", "ΒΏcΓ³mo estΓ‘s?")
- mango: preguntas su veneficios, vitaminas, salud, calorΓ­as, alimentos, etc.
- OTRO: preguntas no relacionadas con nutriciΓ³n ni saludos
Responde SOLO con la categorΓ­a, sin explicaciΓ³n.
Mensaje: {query}
CategorΓ­a:""",
input_variables=["query"]
)
SALUDO_PROMPT = PromptTemplate(
template="""hola soy tu ayudante para entender todo sobre los mangosπŸ₯­.
Mensaje: {query}
Respuesta:""",
input_variables=["query"]
)
RAG_PROMPT = PromptTemplate(
template="""eres un asistente personal centrado en enseΓ±ar sobre el mango y sus veneficios entonces .
Contexto de la base de datos: {context}
Pregunta del usuario: {question}
Respuesta:""",
input_variables=["context", "question"]
)
# --------------------------------------------------------
# 3. FUNCIONES DE DESCARGA Y CARGA
# --------------------------------------------------------
class QueryRequest(BaseModel):
query: str
def download_file(url, local_path):
file_name = os.path.basename(local_path)
print(f"Descargando: {file_name}...")
headers = {'User-Agent': 'Mozilla/5.0'}
try:
response = requests.get(url, stream=True, headers=headers, timeout=30)
if response.status_code == 403:
raise PermissionError(f"Error 403: {file_name} no es pΓΊblico.")
response.raise_for_status()
os.makedirs(os.path.dirname(local_path), exist_ok=True)
with open(local_path, 'wb') as f:
shutil.copyfileobj(response.raw, f)
print(f"βœ“ {file_name} descargado.")
except requests.exceptions.RequestException as e:
raise RuntimeError(f"Fallo al descargar {file_name}: {e}")
def load_and_configure_rag():
try:
download_file(URL_FAISS, os.path.join(DOWNLOAD_DIR, 'index.faiss'))
download_file(URL_PKL, os.path.join(DOWNLOAD_DIR, 'index.pkl'))
print("Cargando embeddings...")
embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-MiniLM-L6-v2",
model_kwargs={'device': 'cpu'},
cache_folder=TEMP_CACHE_DIR
)
print("Cargando FAISS...")
vectorstore = FAISS.load_local(
DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True
)
llm = ChatGroq(temperature=0.3, model_name="llama-3.3-70b-versatile")
# Cadena clasificadora de intenciΓ³n
intent_chain = INTENT_PROMPT | llm
# Cadena para saludos
saludo_chain = SALUDO_PROMPT | llm
# Cadena RAG principal
retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
rag_chain = (
{"context": retriever, "question": RunnablePassthrough()}
| RAG_PROMPT
| llm
)
return intent_chain, saludo_chain, rag_chain, retriever
except Exception as e:
print(f"Error CRÍTICO al inicializar: {type(e).__name__}: {e}")
raise RuntimeError(f"Falla al cargar RAG: {e}")
# --------------------------------------------------------
# 4. FASTAPI
# --------------------------------------------------------
app = FastAPI(title="NutriActive RAG API")
intent_chain = saludo_chain = qa_chain = retriever = None
try:
intent_chain, saludo_chain, qa_chain, retriever = load_and_configure_rag()
except RuntimeError:
pass
@app.get("/")
def home():
if qa_chain is None:
return {"error": "RAG no inicializado. Revisa los logs."}
return {"message": "API mango ACTIVA!!!. Usa /query."}
@app.post("/query")
async def process_query(request: QueryRequest):
if qa_chain is None:
return {"error": "El sistema RAG no se pudo cargar."}
try:
# ── 1. Clasificar intenciΓ³n ──────────────────────────────
intent_result = intent_chain.invoke({"query": request.query})
intent = intent_result.content.strip().upper()
print(f"[Intent] '{request.query}' β†’ {intent}")
# ── 2. Ruta segΓΊn intenciΓ³n ──────────────────────────────
if "SALUDO" in intent:
respuesta = saludo_chain.invoke({"query": request.query})
return {
"query": request.query,
"response": respuesta.content,
"intent": "SALUDO",
"sources": []
}
elif "OTRO" in intent:
return {
"query": request.query,
"response": "Soy tu ayudante para darte todo sobre los mangos y su conocimiento frutal πŸ₯­",
"intent": "OTRO",
"sources": []
}
else:
# NUTRICION β†’ RAG completo
respuesta = qa_chain.invoke(request.query)
docs = retriever.invoke(request.query)
sources = [doc.metadata.get("source", "N/A") for doc in docs]
return {
"query": request.query,
"response": respuesta.content,
"intent": "NUTRICION",
"sources": sources
}
except Exception as e:
return {"error": f"Error al procesar la consulta: {e}"}