lybic-guiagents 0.1.0__py3-none-any.whl

gui_agents/core/engine.py

@@ -0,0 +1,1496 @@
+import os
+import json
+import backoff
+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"]:
+            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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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"""
+        response = self.llm_client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            max_tokens=max_new_tokens if max_new_tokens else 4096,
+            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)
+        
+        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 4096,
+                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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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 4096,
+            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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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_tokens=max_new_tokens if max_new_tokens else 4096,
+            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_tokens": max_new_tokens if max_new_tokens else 4096,
+            "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]:
+        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)
+        
+        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