orchestrator.py

import logging
from datetime import datetime
from typing import List
import os

class TaskType:
    CONVERSATION = "conversación"
    CODE_GENERATION = "código"
    CODE_EXPLANATION = "explicación_código"
    CREATIVE_WRITING = "escritura_creativa"
    TECHNICAL_EXPLANATION = "explicación_técnica"
    TRANSLATION = "traducción"
    SUMMARIZATION = "resumen"
    MATH_REASONING = "razonamiento_matemático"
    IMAGE_GENERATION = "generación_imagen"
    RESEARCH = "investigación"
    DEBUGGING = "depuración"

class SuperConfig:
    def __init__(self, deepseek_api_key="", openai_api_key="", hf_token="", enable_local_models=True, debug_mode=False, max_history=20):
        self.deepseek_api_key = deepseek_api_key
        self.openai_api_key = openai_api_key
        self.hf_token = hf_token
        self.enable_local_models = enable_local_models
        self.debug_mode = debug_mode
        self.max_history = max_history

class SuperResponse:
    def __init__(self, content, task_type, source, confidence, processing_time):
        self.content = content
        self.task_type = task_type
        self.source = source
        self.confidence = confidence
        self.processing_time = processing_time

try:
    from models.local_chat_model import ChatLocal
    from models.local_code_model import CodeLocal
    from models.local_image_model import ImageLocal
except ImportError:
    ChatLocal = None
    CodeLocal = None
    ImageLocal = None

class SuperOrchestrator:
    def __init__(self, config: SuperConfig):
        self.config = config
        self.logger = logging.getLogger("super_chatbot")
        self._setup_logger()
        self.history: List[str] = []
        self.local_chat = ChatLocal() if (ChatLocal and config.enable_local_models) else None
        self.local_code = CodeLocal() if (CodeLocal and config.enable_local_models) else None
        self.local_image = ImageLocal() if (ImageLocal and config.enable_local_models) else None

    def _setup_logger(self):
        if not self.logger.handlers:
            self.logger.setLevel(logging.DEBUG if self.config.debug_mode else logging.INFO)
            ch = logging.StreamHandler()
            ch.setLevel(logging.DEBUG)
            self.logger.addHandler(ch)

    def process(self, message: str, task_type: str = TaskType.CONVERSATION) -> SuperResponse:
        start = datetime.now()
        self.history.append(message)
        if len(self.history) > self.config.max_history:
            self.history = self.history[-self.config.max_history:]

        content = None
        source = "local"

        if task_type in [TaskType.CONVERSATION, TaskType.CREATIVE_WRITING]:
            if self.local_chat:
                try:
                    content = self.local_chat.generate(message, task_type)
                    source = "local_chat"
                except Exception as e:
                    self.logger.error(f"Fallo local_chat: {str(e)}")
                    content = None

        if not content and self.local_code:
            try:
                content = self.local_code.generate_code(message, task_type)
                source = "local_code"
            except Exception as e:
                self.logger.error(f"Fallo local_code: {str(e)}")
                content = None

        if not content and self.local_image and task_type == TaskType.IMAGE_GENERATION:
            try:
                content = self.local_image.generate(message)
                source = "local_image"
            except Exception as e:
                self.logger.error(f"Fallo local_image: {str(e)}")
                content = None

        if not content:
            content = "Lo siento, no pude generar una respuesta con los modelos locales disponibles."
            source = "fallback"

        duration = (datetime.now() - start).total_seconds()
        return SuperResponse(content, task_type, source, 0.85, duration)