rag_api.py

import os
import requests
import shutil
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)

# --------------------------------------------------------
# 1. CONFIGURACIÓN
# --------------------------------------------------------
URL_FAISS = "https://drive.google.com/uc?export=download&id=1pFE0RqM5QAKDkRqj2FHUoEglXDhr48nj"
URL_PKL   = "https://drive.google.com/uc?export=download&id=1md--JucisjwlCCarE-HQqst9K73Pom0J"
DOWNLOAD_DIR  = "/tmp/db_faiss"
DB_FAISS_PATH = DOWNLOAD_DIR

# --------------------------------------------------------
# 2. CLASIFICADOR DE INTENCIÓN  ← NUEVO
# --------------------------------------------------------
from langchain_core.prompts import PromptTemplate

# 1. CLASIFICADOR DE INTENCIONES
# Ahora distingue entre saludos, temas específicos de robótica o temas ajenos.
INTENT_PROMPT = PromptTemplate(
    template="""Eres un clasificador de intenciones para un asistente experto en la Enciclopedia de Robótica.
Analiza el mensaje del usuario y clasifícalo en UNA de estas categorías:
- SALUDO: saludos, despedidas, conversación casual ("hola", "gracias", "adiós", "¿qué tal?").
- ROBOTICA: preguntas técnicas sobre robots, sensores, actuadores, historia de la robótica, programación, cinemática y cualquier contenido presente en la enciclopedia.
- OTRO: preguntas claramente NO relacionadas con la robótica (clima, deportes, política, recetas de cocina, etc.).

IMPORTANTE: Ante la duda, clasifica como ROBOTICA. Solo usa OTRO cuando estés completamente seguro de que no tiene relación con el área técnica.
Responde SOLO con la categoría, sin explicación.

Mensaje: {query}
Categoría:""",
    input_variables=["query"]
)

# 2. PROMPT DE SALUDO
# Define la identidad del bot como un tutor especializado.
SALUDO_PROMPT = PromptTemplate(
    template="""Eres Robotech, el Asistente Virtual experto de la Enciclopedia de Robótica. 
Estás aquí para ayudar a entender conceptos complejos, componentes y la historia de la robótica. 
Si el usuario se despide o agradece, invítalo a profundizar en algún tema técnico del manual.

Mensaje: {query}
Respuesta:""",
    input_variables=["query"]
)

# 3. PROMPT RAG (PARA LA TRADUCCIÓN Y CONSULTA)
# Este es el más importante, ya que usa la base de datos FAISS que creamos.
RAG_PROMPT = PromptTemplate(
    template="""Eres RoboGuide, un Asistente Virtual experto basado en la Enciclopedia de Robótica. 
Tu tarea es responder preguntas técnicas de forma amigable, clara y precisa, utilizando EXCLUSIVAMENTE el contexto proporcionado.

Si el contexto no contiene la información necesaria para responder, dile amablemente al usuario que ese tema específico no está cubierto en la Enciclopedia de Robótica actual.

Contexto de la base de datos (Fragmentos del PDF): {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.150, model_name="openai/gpt-oss-120b")

        # 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
# --------------------------------------------------------
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI(title="RoboGuide RAG API - Enciclopedia de Robótica")

# Variable para la solicitud
class QueryRequest(BaseModel):
    query: str

intent_chain = saludo_chain = qa_chain = retriever = None

try:
    # Esta función debe cargar tus nuevos prompts de Robótica y la base FAISS
    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": "Sistema RoboGuide no inicializado. Revisa los logs y la base FAISS."}
    return {"message": "API RoboGuide (Enciclopedia de Robótica) operativa. Usa /query."}

@app.post("/query")
async def process_query(request: QueryRequest):
    if qa_chain is None:
        return {"error": "El sistema RAG de robótica 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 RoboGuide, experto en la Enciclopedia de Robótica. Mi especialidad son los sensores, actuadores y la historia de la robótica. ¿Tienes alguna duda técnica sobre el manual? 🤖",
                "intent": "OTRO",
                "sources": []
            }

        else:
            # ROBOTICA o cualquier duda técnica → Uso de RAG
            # Invocamos la cadena de pregunta-respuesta
            respuesta = qa_chain.invoke(request.query)
            
            # Recuperamos los fragmentos para mostrar las fuentes (opcional)
            docs = retriever.invoke(request.query)
            # En el PDF extraído, la fuente suele ser el nombre del archivo
            sources = list(set([doc.metadata.get("source", "Enciclopedia_Robotica.pdf") for doc in docs]))
            
            return {
                "query": request.query,
                "response": respuesta.content,
                "intent": "ROBOTICA",
                "sources": sources
            }

    except Exception as e:
        return {"error": f"Error al procesar la consulta técnica: {e}"}