llm_manager.py

import os
import json
import re
import traceback
from typing import Any, Dict, List, Type, Union
from pydantic import BaseModel
from dotenv import load_dotenv

# LangChain imports
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage
from langchain_core.runnables import Runnable

load_dotenv()

# List of models to try in order
FALLBACK_MODELS = [
    {"provider": "google", "model": "gemini-2.5-flash"},
    {"provider": "groq", "model": "llama-3.3-70b-versatile"},
    {"provider": "groq", "model": "llama-3.1-8b-instant"},
    {"provider": "openrouter", "model": "inclusionai/ling-3.0-flash:free"},
    {"provider": "openrouter", "model": "poolside/laguna-s-2.1:free"},
    {"provider": "openrouter", "model": "poolside/laguna-xs-2.1:free"}
]

class FallbackLLMManager:
    def __init__(self):
        self.current_model_idx = 0
        self.google_api_key = os.getenv("GOOGLE_API_KEY")
        self.groq_api_key = os.getenv("GROQ_API_KEY")
        self.openrouter_api_key = os.getenv("OPENROUTER_API_KEY")

    def _get_llm(self, provider: str, model_name: str) -> Union[ChatGoogleGenerativeAI, ChatOpenAI]:
        if provider == "google":
            if not self.google_api_key:
                raise ValueError("GOOGLE_API_KEY is not set")
            return ChatGoogleGenerativeAI(
                model=model_name,
                google_api_key=self.google_api_key,
                temperature=0.7
            )
        elif provider == "groq":
            if not self.groq_api_key:
                raise ValueError("GROQ_API_KEY is not set")
            return ChatOpenAI(
                model=model_name,
                openai_api_key=self.groq_api_key,
                openai_api_base="https://api.groq.com/openai/v1",
                temperature=0.7
            )
        elif provider == "openrouter":
            if not self.openrouter_api_key:
                raise ValueError("OPENROUTER_API_KEY is not set")
            return ChatOpenAI(
                model=model_name,
                openai_api_key=self.openrouter_api_key,
                openai_api_base="https://openrouter.ai/api/v1",
                temperature=0.7
            )
        else:
            raise ValueError(f"Unknown provider: {provider}")

    def invoke(self, messages: List[BaseMessage], **kwargs) -> AIMessage:
        attempts = 0
        max_attempts = len(FALLBACK_MODELS)
        
        while attempts < max_attempts:
            model_info = FALLBACK_MODELS[self.current_model_idx]
            provider = model_info["provider"]
            model_name = model_info["model"]
            
            try:
                print(f"[LLM] Attempting invoke with model: {provider}/{model_name}")
                llm = self._get_llm(provider, model_name)
                response = llm.invoke(messages, **kwargs)
                return response
            except Exception as e:
                print(f"[LLM] Error using {provider}/{model_name}: {e}")
                # Move to next model
                self.current_model_idx = (self.current_model_idx + 1) % len(FALLBACK_MODELS)
                attempts += 1
                
        raise RuntimeError("All fallback models failed to respond to invoke.")

    def with_structured_output(self, schema: Type[BaseModel]) -> 'StructuredFallbackLLM':
        return StructuredFallbackLLM(self, schema)


class StructuredFallbackLLM:
    def __init__(self, manager: FallbackLLMManager, schema: Type[BaseModel]):
        self.manager = manager
        self.schema = schema

    def _clean_json_string(self, text: str) -> str:
        # Find first '{' and last '}'
        start_idx = text.find('{')
        end_idx = text.rfind('}')
        if start_idx != -1 and end_idx != -1:
            return text[start_idx:end_idx + 1]
        return text

    def invoke(self, messages: List[BaseMessage], **kwargs) -> BaseModel:
        attempts = 0
        max_attempts = len(FALLBACK_MODELS)
        
        while attempts < max_attempts:
            model_info = FALLBACK_MODELS[self.manager.current_model_idx]
            provider = model_info["provider"]
            model_name = model_info["model"]
            
            try:
                print(f"[Structured LLM] Attempting structured output with: {provider}/{model_name}")
                llm = self.manager._get_llm(provider, model_name)
                
                # We try native structured output first
                try:
                    structured_llm = llm.with_structured_output(self.schema)
                    res = structured_llm.invoke(messages, **kwargs)
                    return res
                except Exception as structured_err:
                    print(f"[Structured LLM] Native structured output failed for {provider}/{model_name}: {structured_err}")
                    print("[Structured LLM] Falling back to manual JSON prompt parsing...")
                    
                    # Fallback to manual JSON formatting and parsing
                    schema_json = json.dumps(self.schema.model_json_schema(), indent=2)
                    json_instruction = (
                        f"\n\nYou MUST return a JSON object that strictly conforms to the following schema:\n"
                        f"{schema_json}\n\n"
                        f"CRITICAL: Return ONLY valid JSON. Do not include any explanations, markdown code blocks, "
                        f"introductory text, or trailing content. Your entire response must be parseable as a single JSON object."
                    )
                    
                    # Clone messages and modify system message or add instruction
                    fallback_messages = list(messages)
                    if len(fallback_messages) > 0 and isinstance(fallback_messages[0], SystemMessage):
                        fallback_messages[0] = SystemMessage(content=fallback_messages[0].content + json_instruction)
                    else:
                        fallback_messages.append(HumanMessage(content=json_instruction))
                        
                    res_raw = llm.invoke(fallback_messages, **kwargs)
                    raw_text = res_raw.content
                    if isinstance(raw_text, list):
                        raw_text = "".join(item if isinstance(item, str) else item.get("text", "") for item in raw_text)
                    
                    cleaned_json = self._clean_json_string(raw_text)
                    parsed_obj = json.loads(cleaned_json)
                    validated_obj = self.schema.model_validate(parsed_obj)
                    return validated_obj
                    
            except Exception as e:
                print(f"[Structured LLM] Error using {provider}/{model_name}: {e}")
                # Move to next model
                self.manager.current_model_idx = (self.manager.current_model_idx + 1) % len(FALLBACK_MODELS)
                attempts += 1
                
        raise RuntimeError("All fallback models failed to respond to structured invoke.")