from __future__ import annotations
import os
import json
import base64
from copy import deepcopy
from dotenv import load_dotenv
load_dotenv()
from pathlib import Path
from docx import Document
from typing import Optional, Tuple, Dict, Any, Generator, Union
from dataclasses import dataclass
from pydantic import BaseModel
#from pydantic_processor import PydanticGuidelinesProcessor # experimental, might be removed
from openai import OpenAI
from xai_sdk import Client as grok_client
from xai_sdk.chat import system, user, image
from google.genai import Client as google_client
from google.genai import types
from groq import Groq
import anthropic
from exa_py import Exa
import pymupdf
import httpx
import requests
from serpapi import GoogleSearch
import time
import traceback
import warnings
warnings.filterwarnings("ignore")
import pprint


@dataclass
class StreamChunk:
    """Represents a single chunk during streaming."""
    content: str      # Accumulated content so far
    delta: str        # New token(s) in this chunk
    token_usage: Optional[dict] = None
    cost: Optional[float] = None
    done: bool = False


@dataclass
class StreamResult:
    """Final result returned after streaming completes."""
    content: str
    execution_time: float
    token_usage: dict
    cost: float


# Define model-specific configurations
MODEL_CONFIGS = {
    "gpt-4o-mini": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.15 * 1e-6,
        "price_per_completion_tokens": 0.6 * 1e-6,
        "context_length": 128000,
        "max_tokens": 16000,
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-4o": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 2.50 * 1e-6,
        "price_per_completion_tokens": 10.00 * 1e-6,
        "context_length": 128000,
        "max_tokens": 16000,
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-4.1": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 2.00 * 1e-6,
        "price_per_completion_tokens": 8.00 * 1e-6,
        "context_length": 1000000,
        "max_tokens": 32000,
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-4.1-mini": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.40 * 1e-6,
        "price_per_completion_tokens": 1.60 * 1e-6,
        "context_length": 1000000,
        "max_tokens": 32000,
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-4.1-nano": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.1 * 1e-6,
        "price_per_completion_tokens": 0.4 * 1e-6,
        "context_length": 1000000,
        "max_tokens": 32000,
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-5.1-chat-latest": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 5.0 * 1e-6,
        "price_per_completion_tokens": 15.00 * 1e-6,
        "context_length": 128000,
        "max_tokens": 16000,
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "o3": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 2.00 * 1e-6,
        "price_per_completion_tokens": 8.00 * 1e-6,
        "context_length": 200000,
        "max_tokens": 100000,
        "reasoning_effort": "high",
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "o3-mini": {
        "structured_output": True,
        "vision": False,
        "price_per_prompt_tokens": 1.1 * 1e-6,
        "price_per_completion_tokens": 4.40 * 1e-6,
        "context_length": 200000,
        "max_tokens": 100000,
        "reasoning_effort": "medium",
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "o4-mini": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 1.1 * 1e-6,
        "price_per_completion_tokens": 4.40 * 1e-6,
        "context_length": 200000,
        "max_tokens": 100000,
        "reasoning_effort": "high",
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-5": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 1.25 * 1e-6,                 
        "price_per_completion_tokens": 10.00 * 1e-6,         
        "context_length": 256000,                            
        "max_tokens": 128000,
        "reasoning_effort": "medium",                  
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-5-nano": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.05 * 1e-6,                 
        "price_per_completion_tokens": 0.4 * 1e-6,         
        "context_length": 400000,                            
        "max_tokens": 128000,
        "reasoning_effort": "medium",                  
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-5-mini": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.25 * 1e-6,                 
        "price_per_completion_tokens": 2.00 * 1e-6,         
        "context_length": 400000,                            
        "max_tokens": 128000,
        "reasoning_effort": "medium",                  
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-5.2": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 1.75 * 1e-6,                 
        "price_per_completion_tokens": 14.00 * 1e-6,         
        "context_length": 400000,                            
        "max_tokens": 128000,
        "reasoning_effort": "medium",                  
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-5.2-pro": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 21 * 1e-6,                 
        "price_per_completion_tokens": 168.00 * 1e-6,         
        "context_length": 400000,                            
        "max_tokens": 128000,
        "reasoning_effort": "medium",                  
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-5.4": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 1.75 * 1e-6,
        "price_per_completion_tokens": 14.00 * 1e-6,
        "context_length": 400000,
        "max_tokens": 128000,
        "reasoning_effort": "medium",
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-5.4-pro": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 21.00 * 1e-6,
        "price_per_completion_tokens": 168.00 * 1e-6,         
        "context_length": 400000,                            
        "max_tokens": 128000,
        "reasoning_effort": "medium",                  
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-5.4": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 1.75 * 1e-6,
        "price_per_completion_tokens": 14.00 * 1e-6,
        "context_length": 400000,
        "max_tokens": 128000,
        "reasoning_effort": "medium",
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "gpt-5.4-pro": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 21.00 * 1e-6,
        "price_per_completion_tokens": 168.00 * 1e-6,
        "context_length": 400000,
        "max_tokens": 128000,
        "reasoning_effort": "medium",
        "client_config": {
            "api_key_env": "OPENAI_API_KEY",
            "base_url": "https://api.openai.com/v1",
            "client_class": "OpenAI"
        }
    },
    "deepseek-chat": {
        "structured_output": True,
        "vision": False,
        "price_per_prompt_tokens": 0.56 * 1e-6,
        "price_per_completion_tokens": 1.68 * 1e-6,
        "context_length": 64000,
        "max_tokens": 8096,
        "client_config": {
            "api_key_env": "DEEPSEEK_API_KEY",
            "base_url": "https://api.deepseek.com",
            "client_class": "DeepSeek"
        },
        "prompt_limit": 70000
    },
    "deepseek-reasoner": {
        "structured_output": True,
        "vision": False,
        "price_per_prompt_tokens": 0.56 * 1e-6,
        "price_per_completion_tokens": 1.68 * 1e-6,
        "context_length": 64000,
        "max_tokens": 8096,
        "client_config": {
            "api_key_env": "DEEPSEEK_API_KEY",
            "base_url": "https://api.deepseek.com",
            "client_class": "DeepSeek"
        },
        "prompt_limit": 70000
    },
    "llama-3.1-8b-instant": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.05 * 1e-6,
        "price_per_completion_tokens": 0.08 * 1e-6,
        "context_length": 131000,
        "max_tokens": 131000,
        "client_config": {
            "api_key_env": "GROQ_API_KEY",
            "client_class": "Groq"
        },
        "free_tier": {
            "max_tokens": 7999,
            "prompt_limit": 6000
        }
    },
    "llama-3.3-70b-versatile": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.59 * 1e-6,
        "price_per_completion_tokens": 0.79 * 1e-6,
        "context_length": 131000,
        "max_tokens": 32000,
        "client_config": {
            "api_key_env": "GROQ_API_KEY",
            "client_class": "Groq"
        },
        "free_tier": {
            "max_tokens": 7999,
            "prompt_limit": 6000
        }
    },
    "moonshotai/kimi-k2-instruct-0905": {
        "structured_output": True,
        "vision": False,
        "price_per_prompt_tokens": 1.00 * 1e-6,
        "price_per_completion_tokens": 3.00 * 1e-6,
        "context_length": 131072,
        "max_tokens": 16384,
        "client_config": {
            "api_key_env": "GROQ_API_KEY",
            "client_class": "Groq"
        },
        "free_tier": {
            "max_tokens": 7999,
            "prompt_limit": 6000
        }
    },
    "openai/gpt-oss-120b": {
        "structured_output": True,
        "vision": False,
        "price_per_prompt_tokens": 0.15 * 1e-6,
        "price_per_completion_tokens": 0.75 * 1e-6,
        "context_length": 131072,
        "max_tokens": 65526,
        "client_config": {
            "api_key_env": "GROQ_API_KEY",
            "client_class": "Groq"
        },
        "free_tier": {
            "max_tokens": 7999,
            "prompt_limit": 6000
        }
    },
    "openai/gpt-oss-20b": {
        "structured_output": True,
        "vision": False,
        "price_per_prompt_tokens": 0.10 * 1e-6,
        "price_per_completion_tokens": 0.50 * 1e-6,
        "context_length": 131072,
        "max_tokens": 65526,
        "client_config": {
            "api_key_env": "GROQ_API_KEY",
            "client_class": "Groq"
        },
        "free_tier": {
            "max_tokens": 7999,
            "prompt_limit": 6000
        }
    },
    "meta-llama/llama-4-scout-17b-16e-instruct": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.11 * 1e-6,
        "price_per_completion_tokens": 0.34 * 1e-6,
        "context_length": 	131072,
        "max_tokens": 8192,
        "client_config": {
            "api_key_env": "GROQ_API_KEY",
            "client_class": "Groq"
        },
        "free_tier": {
            "max_tokens": 7999,
            "prompt_limit": 6000
        }
    },
    "openai/gpt-oss-safeguard-20b": {
        "structured_output": True,
        "vision": False,
        "price_per_prompt_tokens": 0.10 * 1e-6,
        "price_per_completion_tokens": 0.50 * 1e-6,
        "context_length": 131072,
        "max_tokens": 65526,
        "client_config": {
            "api_key_env": "GROQ_API_KEY",
            "client_class": "Groq"
        },
        "free_tier": {
            "max_tokens": 7999,
            "prompt_limit": 6000
        }
    },
    "qwen-max-latest": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 1.6 * 1e-6,
        "price_per_completion_tokens": 6.4 * 1e-6,
        "context_length": 30720,
        "max_tokens": 8192,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen-turbo-latest": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.05 * 1e-6,
        "price_per_completion_tokens": 0.2 * 1e-6,
        "context_length": 1000000,
        "max_tokens": 16384,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3.5-plus": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.4 * 1e-6,
        "price_per_completion_tokens": 2.4 * 1e-6,
        "context_length": 983616,
        "max_tokens": 65536,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3.5-plus-thinking": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.4 * 1e-6,
        "price_per_completion_tokens": 2.4 * 1e-6,
        "context_length": 983616,
        "max_tokens": 65536,
        "thinking": {
            "enable_thinking": True
        },
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen-plus-latest": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.4 * 1e-6,
        "price_per_completion_tokens": 1.2 * 1e-6,
        "context_length": 995904,
        "max_tokens": 32768,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen-plus-latest-thinking": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.4 * 1e-6,
        "price_per_completion_tokens": 1.2 * 1e-6,
        "context_length": 995904,
        "max_tokens": 32768,
        "thinking": {
            "enable_thinking": True
        },
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen-flash": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.05 * 1e-6,
        "price_per_completion_tokens": 0.40 * 1e-6,
        "context_length": 1000000,
        "max_tokens": 32768,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3-max": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 1.20 * 1e-6,
        "price_per_completion_tokens": 6.00 * 1e-6,
        "context_length": 262144,
        "max_tokens": 65536,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3.5-flash": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.10 * 1e-6,
        "price_per_completion_tokens": 0.40 * 1e-6,
        "context_length": 1000000,
        "max_tokens": 65536,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen-flash-thinking": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.05 * 1e-6,
        "price_per_completion_tokens": 0.40 * 1e-6,
        "context_length": 1000000,
        "max_tokens": 32768,
        "thinking": {
            "enable_thinking": True
        },
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3-max": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 1.20 * 1e-6,
        "price_per_completion_tokens": 6.00 * 1e-6,
        "context_length": 262144,
        "max_tokens": 65536,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3-max-thinking": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 1.20 * 1e-6,
        "price_per_completion_tokens": 6.00 * 1e-6,
        "context_length": 262144,
        "max_tokens": 65536,
        "thinking": {
            "enable_thinking": True
        },
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3.5-flash": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.10 * 1e-6,
        "price_per_completion_tokens": 0.40 * 1e-6,
        "context_length": 1000000,
        "max_tokens": 65536,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3.5-flash-thinking": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.10 * 1e-6,
        "price_per_completion_tokens": 0.40 * 1e-6,
        "context_length": 1000000,
        "max_tokens": 65536,
        "thinking": {
            "enable_thinking": True
        },
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3-235b-a22b-thinking-2507": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.7 * 1e-6,
        "price_per_completion_tokens": 8.4 * 1e-6,
        "context_length": 126976,
        "max_tokens": 32768,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3-235b-a22b-instruct-2507": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.7 * 1e-6,
        "price_per_completion_tokens": 2.8 * 1e-6,
        "context_length": 129024,
        "max_tokens": 32768,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3-30b-a3b-thinking-2507": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.2 * 1e-6,
        "price_per_completion_tokens": 2.4 * 1e-6,
        "context_length": 126976,
        "max_tokens": 32768,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "qwen3-30b-a3b-instruct-2507": {
        "structured_output": False,
        "vision": False,
        "price_per_prompt_tokens": 0.20 * 1e-6,
        "price_per_completion_tokens": 0.80 * 1e-6,
        "context_length": 126976,
        "max_tokens": 32768,
        "client_config": {
            "api_key_env": "QWEN_API_KEY",
            "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
            "client_class": "Alibaba"
        }
    },
    "grok-3-mini-latest": {
        "structured_output": True,
        "vision": False,
        "price_per_prompt_tokens": 0.30 * 1e-6,
        "price_per_completion_tokens": 0.50 * 1e-6,
        "context_length": 131072,
        "max_tokens": 131072,
        "client_config": {
            "api_key_env": "GROK_API_KEY",
            "base_url": "https://api.x.ai/v1",
            "client_class": "Grok"
        }
    },
    "grok-3-latest": {
        "structured_output": True,
        "vision": False,
        "price_per_prompt_tokens": 3.00 * 1e-6,
        "price_per_completion_tokens": 15.00 * 1e-6,
        "context_length": 131072,
        "max_tokens": 131072,
        "client_config": {
            "api_key_env": "GROK_API_KEY",
            "base_url": "https://api.x.ai/v1",
            "client_class": "Grok"
        }
    },
    "grok-3-mini-fast-latest": {
        "structured_output": True,
        "vision": False,
        "price_per_prompt_tokens": 0.60 * 1e-6,
        "price_per_completion_tokens": 4.00 * 1e-6,
        "context_length": 131072,
        "max_tokens": 131072,
        "client_config": {
            "api_key_env": "GROK_API_KEY",
            "base_url": "https://api.x.ai/v1",
            "client_class": "Grok"
        }
    },
    "grok-4-latest": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 3.00 * 1e-6,
        "price_per_completion_tokens": 15.00 * 1e-6,
        "context_length": 256000,
        "max_tokens": 256000,
        "client_config": {
            "api_key_env": "GROK_API_KEY",
            "base_url": "https://api.x.ai/v1",
            "client_class": "Grok"
        }
    },
    "grok-4-fast-non-reasoning": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.20 * 1e-6,
        "price_per_completion_tokens": 0.50 * 1e-6,
        "context_length": 2000000,
        "max_tokens": 256000,
        "client_config": {
            "api_key_env": "GROK_API_KEY",
            "base_url": "https://api.x.ai/v1",
            "client_class": "Grok"
        }
    },
    "grok-4-fast-reasoning": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.20 * 1e-6,
        "price_per_completion_tokens": 0.50 * 1e-6,
        "context_length": 2000000,
        "max_tokens": 256000,
        "client_config": {
            "api_key_env": "GROK_API_KEY",
            "base_url": "https://api.x.ai/v1",
            "client_class": "Grok"
        }
    },
    "grok-4-1-fast-non-reasoning": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.20 * 1e-6,
        "price_per_completion_tokens": 0.50 * 1e-6,
        "context_length": 2000000,
        "max_tokens": 256000,
        "client_config": {
            "api_key_env": "GROK_API_KEY",
            "base_url": "https://api.x.ai/v1",
            "client_class": "Grok"
        }
    },
    "grok-4-1-fast-reasoning": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.20 * 1e-6,
        "price_per_completion_tokens": 0.50 * 1e-6,
        "context_length": 2000000,
        "max_tokens": 256000,
        "client_config": {
            "api_key_env": "GROK_API_KEY",
            "base_url": "https://api.x.ai/v1",
            "client_class": "Grok"
        }
    },
    "grok-4.20-beta": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 3.00 * 1e-6,
        "price_per_completion_tokens": 15.00 * 1e-6,
        "context_length": 2000000,
        "max_tokens": 256000,
        "client_config": {
            "api_key_env": "GROK_API_KEY",
            "base_url": "https://api.x.ai/v1",
            "client_class": "Grok"
        }
    },
    "grok-4.20-beta": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 3.00 * 1e-6,
        "price_per_completion_tokens": 15.00 * 1e-6,
        "context_length": 2000000,
        "max_tokens": 256000,
        "client_config": {
            "api_key_env": "GROK_API_KEY",
            "base_url": "https://api.x.ai/v1",
            "client_class": "Grok"
        }
    },
    "claude-haiku-4-5": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 1.00 * 1e-6,
        "price_per_completion_tokens": 5.00 * 1e-6,
        "context_length": 200000,
        "max_tokens": 64000,
        "client_config": {
            "api_key_env": "ANTHROPIC_API_KEY",
            "base_url": "https://api.anthropic.com/v1",
            "client_class": "Anthropic"
        }
    },
    "claude-haiku-4-5-thinking": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 1.00 * 1e-6,
        "price_per_completion_tokens": 5.00 * 1e-6,
        "context_length": 200000,
        "max_tokens": 64000,
        "thinking": {
            "type": "enabled",
            "budget_tokens": 4096
        },
        "client_config": {
            "api_key_env": "ANTHROPIC_API_KEY",
            "base_url": "https://api.anthropic.com/v1",
            "client_class": "Anthropic"
        }
    },
    "claude-sonnet-4-6": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 3.00 * 1e-6,
        "price_per_completion_tokens": 3.75 * 1e-6,
        "context_length": 200000,
        "max_tokens": 64000,
        "client_config": {
            "api_key_env": "ANTHROPIC_API_KEY",
            "base_url": "https://api.anthropic.com/v1",
            "client_class": "Anthropic"
        }
    },
    "claude-sonnet-4-6-thinking": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 3.00 * 1e-6,
        "price_per_completion_tokens": 3.75 * 1e-6,
        "context_length": 200000,
        "max_tokens": 64000,
        "thinking": {
            "type": "enabled",
            "budget_tokens": 4096
        },
        "client_config": {
            "api_key_env": "ANTHROPIC_API_KEY",
            "base_url": "https://api.anthropic.com/v1",
            "client_class": "Anthropic"
        }
    },
    "claude-sonnet-4-5-thinking": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 3.00 * 1e-6,
        "price_per_completion_tokens": 15.00 * 1e-6,
        "context_length": 200000,
        "max_tokens": 64000,
        "thinking": {
            "type": "enabled",
            "budget_tokens": 4096
        },
        "client_config": {
            "api_key_env": "ANTHROPIC_API_KEY",
            "base_url": "https://api.anthropic.com/v1",
            "client_class": "Anthropic"
        }
    },
    "claude-opus-4-5": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 5.00 * 1e-6,
        "price_per_completion_tokens": 25.00 * 1e-6,
        "context_length": 200000,
        "max_tokens": 64000,
        "client_config": {
            "api_key_env": "ANTHROPIC_API_KEY",
            "base_url": "https://api.anthropic.com/v1",
            "client_class": "Anthropic"
        }
    },
    "claude-opus-4-5-thinking": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 5.00 * 1e-6,
        "price_per_completion_tokens": 25.00 * 1e-6,
        "context_length": 200000,
        "max_tokens": 64000,
        "thinking": {
            "type": "enabled",
            "budget_tokens": 4096
        },
        "client_config": {
            "api_key_env": "ANTHROPIC_API_KEY",
            "base_url": "https://api.anthropic.com/v1",
            "client_class": "Anthropic"
        }
    },
    "claude-opus-4-6": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 5.00 * 1e-6,
        "price_per_completion_tokens": 6.25 * 1e-6,
        "context_length": 200000,
        "max_tokens": 64000,
        "client_config": {
            "api_key_env": "ANTHROPIC_API_KEY",
            "base_url": "https://api.anthropic.com/v1",
            "client_class": "Anthropic"
        }
    },
    "claude-opus-4-6-thinking": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 5.00 * 1e-6,
        "price_per_completion_tokens": 6.25 * 1e-6,
        "context_length": 200000,
        "max_tokens": 64000,
        "thinking": {
            "type": "enabled",
            "budget_tokens": 4096
        },
        "client_config": {
            "api_key_env": "ANTHROPIC_API_KEY",
            "base_url": "https://api.anthropic.com/v1",
            "client_class": "Anthropic"
        }
    },
    "claude-sonnet-4-5": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 3.00 * 1e-6,
        "price_per_completion_tokens": 15.00 * 1e-6,
        "context_length": 200000,
        "max_tokens": 64000,
        "client_config": {
            "api_key_env": "ANTHROPIC_API_KEY",
            "base_url": "https://api.anthropic.com/v1",
            "client_class": "Anthropic"
        }
    },
    "gemini-2.0-flash": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.10 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 0.40 * 1e-6, 
        "context_length": 500000,
        "max_tokens": 500000,
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    }, 
    "gemini-2.5-flash": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.3 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 2.50 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65000,
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
    "gemini-3-flash-preview": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.5 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 3.0 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65000,
        "thinking": {
            "thinking_level": "medium"
        },
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
    "gemini-3-flash-preview-thinking-minimal": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.5 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 3.0 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65000,
        "thinking": {
            "thinking_level": "minimal"
        },
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
    "gemini-3-flash-preview-thinking-high": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.5 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 3.0 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65000,
        "thinking": {
            "thinking_level": "high"
        },
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
    "gemini-2.5-pro": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 1.25 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 10.00 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65000,
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
    "gemini-3.1-pro-preview": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 2.00 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 12.00 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65000,
        "thinking": {
            "thinking_level": "medium"
        },
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
    "gemini-3.1-pro-preview-thinking-low": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 2.00 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 12.00 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65000,
        "thinking": {
            "thinking_level": "low"
        },
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
    "gemini-3.1-pro-preview-thinking-high": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 2.00 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 12.00 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65000,
        "thinking": {
            "thinking_level": "high"
        },
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
    "gemini-3.1-flash-lite-preview": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.10 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 0.40 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65536,
        "thinking": {
            "thinking_level": "medium"
        },
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
    "gemini-3.1-flash-lite-preview-thinking-minimal": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.10 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 0.40 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65536,
        "thinking": {
            "thinking_level": "minimal"
        },
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
    "gemini-3.1-flash-lite-preview-thinking-high": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.10 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 0.40 * 1e-6, 
        "context_length": 1000000,
        "max_tokens": 65536,
        "thinking": {
            "thinking_level": "high"
        },
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    }
    "gemini-2.5-flash-lite": {
        "structured_output": True,
        "vision": True,
        "price_per_prompt_tokens": 0.1 * 1e-6, # (text / image / video)
        "price_per_completion_tokens": 0.40 * 1e-6,
        "context_length": 1000000,
        "max_tokens": 65536,
        "client_config": {
            "api_key_env": "GEMINI_API_KEY",
            "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
            "client_class": "Google"
        }
    },
}

class QuerySchema(BaseModel):
    query: str

class SearchSchema(BaseModel):
    result: str
    sources: list[str]

class SimpleLogger:
    """Simple colored logger with 4 methods: info, warning, data, error"""
    
    # ANSI color codes
    CYAN = '\033[96m'
    YELLOW = '\033[93m'
    PURPLE = '\033[95m'
    RED = '\033[91m'
    GREEN = '\033[92m'
    RESET = '\033[0m'
    
    def __init__(self, prefix="LLMProxy", verbose=1):
        self.prefix = prefix
        self.verbose = verbose
    
    def success(self, message):
        """Log success messages in green"""
        if self.verbose >= 1:
            print(f"{self.GREEN}{self.prefix} SUCCESS: {message}{self.RESET}")
    
    def info(self, message):
        """Log info messages in light blue"""
        if self.verbose >= 1:
            print(f"{self.CYAN}{self.prefix} {message}{self.RESET}")
    
    def warning(self, message):
        """Log warning messages in yellow"""
        if self.verbose >= 1:
            print(f"{self.YELLOW}{self.prefix} WARNING: {message}{self.RESET}")
    
    def data(self, message):
        """Log data messages in purple"""
        if self.verbose >= 2:
            print(f"{self.PURPLE}{self.prefix} {message}{self.RESET}")
    
    def error(self, message):
        """Log error messages in red"""
        print(f"{self.RED}{self.prefix} ERROR: {message}{self.RESET}")

class LLMProxy:

    def __init__(self, verbose=1, api_key: Optional[str] = None, enable_timeouts= False, timeouts_options= None, enable_tracing=False):
        
        # Create simple colored logger
        self.logger = SimpleLogger(verbose)
        
        # Set up timeout configuration
        if enable_timeouts and timeouts_options is None:
            timeouts_options = {"total": 120, "read": 60.0, "write": 60.0, "connect": 10.0}
        self.timeout_config = httpx.Timeout(
            timeouts_options["total"],
            read=timeouts_options["read"],
            write=timeouts_options["write"],
            connect=timeouts_options["connect"]
        ) if enable_timeouts else None
        
        # Initialize OpenAI client
        openai_api_key = os.getenv("OPENAI_API_KEY")
        if not openai_api_key:
            self.logger.warning("OpenAI API key not found. Using placeholder for now. Please set OPENAI_API_KEY in your .env file.")
            openai_api_key = "placeholder-key"
        self.openai_client = OpenAI(api_key=openai_api_key,
            timeout=self.timeout_config if enable_timeouts else None
        )
        self.default_openai_base_url = "https://api.openai.com/v1"
        
        # Initialize Anthropic client
        anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
        if not anthropic_api_key:
            self.logger.warning("Anthropic API key not found. Using placeholder for now. Please set ANTHROPIC_API_KEY in your .env file.")
            anthropic_api_key = "placeholder-key"
        self.anthropic_client = anthropic.Anthropic(api_key=anthropic_api_key)
        
        # Initialize Groq client
        groq_api_key = os.getenv("GROQ_API_KEY")
        if not groq_api_key:
            self.logger.warning("Groq API key not found. Using placeholder for now. Please set GROQ_API_KEY in your .env file.")
            groq_api_key = "placeholder-key"
        self.groq_client = Groq(api_key=groq_api_key)
        
        # Initialize Google client
        gemini_api_key = os.getenv("GEMINI_API_KEY")
        if not gemini_api_key:
            self.logger.warning("Gemini API key not found. Using placeholder for now. Please set GEMINI_API_KEY in your .env file.")
            gemini_api_key = "placeholder-key"
        self.google_client = google_client(api_key=gemini_api_key)
        
        # Initialize Grok client
        grok_api_key = os.getenv("GROK_API_KEY")
        if not grok_api_key:
            self.logger.warning("Grok API key not found. Using placeholder for now. Please set GROK_API_KEY in your .env file.")
            grok_api_key = "placeholder-key"
        self.grok_client = grok_client(api_key=grok_api_key)
        
        # Initialize DeepSeek client
        deepseek_api_key = os.getenv("DEEPSEEK_API_KEY")
        if not deepseek_api_key:
            self.logger.warning("DeepSeek API key not found. Using placeholder for now. Please set DEEPSEEK_API_KEY in your .env file.")
            deepseek_api_key = "placeholder-key"
        self.deepseek_client = OpenAI(base_url="https://api.deepseek.com", api_key=deepseek_api_key)
        
        # Initialize Qwen client
        qwen_api_key = os.getenv("QWEN_API_KEY")
        if not qwen_api_key:
            self.logger.warning("Qwen API key not found. Using placeholder for now. Please set QWEN_API_KEY in your .env file.")
            qwen_api_key = "placeholder-key"
        self.qwen_client = OpenAI(base_url="https://dashscope-intl.aliyuncs.com/compatible-mode/v1", api_key=qwen_api_key)
        
        # Initialize Exa client
        exa_key = os.getenv("EXA_API_KEY")
        if not exa_key:
            self.logger.warning("Exa API key not found. Using placeholder for now. Please set EXA_API_KEY in your .env file.")
            exa_key = "placeholder-key"
        self.exa_key = exa_key

        self.exa_client = Exa(api_key=exa_key)

        # Initialize SERP API key
        serp_api_key = os.getenv("SERP_API_KEY")
        if not serp_api_key:
            self.logger.warning("SERP API key not found. Using placeholder for now. Please set SERP_API_KEY in your .env file.")
            serp_api_key = "placeholder-key"
        self.serp_api_key = serp_api_key
        
        # Initialize RunPod settings
        runpod_token = os.getenv("RUNPOD_TOKEN")
        if not runpod_token:
            self.logger.warning("RunPod token not found. Set RUNPOD_TOKEN to enable RunPod provider.")
        self.runpod_token = runpod_token
        self.runpod_session = requests.Session()
        self.runpod_poll_interval = float(os.getenv("RUNPOD_POLL_INTERVAL", "5"))
        self.runpod_poll_timeout = float(os.getenv("RUNPOD_POLL_TIMEOUT", "300"))
        
        # Cache for dynamically loaded RunPod model configs
        self._runpod_model_cache = {}
        
        # Initialize Ollama client
        ollama_api_key = os.getenv("OLLAMA_API_KEY", "ollama")  # Local Ollama doesn't require a real API key
        ollama_base_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1")
        self.ollama_client = OpenAI(
            api_key=ollama_api_key,
            base_url=ollama_base_url,
            timeout=self.timeout_config if enable_timeouts else None
        )
        
        # Cache for clients with custom base URLs
        self._client_cache = {}

        # Map provider names to client and response function
        self._provider_map = {
            "OpenAI": {"client":self.openai_client, "response_function":self.openai_response},
            "Anthropic": {"client":self.anthropic_client, "response_function":self.claude_response},
            "Groq": {"client":self.groq_client, "response_function":self.groq_response},
            "Google": {"client":self.google_client, "response_function":self.gemini_response},
            "DeepSeek": {"client":self.deepseek_client, "response_function":self.deepseek_response},
            "Alibaba": {"client":self.qwen_client, "response_function":self.qwen_response},
            "Grok": {"client":self.grok_client, "response_function":self.grok_response},
            "Ollama": {"client":self.ollama_client, "response_function":self.ollama_response},
            "RunPod": {"client":self.runpod_session, "response_function":self.runpod_response},
        }

        self.best_model = {
            "OpenAI": "gpt-5.1-chat-latest",
            "Anthropic": "claude-sonnet-4-6",
            "Groq": "moonshotai/kimi-k2-instruct-0905",
            "Google": "gemini-2.5-pro",
            "DeepSeek": "deepseek-reasoner",
            "Alibaba": "qwen-max-latest",
            "Grok": "grok-4-latest",
            "Ollama": "ollama/qwen3:4b",
            "RunPod": None,  # Must be configured via environment variables
        }

        self.fast_model = {
            "OpenAI": "gpt-4.1-nano",
            "Anthropic": "claude-haiku-4-5",
            "Groq": "meta-llama/llama-4-scout-17b-16e-instruct",
            "Google": "gemini-2.5-flash-lite",
            "DeepSeek": "deepseek-chat",
            "Alibaba": "qwen-flash",
            "Grok": "grok-3-mini-latest",
            "Ollama": "gemini-2.5-flash-lite",
            "RunPod": None,  # Must be configured via environment variables
        }

        self.default_vision_models = {
            "OpenAI": "gpt-4.1-mini",
            "Anthropic": "claude-sonnet-4-6",
            "Groq": "meta-llama/llama-4-scout-17b-16e-instruct",
            "Google": "gemini-2.5-flash",
            "DeepSeek": "gemini-2.5-flash",
            "Alibaba": "gemini-2.5-flash",
            "Grok": "grok-4-latest",
            "Ollama": "gemini-2.5-flash",
            "RunPod": "gemini-2.5-flash",
        }

        #self.schema_modification_model = "moonshotai/kimi-k2-instruct-0905"
        self.structured_output_fallback_model = "gpt-5.1-chat-latest"
        self.fallback_to_provider_best_model = True
        self.fallback_to_standard_model = True

        """# load guidelines.txt
        with open("guidelines.txt", "r") as f:
            self.guidelines = f.read()"""

        self.json_instructions = """
        Make sure your output is a valid JSON object as per the schema.
        Do not produce a schema, provide a json object that adheres to the pydantic schema and can be parsed into a python dict.
        Do not include any additional text outside the json format.
        Also do not use json markers and backticks like ```json and ```
        This is important because your output will be automatically parsed.
        """
        
        # LangSmith tracing (optional — requires `pip install langsmith`)
        self._traceable = None
        self._get_current_run_tree = None
        self._in_traced_call = False
        if enable_tracing:
            try:
                from langsmith import traceable
                from langsmith import get_current_run_tree
                self._traceable = traceable
                self._get_current_run_tree = get_current_run_tree
                self.logger.success("LangSmith tracing enabled.")
            except ImportError:
                self.logger.warning("enable_tracing=True but 'langsmith' is not installed. Run: pip install langsmith")
    
    def get_model_config(self, model_name: str):
        """Return model configuration, supporting dynamic providers such as Ollama."""
        if model_name in MODEL_CONFIGS:
            return MODEL_CONFIGS[model_name]

        if model_name.startswith("runpod/"):
            # Try to load RunPod model dynamically from environment variables
            dynamic_config = self._load_runpod_config_from_env(model_name)
            if dynamic_config:
                return dynamic_config
            
        if model_name.startswith("ollama/"):
            engine_model = model_name.split("/", 1)[1]
            return {
                "structured_output": True,
                "vision": False,
                "price_per_prompt_tokens": 0.0,
                "price_per_completion_tokens": 0.0,
                "context_length": 1_000_000,
                "max_tokens": 100_000,
                "client_config": {
                    "api_key_env": "OLLAMA_API_KEY",
                    "base_url": os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1"),
                    "client_class": "Ollama"
                },
                "engine_model_name": engine_model
            }

        raise ValueError(f"Unsupported model: {model_name}")

    def _load_runpod_config_from_env(self, model_name: str) -> Optional[dict]:
        """
        Dynamically load RunPod model configuration from environment variables.
        
        Environment variable format:
        RUNPOD_MODEL_<MODEL_NAME_UPPER>_RUN_ENDPOINT
        RUNPOD_MODEL_<MODEL_NAME_UPPER>_STATUS_ENDPOINT
        RUNPOD_MODEL_<MODEL_NAME_UPPER>_CONTEXT_LENGTH
        RUNPOD_MODEL_<MODEL_NAME_UPPER>_MAX_TOKENS
        
        Example for "runpod/openai/gpt-oss-20b":
        RUNPOD_MODEL_OPENAI_GPT_OSS_20B_RUN_ENDPOINT
        """
        # Check cache first
        if model_name in self._runpod_model_cache:
            return self._runpod_model_cache[model_name]
        
        # Convert model name to env var format
        # e.g., "runpod/openai/gpt-oss-20b" -> "OPENAI_GPT_OSS_20B"
        model_part = model_name.replace("runpod/", "").replace("-", "_").replace("/", "_").upper()
        
        # Required env vars
        run_endpoint = os.getenv(f"RUNPOD_MODEL_{model_part}_RUN_ENDPOINT")
        status_endpoint = os.getenv(f"RUNPOD_MODEL_{model_part}_STATUS_ENDPOINT")
        
        if not run_endpoint or not status_endpoint:
            self.logger.warning(f"[RunPod] Missing env vars for {model_name}. Expected: RUNPOD_MODEL_{model_part}_RUN_ENDPOINT and RUNPOD_MODEL_{model_part}_STATUS_ENDPOINT")
            return None
        
        # Optional env vars with defaults
        context_length = int(os.getenv(f"RUNPOD_MODEL_{model_part}_CONTEXT_LENGTH", "8192"))
        max_tokens = int(os.getenv(f"RUNPOD_MODEL_{model_part}_MAX_TOKENS", "4096"))
        
        config = {
            "structured_output": False,
            "vision": False,
            "price_per_prompt_tokens": 0.0,
            "price_per_completion_tokens": 0.0,
            "context_length": context_length,
            "max_tokens": max_tokens,
            "client_config": {
                "client_class": "RunPod",
                "run_endpoint": run_endpoint,
                "status_endpoint": status_endpoint
            }
        }
        
        # Cache and return
        self._runpod_model_cache[model_name] = config
        self.logger.info(f"[RunPod] Loaded dynamic config for {model_name}")
        return config

    def list_available_models(self) -> list[dict]:
        """Return all available model definitions and metadata."""
        models: list[dict] = []
        for model_name, config in MODEL_CONFIGS.items():
            model_info = deepcopy(config)
            model_info["model_name"] = model_name
            model_info["provider"] = config.get("client_config", {}).get("client_class")
            models.append(model_info)
        models.extend(self._discover_runpod_models())
        return sorted(models, key=lambda item: item["model_name"])

    def _discover_runpod_models(self) -> list[dict]:
        """Discover RunPod models configured via environment variables."""
        discovered: list[dict] = []
        prefix = "RUNPOD_MODEL_"
        suffix = "_RUN_ENDPOINT"

        for key in os.environ:
            if not (key.startswith(prefix) and key.endswith(suffix)):
                continue

            token = key[len(prefix):-len(suffix)]
            model_path = self._reconstruct_runpod_model_name(token)
            if not model_path:
                continue

            full_name = f"runpod/{model_path}"
            config = self._load_runpod_config_from_env(full_name)
            if not config:
                continue

            model_info = deepcopy(config)
            model_info["model_name"] = full_name
            model_info["provider"] = "RunPod"
            discovered.append(model_info)

        return discovered

    def _reconstruct_runpod_model_name(self, token: str) -> Optional[str]:
        """Reconstruct runpod/<provider>/<model> from env key token."""
        if not token:
            return None

        normalized = token.strip("_").lower()
        if not normalized:
            return None

        if "_" not in normalized:
            return normalized

        provider, rest = normalized.split("_", 1)
        rest = rest.replace("_", "-")
        return f"{provider}/{rest}"

    def get_client_for_model(self, model_name: str, model_config: dict):
        """Return an API client instance for the given model, honoring custom base URLs."""
        client_cfg = model_config["client_config"]
        provider = client_cfg["client_class"]
        base_url = client_cfg.get("base_url")

        if provider != "OpenAI":
            if provider not in self._provider_map:
                raise ValueError(f"Unknown provider '{provider}' for model '{model_name}'. Available providers: {list(self._provider_map.keys())}")
            return self._provider_map[provider]["client"]

        if not base_url or base_url == self.default_openai_base_url:
            return self.openai_client

        cache_key = (provider, base_url)
        if cache_key in self._client_cache:
            return self._client_cache[cache_key]

        api_key_env = client_cfg.get("api_key_env", "OPENAI_API_KEY")
        api_key = os.getenv(api_key_env)
        if not api_key:
            self.logger.warning(f"{provider} API key not found for {model_name}. Using placeholder. Please set {api_key_env} if required.")
            api_key = "placeholder-key"

        client = OpenAI(
            api_key=api_key,
            base_url=base_url,
            timeout=self.timeout_config
        )
        self._client_cache[cache_key] = client
        return client

    def generate_query(self, user_input, model_name):
        system_prompt = """
        You are a helpful assistant.
        Your task is to generate a query for a search engine, given a user input.
        """
        schema = QuerySchema
        model_config = self.get_model_config(model_name)
        client = model_config["client_config"]["client_class"]
        model_name = self.fast_model[client]
        user_prompt = f"User Input:\n{user_input}\n\nGenerate a suitable query for a search engine:\n"
        try:
            response, execution_time, token_usage, price = self.ask_llm(model_name=model_name, system_prompt=system_prompt, user_prompt=user_prompt, schema=schema)
            response = json.loads(response)
            return response["query"], token_usage, price
        except Exception as e:
            self.logger.error(f"[LLMProxy] Query generation failed using model '{model_name}'. Ensure API keys are configured. Error: {e}")
            raise

    def describe_image(self, image_path, model_name):
        """
        Describe image content for models that don't support vision.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            str: Detailed description of the image content
        """
        self.logger.info(f"[LLMProxy] Using fallback for image description...")
        
        system_prompt = """
        You are a computer vision expert. Provide a detailed, accurate description of the image.
        Focus on objects, text, colors, scene context, and any relevant details that would be useful for analysis.
        """
        
        user_prompt = f"Describe this image in detail, including any text, objects, colors, and context."
        
        model_config = self.get_model_config(model_name)
        provider = model_config["client_config"]["client_class"]
        model_name = self.default_vision_models[provider]
        
        try:
            response, execution_time, token_usage, price = self.ask_llm(model_name=model_name, system_prompt=system_prompt, user_prompt=user_prompt, image_path=image_path)
                
            self.logger.info(f"[LLMProxy] Image description:\n{response}")
            
            return response, token_usage, price
            
        except Exception as e:
            self.logger.error(f"[LLMProxy] Vision fallback failed using model '{model_name}'. Ensure API keys are configured. Error: {e}")
            return f"Error describing image: {str(e)}", {}, 0

    def ask_llm(self, model_name="gpt-4o-mini", 
                system_prompt="", 
                user_prompt="",
                temperature=None,
                schema=None,
                image_path=None,
                image_paths=None,
                file_path=None,
                websearch=False,
                use_query_generator=False,
                max_search_results=12,
                extracted_search_content_limit=0,
                search_tool="exa", # exa, serp, both (exa + serp)
                serp_engine="auto",
                search_domain=None,
                max_tokens=None,
                fallback_to_provider_best_model=False,
                fallback_to_standard_model=False,
                retry_limit=1,
                stream=False):
        """
        Makes a request to the model using the specified parameters.

        :param model_name: The name of the model to use. Default is 'gpt-4o-mini'.
        :param system_prompt: The system prompt for the model. Default is an empty string.
        :param user_prompt: The user prompt for the model. Default is an empty string.
        :param temperature: Controls randomness in the response. Default is 0.7.
        :param schema: The schema for the structured response. Default is None.
        :param image_path: Path to an image file to be used as input. Default is None.
        :param image_paths: List of paths to image files for multi-image input. Default is None.
        :param file_path: Path to a file to be used as input. Default is None.
        :param websearch: Whether to perform a web search. Default is False.
        :param use_query_generator: Whether to use query generator. Default is False.
        :param max_search_results: Maximum number of search results to return. Default is 12.
        :param extracted_search_content_limit: Max characters of extracted content per Exa result (0 = no limit). Default is 0.
        :param search_tool: Whether to use exa, serp or both (exa + serp). Default is exa.
        :param serp_engine: Search engine type for SERP ("auto", "google", "bing", etc.). Default is "auto".
        :param fallback_to_provider_best_model: Whether to fallback to provider best model. Default is False.
        :param fallback_to_standard_model: Whether to fallback to standard model. Default is False.
        :param retry_limit: The number of times to retry the request if an exception occurs.
        :param stream: Whether to stream the response. Default is False. When True, returns a generator
                       yielding StreamChunk objects. Incompatible with schema parameter.

        :return: A tuple containing:
                - response: The API response object
                - execution_time: Time taken to execute the request in seconds
                - token_usage: Dictionary containing token usage information
                - price: Calculated price for the request
                When stream=True, returns a Generator[StreamChunk, None, None] that yields
                StreamChunk objects. After iteration, call the generator's final StreamResult.
        :raises: ValueError if model is not supported
                Exception for API errors after retry limit is reached
        """
        # LangSmith tracing: wrap this call if tracing is enabled
        if self._traceable and not self._in_traced_call:
            # Resolve provider name for LangSmith metadata
            _ls_provider = "unknown"
            try:
                _cfg = self.get_model_config(model_name)
                if _cfg:
                    _ls_provider = _cfg.get("client_config", {}).get("client_class", "unknown")
            except Exception:
                pass

            def _process_inputs(inputs: dict) -> dict:
                """Convert Pydantic schema class to its JSON schema for LangSmith logging."""
                processed = dict(inputs)
                _schema = processed.get("schema")
                if _schema is not None and hasattr(_schema, "model_json_schema"):
                    processed["schema"] = _schema.model_json_schema()
                return processed

            @self._traceable(
                name="LLMProxy.ask_llm",
                run_type="llm",
                process_inputs=_process_inputs,
                metadata={
                    "ls_provider": _ls_provider.lower(),
                    "ls_model_name": model_name,
                },
            )
            def _traced_ask_llm(**kwargs):
                self._in_traced_call = True
                try:
                    result = self.ask_llm(**kwargs)
                    # Attach token usage and cost to the LangSmith run
                    if self._get_current_run_tree and not kwargs.get("stream"):
                        try:
                            _response, _exec_time, _token_usage, _cost = result
                            run = self._get_current_run_tree()
                            run.set(usage_metadata={
                                "input_tokens": _token_usage.get("prompt_tokens", 0),
                                "output_tokens": _token_usage.get("completion_tokens", 0),
                                "total_tokens": _token_usage.get("total_tokens", 0),
                            })
                            run.extra = run.extra or {}
                            run.extra.setdefault("metadata", {}).update({
                                "execution_time_seconds": round(_exec_time, 4),
                                "cost_usd": round(_cost, 8),
                            })
                        except Exception:
                            pass
                    return result
                finally:
                    self._in_traced_call = False
            return _traced_ask_llm(
                model_name=model_name, system_prompt=system_prompt, user_prompt=user_prompt,
                temperature=temperature, schema=schema, image_path=image_path, image_paths=image_paths,
                file_path=file_path, websearch=websearch, use_query_generator=use_query_generator,
                max_search_results=max_search_results, extracted_search_content_limit=extracted_search_content_limit,
                search_tool=search_tool, serp_engine=serp_engine,
                search_domain=search_domain, max_tokens=max_tokens,
                fallback_to_provider_best_model=fallback_to_provider_best_model,
                fallback_to_standard_model=fallback_to_standard_model,
                retry_limit=retry_limit, stream=stream
            )

        start_time = time.time()
        retry_count = 0
        last_error = None

        self.fallback_to_provider_best_model = fallback_to_provider_best_model
        self.fallback_to_standard_model = fallback_to_standard_model
        
        if websearch and not schema and not stream:
            schema = SearchSchema

        # Get model configuration
        model_config = self.get_model_config(model_name)
        
        # Streaming validation
        if stream:
            if schema is not None:
                self.logger.warning("[LLMProxy] stream=True was requested with schema. Streaming is incompatible with structured outputs, so streaming will be disabled and the request will proceed in non-streaming structured mode.")
                stream = False
            provider = model_config["client_config"]["client_class"]
            if provider in ["Ollama", "RunPod"]:
                self.logger.warning(f"[LLMProxy] Streaming not supported for provider {provider}. Disabling streaming.")
                stream = False
        
        if max_tokens and max_tokens>model_config["max_tokens"]:
            self.logger.warning(f"[ask_llm] The specified max tokens exceeds model limit, setting max_tokens to default")
            max_tokens = None
        
        # Handle vision fallback for models that don't support vision
        if image_path and not model_config["vision"]:
            self.logger.warning(f"[LLMProxy] Model {model_name} doesn't support vision, using fallback...")
            image_description, image_token_usage, image_price = self.describe_image(image_path, model_name)
            user_prompt = f"{user_prompt}\n\nImage Description (provided by fallback):\n{image_description}"
            image_path = None  # Prevent passing image to non-vision model
        if image_paths and not model_config["vision"]:
            self.logger.warning(f"[LLMProxy] Model {model_name} doesn't support vision, using fallback for multiple images...")
            for idx, img_path in enumerate(image_paths):
                img_description, img_token_usage, img_price = self.describe_image(img_path, model_name)
                user_prompt = f"{user_prompt}\n\nImage {idx+1} Description (provided by fallback):\n{img_description}"
            image_paths = None
        
        if file_path:
            provider_name = model_config["client_config"]["client_class"]
            needs_manual_embedding = provider_name in ["Groq", "Grok", "DeepSeek", "Qwen", "Ollama", "RunPod"]
            if needs_manual_embedding:
                try:
                    file_contents = self.read_file_content(file_path)
                    user_prompt = f"{user_prompt}\nDocument Content: {file_contents}"
                    if provider_name in ["Ollama", "RunPod"]:
                        self.logger.warning(f"[LLMProxy] {provider_name} models do not support direct file uploads. Inlining file content into the prompt.")
                        file_path = None
                except Exception as exc:
                    self.logger.error(f"[LLMProxy] Failed to read file {file_path}: {exc}")
                    raise

        if websearch:
            self.logger.info("[ask_llm] Using web search...")
            if use_query_generator:
                self.logger.info(f"[ask_llm] Generating Search Query...")
                search_query, query_token_usage, query_price = self.generate_query(user_input=user_prompt, model_name=model_name)
                self.logger.info(f"[ask_llm] Search Query: {search_query}")
            else:
                search_query = user_prompt
                query_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
                query_price = 0
            
            if search_tool == "exa":
                search_results, search_cost = self.search_web_with_exa(
                    search_query,
                    max_results=max_search_results,
                    search_domain=search_domain,
                    extracted_search_content_limit=extracted_search_content_limit
                )
            elif search_tool == "serp":
                search_results, search_cost = self.search_web_with_serp(
                    search_query,
                    max_results=max_search_results,
                    search_domain=search_domain,
                    engine_type=serp_engine
                )
            else: # search_tool == "both"   
                search_results_exa, search_cost_exa = self.search_web_with_exa(
                    search_query,
                    max_results=max_search_results,
                    search_domain=search_domain,
                    extracted_search_content_limit=extracted_search_content_limit
                )
                search_results_serp, search_cost_serp = self.search_web_with_serp(
                    search_query,
                    max_results=max_search_results,
                    search_domain=search_domain,
                    engine_type=serp_engine
                )
                search_results = f">> EXA Search Results:\n{search_results_exa}\n\n\n>> SERP Search Results:\n{search_results_serp}"
                search_cost = search_cost_exa + search_cost_serp
            print(f"\n\nSEARCH RESULTS:\n{search_results}\n\n")
            user_prompt = f"{user_prompt}\n\nRelevant Information from {search_tool.upper()} Web Search:\n{search_results}"

        # Extract pricing information
        price_per_prompt_tokens = model_config["price_per_prompt_tokens"]
        price_per_completion_tokens = model_config["price_per_completion_tokens"]

        # Print parameters for debugging
        self.logger.info(f"[ask_llm] Model Name: {model_name}")
        self.logger.info(f"[ask_llm] Temperature: {temperature}")
        self.logger.info(f"[ask_llm] Schema: {True if schema else False}")
        self.logger.info(f"[ask_llm] Max Tokens: {max_tokens}")
        self.logger.info(f"[ask_llm] Vision: {True if image_path else False}")
        self.logger.info(f"[ask_llm] File: {True if file_path else False}")
        self.logger.info(f"[ask_llm] Websearch: {websearch}")
        if websearch:
            self.logger.info(f"[ask_llm] Use Query Generator: {use_query_generator}")
            self.logger.info(f"[ask_llm] Max Search Results: {max_search_results}")
            self.logger.info(f"[ask_llm] Search Tool: {search_tool}")
        self.logger.info(f"[ask_llm] Retry Limit: {retry_limit}")
        self.logger.info(f"[ask_llm] Stream: {stream}")

        # Handle streaming mode - return generator immediately
        if stream:
            extra_cost = 0
            extra_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
            if "query_price" in locals():
                extra_cost += query_price
            if "image_price" in locals():
                extra_cost += image_price
            if "search_cost" in locals():
                extra_cost += search_cost
            if "query_token_usage" in locals():
                extra_token_usage["prompt_tokens"] += query_token_usage.get("prompt_tokens", 0)
                extra_token_usage["completion_tokens"] += query_token_usage.get("completion_tokens", 0)
                extra_token_usage["total_tokens"] += query_token_usage.get("total_tokens", 0)
            if "image_token_usage" in locals():
                extra_token_usage["prompt_tokens"] += image_token_usage.get("prompt_tokens", 0)
                extra_token_usage["completion_tokens"] += image_token_usage.get("completion_tokens", 0)
                extra_token_usage["total_tokens"] += image_token_usage.get("total_tokens", 0)
            provider = model_config["client_config"]["client_class"]
            client = self.get_client_for_model(model_name, model_config)
            return self._stream_response(
                client=client,
                provider=provider,
                model_name=model_name,
                system_prompt=system_prompt,
                user_prompt=user_prompt,
                temperature=temperature,
                image_path=image_path,
                image_paths=image_paths,
                file_path=file_path,
                max_tokens=max_tokens,
                price_per_prompt_tokens=price_per_prompt_tokens,
                price_per_completion_tokens=price_per_completion_tokens,
                extra_cost=extra_cost,
                extra_token_usage=extra_token_usage,
                start_time=start_time,
            )

        while retry_count < retry_limit:
            try:
                provider = model_config["client_config"]["client_class"]
                response_function = self._provider_map[provider]["response_function"]
                client = self.get_client_for_model(model_name, model_config)
                
                # Generate response with fallback logic
                response, token_usage = self._generate_response_with_fallbacks(
                    client, response_function, model_name, system_prompt, user_prompt, 
                    schema, temperature, provider, image_path, image_paths, file_path, max_tokens
                )

                end_time = time.time()
                execution_time = end_time - start_time

                # Calculate price based on token usage
                price = (token_usage["prompt_tokens"] * price_per_prompt_tokens + 
                        token_usage["completion_tokens"] * price_per_completion_tokens)
                if "query_price" in locals():
                    price += query_price
                
                if "image_price" in locals():
                    price += image_price
                
                if "search_cost" in locals():
                    price += search_cost
                
                if "query_token_usage" in locals():
                    token_usage["prompt_tokens"] += query_token_usage["prompt_tokens"]
                    token_usage["completion_tokens"] += query_token_usage["completion_tokens"]
                    token_usage["total_tokens"] += query_token_usage["total_tokens"]
                
                if "image_token_usage" in locals():
                    token_usage["prompt_tokens"] += image_token_usage["prompt_tokens"]
                    token_usage["completion_tokens"] += image_token_usage["completion_tokens"]
                    token_usage["total_tokens"] += image_token_usage["total_tokens"]
                
                self.logger.success("SUCCESS")
                self.logger.data(f"Response:\n{response}")
                self.logger.success(f"Execution Time: {execution_time:.2f} seconds")
                self.logger.success(f"Token Usage:\n{json.dumps(token_usage, indent=2)}")
                self.logger.success(f"Price: {price:.4f} USD")
                
                return response, execution_time, token_usage, price

            except Exception as e:
                last_error = e
                self.logger.error(f"[LLMProxy] Exception: {e}")
                self.logger.warning(f'Failed to generate response ({retry_count}/{retry_limit}). Retrying...')
                if retry_count == retry_limit-1:
                    self.logger.error(traceback.format_exc())
                retry_count += 1
                time.sleep(2)
                if retry_count >= retry_limit:
                    break
            
        # If we've exhausted all retries, raise the last error with context
        error_msg = f"[LLMProxy] Gave up after {retry_limit} attempts. Last error: {str(last_error)}"
        self.logger.error(error_msg)
        raise Exception(error_msg)

    def _generate_response_with_fallbacks(self, client, response_function, model_name, 
                                        system_prompt, user_prompt, schema, temperature, provider, image_path, image_paths, file_path, max_tokens):
        """
        Generate response with structured fallback logic for validation failures.
        """
        model_config = self.get_model_config(model_name)
        provider_name = model_config["client_config"]["client_class"]
        
        if schema is None:
            return self._generate_unstructured_response(
                client, response_function, model_name, system_prompt, user_prompt, temperature, image_path, image_paths, file_path, max_tokens
            )
        
        if provider_name == "RunPod":
            return self._generate_pseudo_structured_response(
                client, response_function, model_name, system_prompt, user_prompt, 
                schema, temperature, provider, image_path, image_paths, file_path, max_tokens, allow_fallbacks=False
            )
        
        if model_config["structured_output"]:
            return self._generate_structured_response(
                client, response_function, model_name, system_prompt, user_prompt, 
                schema, temperature, provider, image_path, image_paths, file_path, max_tokens
            )
        else:
            return self._generate_pseudo_structured_response(
                client, response_function, model_name, system_prompt, user_prompt, 
                schema, temperature, provider, image_path, image_paths, file_path, max_tokens
            )

    def _generate_unstructured_response(self, client, response_function, model_name, 
                                    system_prompt, user_prompt, temperature, image_path, image_paths, file_path, max_tokens):
        """Generate response without schema validation."""
        self.logger.info(f"[LLMProxy] Generating unstructured output...")
        response, token_usage = response_function(
            client, model_name, system_prompt, user_prompt, temperature=temperature, image_path=image_path, image_paths=image_paths, file_path=file_path, max_tokens=max_tokens
        )
        #self.logger.info(f"RESPONSE:\n{response}")
        return response, token_usage

    def _generate_structured_response(self, client, response_function, model_name, 
                                    system_prompt, user_prompt, schema, temperature, provider, image_path, image_paths, file_path, max_tokens):
        """Generate response with native structured output support."""
        self.logger.info(f"[LLMProxy] Generating structured output...")
        fallback_models = [model_name]
        allow_fallbacks = provider not in ["Ollama", "RunPod"]

        if not allow_fallbacks and (self.fallback_to_provider_best_model or self.fallback_to_standard_model):
            self.logger.warning("[LLMProxy] Fallbacks are disabled for Ollama models. Ignoring fallback settings.")

        if allow_fallbacks and self.fallback_to_provider_best_model:
            fallback_models.append(self.best_model[provider])
        
        if allow_fallbacks and self.fallback_to_standard_model:
            fallback_models.append(self.structured_output_fallback_model)
        
        for i, current_model in enumerate(fallback_models):
            try:
                # Get the correct client and response function for the current model
                if i < 2:
                    current_client = client
                    current_response_function = response_function
                else:
                    # For fallback model, get the correct provider and response function
                    current_client = self.openai_client
                    fallback_provider = MODEL_CONFIGS[self.structured_output_fallback_model]["client_config"]["client_class"]
                    current_response_function = self._provider_map[fallback_provider]["response_function"]
                
                response, token_usage = current_response_function(
                    current_client, current_model, system_prompt, user_prompt, 
                    schema, temperature=temperature, image_path=image_path, image_paths=image_paths, file_path=file_path, max_tokens=max_tokens
                )
                #self.logger.info(f"RESPONSE:\n{response}")
                schema.model_validate_json(response)
                self.logger.success(f"[LLMProxy] Pydantic validation successful")
                return response, token_usage
                
            except Exception as e:
                self._log_validation_error(e, current_model, i, fallback_models)
                self.logger.error(traceback.format_exc())
                if i == len(fallback_models) - 1:  # Last attempt failed
                    self.logger.error("Failed to validate response.")
                    raise
        
    def _generate_pseudo_structured_response(self, client, response_function, model_name, 
                                    system_prompt, user_prompt, schema, temperature, provider, image_path, image_paths, file_path, max_tokens, allow_fallbacks=None):
        """Generate response for models without native structured output."""
        self.logger.info(f"[LLMProxy] Generating pseudo-structured output...")
        self.logger.warning(f"Model {model_name} does not natively support structured responses. Exact schema not guaranteed.")
        
        fallback_models = [model_name]
        if allow_fallbacks is None:
            allow_fallbacks = provider not in ["Ollama", "RunPod"]

        if not allow_fallbacks and (self.fallback_to_provider_best_model or self.fallback_to_standard_model):
            self.logger.warning("[LLMProxy] Fallbacks are disabled for this provider. Ignoring fallback settings.")

        if allow_fallbacks and self.fallback_to_provider_best_model:
            fallback_models.append(self.best_model[provider])
        
        if allow_fallbacks and self.fallback_to_standard_model:
            fallback_models.append(self.structured_output_fallback_model)
        
        modified_system_prompt = self.prompt_with_structure(system_prompt, schema)
        
        for i, current_model in enumerate(fallback_models):
            try:
                # Get the correct client and response function for the current model
                if i < 2:
                    current_client = client
                    current_response_function = response_function
                else:
                    # For fallback model, get the correct provider and response function
                    current_client = self.openai_client
                    fallback_provider = MODEL_CONFIGS[self.structured_output_fallback_model]["client_config"]["client_class"]
                    current_response_function = self._provider_map[fallback_provider]["response_function"]
                
                response, token_usage = current_response_function(
                    current_client, current_model, modified_system_prompt, user_prompt, 
                    temperature=temperature, image_path=image_path, image_paths=image_paths, file_path=file_path, max_tokens=max_tokens
                )
                #self.logger.info(f"RESPONSE:\n{response}")
                cleaned_response = response.replace("```json", "").replace("```", "")
                schema.model_validate_json(cleaned_response)
                self.logger.success(f"[LLMProxy] Pydantic validation successful")
                return cleaned_response, token_usage
                
            except Exception as e:
                self._log_validation_error(e, current_model, i, fallback_models)
                if i == len(fallback_models) - 1:  # Last attempt failed
                    self.logger.error("Failed to validate response.")
                    raise

    def _log_validation_error(self, error, model_name, attempt_index, fallback_models):
        """Log validation errors with appropriate fallback messaging."""
        self.logger.error(f"[LLMProxy] Pydantic validation error with model {model_name}:\n{error}")
        
        if attempt_index == 0 and len(fallback_models) > 1:
            next_model = fallback_models[1]
            self.logger.warning(f"Will attempt with a different model from the same provider: {next_model}")
        elif attempt_index == 1 and len(fallback_models) > 2:
            self.logger.warning(f"Will attempt with {self.structured_output_fallback_model}")
        else:
            self.logger.error(f"[LLMProxy] Pydantic validation error with model {model_name}: {error}")
            self.logger.error(traceback.format_exc())
    
    def prompt_with_structure(self, system_prompt, schema):
        return system_prompt + "\n\nFormat your response as a JSON object:\n\n"+json.dumps(schema.model_json_schema(), indent=2) + "\n\n" + self.json_instructions
    
    def create_image_openai(self, file_path):
        with open(file_path, "rb") as file_content:
            result = self.openai_client.files.create(
                file=file_content,
                purpose="vision",
            )
            return result.id
    
    def create_file_openai(self, file_path):
        with open(file_path, "rb") as file_content:
            result = self.openai_client.files.create(
                file=file_content,
                purpose="user_data",
            )
            return result.id
    
    def encode_image(self,image_path):
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode("utf-8")

    def ollama_response(self, client, model, system_prompt, user_prompt, schema=None, temperature=None, image_path=None, image_paths=None, file_path=None, max_tokens=None):
        """
        Handle Ollama model responses using OpenAI-compatible chat completions API.
        
        Note: File and image inputs are handled upstream in ask_llm via content inlining
        and vision fallback, so file_path and image_path will typically be None here.
        """
        # Extract the actual model name (e.g., "llama3" from "ollama/llama3")
        engine_model = model.split("/", 1)[1] if model.startswith("ollama/") else model
        
        if image_path:
            self.logger.warning(f"[LLMProxy] Ollama models do not support direct image inputs. Image should have been handled via vision fallback upstream.")
        
        if file_path:
            self.logger.warning(f"[LLMProxy] Ollama models do not support direct file inputs. File should have been inlined upstream.")
        
        # Build messages
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": user_prompt})
        
        # Build request parameters
        chat_kwargs = {
            "model": engine_model,
            "messages": messages,
        }
        if temperature is not None:
            chat_kwargs["temperature"] = temperature
        if max_tokens is not None:
            chat_kwargs["max_tokens"] = max_tokens
        if schema is not None:
            chat_kwargs["response_format"] = {
                "type": "json_schema",
                "json_schema": {
                    "name": schema.__name__,
                    "schema": schema.model_json_schema(),
                },
            }
        
        # Make request
        response = client.chat.completions.create(**chat_kwargs)
        
        # Extract text
        message_content = response.choices[0].message.content
        if isinstance(message_content, list):
            text = ''.join(
                part.get("text", "") if isinstance(part, dict) else str(part)
                for part in message_content
            )
        else:
            text = message_content or ""
        
        # Validate structured output if schema provided
        if schema is not None:
            try:
                parsed = schema.model_validate_json(text)
                text = parsed.model_dump_json(indent=2)
            except Exception as exc:
                self.logger.error(f"[LLMProxy] Failed to validate Ollama structured response: {exc}")
                raise
        
        # Extract token usage safely
        usage = response.usage
        def _usage_value(obj, name):
            if obj is None:
                return 0
            if isinstance(obj, dict):
                return obj.get(name, 0)
            return getattr(obj, name, 0)
        
        token_usage = {
            "prompt_tokens": _usage_value(usage, "prompt_tokens"),
            "completion_tokens": _usage_value(usage, "completion_tokens"),
            "total_tokens": _usage_value(usage, "total_tokens"),
        }
        
        return text, token_usage

    def openai_response(self, client, model, system_prompt, user_prompt, schema=None, temperature=None, image_path=None, image_paths=None, file_path=None, max_tokens=None):
        # All models reaching this function should be in MODEL_CONFIGS
        # (Ollama now has its own ollama_response function)
        model_cfg = MODEL_CONFIGS.get(model)
        
        user_content = [
            {"type": "input_text", "text": user_prompt}
        ]
        if image_path:
            user_content.append({
                "type": "input_image",
                "file_id": self.create_image_openai(image_path)
            })
        if image_paths:
            for img_path in image_paths:
                user_content.append({
                    "type": "input_image",
                    "file_id": self.create_image_openai(img_path)
                })
        if file_path:
            if file_path.split('.')[-1] == "pdf":
                user_content.append({
                    "type": "input_file",
                    "file_id": self.create_file_openai(file_path)
                })
            else:
                user_content.append({
                    "type": "input_text",
                    "text": self.read_file_content(file_path)
                })
        completion_params = {
            "model": model,
            "input": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_content}
            ]
        }
        if max_tokens:
            completion_params["max_output_tokens"] = max_tokens
        if model_cfg and "reasoning_effort" in model_cfg:
            completion_params["reasoning"] = {"effort": model_cfg["reasoning_effort"]}
        elif temperature:
            completion_params["temperature"] = temperature
        if schema:
            completion_params["text_format"] = schema
            response = client.responses.parse(**completion_params)
            text = response.output[-1].content[0].parsed
            text = text.model_dump_json(indent=2)
        else:
            response = client.responses.create(**completion_params)
            text = response.output_text
        token_usage = {
            "prompt_tokens": response.usage.input_tokens,
            "completion_tokens": response.usage.output_tokens,
            "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
        }
        return text, token_usage

    def grok_response(self, client, model, system_prompt, user_prompt, schema=None, temperature=None, image_path=None, image_paths=None, file_path=None, max_tokens=None):
        user_input = [user_prompt]
        if image_path:
            base64_image = self.encode_image(image_path)
            image_input = image(image_url=f"data:image/jpeg;base64,{base64_image}", detail="auto")
            user_input.append(image_input)
        if image_paths:
            for img_path in image_paths:
                base64_img = self.encode_image(img_path)
                user_input.append(image(image_url=f"data:image/jpeg;base64,{base64_img}", detail="auto"))
        if max_tokens:
            chat = client.chat.create(model=model, temperature=temperature, max_tokens=max_tokens)
        else:
            chat = client.chat.create(model=model, temperature=temperature)
        chat.append(system(system_prompt))
        chat.append(user(*user_input))
        if schema:
            response, structure = chat.parse(schema)
            text = structure.model_dump_json(indent=2)
        else:
            response = chat.sample()
            text = response.content
        prompt_tokens = response.usage.prompt_tokens
        completion_tokens = response.usage.completion_tokens
        total_tokens = response.usage.total_tokens
        reasoning_tokens = response.usage.reasoning_tokens
        token_usage = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens,
            "reasoning_tokens": reasoning_tokens
        }
        return text, token_usage
    
    def gemini_response(self, client, model, system_prompt, user_prompt, schema=None, temperature=None, image_path=None, image_paths=None, file_path=None, max_tokens=None):
        # Extract thinking level from model name or config
        model_name = model
        thinking_level = None
        
        # Check for explicit thinking level in model name
        if model.endswith("-thinking-high"):
            thinking_level = "high"
            model_name = model.replace("-thinking-high", "")
        elif model.endswith("-thinking-minimal"):
            thinking_level = "minimal"
            model_name = model.replace("-thinking-minimal", "")
        elif model.endswith("-thinking-low"):
            thinking_level = "low"
            model_name = model.replace("-thinking-low", "")
        else:
            # For base models, check if thinking config exists
            model_config = MODEL_CONFIGS.get(model)
            if model_config and "thinking" in model_config:
                thinking_level = model_config["thinking"]["thinking_level"]
        
        prompt = system_prompt + "\n\n" + user_prompt
        completion_params = {
            "model": model_name,
            "contents": [prompt]
        }
        config_params = {}
        
        # Handle thinking mode for Gemini 3 models
        # Note: Google Genai SDK doesn't support thinking_level directly in ThinkingConfig
        # We need to map thinking levels to thinking_budget values
        if thinking_level:
            # Map thinking levels to budget values
            # minimal/low = small budget, medium = moderate, high = large budget
            budget_map = {
                "minimal": 1000,
                "low": 2000,
                "medium": 5000,
                "high": -1  # -1 means automatic/unlimited
            }
            thinking_budget = budget_map.get(thinking_level, -1)
            config_params["thinking_config"] = types.ThinkingConfig(
                thinking_budget=thinking_budget,
                include_thoughts=True
            )
        
        if temperature:
            config_params["temperature"] = temperature
        if max_tokens:
            config_params["maxOutputTokens"] = max_tokens
        if schema:
            config_params["response_mime_type"] = "application/json"
            config_params["response_schema"] = schema
        if image_path:
            with open(image_path, "rb") as image:
                image_bytes = image.read()
            completion_params["contents"].append(
                types.Part.from_bytes(
                    data=image_bytes,
                    mime_type=f"image/{image_path.split('.')[-1]}"
                )
            )
        if image_paths:
            for img_path in image_paths:
                with open(img_path, "rb") as img_file:
                    img_bytes = img_file.read()
                completion_params["contents"].append(
                    types.Part.from_bytes(
                        data=img_bytes,
                        mime_type=f"image/{img_path.split('.')[-1]}"
                    )
                )
        if file_path:
            if file_path.split('.')[-1] in ["pdf", "txt"]: 
                myfile = client.files.upload(file=file_path)
                completion_params["contents"].append(myfile)
            else:
                completion_params["contents"].append(f"Document Content:\n{self.read_file_content(file_path)}")
        completion_params["config"] = types.GenerateContentConfig(**config_params)
        response = client.models.generate_content(**completion_params)
        if schema:
            text = response.text #response.parsed
            if type(text) == dict:
                text = json.dumps(text, indent=2)
            elif type(text) == str:
                text = json.loads(text)
                text = json.dumps(text, indent=2)
            else:
                text = text.model_dump_json(indent=2)
        else:
            text = response.candidates[0].content.parts[0].text
        prompt_tokens = response.usage_metadata.prompt_token_count
        completion_tokens = response.usage_metadata.candidates_token_count
        total_tokens = response.usage_metadata.total_token_count
        token_usage = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens
        }
        return text, token_usage
    
    def groq_response(self, client, model, system_prompt, user_prompt, schema=None, temperature=None, image_path=None, image_paths=None, file_path=None, max_tokens=None):
        user_prompt = {
            "type": "text",
            "text": user_prompt
        }
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": [user_prompt]}
        ]
        request_params = {
            "model": model,
            "messages": messages
        }
        if schema:
            request_params["response_format"] = {
                "type": "json_schema",
                "json_schema": {
                    "name": schema.__name__,
                    "schema": schema.model_json_schema()
                }
            }
        if temperature:
            request_params["temperature"] = temperature
        if max_tokens:
            request_params["max_completion_tokens"] = max_tokens
        if image_path:
            base64_image = self.encode_image(image_path)
            request_params["messages"][1]["content"].append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/jpeg;base64,{base64_image}"
                }
            })
        if image_paths:
            for img_path in image_paths:
                base64_img = self.encode_image(img_path)
                request_params["messages"][1]["content"].append({
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_img}"}
                })
        response = client.chat.completions.create(**request_params)
        text = response.choices[0].message.content
        if schema:
            text = text.replace("```json", "").replace("```", "")
            text = json.loads(text)
            text = json.dumps(text, indent=2)
        usage_data = response.usage
        token_usage = {
            "prompt_tokens": usage_data.prompt_tokens,
            "completion_tokens": usage_data.completion_tokens,
            "total_tokens": usage_data.total_tokens
        }
        return text, token_usage  

    def claude_response(self, client, model, system_prompt, user_prompt, schema=None, temperature=None, image_path=None, image_paths=None, file_path=None, max_tokens=None):
        # Clean model name and prepare prompts
        model_name = model.replace("-thinking", "")
        system_prompt = system_prompt.replace('%', 'percent')
        user_prompt = user_prompt.replace('%', 'percent')
        model_config = MODEL_CONFIGS.get(model) or MODEL_CONFIGS.get(model_name)
        content = []
        content.append({
            "type": "text",
            "text": user_prompt
        })
        if image_path:
            base64_image = self.encode_image(image_path)
            media_type = f"image/{image_path.split('.')[-1]}"
            content.append({
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": media_type,
                    "data": base64_image
                }
            })
        if image_paths:
            for img_path in image_paths:
                base64_img = self.encode_image(img_path)
                img_media_type = f"image/{img_path.split('.')[-1]}"
                content.append({
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": img_media_type,
                        "data": base64_img
                    }
                })
        if file_path:
            if file_path.split('.')[-1] in ['pdf','txt']:
                type_mapper = {
                    "pdf": "application/pdf",
                    "txt": "text/plain"
                }
                file = client.beta.files.upload(
                    file=(file_path, open(file_path, "rb"), type_mapper[file_path.split('.')[-1]])
                )
                content.append({
                    "type": "document",
                    "source": {
                        "type": "file",
                        "file_id": file.id
                    }
                })
            else:
                content.append({
                    "type": "text",
                    "text": f"\n\nDocument Content:\n{self.read_file_content(file_path)}\n\n"
                })
        # Build message parameters
        params = {
            "model": model_name,
            "system": system_prompt,
            "messages": [{"role": "user", "content": content}],
        }
        supports_structured_outputs = bool(
            schema is not None
            and model_config
            and model_config.get("structured_output")
        )
        if max_tokens:
            params["max_tokens"] = max_tokens
        else:
            params["max_tokens"] = 16000 if "3-5" not in model else 8000
        if file_path:
            params["betas"]=["files-api-2025-04-14"]
        # Add conditional parameters
        if temperature:
            params["temperature"] = temperature
        if "thinking" in model:
            params["thinking"] = MODEL_CONFIGS[model]["thinking"]
            params["temperature"] = 1
        if schema and not supports_structured_outputs:
            self.logger.warning("[LLMProxy] Claude structured outputs require a model with structured_output=True. Falling back to pseudo-structured output.")
        
        try:
            if supports_structured_outputs:
                parse_fn = None
                # Prefer GA parse if present, otherwise fall back to beta.parse
                if hasattr(client, "messages") and hasattr(client.messages, "parse"):
                    parse_fn = client.messages.parse
                elif hasattr(client, "beta") and hasattr(client.beta, "messages") and hasattr(client.beta.messages, "parse"):
                    parse_fn = client.beta.messages.parse

                if parse_fn is None:
                    raise AttributeError("Anthropic client has no messages.parse or beta.messages.parse available")

                try:
                    response = parse_fn(
                        **params,
                        output_format=schema,
                    )
                except Exception as parse_err:
                    err_msg = str(parse_err)
                    if "Streaming is required" in err_msg:
                        self.logger.warning(f"[LLMProxy] Claude non-streaming blocked for {model}, falling back to streaming")
                        response = self._claude_streaming_fallback(client, params)
                    elif "compiled grammar is too large" in err_msg:
                        self.logger.warning(f"[LLMProxy] Claude schema too complex for native structured output, falling back to pseudo-structured")
                        return self._claude_pseudo_structured_fallback(client, model, model_name, system_prompt, user_prompt, content, schema, params, max_tokens)
                    else:
                        raise
                if hasattr(response, 'parsed_output') and response.parsed_output is not None:
                    parsed_output = response.parsed_output
                    text = parsed_output.model_dump_json(indent=2)
                else:
                    if "thinking" in model:
                        text = response.content[1].text
                    else:
                        text = response.content[0].text
            else:
                try:
                    if file_path:
                        response = client.beta.messages.create(**params)
                    else:
                        response = client.messages.create(**params)
                except Exception as create_err:
                    if "Streaming is required" in str(create_err):
                        self.logger.warning(f"[LLMProxy] Claude non-streaming blocked for {model}, falling back to streaming")
                        response = self._claude_streaming_fallback(client, params)
                    else:
                        raise
                if "thinking" in model:
                    text = response.content[1].text
                else:
                    text = response.content[0].text
            usage = response.usage
            
            token_usage = {
                "prompt_tokens": usage.input_tokens,
                "completion_tokens": usage.output_tokens,
                "total_tokens": usage.input_tokens + usage.output_tokens
            }
            
            return text, token_usage
            
        except Exception as e:
            traceback.print_exc()
            self.logger.error(f"[LLMProxy] Claude request failed for model {model}: {e}")
            raise

    def _claude_streaming_fallback(self, client, params):
        """Use streaming to collect a complete response when non-streaming is blocked."""
        stream_params = {k: v for k, v in params.items() if k != "betas"}
        accumulated = ""
        final_response = None
        with client.messages.stream(**stream_params) as stream:
            for text in stream.text_stream:
                accumulated += text
            final_response = stream.get_final_message()
        return final_response

    def _claude_pseudo_structured_fallback(self, client, model, model_name, system_prompt, user_prompt, content, schema, params, max_tokens):
        """Fall back to pseudo-structured output when Claude's native structured output rejects the schema."""
        schema_json = json.dumps(schema.model_json_schema(), indent=2)
        pseudo_system = system_prompt + "\n\nFormat your response as a JSON object:\n\n" + schema_json + "\n\n" + self.json_instructions
        pseudo_params = {
            "model": model_name,
            "system": pseudo_system,
            "messages": [{"role": "user", "content": content}],
            "max_tokens": max_tokens or params.get("max_tokens", 16000),
        }
        if params.get("temperature"):
            pseudo_params["temperature"] = params["temperature"]
        if "thinking" in model:
            pseudo_params["thinking"] = MODEL_CONFIGS[model]["thinking"]
            pseudo_params["temperature"] = 1

        try:
            response = client.messages.create(**pseudo_params)
        except Exception as create_err:
            if "Streaming is required" in str(create_err):
                self.logger.warning(f"[LLMProxy] Claude pseudo-structured also blocked non-streaming, using streaming")
                response = self._claude_streaming_fallback(client, pseudo_params)
            else:
                raise

        if "thinking" in model:
            text = response.content[1].text
        else:
            text = response.content[0].text
        text = text.replace("```json", "").replace("```", "").strip()

        usage = response.usage
        token_usage = {
            "prompt_tokens": usage.input_tokens,
            "completion_tokens": usage.output_tokens,
            "total_tokens": usage.input_tokens + usage.output_tokens,
        }
        return text, token_usage

    def deepseek_response(self, client, model, system_prompt, user_prompt, schema=None, temperature=None, image_path=None, image_paths=None, file_path=None, max_tokens=None):
        if image_path:
            self.logger.error("[LLMProxy] DeepSeek does not support image inputs. Image input will be ignored.")
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        request_params = {
            "model": model,
            "messages": messages
        }
        if temperature:
            request_params["temperature"] = temperature
        if max_tokens:
            request_params["max_tokens"] = max_tokens
        if schema:
            messages[1]["content"] += f"\n\nFormat your response as a JSON object:\n\n{json.dumps(schema.model_json_schema(), indent=2)}"
            request_params["response_format"] = {"type": "json_object"}
        
        response = client.chat.completions.create(**request_params)
        
        text = response.choices[0].message.content
        
        if schema:
            text = text.replace("```json", "").replace("```", "")
            text = json.loads(text)
            text = json.dumps(text, indent=2)
        
        token_usage = {
            "prompt_tokens": response.usage.prompt_tokens,
            "completion_tokens": response.usage.completion_tokens,
            "total_tokens": response.usage.total_tokens,
        }
        
        return text, token_usage
    
    def runpod_response(self, client, model, system_prompt, user_prompt, schema=None, temperature=None, image_path=None, image_paths=None, file_path=None, max_tokens=None):
        if not self.runpod_token:
            raise ValueError("RUNPOD_TOKEN is not configured. Please set it to use RunPod models.")
        
        # Strip "runpod/" prefix if present
        actual_model = model.split("/", 1)[1] if model.startswith("runpod/") else model

        if image_path:
            self.logger.warning("[LLMProxy] RunPod models do not support direct image inputs. Using image description fallback.")
        if schema:
            self.logger.warning("[LLMProxy] RunPod models do not support native structured outputs. Falling back to pseudo-structured enforcement.")
        model_config = self.get_model_config(model)
        run_endpoint, status_endpoint = self._resolve_runpod_endpoints(model_config)
        self.logger.info(f"[RunPod] Using run endpoint: {run_endpoint}")
        self.logger.info(f"[RunPod] Using status endpoint: {status_endpoint}")
        headers = {
            "Authorization": f"Bearer {self.runpod_token}",
            "Content-Type": "application/json"
        }
        payload = self._build_runpod_payload(system_prompt, user_prompt, temperature, max_tokens, model_config, actual_model)
        self.logger.info(f"[RunPod] Sending request with timeout: {self.runpod_poll_timeout}s")
        try:
            response = client.post(run_endpoint, headers=headers, json=payload, timeout=self.runpod_poll_timeout)
            response.raise_for_status()
        except requests.RequestException as exc:
            raise ValueError(f"RunPod request failed: {exc}") from exc
        
        data = response.json()
        final_payload = self._finalize_runpod_job(data, client, status_endpoint, headers)
        output_text, token_usage = self._parse_runpod_output(final_payload.get("output"))
        return output_text, token_usage

    def _resolve_runpod_endpoints(self, model_config: dict) -> Tuple[str, str]:
        """Resolve RunPod endpoints from model config (supports both old and new format)."""
        client_cfg = model_config["client_config"]
        
        # New format: direct endpoint URLs in config
        if "run_endpoint" in client_cfg:
            run_endpoint = client_cfg["run_endpoint"]
            status_endpoint = client_cfg["status_endpoint"]
        else:
            # Old format: env var names (for backward compatibility if needed)
            run_endpoint = os.getenv(client_cfg.get("run_endpoint_env", ""), "")
            status_endpoint = os.getenv(client_cfg.get("status_endpoint_env", ""), "")
            if not run_endpoint:
                run_endpoint = client_cfg.get("default_run_endpoint")
            if not status_endpoint:
                status_endpoint = client_cfg.get("default_status_endpoint")
        
        if not run_endpoint or not status_endpoint:
            raise ValueError("RunPod endpoints are not configured correctly.")
        return run_endpoint, status_endpoint

    def _build_runpod_payload(self, system_prompt: str, user_prompt: str, temperature: Optional[float], max_tokens: Optional[int], model_config: dict, model_name: str) -> Dict[str, Any]:
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": user_prompt})
        sampling_params: Dict[str, Any] = {}
        sampling_params["temperature"] = temperature if temperature is not None else 0.7
        max_tokens_value = max_tokens if max_tokens else min(model_config.get("max_tokens", 1024), 2048)
        sampling_params["max_tokens"] = max_tokens_value
        payload = {
            "input": {
                "model": model_name,
                "messages": messages,
                "sampling_params": sampling_params
            }
        }
        return payload

    def _finalize_runpod_job(self, payload: dict, client, status_endpoint: str, headers: dict) -> dict:
        if self._runpod_payload_has_output(payload):
            self.logger.info(f"[RunPod] Job completed immediately with output.")
            return payload

        job_id = payload.get("id")
        if not job_id:
            raise ValueError("RunPod response missing job id and output.")

        status_url = status_endpoint.rstrip("/") + f"/{job_id}"
        initial_status = payload.get("status", "UNKNOWN")
        self.logger.info(f"[RunPod] Job {job_id} submitted. Initial status: {initial_status}")
        self.logger.info(f"[RunPod] Polling status endpoint: {status_url}")
        
        deadline = time.time() + self.runpod_poll_timeout
        poll_count = 0
        last_status = initial_status
        
        while time.time() < deadline:
            poll_count += 1
            elapsed = int(time.time() - (deadline - self.runpod_poll_timeout))
            remaining = int(deadline - time.time())
            
            self.logger.info(f"[RunPod] Poll #{poll_count} - Elapsed: {elapsed}s, Remaining: {remaining}s, Status: {last_status}")
            
            time.sleep(self.runpod_poll_interval)
            try:
                status_response = client.get(status_url, headers=headers, timeout=self.runpod_poll_interval + 5)
                status_response.raise_for_status()
                latest_payload = status_response.json()
            except requests.RequestException as exc:
                self.logger.error(f"[RunPod] Status check failed: {exc}")
                raise ValueError(f"RunPod status check failed: {exc}") from exc
            
            if self._runpod_payload_has_output(latest_payload):
                self.logger.success(f"[RunPod] Job {job_id} completed with output after {elapsed}s")
                return latest_payload
            
            status_value = latest_payload.get("status")
            if isinstance(status_value, str):
                status_upper = status_value.upper()
                last_status = status_upper
                
                if status_upper == "FAILED":
                    raise ValueError(f"RunPod job {job_id} failed: {latest_payload}")
                elif status_upper in ["IN_QUEUE", "IN_PROGRESS"]:
                    self.logger.info(f"[RunPod] Job {job_id} is {status_upper}, continuing to poll...")
                else:
                    self.logger.warning(f"[RunPod] Job {job_id} has unexpected status: {status_value}")
            else:
                self.logger.warning(f"[RunPod] Job {job_id} has non-string status: {status_value}")

        raise TimeoutError(f"RunPod job {job_id} did not complete within {self.runpod_poll_timeout} seconds. Last status: {last_status}")

    def _runpod_payload_has_output(self, payload: dict) -> bool:
        if not payload:
            return False
        if payload.get("output"):
            return True
        status_field = payload.get("status")
        if isinstance(status_field, dict) and status_field.get("output"):
            payload["output"] = status_field.get("output")
            return True
        return False

    def _parse_runpod_output(self, output: Optional[list]) -> Tuple[str, Dict[str, int]]:
        if not output:
            raise ValueError("RunPod job returned no output.")
        segment = output[0]
        choices = segment.get("choices", [])
        text_content = ""
        if choices:
            tokens = choices[0].get("tokens", [])
            if isinstance(tokens, list):
                text_content = "".join(str(token) for token in tokens)
            else:
                text_content = str(tokens)
        analysis, final = self._split_runpod_sections(text_content)
        final_text = final or analysis or ""
        usage = segment.get("usage", {})
        prompt_tokens = usage.get("input", 0)
        completion_tokens = usage.get("output", 0)
        token_usage = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens
        }
        return final_text, token_usage

    def _split_runpod_sections(self, text: str) -> Tuple[str, str]:
        marker = "assistantfinal"
        idx = text.find(marker)
        if idx == -1:
            return "", text.strip()
        analysis = text[:idx].strip()
        final = text[idx + len(marker):].strip()
        return analysis, final
    
    def qwen_response(self, client, model, system_prompt, user_prompt, schema=None, temperature=None, image_path=None, image_paths=None, file_path=None, max_tokens=None):
        if image_path:
            self.logger.error("[LLMProxy] Qwen does not support image inputs. Image input will be ignored.")
        
        # Clean model name and get config
        model_name = model.replace("-thinking", "")
        model_config = MODEL_CONFIGS.get(model) or MODEL_CONFIGS.get(model_name)
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        request_params = {
            "model": model_name,
            "messages": messages
        }
        
        # Handle thinking mode for Qwen models
        if model_config and "thinking" in model_config:
            thinking_config = model_config["thinking"]
            if "enable_thinking" in thinking_config:
                request_params["extra_body"] = {"enable_thinking": thinking_config["enable_thinking"]}
        
        if schema:
            request_params["messages"][1]["content"] += f"\n\nFormat your response as a JSON object:\n\n{json.dumps(schema.model_json_schema(), indent=2)}"
            request_params["response_format"] = {"type": "json_object"}
        if temperature:
            request_params["temperature"] = temperature
        if max_tokens:
            request_params["max_tokens"] = max_tokens
        response = client.chat.completions.create(**request_params)
        text = response.choices[0].message.content
        if schema:
            text = text.replace("```json", "").replace("```", "")
        token_usage = {
            "prompt_tokens": response.usage.prompt_tokens,
            "completion_tokens": response.usage.completion_tokens,
            "total_tokens": response.usage.total_tokens,
        }
        return text, token_usage
    
    def read_text_file(self, file_path: str) -> str:
        """Read content from a plain text file."""
        self.logger.info("Reading text file...")
        with open(file_path, 'r', encoding='utf-8') as file:
            return file.read()

    def read_docx_file(self, file_path: str) -> str:
        """Read content from a DOCX file."""
        self.logger.info("Reading DOCX file...")
        doc = Document(file_path)
        content = []
        for paragraph in doc.paragraphs:
            content.append(paragraph.text)
        return '\n'.join(content)


    def read_json_file(self, file_path: str) -> str:
        """Read content from a JSON file and return as formatted string."""
        self.logger.info("Reading JSON file...")
        with open(file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)
        return json.dumps(data, indent=2)


    def extract_text_with_pymupdf(self, file_path: str) -> tuple[str, int]:
        """
        Extract text from PDF using PyMuPDF.
        Returns tuple of (text_content, page_count).
        """
        self.logger.info("Extracting text from PDF using PyMuPDF...")
        doc = pymupdf.open(file_path)
        content = []
        
        for page in doc:
            content.append(page.get_text())
        
        doc.close()
        return '\n'.join(content), len(content)

    def is_text_based_pdf(self, file_path: str, min_words_per_page: int = 20) -> bool:
        """
        Check if PDF is text-based by extracting text and checking word density.
        Returns True if it's text-based with sufficient content.
        """
        self.logger.info("Checking if PDF is text-based...")
        try:
            text_content, page_count = self.extract_text_with_pymupdf(file_path)
            
            if not text_content.strip():
                return False
            
            word_count = len(text_content.split())
            avg_words_per_page = word_count / page_count if page_count > 0 else 0
            
            self.logger.info(f"Average words per page: {avg_words_per_page}")
            return avg_words_per_page >= min_words_per_page
        except Exception:
            traceback.print_exc()
            return False


    def extract_with_gemini(self, file_path: str) -> str:
        """Extract text from file using Gemini OCR."""
        self.logger.info("Extracting text from file using Gemini OCR...")
        client = self.google_client
        uploaded_file = client.files.upload(file=file_path)
        response = client.models.generate_content(
            model="gemini-2.0-flash-lite",
            contents=["extract all the text from this file.", uploaded_file],
        )
        return response.text


    def read_pdf_file(self, file_path: str) -> str:
        """
        Read content from PDF file.
        Uses PyMuPDF for text-based PDFs, Gemini OCR for image-based or sparse PDFs.
        """
        self.logger.info("Reading PDF file...")
        if self.is_text_based_pdf(file_path):
            self.logger.info("PDF is text-based.")
            text_content, _ = self.extract_text_with_pymupdf(file_path)
            return text_content
        else:
            self.logger.info("PDF is not text-based.")
            return self.extract_with_gemini(file_path)

    def read_file_content(self, file_path: str) -> str:
        """
        Main function to read content from various file types.
        
        Args:
            file_path: Path to the file to read
            
        Returns:
            str: Extracted text content
            
        Raises:
            FileNotFoundError: If file doesn't exist
            ValueError: If file type is not supported
        """
        self.logger.info("Reading file content...")
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found: {file_path}")
        
        extension = Path(file_path).suffix.lower()
        
        if extension == '.txt':
            return self.read_text_file(file_path)
        elif extension == '.docx':
            return self.read_docx_file(file_path)
        elif extension == '.json':
            return self.read_json_file(file_path)
        elif extension == '.pdf':
            return self.read_pdf_file(file_path)
        else:
            raise ValueError(f"Unsupported file type: {extension}")

    def search_web_with_exa(self, query: str, max_results: int = 8, search_domain: Optional[str] = None, extracted_search_content_limit: int = 0) -> tuple:
        """
        Search the web using Exa API.
        
        Args:
            query: Search query string
            max_results: Maximum number of results to return (default: 8)
            
        Returns:
            tuple: (context, cost)
                - context: Search results including organic results, knowledge graph, and related searches
                - cost: Cost of the search in USD
            
        Raises:
            ValueError: If Exa_KEY is not configured
            requests.RequestException: If API request fails
        """
        self.logger.info(f"[search_web] Searching the web with Exa")
        exa_kwargs = {
            "type": "auto",
            "context": True,
            "num_results": max_results,
            "summary": True,
        }
        if extracted_search_content_limit > 0:
            exa_kwargs["text"] = {"max_characters": extracted_search_content_limit}
        if search_domain:
            exa_kwargs["include_domains"] = [search_domain]
        result = self.exa_client.search_and_contents(query, **exa_kwargs)
        context = result.context
        cost = result.cost_dollars.total
        self.logger.success(f"[search_web] Exa Search completed with cost: {cost}")
        self.logger.data(f"[search_web] Exa Search results:\n{context}")
        return context, cost
    
    def search_web_with_serp(self, query: str, max_results: int = 10, search_domain: Optional[str] = None, engine_type: str = "auto") -> tuple:
        """
        Search the web using SERP API (Google search engine).
        
        Args:
            query: Search query string
            max_results: Maximum number of results to return (default: 10)
            
        Returns:
            tuple: (search_results_dict, estimated_cost)
                - search_results_dict: Raw dictionary of search results from SERP API
                - estimated_cost: Estimated cost in USD for the search operation
            
        Raises:
            ValueError: If SERP_API_KEY is not configured
            requests.RequestException: If API request fails
        """
        self.logger.info(f"[search_web] Searching the web with SERP")
        if self.serp_api_key == "placeholder-key":
            raise ValueError("SERP API key not configured. Please set SERP_API_KEY in your .env file.")
        
        
        try:
            # Configure search parameters based on engine
            search_params = {
                "q": query,
                "api_key": self.serp_api_key,
                "engine": "google" if engine_type == "auto" else engine_type,
                "num": min(max_results, 20)  # SERP API typically limits to 20 results
            }
            if search_domain:
                search_params["as_sitesearch"] = search_domain
            
            # Perform the search
            search = GoogleSearch(search_params)
            results = search.get_dict()

            results = self.format_serp_results(results)
            
            # Estimate cost (SERP API typically charges per search)
            estimated_cost = 0.015  # Approximate cost per search, adjust based on actual pricing

            self.logger.success(f"[search_web] SERP Search completed with estimated cost: {estimated_cost}\nNote: This is an estimated cost. Adjust based on plan.")
            self.logger.data(f"[search_web] SERP Search results:\n{results}")
            
            return results, estimated_cost
            
        except Exception as e:
            self.logger.error(f"[search_web_serp] Error during search: {str(e)}")
            raise requests.RequestException(f"SERP API search failed: {str(e)}")

    def format_serp_results(self, results: dict) -> str:
        """
        Format SERP search results into a readable string.
        
        Args:
            results: Raw dictionary of search results from SERP API
            
        Returns:
            str: Formatted string of search results
        """
        formatted_results = ""
        if 'organic_results' in results:
            for result in results['organic_results']:
                if 'title' in result:
                    formatted_results += f"Title: {result['title']}\n"
                if 'link' in result:
                    formatted_results += f"Link: {result['link']}\n"
                if 'snippet' in result:
                    formatted_results += f"Snippet: {result['snippet']}\n"
                formatted_results += "\n"
        if 'related_questions' in results:
            for result in results['related_questions']:
                if 'question' in result:
                    formatted_results += f"Question: {result['question']}\n"
                if 'title' in result:
                    formatted_results += f"Title: {result['title']}\n"
                if 'link' in result:
                    formatted_results += f"Link: {result['link']}\n"
                if 'snippet' in result:
                    formatted_results += f"Snippet: {result['snippet']}\n"
                formatted_results += "\n"
        return formatted_results

    def _stream_response(
        self,
        client,
        provider: str,
        model_name: str,
        system_prompt: str,
        user_prompt: str,
        temperature: float,
        image_path: str,
        image_paths: list,
        file_path: str,
        max_tokens: int,
        price_per_prompt_tokens: float,
        price_per_completion_tokens: float,
        extra_cost: float,
        extra_token_usage: dict,
        start_time: float,
    ) -> Generator[StreamChunk, None, StreamResult]:
        """
        Wrapper that dispatches to provider-specific streaming functions
        and yields StreamChunk objects. Returns StreamResult when complete.
        """
        stream_functions = {
            "OpenAI": self._openai_stream,
            "Anthropic": self._claude_stream,
            "Grok": self._grok_stream,
            "Groq": self._groq_stream,
            "Google": self._gemini_stream,
            "DeepSeek": self._deepseek_stream,
            "Alibaba": self._qwen_stream,
        }
        
        stream_func = stream_functions.get(provider)
        if not stream_func:
            raise ValueError(f"Streaming not supported for provider: {provider}")
        
        accumulated = ""
        token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
        
        for chunk, usage in stream_func(
            client, model_name, system_prompt, user_prompt,
            temperature, image_path, image_paths, file_path, max_tokens
        ):
            accumulated = chunk.content
            if usage:
                token_usage = usage
            yield chunk
        
        execution_time = time.time() - start_time
        token_usage["prompt_tokens"] += extra_token_usage.get("prompt_tokens", 0)
        token_usage["completion_tokens"] += extra_token_usage.get("completion_tokens", 0)
        token_usage["total_tokens"] += extra_token_usage.get("total_tokens", 0)
        cost = (
            token_usage["prompt_tokens"] * price_per_prompt_tokens +
            token_usage["completion_tokens"] * price_per_completion_tokens
        ) + extra_cost
        
        self.logger.success("SUCCESS (streaming)")
        self.logger.success(f"Execution Time: {execution_time:.2f} seconds")
        self.logger.success(f"Token Usage:\n{json.dumps(token_usage, indent=2)}")
        self.logger.success(f"Price: {cost:.4f} USD")
        
        yield StreamChunk(content=accumulated, delta="", token_usage=token_usage, cost=cost, done=True)

    def _openai_stream(
        self, client, model, system_prompt, user_prompt,
        temperature, image_path, image_paths, file_path, max_tokens
    ) -> Generator[tuple, None, None]:
        """OpenAI streaming implementation."""
        model_cfg = MODEL_CONFIGS.get(model)
        
        user_content = [{"type": "input_text", "text": user_prompt}]
        if image_path:
            user_content.append({
                "type": "input_image",
                "file_id": self.create_image_openai(image_path)
            })
        if image_paths:
            for img_path in image_paths:
                user_content.append({
                    "type": "input_image",
                    "file_id": self.create_image_openai(img_path)
                })
        if file_path:
            if file_path.split('.')[-1] == "pdf":
                user_content.append({
                    "type": "input_file",
                    "file_id": self.create_file_openai(file_path)
                })
            else:
                user_content.append({
                    "type": "input_text",
                    "text": self.read_file_content(file_path)
                })
        
        completion_params = {
            "model": model,
            "input": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_content}
            ],
            "stream": True,
        }
        if max_tokens:
            completion_params["max_output_tokens"] = max_tokens
        if model_cfg and "reasoning_effort" in model_cfg:
            completion_params["reasoning"] = {"effort": model_cfg["reasoning_effort"]}
        elif temperature:
            completion_params["temperature"] = temperature
        
        accumulated = ""
        token_usage = None
        
        for event in client.responses.create(**completion_params):
            if hasattr(event, 'type'):
                if event.type == 'response.output_text.delta':
                    delta = event.delta if hasattr(event, 'delta') else ""
                    accumulated += delta
                    yield StreamChunk(content=accumulated, delta=delta), None
                elif event.type == 'response.completed':
                    if hasattr(event, 'response') and hasattr(event.response, 'usage'):
                        usage = event.response.usage
                        token_usage = {
                            "prompt_tokens": usage.input_tokens,
                            "completion_tokens": usage.output_tokens,
                            "total_tokens": usage.input_tokens + usage.output_tokens,
                        }
                        yield StreamChunk(content=accumulated, delta=""), token_usage

    def _claude_stream(
        self, client, model, system_prompt, user_prompt,
        temperature, image_path, image_paths, file_path, max_tokens
    ) -> Generator[tuple, None, None]:
        """Anthropic Claude streaming implementation."""
        model_name = model.replace("-thinking", "")
        system_prompt = system_prompt.replace('%', 'percent')
        user_prompt = user_prompt.replace('%', 'percent')
        
        content = [{"type": "text", "text": user_prompt}]
        if image_path:
            base64_image = self.encode_image(image_path)
            media_type = f"image/{image_path.split('.')[-1]}"
            content.append({
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": media_type,
                    "data": base64_image
                }
            })
        if image_paths:
            for img_path in image_paths:
                base64_img = self.encode_image(img_path)
                img_media_type = f"image/{img_path.split('.')[-1]}"
                content.append({
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": img_media_type,
                        "data": base64_img
                    }
                })
        if file_path:
            if file_path.split('.')[-1] in ['pdf', 'txt']:
                type_mapper = {"pdf": "application/pdf", "txt": "text/plain"}
                file = client.beta.files.upload(
                    file=(file_path, open(file_path, "rb"), type_mapper[file_path.split('.')[-1]])
                )
                content.append({
                    "type": "document",
                    "source": {"type": "file", "file_id": file.id}
                })
            else:
                content.append({
                    "type": "text",
                    "text": f"\n\nDocument Content:\n{self.read_file_content(file_path)}\n\n"
                })
        
        params = {
            "model": model_name,
            "system": system_prompt,
            "messages": [{"role": "user", "content": content}],
            "max_tokens": max_tokens or (16000 if "3-5" not in model else 8000),
        }
        if temperature:
            params["temperature"] = temperature
        if "thinking" in model:
            params["thinking"] = MODEL_CONFIGS[model]["thinking"]
            params["temperature"] = 1
        
        accumulated = ""
        token_usage = None
        
        with client.messages.stream(**params) as stream:
            for text in stream.text_stream:
                accumulated += text
                yield StreamChunk(content=accumulated, delta=text), None
            
            response = stream.get_final_message()
            token_usage = {
                "prompt_tokens": response.usage.input_tokens,
                "completion_tokens": response.usage.output_tokens,
                "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
            }
            yield StreamChunk(content=accumulated, delta=""), token_usage

    def _grok_stream(
        self, client, model, system_prompt, user_prompt,
        temperature, image_path, image_paths, file_path, max_tokens
    ) -> Generator[tuple, None, None]:
        """Grok streaming implementation."""
        user_input = [user_prompt]
        if image_path:
            base64_image = self.encode_image(image_path)
            image_input = image(image_url=f"data:image/jpeg;base64,{base64_image}", detail="auto")
            user_input.append(image_input)
        if image_paths:
            for img_path in image_paths:
                base64_img = self.encode_image(img_path)
                user_input.append(image(image_url=f"data:image/jpeg;base64,{base64_img}", detail="auto"))
        
        if max_tokens:
            chat = client.chat.create(model=model, temperature=temperature, max_tokens=max_tokens)
        else:
            chat = client.chat.create(model=model, temperature=temperature)
        
        chat.append(system(system_prompt))
        chat.append(user(*user_input))
        
        accumulated = ""
        token_usage = None
        
        for response, chunk in chat.stream():
            delta = chunk.content if hasattr(chunk, 'content') else ""
            accumulated += delta
            yield StreamChunk(content=accumulated, delta=delta), None
        
        if hasattr(response, 'usage'):
            token_usage = {
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            }
        yield StreamChunk(content=accumulated, delta=""), token_usage

    def _groq_stream(
        self, client, model, system_prompt, user_prompt,
        temperature, image_path, image_paths, file_path, max_tokens
    ) -> Generator[tuple, None, None]:
        """Groq streaming implementation."""
        user_content = [{"type": "text", "text": user_prompt}]
        if image_path:
            base64_image = self.encode_image(image_path)
            user_content.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}
            })
        if image_paths:
            for img_path in image_paths:
                base64_img = self.encode_image(img_path)
                user_content.append({
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_img}"}
                })
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content}
        ]
        request_params = {"model": model, "messages": messages, "stream": True}
        if temperature:
            request_params["temperature"] = temperature
        if max_tokens:
            request_params["max_completion_tokens"] = max_tokens
        
        accumulated = ""
        token_usage = None
        
        for chunk in client.chat.completions.create(**request_params):
            if chunk.choices and chunk.choices[0].delta.content:
                delta = chunk.choices[0].delta.content
                accumulated += delta
                yield StreamChunk(content=accumulated, delta=delta), None
            if hasattr(chunk, 'x_groq') and chunk.x_groq and hasattr(chunk.x_groq, 'usage'):
                usage = chunk.x_groq.usage
                if usage is not None:
                    token_usage = {
                        "prompt_tokens": usage.prompt_tokens,
                        "completion_tokens": usage.completion_tokens,
                        "total_tokens": usage.total_tokens,
                    }
        
        yield StreamChunk(content=accumulated, delta=""), token_usage

    def _gemini_stream(
        self, client, model, system_prompt, user_prompt,
        temperature, image_path, image_paths, file_path, max_tokens
    ) -> Generator[tuple, None, None]:
        """Google Gemini streaming implementation."""
        prompt = system_prompt + "\n\n" + user_prompt
        contents = [prompt]
        
        if image_path:
            with open(image_path, "rb") as img:
                image_bytes = img.read()
            contents.append(
                types.Part.from_bytes(
                    data=image_bytes,
                    mime_type=f"image/{image_path.split('.')[-1]}"
                )
            )
        if image_paths:
            for img_path in image_paths:
                with open(img_path, "rb") as img_file:
                    img_bytes = img_file.read()
                contents.append(
                    types.Part.from_bytes(
                        data=img_bytes,
                        mime_type=f"image/{img_path.split('.')[-1]}"
                    )
                )
        if file_path:
            if file_path.split('.')[-1] in ["pdf", "txt"]:
                myfile = client.files.upload(file=file_path)
                contents.append(myfile)
            else:
                contents.append(f"Document Content:\n{self.read_file_content(file_path)}")
        
        config_params = {}
        if temperature:
            config_params["temperature"] = temperature
        if max_tokens:
            config_params["maxOutputTokens"] = max_tokens
        
        accumulated = ""
        token_usage = None
        
        for chunk in client.models.generate_content_stream(
            model=model,
            contents=contents,
            config=types.GenerateContentConfig(**config_params) if config_params else None
        ):
            if hasattr(chunk, 'text') and chunk.text:
                delta = chunk.text
                accumulated += delta
                yield StreamChunk(content=accumulated, delta=delta), None
            
            if hasattr(chunk, 'usage_metadata') and chunk.usage_metadata:
                token_usage = {
                    "prompt_tokens": chunk.usage_metadata.prompt_token_count or 0,
                    "completion_tokens": chunk.usage_metadata.candidates_token_count or 0,
                    "total_tokens": chunk.usage_metadata.total_token_count or 0,
                }
        
        yield StreamChunk(content=accumulated, delta=""), token_usage

    def _deepseek_stream(
        self, client, model, system_prompt, user_prompt,
        temperature, image_path, image_paths, file_path, max_tokens
    ) -> Generator[tuple, None, None]:
        """DeepSeek streaming implementation (OpenAI-compatible)."""
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        request_params = {"model": model, "messages": messages, "stream": True}
        if temperature:
            request_params["temperature"] = temperature
        if max_tokens:
            request_params["max_tokens"] = max_tokens
        
        accumulated = ""
        token_usage = None
        
        for chunk in client.chat.completions.create(**request_params):
            if chunk.choices and chunk.choices[0].delta.content:
                delta = chunk.choices[0].delta.content
                accumulated += delta
                yield StreamChunk(content=accumulated, delta=delta), None
            if hasattr(chunk, 'usage') and chunk.usage:
                token_usage = {
                    "prompt_tokens": chunk.usage.prompt_tokens or 0,
                    "completion_tokens": chunk.usage.completion_tokens or 0,
                    "total_tokens": chunk.usage.total_tokens or 0,
                }
        
        yield StreamChunk(content=accumulated, delta=""), token_usage

    def _qwen_stream(
        self, client, model, system_prompt, user_prompt,
        temperature, image_path, image_paths, file_path, max_tokens
    ) -> Generator[tuple, None, None]:
        """Qwen streaming implementation (OpenAI-compatible)."""
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        request_params = {
            "model": model,
            "messages": messages,
            "stream": True,
            "stream_options": {"include_usage": True},
        }
        if temperature:
            request_params["temperature"] = temperature
        if max_tokens:
            request_params["max_tokens"] = max_tokens
        
        accumulated = ""
        token_usage = None
        
        for chunk in client.chat.completions.create(**request_params):
            if chunk.choices and chunk.choices[0].delta.content:
                delta = chunk.choices[0].delta.content
                accumulated += delta
                yield StreamChunk(content=accumulated, delta=delta), None
            if hasattr(chunk, 'usage') and chunk.usage:
                token_usage = {
                    "prompt_tokens": chunk.usage.prompt_tokens or 0,
                    "completion_tokens": chunk.usage.completion_tokens or 0,
                    "total_tokens": chunk.usage.total_tokens or 0,
                }
        
        yield StreamChunk(content=accumulated, delta=""), token_usage