import os
import json
import logging
import backoff

logger = logging.getLogger()
doubao_logger = logging.getLogger("doubao_api")
import requests
from typing import List, Dict, Any, Optional, Union
import numpy as np
from anthropic import Anthropic
from openai import (
    AzureOpenAI,
    APIConnectionError,
    APIError,
    AzureOpenAI,
    OpenAI,
    RateLimitError,
)
from google import genai
from google.genai import types
from zhipuai import ZhipuAI
from groq import Groq
import boto3
import exa_py
from typing import List, Dict, Any, Optional, Union, Tuple

class ModelPricing:
    def __init__(self, pricing_file: str = "model_pricing.json"):
        self.pricing_file = pricing_file
        self.pricing_data = self._load_pricing()
    
    def _load_pricing(self) -> Dict:
        if os.path.exists(self.pricing_file):
            try:
                with open(self.pricing_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                print(f"Warning: Failed to load pricing file {self.pricing_file}: {e}")
        
        return {
            "default": {"input": 0, "output": 0}
        }
    
    def get_price(self, model: str) -> Dict[str, float]:
        # Handle nested pricing data structure
        if "llm_models" in self.pricing_data:
            # Iterate through all LLM model categories
            for category, models in self.pricing_data["llm_models"].items():
                # Direct model name matching
                if model in models:
                    pricing = models[model]
                    return self._parse_pricing(pricing)
                
                # Fuzzy matching for model names
                for model_name in models:
                    if model_name in model or model in model_name:
                        pricing = models[model_name]
                        return self._parse_pricing(pricing)
        
        # Handle embedding models
        if "embedding_models" in self.pricing_data:
            for category, models in self.pricing_data["embedding_models"].items():
                if model in models:
                    pricing = models[model]
                    return self._parse_pricing(pricing)
                
                for model_name in models:
                    if model_name in model or model in model_name:
                        pricing = models[model_name]
                        return self._parse_pricing(pricing)
        
        # Default pricing
        return {"input": 0, "output": 0}
    
    def _parse_pricing(self, pricing: Dict[str, str]) -> Dict[str, float]:
        """Parse pricing strings and convert to numeric values"""
        result = {}
        
        for key, value in pricing.items():
            if isinstance(value, str):
                # Remove currency symbols and units, convert to float
                clean_value = value.replace('$', '').replace('￥', '').replace(',', '')
                try:
                    result[key] = float(clean_value)
                except ValueError:
                    result[key] = 0.0
            else:
                result[key] = float(value) if value else 0.0
        
        return result
    
    def calculate_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        pricing = self.get_price(model)
        input_cost = (input_tokens / 1000000) * pricing["input"]
        output_cost = (output_tokens / 1000000) * pricing["output"]
        return input_cost + output_cost

# Initialize pricing manager with correct pricing file path
pricing_file = os.path.join(os.path.dirname(__file__), 'model_pricing.json')
pricing_manager = ModelPricing(pricing_file)

def extract_token_usage(response, provider: str) -> Tuple[int, int]:
    if "-" in provider:
        api_type, vendor = provider.split("-", 1)
    else:
        api_type, vendor = "llm", provider

    if api_type == "llm":
        if vendor in ["openai", "qwen", "deepseek", "doubao", "siliconflow", "monica", "vllm", "groq", "zhipu", "gemini", "openrouter", "azureopenai", "huggingface", "exa", "lybic"]:
            if hasattr(response, 'usage') and response.usage:
                return response.usage.prompt_tokens, response.usage.completion_tokens
        
        elif vendor == "anthropic":
            if hasattr(response, 'usage') and response.usage:
                return response.usage.input_tokens, response.usage.output_tokens
        
        elif vendor == "bedrock":
            if isinstance(response, dict) and "usage" in response:
                usage = response["usage"]
                return usage.get("input_tokens", 0), usage.get("output_tokens", 0)
    
    elif api_type == "embedding":
        if vendor in ["openai", "azureopenai", "qwen", "doubao"]:
            if hasattr(response, 'usage') and response.usage:
                return response.usage.prompt_tokens, 0
        
        elif vendor == "jina":
            if isinstance(response, dict) and "usage" in response:
                total_tokens = response["usage"].get("total_tokens", 0)
                return total_tokens, 0
        
        elif vendor == "gemini":
            if hasattr(response, 'usage') and response.usage:
                return response.usage.prompt_tokens, 0

    return 0, 0

def calculate_tokens_and_cost(response, provider: str, model: str) -> Tuple[List[int], float]:
    input_tokens, output_tokens = extract_token_usage(response, provider)
    total_tokens = input_tokens + output_tokens
    cost = pricing_manager.calculate_cost(model, input_tokens, output_tokens)
    
    return [input_tokens, output_tokens, total_tokens], cost

class LMMEngine:
    pass

# ==================== LLM ====================

class LMMEngineOpenAI(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-openai"

        api_key = api_key or os.getenv("OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENAI_API_KEY"
            )

        self.base_url = base_url

        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        if not self.base_url:
            self.llm_client = OpenAI(api_key=self.api_key)
        else:
            self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            **({} if self.model in ["o3", "o3-pro"] else {"temperature": temperature}),
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        return content, total_tokens, cost


class LMMEngineLybic(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-lybic"

        api_key = api_key or os.getenv("LYBIC_LLM_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named LYBIC_LLM_API_KEY"
            )

        self.base_url = base_url or "https://aigw.lybicai.com/v1"
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=1, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            # temperature=temperature,
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        return content, total_tokens, cost


class LMMEngineQwen(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, enable_thinking=False, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.enable_thinking = enable_thinking
        self.provider = "llm-qwen"

        api_key = api_key or os.getenv("DASHSCOPE_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named DASHSCOPE_API_KEY"
            )

        self.base_url = base_url or "https://dashscope.aliyuncs.com/compatible-mode/v1"
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        # For Qwen3 models, we need to handle thinking mode
        extra_body = {}
        if self.model.startswith("qwen3") and not self.enable_thinking:
            extra_body["enable_thinking"] = False

        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            **extra_body,
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        return content, total_tokens, cost


class LMMEngineDoubao(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-doubao"

        api_key = api_key or os.getenv("ARK_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named ARK_API_KEY"
            )

        self.base_url = base_url or "https://ark.cn-beijing.volces.com/api/v3"
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        
        # doubao_logger.info(f"Doubao API Call - Model: {self.model}, Temperature: {temperature}, Max Tokens: {max_new_tokens}")
        # doubao_logger.info(f"Doubao API Input - Messages count: {len(messages)}")
        # doubao_logger.info(f"Doubao API Input - messages: {messages}")
        
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            extra_body={
                "thinking": {
                    "type": "disabled",
                    # "type": "enabled",
                    # "type": "auto",
                }
            },
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        # doubao_logger.info(f"Doubao API Response - Content length: {len(content) if content else 0}, Tokens: {total_tokens}, Cost: {cost}")

        # doubao_logger.info(f"Doubao API Response - Content: {content}")
        
        return content, total_tokens, cost


class LMMEngineAnthropic(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, thinking=False, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.thinking = thinking
        self.provider = "llm-anthropic"

        api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named ANTHROPIC_API_KEY"
            )

        self.api_key = api_key

        self.llm_client = Anthropic(api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        if self.thinking:
            response = self.llm_client.messages.create(
                system=messages[0]["content"][0]["text"],
                model=self.model,
                messages=messages[1:],
                max_tokens=8192,
                thinking={"type": "enabled", "budget_tokens": 4096},
                **kwargs,
            )
            thoughts = response.content[0].thinking
            print("CLAUDE 3.7 THOUGHTS:", thoughts)
            content = response.content[1].text
        else:
            response = self.llm_client.messages.create(
                system=messages[0]["content"][0]["text"],
                model=self.model,
                messages=messages[1:],
                max_tokens=max_new_tokens if max_new_tokens else 8192,
                temperature=temperature,
                **kwargs,
            )
            content = response.content[0].text
        
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        return content, total_tokens, cost


class LMMEngineGemini(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-gemini"

        api_key = api_key or os.getenv("GEMINI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named GEMINI_API_KEY"
            )

        self.base_url = base_url or os.getenv("GEMINI_ENDPOINT_URL")
        if self.base_url is None:
            raise ValueError(
                "An endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named GEMINI_ENDPOINT_URL"
            )

        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            # reasoning_effort="low",
            extra_body={
                'extra_body': {
                    "google": {
                        "thinking_config": {
                            "thinking_budget": 128,
                            "include_thoughts": True
                        }
                    }
                }
            },
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        return content, total_tokens, cost



class LMMEngineOpenRouter(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-openrouter"

        api_key = api_key or os.getenv("OPENROUTER_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENROUTER_API_KEY"
            )

        self.base_url = base_url or os.getenv("OPEN_ROUTER_ENDPOINT_URL")
        if self.base_url is None:
            raise ValueError(
                "An endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named OPEN_ROUTER_ENDPOINT_URL"
            )

        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        return content, total_tokens, cost


class LMMEngineAzureOpenAI(LMMEngine):
    def __init__(
        self,
        base_url=None,
        api_key=None,
        azure_endpoint=None,
        model=None,
        api_version=None,
        rate_limit=-1,
        **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-azureopenai"

        assert api_version is not None, "api_version must be provided"
        self.api_version = api_version

        api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named AZURE_OPENAI_API_KEY"
            )

        self.api_key = api_key

        azure_endpoint = azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT")
        if azure_endpoint is None:
            raise ValueError(
                "An Azure API endpoint needs to be provided in either the azure_endpoint parameter or as an environment variable named AZURE_OPENAI_ENDPOINT"
            )

        self.azure_endpoint = azure_endpoint
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = AzureOpenAI(
            azure_endpoint=self.azure_endpoint,
            api_key=self.api_key,
            api_version=self.api_version,
        )
        self.cost = 0.0

    # @backoff.on_exception(backoff.expo, (APIConnectionError, APIError, RateLimitError), max_tries=10)
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            **kwargs,
        )
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        return content, total_tokens, cost


class LMMEnginevLLM(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.api_key = api_key
        self.provider = "llm-vllm"

        self.base_url = base_url or os.getenv("vLLM_ENDPOINT_URL")
        if self.base_url is None:
            raise ValueError(
                "An endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named vLLM_ENDPOINT_URL"
            )

        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    # @backoff.on_exception(backoff.expo, (APIConnectionError, APIError, RateLimitError), max_tries=10)
    # TODO: Default params chosen for the Qwen model
    def generate(
        self,
        messages,
        temperature=0.0,
        top_p=0.8,
        repetition_penalty=1.05,
        max_new_tokens=512,
        **kwargs
    ):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            top_p=top_p,
            extra_body={"repetition_penalty": repetition_penalty},
        )
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        return content, total_tokens, cost


class LMMEngineHuggingFace(LMMEngine):
    def __init__(self, base_url=None, api_key=None, rate_limit=-1, **kwargs):
        assert base_url is not None, "HuggingFace endpoint must be provided"
        self.base_url = base_url
        self.model = base_url.split('/')[-1] if base_url else "huggingface-tgi"
        self.provider = "llm-huggingface"

        api_key = api_key or os.getenv("HF_TOKEN")
        if api_key is None:
            raise ValueError(
                "A HuggingFace token needs to be provided in either the api_key parameter or as an environment variable named HF_TOKEN"
            )

        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model="tgi",
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        return content, total_tokens, cost


class LMMEngineDeepSeek(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-deepseek"

        api_key = api_key or os.getenv("DEEPSEEK_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named DEEPSEEK_API_KEY"
            )

        self.base_url = base_url or "https://api.deepseek.com"
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        return content, total_tokens, cost


class LMMEngineZhipu(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-zhipu"

        api_key = api_key or os.getenv("ZHIPU_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named ZHIPU_API_KEY"
            )

        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        # Use ZhipuAI client directly instead of OpenAI compatibility layer
        self.llm_client = ZhipuAI(api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            **kwargs,
        )
        
        content = response.choices[0].message.content # type: ignore
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        return content, total_tokens, cost



class LMMEngineGroq(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-groq"

        api_key = api_key or os.getenv("GROQ_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named GROQ_API_KEY"
            )

        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        # Use Groq client directly
        self.llm_client = Groq(api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        return content, total_tokens, cost


class LMMEngineSiliconflow(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-siliconflow"

        api_key = api_key or os.getenv("SILICONFLOW_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named SILICONFLOW_API_KEY"
            )

        self.base_url = base_url or "https://api.siliconflow.cn/v1"
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        return content, total_tokens, cost


class LMMEngineMonica(LMMEngine):
    def __init__(
        self, base_url=None, api_key=None, model=None, rate_limit=-1, **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-monica"

        api_key = api_key or os.getenv("MONICA_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named MONICA_API_KEY"
            )

        self.base_url = base_url or "https://openapi.monica.im/v1"
        self.api_key = api_key
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        self.llm_client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""
        response = self.llm_client.chat.completions.create(
            model=self.model,
            messages=messages,
            max_completion_tokens=max_new_tokens if max_new_tokens else 8192,
            temperature=temperature,
            **kwargs,
        )
        
        content = response.choices[0].message.content
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        return content, total_tokens, cost


class LMMEngineAWSBedrock(LMMEngine):
    def __init__(
        self,
        aws_access_key=None,
        aws_secret_key=None,
        aws_region=None,
        model=None,
        rate_limit=-1,
        **kwargs
    ):
        assert model is not None, "model must be provided"
        self.model = model
        self.provider = "llm-bedrock"

        # Claude model mapping for AWS Bedrock
        self.claude_model_map = {
            "claude-opus-4": "anthropic.claude-opus-4-20250514-v1:0",
            "claude-sonnet-4": "anthropic.claude-sonnet-4-20250514-v1:0",
            "claude-3-7-sonnet": "anthropic.claude-3-7-sonnet-20250219-v1:0",
            "claude-3-5-sonnet": "anthropic.claude-3-5-sonnet-20241022-v2:0",
            "claude-3-5-sonnet-20241022": "anthropic.claude-3-5-sonnet-20241022-v2:0",
            "claude-3-5-sonnet-20240620": "anthropic.claude-3-5-sonnet-20240620-v1:0",
            "claude-3-5-haiku": "anthropic.claude-3-5-haiku-20241022-v1:0",
            "claude-3-haiku": "anthropic.claude-3-haiku-20240307-v1:0",
            "claude-3-sonnet": "anthropic.claude-3-sonnet-20240229-v1:0",
            "claude-3-opus": "anthropic.claude-3-opus-20240229-v1:0",
        }

        # Get the actual Bedrock model ID
        self.bedrock_model_id = self.claude_model_map.get(model, model)

        # AWS credentials
        aws_access_key = aws_access_key or os.getenv("AWS_ACCESS_KEY_ID")
        aws_secret_key = aws_secret_key or os.getenv("AWS_SECRET_ACCESS_KEY")
        aws_region = aws_region or os.getenv("AWS_DEFAULT_REGION") or "us-west-2"

        if aws_access_key is None:
            raise ValueError(
                "AWS Access Key needs to be provided in either the aws_access_key parameter or as an environment variable named AWS_ACCESS_KEY_ID"
            )
        if aws_secret_key is None:
            raise ValueError(
                "AWS Secret Key needs to be provided in either the aws_secret_key parameter or as an environment variable named AWS_SECRET_ACCESS_KEY"
            )

        self.aws_region = aws_region
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        # Initialize Bedrock client
        self.bedrock_client = boto3.client(
            service_name="bedrock-runtime",
            region_name=aws_region,
            aws_access_key_id=aws_access_key,
            aws_secret_access_key=aws_secret_key
        )

    @backoff.on_exception(
        backoff.expo, (APIConnectionError, APIError, RateLimitError), max_time=60
    )
    def generate(self, messages, temperature=0.0, max_new_tokens=None, **kwargs):
        """Generate the next message based on previous messages"""

        # Convert messages to Bedrock format
        # Extract system message if present
        system_message = None
        user_messages = []

        for message in messages:
            if message["role"] == "system":
                if isinstance(message["content"], list):
                    system_message = message["content"][0]["text"]
                else:
                    system_message = message["content"]
            else:
                # Handle both list and string content formats
                if isinstance(message["content"], list):
                    content = message["content"][0]["text"] if message["content"] else ""
                else:
                    content = message["content"]

                user_messages.append({
                    "role": message["role"],
                    "content": content
                })

        # Prepare the body for Bedrock
        body = {
            "max_completion_tokens": max_new_tokens if max_new_tokens else 8192,
            "messages": user_messages,
            "anthropic_version": "bedrock-2023-05-31"
        }

        if temperature > 0:
            body["temperature"] = temperature

        if system_message:
            body["system"] = system_message

        try:
            response = self.bedrock_client.invoke_model(
                body=json.dumps(body),
                modelId=self.bedrock_model_id
            )

            response_body = json.loads(response.get("body").read())

            if "content" in response_body and response_body["content"]:
                content = response_body["content"][0]["text"]
            else:
                raise ValueError("No content in response")
            
            total_tokens, cost = calculate_tokens_and_cost(response_body, self.provider, self.model)
            return content, total_tokens, cost

        except Exception as e:
            print(f"AWS Bedrock error: {e}")
            raise

# ==================== Embedding ====================

class OpenAIEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "text-embedding-3-small",
        api_key=None,
        **kwargs
    ):
        """Init an OpenAI Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "text-embedding-3-small".
            api_key (_type_, optional): Auth key from OpenAI. Defaults to None.
        """
        self.model = embedding_model
        self.provider = "embedding-openai"

        api_key = api_key or os.getenv("OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named OPENAI_API_KEY"
            )
        self.api_key = api_key

    @backoff.on_exception(
        backoff.expo,
        (
            APIError,
            RateLimitError,
            APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> Tuple[np.ndarray, List[int], float]:
        client = OpenAI(api_key=self.api_key)
        response = client.embeddings.create(model=self.model, input=text)
        
        embeddings = np.array([data.embedding for data in response.data])
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        return embeddings, total_tokens, cost



class GeminiEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "text-embedding-004",
        api_key=None,
        **kwargs
    ):
        """Init an Gemini Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "text-embedding-004".
            api_key (_type_, optional): Auth key from Gemini. Defaults to None.
        """
        self.model = embedding_model
        self.provider = "embedding-gemini"

        api_key = api_key or os.getenv("GEMINI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named GEMINI_API_KEY"
            )
        self.api_key = api_key

    @backoff.on_exception(
        backoff.expo,
        (
            APIError,
            RateLimitError,
            APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> Tuple[np.ndarray, List[int], float]:
        client = genai.Client(api_key=self.api_key)

        result = client.models.embed_content(
            model=self.model,
            contents=text,
            config=types.EmbedContentConfig(task_type="SEMANTIC_SIMILARITY"),
        )

        embeddings = np.array([i.values for i in result.embeddings]) # type: ignore
        total_tokens, cost = calculate_tokens_and_cost(result, self.provider, self.model)
        
        return embeddings, total_tokens, cost



class AzureOpenAIEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "text-embedding-3-small",
        api_key=None,
        api_version=None,
        endpoint_url=None,
        **kwargs
    ):
        """Init an Azure OpenAI Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "text-embedding-3-small".
            api_key (_type_, optional): Auth key from Azure OpenAI. Defaults to None.
            api_version (_type_, optional): API version. Defaults to None.
            endpoint_url (_type_, optional): Endpoint URL. Defaults to None.
        """
        self.model = embedding_model
        self.provider = "embedding-azureopenai"

        api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named AZURE_OPENAI_API_KEY"
            )
        self.api_key = api_key

        api_version = api_version or os.getenv("OPENAI_API_VERSION")
        if api_version is None:
            raise ValueError(
                "An API Version needs to be provided in either the api_version parameter or as an environment variable named OPENAI_API_VERSION"
            )
        self.api_version = api_version

        endpoint_url = endpoint_url or os.getenv("AZURE_OPENAI_ENDPOINT")
        if endpoint_url is None:
            raise ValueError(
                "An Endpoint URL needs to be provided in either the endpoint_url parameter or as an environment variable named AZURE_OPENAI_ENDPOINT"
            )
        self.endpoint_url = endpoint_url

    @backoff.on_exception(
        backoff.expo,
        (
            APIError,
            RateLimitError,
            APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> Tuple[np.ndarray, List[int], float]:
        client = AzureOpenAI(
            api_key=self.api_key,
            api_version=self.api_version,
            azure_endpoint=self.endpoint_url,
        )
        response = client.embeddings.create(input=text, model=self.model)
        
        embeddings = np.array([data.embedding for data in response.data])
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        return embeddings, total_tokens, cost


class DashScopeEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "text-embedding-v4",
        api_key=None,
        dimensions: int = 1024,
        **kwargs
    ):
        """Init a DashScope Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "text-embedding-v4".
            api_key (_type_, optional): Auth key from DashScope. Defaults to None.
            dimensions (int, optional): Embedding dimensions. Defaults to 1024.
        """
        self.model = embedding_model
        self.dimensions = dimensions
        self.provider = "embedding-qwen"

        api_key = api_key or os.getenv("DASHSCOPE_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named DASHSCOPE_API_KEY"
            )
        self.api_key = api_key

        # Initialize OpenAI client with DashScope base URL
        self.client = OpenAI(
            api_key=self.api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )

    @backoff.on_exception(
        backoff.expo,
        (
                APIError,
                RateLimitError,
                APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> Tuple[np.ndarray, List[int], float]:
        response = self.client.embeddings.create(
            model=self.model,
            input=text,
            dimensions=self.dimensions,
            encoding_format="float"
        )
        
        embeddings = np.array([data.embedding for data in response.data])
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        return embeddings, total_tokens, cost



class DoubaoEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "doubao-embedding-256",
        api_key=None,
        **kwargs
    ):
        """Init a Doubao Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "doubao-embedding-256".
            api_key (_type_, optional): Auth key from Doubao. Defaults to None.
        """
        self.model = embedding_model
        self.provider = "embedding-doubao"

        api_key = api_key or os.getenv("ARK_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named ARK_API_KEY"
            )
        self.api_key = api_key
        self.base_url = "https://ark.cn-beijing.volces.com/api/v3"

        # Use OpenAI-compatible client for text embeddings
        self.client = OpenAI(
            api_key=self.api_key,
            base_url=self.base_url
        )

    @backoff.on_exception(
        backoff.expo,
        (
                APIError,
                RateLimitError,
                APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> Tuple[np.ndarray, List[int], float]:
        # Log embedding request
        logger.info(f"Doubao Embedding API Call - Model: {self.model}, Text length: {len(text)}")
        doubao_logger.info(f"Doubao Embedding API Call - Model: {self.model}, Text length: {len(text)}")
        
        response = self.client.embeddings.create(
            model=self.model,
            input=text,
            encoding_format="float"
        )
        
        embeddings = np.array([data.embedding for data in response.data])
        total_tokens, cost = calculate_tokens_and_cost(response, self.provider, self.model)
        
        # Log embedding response
        logger.info(f"Doubao Embedding API Response - Embedding dimension: {embeddings.shape}, Tokens: {total_tokens}, Cost: {cost}")
        doubao_logger.info(f"Doubao Embedding API Response - Embedding dimension: {embeddings.shape}, Tokens: {total_tokens}, Cost: {cost}")
        
        return embeddings, total_tokens, cost


class JinaEmbeddingEngine(LMMEngine):
    def __init__(
        self,
        embedding_model: str = "jina-embeddings-v4",
        api_key=None,
        task: str = "retrieval.query",
        **kwargs
    ):
        """Init a Jina AI Embedding engine

        Args:
            embedding_model (str, optional): Model name. Defaults to "jina-embeddings-v4".
            api_key (_type_, optional): Auth key from Jina AI. Defaults to None.
            task (str, optional): Task type. Options: "retrieval.query", "retrieval.passage", "text-matching". Defaults to "retrieval.query".
        """
        self.model = embedding_model
        self.task = task
        self.provider = "embedding-jina" 

        api_key = api_key or os.getenv("JINA_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named JINA_API_KEY"
            )
        self.api_key = api_key
        self.base_url = "https://api.jina.ai/v1"

    @backoff.on_exception(
        backoff.expo,
        (
                APIError,
                RateLimitError,
                APIConnectionError,
        ),
    )
    def get_embeddings(self, text: str) -> Tuple[np.ndarray, List[int], float]:
        import requests

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        data = {
            "model": self.model,
            "task": self.task,
            "input": [
                {
                    "text": text
                }
            ]
        }

        response = requests.post(
            f"{self.base_url}/embeddings",
            headers=headers,
            json=data
        )

        if response.status_code != 200:
            raise Exception(f"Jina AI API error: {response.text}")

        result = response.json()
        embeddings = np.array([data["embedding"] for data in result["data"]])
        
        total_tokens, cost = calculate_tokens_and_cost(result, self.provider, self.model)
    
        return embeddings, total_tokens, cost


# ==================== webSearch ====================
class SearchEngine:
    """Base class for search engines"""
    pass

class BochaAISearchEngine(SearchEngine):
    def __init__(
            self,
            api_key: str|None = None,
            base_url: str = "https://api.bochaai.com/v1",
            rate_limit: int = -1,
            **kwargs
    ):
        """Init a Bocha AI Search engine

        Args:
            api_key (str, optional): Auth key from Bocha AI. Defaults to None.
            base_url (str, optional): Base URL for the API. Defaults to "https://api.bochaai.com/v1".
            rate_limit (int, optional): Rate limit per minute. Defaults to -1 (no limit).
        """
        api_key = api_key or os.getenv("BOCHA_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named BOCHA_API_KEY"
            )

        self.api_key = api_key
        self.base_url = base_url
        self.endpoint = f"{base_url}/ai-search"
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

    @backoff.on_exception(
        backoff.expo,
        (
                APIConnectionError,
                APIError,
                RateLimitError,
                requests.exceptions.RequestException,
        ),
        max_time=60
    )
    def search(
            self,
            query: str,
            freshness: str = "noLimit",
            answer: bool = True,
            stream: bool = False,
            **kwargs
    ) -> Union[Dict[str, Any], Any]:
        """Search with AI and return intelligent answer

        Args:
            query (str): Search query
            freshness (str, optional): Freshness filter. Defaults to "noLimit".
            answer (bool, optional): Whether to return answer. Defaults to True.
            stream (bool, optional): Whether to stream response. Defaults to False.

        Returns:
            Union[Dict[str, Any], Any]: AI search results with sources and answer
        """
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }

        payload = {
            "query": query,
            "freshness": freshness,
            "answer": answer,
            "stream": stream,
            **kwargs
        }

        if stream:
            result = self._stream_search(headers, payload)
            return result, [0, 0, 0], 0.06
        else:
            result = self._regular_search(headers, payload)
            return result, [0, 0, 0], 0.06


    def _regular_search(self, headers: Dict[str, str], payload: Dict[str, Any]) -> Dict[str, Any]:
        """Regular non-streaming search"""
        response = requests.post(
            self.endpoint,
            headers=headers,
            json=payload
        )

        if response.status_code != 200:
            raise APIError(f"Bocha AI Search API error: {response.text}") # type: ignore

        return response.json()

    def _stream_search(self, headers: Dict[str, str], payload: Dict[str, Any]):
        """Streaming search response"""
        response = requests.post(
            self.endpoint,
            headers=headers,
            json=payload,
            stream=True
        )

        if response.status_code != 200:
            raise APIError(f"Bocha AI Search API error: {response.text}") # type: ignore

        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data:'):
                    data = line[5:].strip()
                    if data and data != '{"event":"done"}':
                        try:
                            yield json.loads(data)
                        except json.JSONDecodeError:
                            continue

    def get_answer(self, query: str, **kwargs) -> Tuple[str, int, float]:
        """Get AI generated answer only"""
        result, _, remaining_balance = self.search(query, answer=True, **kwargs)

        # Extract answer from messages
        messages = result.get("messages", []) # type: ignore
        answer = ""
        for message in messages:
            if message.get("type") == "answer":
                answer = message.get("content", "")
                break

        return answer, [0,0,0], remaining_balance # type: ignore


    def get_sources(self, query: str, **kwargs) -> List[Dict[str, Any]]:
        """Get source materials only"""
        result, _, remaining_balance = self.search(query, **kwargs)

        # Extract sources from messages
        sources = []
        messages = result.get("messages", []) # type: ignore
        for message in messages:
            if message.get("type") == "source":
                content_type = message.get("content_type", "")
                if content_type in ["webpage", "image", "video", "baike_pro", "medical_common"]:
                    sources.append({
                        "type": content_type,
                        "content": json.loads(message.get("content", "{}"))
                    })

        return sources, 0, remaining_balance # type: ignore


    def get_follow_up_questions(self, query: str, **kwargs) -> List[str]:
        """Get follow-up questions"""
        result, _, remaining_balance = self.search(query, **kwargs)

        # Extract follow-up questions from messages
        follow_ups = []
        messages = result.get("messages", []) # type: ignore
        for message in messages:
            if message.get("type") == "follow_up":
                follow_ups.append(message.get("content", ""))

        return follow_ups, 0, remaining_balance # type: ignore


class ExaResearchEngine(SearchEngine):
    def __init__(
            self,
            api_key: str|None = None,
            base_url: str = "https://api.exa.ai",
            rate_limit: int = -1,
            **kwargs
    ):
        """Init an Exa Research engine

        Args:
            api_key (str, optional): Auth key from Exa AI. Defaults to None.
            base_url (str, optional): Base URL for the API. Defaults to "https://api.exa.ai".
            rate_limit (int, optional): Rate limit per minute. Defaults to -1 (no limit).
        """
        api_key = api_key or os.getenv("EXA_API_KEY")
        if api_key is None:
            raise ValueError(
                "An API Key needs to be provided in either the api_key parameter or as an environment variable named EXA_API_KEY"
            )

        self.api_key = api_key
        self.base_url = base_url
        self.request_interval = 0 if rate_limit == -1 else 60.0 / rate_limit

        # Initialize OpenAI-compatible client for chat completions
        self.chat_client = OpenAI(
            base_url=base_url,
            api_key=api_key
        )

        # Initialize Exa client for research tasks
        try:
            from exa_py import Exa
            self.exa_client = Exa(api_key=api_key)
        except ImportError:
            self.exa_client = None
            print("Warning: exa_py not installed. Research tasks will not be available.")

    @backoff.on_exception(
        backoff.expo,
        (
                APIConnectionError,
                APIError,
                RateLimitError,
        ),
        max_time=60
    )
    def search(self, query: str, **kwargs):
        """Standard Exa search with direct cost from API

        Args:
            query (str): Search query
            **kwargs: Additional search parameters

        Returns:
            tuple: (result, tokens, cost) where cost is actual API cost
        """
        headers = {
            'x-api-key': self.api_key,
            'Content-Type': 'application/json'
        }
        
        payload = {
            "query": query,
            **kwargs
        }
        
        response = requests.post(
            f"{self.base_url}/search",
            headers=headers,
            json=payload
        )
        
        if response.status_code != 200:
            raise APIError(f"Exa Search API error: {response.text}") # type: ignore
        
        result = response.json()
        
        cost = 0.0
        if "costDollars" in result:
            cost = result["costDollars"].get("total", 0.0)
        
        return result, [0, 0, 0], cost
    
    def chat_research(
            self,
            query: str,
            model: str = "exa",
            stream: bool = False,
            **kwargs
    ) -> Union[str, Any]:
        """Research using chat completions interface

        Args:
            query (str): Research query
            model (str, optional): Model name. Defaults to "exa".
            stream (bool, optional): Whether to stream response. Defaults to False.

        Returns:
            Union[str, Any]: Research result or stream
        """
        messages = [
            {"role": "user", "content": query}
        ]

        if stream:
            completion = self.chat_client.chat.completions.create(
                model=model,
                messages=messages, # type: ignore
                stream=True,
                **kwargs
            )
            return completion
        else:
            completion = self.chat_client.chat.completions.create(
                model=model,
                messages=messages, # type: ignore
                **kwargs
            )
            result = completion.choices[0].message.content # type: ignore
            return result,[0,0,0],0.005