src/api/routes.py

from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from src.models.agent_models import QuestionRequest
from src.agents import AGENT_MANAGER
from langchain_core.messages import HumanMessage, AIMessage
import re

router = APIRouter()

@router.post("/ask-stream")
def ask_stream(req: QuestionRequest):
    agent = AGENT_MANAGER.get_agent(req.agent_id)
    if not agent:
        return {"error": "Agente no encontrado"}
    if "get_context" in agent and "prompt" in agent:
        chat_history = "\n".join([
            f"Usuario: {msg.content}" if isinstance(msg, HumanMessage) else f"Asistente: {msg.content}"
            for msg in agent["history"].messages
        ])
        context_result = agent["get_context"](req.question)
        # Handle both old and new context return formats
        if isinstance(context_result, tuple):
            context, language = context_result
        else:
            context = context_result
            language = "ESPAÑOL"  # Default to Spanish if no language detection
            
        prompt_text = agent["prompt"].format(
            context=context, 
            question=req.question, 
            chat_history=chat_history,
            language=language
        )
        def generate():
            full_response = ""
            for chunk in agent["llm"].stream(prompt_text):
                content = chunk.content
                full_response += content
                yield content
            agent["history"].add_user_message(req.question)
            agent["history"].add_ai_message(full_response)
        return StreamingResponse(generate(), media_type="text/plain")
    elif "agent" in agent and "history" in agent:
        def generate():
            inputs = {
                "input": req.question,
                "chat_history": agent["history"].messages,
                "intermediate_steps": []
            }
            response = agent["agent"].invoke(inputs)
            raw_output = getattr(response, "output", str(response)).strip()
            raw_output = re.sub(r"return_values=\{.*?['\"]output['\"]:\s*['\"]", "", raw_output, flags=re.DOTALL)
            raw_output = re.sub(r"['\"]\}\s*(log=.*)?", "", raw_output, flags=re.DOTALL)
            lines = [line.strip() for line in raw_output.splitlines() if line.strip()]
            lines_with_price = [line for line in lines if re.search(r"\$\d{3,6}", line)]
            final_line = lines_with_price[-1] if lines_with_price else lines[-1] if lines else "⚠️ No se pudo generar respuesta."
            print("💬 Respuesta final mostrada al usuario:")
            print(final_line)
            final_msg = f"\n💵 {final_line}"
            agent["history"].add_user_message(req.question)
            agent["history"].add_ai_message(final_msg)
            yield final_msg
        return StreamingResponse(generate(), media_type="text/plain")
    return {"error": "Estructura de agente no válida"}