app.py

# space-orchestrator/app.py
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import requests
import json
import os
import re
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Inicializar la aplicación FastAPI
app = FastAPI()

# --- Middleware de CORS ---
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Permite solicitudes de cualquier origen
    allow_credentials=True,
    allow_methods=["*"],  # Permite todos los métodos HTTP
    allow_headers=["*"],  # Permite todas las cabeceras
)

# --- Configuración del Modelo de Lenguaje ---
# Usamos un modelo instruct para que siga mejor las órdenes
model_id = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float32,  # CPU
    trust_remote_code=True
)

# --- URL del Servicio de Imágenes ---
IMAGE_SERVICE_URL = os.getenv("IMAGE_SERVICE_URL", "http://localhost:8001/generate-image/")

# --- Definición de la Petición ---
class ChatRequest(BaseModel):
    prompt: str

# --- Función para extraer JSON robusto ---
def extract_json(generated_text: str):
    match = re.search(r'\{.*\}', generated_text, re.DOTALL)
    if match:
        try:
            return json.loads(match.group(0))
        except json.JSONDecodeError:
            return {"action": "chat", "response": "Error al parsear JSON"}
    return {"action": "chat", "response": "No pude entender tu petición. ¿Podrías reformularla?"}

# --- Lógica del asistente ---
def generate_llm_response(user_prompt: str):
    """
    Carl IA: asistente de Carley Interactive Studio.
    Clasifica intención del usuario y responde SOLO en JSON.
    """
    system_prompt = """Eres Carl IA, un asistente de Carley Interactive Studio.
Tu tarea es clasificar la intención del usuario y responder SOLO con un objeto JSON válido.
No añadas texto antes ni después.

Ejemplos:
Usuario: hola
Tu JSON: {"action":"chat","response":"¡Hola! Soy Carl IA, asistente de Carley Interactive Studio. ¿En qué puedo ayudarte?"}

Usuario: crea una foto de un perro
Tu JSON: {"action":"generate_image","prompt":"un perro"}

Usuario: quiero hablar
Tu JSON: {"action":"chat","response":"Claro, ¡hablemos! ¿Qué tema te interesa?"}
"""

    full_prompt = f"{system_prompt}\nUsuario: {user_prompt}\nTu JSON:"

    inputs = tokenizer(full_prompt, return_tensors="pt")
    input_ids = inputs.input_ids
    attention_mask = inputs.attention_mask
    input_length = input_ids.shape[1]

    outputs = model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_new_tokens=150,
        pad_token_id=tokenizer.eos_token_id
    )

    generated_tokens = outputs[0, input_length:]
    generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)

    return extract_json(generated_text)

# --- Endpoint de la API ---
@app.post("/chat/")
async def chat_with_carl(request: ChatRequest):
    user_prompt = request.prompt
    print(f"Recibido prompt del usuario: {user_prompt}")

    structured_response = generate_llm_response(user_prompt)
    action = structured_response.get("action")

    if action == "generate_image":
        image_prompt = structured_response.get("prompt", "una imagen aleatoria")
        print(f"LLM decidió generar una imagen con el prompt: '{image_prompt}'")

        try:
            response_from_artist = requests.post(IMAGE_SERVICE_URL, json={"prompt": image_prompt})
            if response_from_artist.status_code == 200:
                return {"type": "image", "content": response_from_artist.content.hex()}
            else:
                raise HTTPException(status_code=500, detail="El servicio de imágenes falló.")
        except requests.exceptions.RequestException as e:
            print(f"Error al conectar con el servicio de imágenes: {e}")
            raise HTTPException(status_code=500, detail="No se pudo conectar con el servicio de imágenes.")

    elif action == "chat":
        chat_response = structured_response.get("response", "No sé qué decir.")
        print(f"LLM decidió chatear con la respuesta: '{chat_response}'")
        return {"type": "text", "content": chat_response}

    else:
        return {"type": "text", "content": "Hubo un error al procesar tu solicitud."}

# --- Endpoint de Bienvenida ---
@app.get("/")
def read_root():
    return {"status": "Carl IA, asistente de Carley Interactive Studio, está en línea."}