src/core/model_manager.py

import os
import time
import json
import asyncio
import requests
from dotenv import load_dotenv

from langchain_core.messages import AIMessage, SystemMessage, HumanMessage
from src.utils.logger import setup_logger

logger = setup_logger("ModelManager")
load_dotenv()

class ReqModel:
    def __init__(self, model: str, temperature: float, base_url: str, api_key: str, headers: dict):
        self.model = model
        self.temperature = temperature
        self.base_url = base_url
        self.api_key = api_key
        self.headers = headers
        self.bound_tools = None

    def bind_tools(self, tools, **kwargs):
        new_model = ReqModel(self.model, self.temperature, self.base_url, self.api_key, self.headers)
        new_model.bound_tools = tools
        return new_model

    def _convert_messages(self, messages):
        req_msgs = []
        for m in messages:
            if isinstance(m, SystemMessage):
                req_msgs.append({"role": "system", "content": m.content})
            elif isinstance(m, HumanMessage):
                req_msgs.append({"role": "user", "content": m.content})
            elif isinstance(m, AIMessage):
                req_msgs.append({"role": "assistant", "content": m.content})
            elif isinstance(m, dict) and "role" in m and "content" in m:
                req_msgs.append(m)
            else:
                req_msgs.append({"role": "user", "content": str(getattr(m, 'content', m))})
        return req_msgs

    def _format_tools(self):
        if not self.bound_tools:
            return None
        tools_list = []
        for tool in self.bound_tools:
            if hasattr(tool, "name") and hasattr(tool, "description") and hasattr(tool, "args_schema"):
                tools_list.append({
                    "type": "function",
                    "function": {
                        "name": tool.name,
                        "description": tool.description,
                        "parameters": tool.args_schema.schema() if tool.args_schema else {"type": "object", "properties": {}}
                    }
                })
        return tools_list

    def _make_request(self, messages, config=None, **kwargs):
        url = f"{self.base_url}/chat/completions"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        headers.update(self.headers)
        
        payload = {
            "model": self.model,
            "messages": self._convert_messages(messages),
            "temperature": self.temperature,
        }
        
        formatted_tools = self._format_tools()
        if formatted_tools:
            payload["tools"] = formatted_tools
            
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()
        data = response.json()
        
        message = data["choices"][0]["message"]
        content = message.get("content", "")
        ai_message = AIMessage(content=content if content else "")
        
        if "tool_calls" in message and message["tool_calls"]:
            tool_calls = []
            for tc in message["tool_calls"]:
                try:
                    args = json.loads(tc["function"]["arguments"])
                except Exception:
                    args = {}
                tool_calls.append({
                    "name": tc["function"]["name"],
                    "args": args,
                    "id": tc["id"]
                })
            ai_message.additional_kwargs["tool_calls"] = message["tool_calls"]
            ai_message.tool_calls = tool_calls

        ai_message.response_metadata = {"token_usage": data.get("usage", {})}
        return ai_message

    def invoke(self, messages, config=None, **kwargs):
        return self._make_request(messages, config, **kwargs)

    async def ainvoke(self, messages, config=None, **kwargs):
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, lambda: self._make_request(messages, config, **kwargs))

    def stream(self, messages, config=None, **kwargs):
        yield self._make_request(messages, config, **kwargs)

    async def astream(self, messages, config=None, **kwargs):
        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(None, lambda: self._make_request(messages, config, **kwargs))
        yield response

class RateLimitFallbackWrapper:
    def __init__(self, main_llm, fallback_llms):
        self.main_llm = main_llm
        self.fallback_llms = fallback_llms
        self.bound_tools = None

    def bind_tools(self, tools, **kwargs):
        new_main = self.main_llm.bind_tools(tools, **kwargs)
        new_falls = [llm.bind_tools(tools, **kwargs) for llm in self.fallback_llms]
        new_wrapper = RateLimitFallbackWrapper(new_main, new_falls)
        new_wrapper.bound_tools = tools
        return new_wrapper

    async def ainvoke(self, messages, config=None, **kwargs):
        try:
            return await self.main_llm.ainvoke(messages, config=config, **kwargs)
        except Exception as e:
            logger.warning(f"LLM Error with main model: {e}. Attempting fallbacks immediately.")

        for idx, fb_llm in enumerate(self.fallback_llms):
            try:
                logger.info(f"Trying fallback model {idx+1} [Model: {fb_llm.model}]")
                return await fb_llm.ainvoke(messages, config=config, **kwargs)
            except Exception as fb_e:
                logger.warning(f"Fallback {idx+1} failed: {fb_e}")

        raise RuntimeError("All models (main and fallbacks) failed.") from None

    def invoke(self, messages, config=None, **kwargs):
        try:
            return self.main_llm.invoke(messages, config=config, **kwargs)
        except Exception as e:
            logger.warning(f"LLM Error with main model: {e}. Attempting fallbacks immediately.")

        for idx, fb_llm in enumerate(self.fallback_llms):
            try:
                logger.info(f"Trying fallback model {idx+1}")
                return fb_llm.invoke(messages, config=config, **kwargs)
            except Exception as fb_e:
                logger.warning(f"Fallback {idx+1} failed: {fb_e}")

        raise RuntimeError("All models failed synchronously.") from None

    def stream(self, messages, config=None, **kwargs):
        try:
            yield from self.main_llm.stream(messages, config=config, **kwargs)
            return
        except Exception as e:
            logger.warning(f"LLM stream error with main model: {e}. Attempting fallbacks immediately.")

        for idx, fb_llm in enumerate(self.fallback_llms):
            try:
                logger.info(f"Trying fallback model {idx+1} for streaming")
                yield from fb_llm.stream(messages, config=config, **kwargs)
                return
            except Exception as fb_e:
                logger.warning(f"Fallback {idx+1} streaming failed: {fb_e}")

        raise RuntimeError("All models failed while streaming.")

    async def astream(self, messages, config=None, **kwargs):
        try:
            async for chunk in self.main_llm.astream(messages, config=config, **kwargs):
                yield chunk
            return
        except Exception as e:
            logger.warning(f"LLM async stream error with main model: {e}. Attempting fallbacks immediately.")

        for idx, fb_llm in enumerate(self.fallback_llms):
            try:
                logger.info(f"Trying fallback model {idx+1} for async streaming")
                async for chunk in fb_llm.astream(messages, config=config, **kwargs):
                    yield chunk
                return
            except Exception as fb_e:
                logger.warning(f"Fallback {idx+1} async streaming failed: {fb_e}")

        raise RuntimeError("All models failed while async streaming.")

class ModelManager:
    def __init__(self, model_name: str = "google/gemma-4-26b-a4b-it:free"):
        self.provider = os.getenv("MODEL_PROVIDER", "openrouter").lower()
        self.model_name = os.getenv("OPENROUTER_MODEL_NAME", model_name)

    def _get_openrouter_api_keys(self):
        primary_api_key = os.getenv("OPENROUTER_API_KEY")
        secondary_api_key = os.getenv("OPENROUTER_SECONDARY_API_KEY")

        if not primary_api_key and secondary_api_key:
            logger.warning("OPENROUTER_API_KEY missing; using OPENROUTER_SECONDARY_API_KEY as the active key.")
            primary_api_key = secondary_api_key
            secondary_api_key = None

        if not primary_api_key:
            raise EnvironmentError(
                "OpenRouter requires OPENROUTER_API_KEY or OPENROUTER_SECONDARY_API_KEY in the environment."
            )

        return primary_api_key, secondary_api_key

    def get_llm(self, temperature: float = 0, model_name: str = None):
        model = model_name or self.model_name
        logger.info(f"Initializing LLM: Provider={self.provider}, Model={model}")
        
        primary_api_key, secondary_api_key = self._get_openrouter_api_keys()
        base_url = "https://openrouter.ai/api/v1"

        main_llm = ReqModel(
            model=model,
            temperature=temperature,
            base_url=base_url,
            api_key=primary_api_key,
            headers={
                "HTTP-Referer": "https://github.com/Sudharshan-3904/dmChatbot",
                "X-Title": "Medical AI Chatbot"
            }
        )

        fallback_llms = []
        if secondary_api_key:
            fallback_llms.append(
                ReqModel(
                    model=model,
                    temperature=temperature,
                    base_url=base_url,
                    api_key=secondary_api_key,
                    headers={
                        "HTTP-Referer": "https://github.com/Sudharshan-3904/dmChatbot",
                        "X-Title": "Medical AI Chatbot"
                    }
                )
            )

        fallback_models = [
            "google/gemma-4-26b-a4b-it:free",
            "google/gemma-4-31b-it:free",
            "openai/gpt-oss-20b:free"
        ]
        
        fallback_llms.extend([
            ReqModel(
                model=m,
                temperature=temperature,
                base_url=base_url,
                api_key=primary_api_key,
                headers={
                    "HTTP-Referer": "https://github.com/Sudharshan-3904/dmChatbot",
                    "X-Title": "Medical AI Chatbot"
                }
            ) for m in fallback_models if m != model
        ])
        
        return RateLimitFallbackWrapper(main_llm, fallback_llms)

model_manager = ModelManager()