"""
This type stub file was generated by pyright.
"""

import anthropic
import anthropic.types as anthropic_types
import backoff
import openai
import openai.types.chat as openai_types
from abc import ABC, abstractmethod
from openai import OpenAI
from tenacity import retry, stop_after_attempt, wait_random_exponential

CLAUDE_OPENAI_MODEL_MAP = ...
ENCODERS = ...
def count_tokens(s: str, model_name: str = ...) -> int:
    """Uses tiktoken"""
    ...

def get_headers(headers, cache_enabled: bool | None = ...): # -> dict[str, str]:
    ...

class AbstractAIHelper(ABC):
    api_base: str
    headers: dict[str, str]
    @abstractmethod
    def __init__(self) -> None:
        ...
    
    @abstractmethod
    def embeddings_with_backoff(self, **kwargs):
        ...
    
    @abstractmethod
    def get_embeddings(self, content_strs: list[str]) -> list[list[float]]:
        ...
    
    @abstractmethod
    def get_embedding(self, content_str: str) -> list[float]:
        ...
    
    @abstractmethod
    def llm_query_with_retry(self, **kwargs):
        ...
    
    @abstractmethod
    def llm_query_no_retry(self, messages: list = ..., model: str = ..., max_tokens: int | None = ...):
        ...
    
    @abstractmethod
    def llm_query_functions_with_retry(self, model: str, messages: list, functions: list[dict], max_tokens: int | None = ...):
        ...
    
    @abstractmethod
    def llm_query_functions(self, model: str, messages: list, functions: list[dict], max_tokens: int | None = ...):
        ...
    
    @abstractmethod
    def llm_response_to_json(response) -> str:
        ...
    


class OpenAIHelper(AbstractAIHelper):
    client: OpenAI = ...
    def __init__(self, openai_key: str, api_base: str = ..., headers=..., cache: bool | None = ...) -> None:
        ...
    
    def set_up_open_ai_key(self) -> None:
        ...
    
    @backoff.on_exception(backoff.expo, openai.RateLimitError)
    def embeddings_with_backoff(self, **kwargs): # -> CreateEmbeddingResponse:
        ...
    
    def get_embeddings(self, content_strs: list[str]) -> list[list[float]]:
        ...
    
    def get_embedding(self, content_str: str) -> list[float]:
        ...
    
    @backoff.on_exception(backoff.expo, openai.RateLimitError)
    def completions_with_backoff(self, **kwargs):
        ...
    
    @retry(wait=wait_random_exponential(min=70, max=600), stop=stop_after_attempt(10))
    def llm_query_with_retry(self, **kwargs):
        ...
    
    def llm_query_no_retry(self, messages: list = ..., model: str = ..., max_tokens: int | None = ..., **kwargs):
        ...
    
    @retry(wait=wait_random_exponential(min=70, max=600), stop=stop_after_attempt(10))
    def llm_query_functions_with_retry(self, model: str, messages: list, functions: list[dict], max_tokens: int | None = ..., **kwargs):
        ...
    
    def llm_query_functions(self, model: str, messages: list, functions: list[dict], max_tokens: int | None = ..., **kwargs):
        ...
    
    @staticmethod
    def llm_response_to_json(response: openai_types.chat_completion.ChatCompletion) -> str:
        ...
    


class AnthropicHelper(AbstractAIHelper):
    def __init__(self, anthropic_key: str, api_base: str = ..., headers=..., openai_anthropic_translation: bool = ..., cache: bool | None = ...) -> None:
        ...
    
    def set_up_claude_key(self) -> None:
        ...
    
    @backoff.on_exception(backoff.expo, anthropic.RateLimitError)
    def embeddings_with_backoff(self, **kwargs):
        ...
    
    def get_embeddings(self, content_strs: list[str]) -> list[list[float]]:
        ...
    
    def get_embedding(self, content_str: str) -> list[float]:
        ...
    
    @backoff.on_exception(backoff.expo, anthropic.RateLimitError)
    def completions_with_backoff(self, **kwargs):
        ...
    
    @backoff.on_exception(backoff.expo, anthropic.RateLimitError)
    def messages_with_backoff(self, **kwargs):
        ...
    
    @retry(wait=wait_random_exponential(min=70, max=600), stop=stop_after_attempt(10))
    def llm_query_with_retry(self, **kwargs): # -> ChatCompletion:
        ...
    
    def llm_query_no_retry(self, messages: list = ..., model: str = ..., max_tokens: int | None = ..., system_prompt: str | anthropic.NotGiven | None = ..., **kwargs): # -> ChatCompletion:
        ...
    
    @retry(wait=wait_random_exponential(min=70, max=600), stop=stop_after_attempt(10))
    def llm_query_functions_with_retry(self, **kwargs): # -> ChatCompletion:
        ...
    
    def llm_query_functions(self, model: str, messages: list, functions: list, max_tokens: int | None = ..., system_prompt: str | anthropic.NotGiven | None = ..., **kwargs): # -> ChatCompletion:
        ...
    
    @staticmethod
    def llm_response_to_json(response: openai_types.chat_completion.ChatCompletion | anthropic_types.Message) -> str:
        ...
    


class MultiProviderAIHelper(AbstractAIHelper):
    def __init__(self, openai_key: str, anthropic_key: str | None = ..., openai_base: str = ..., anthropic_base: str = ..., headers=..., use_openai: bool = ..., use_claude: bool = ..., cache: bool | None = ...) -> None:
        ...
    
    @backoff.on_exception(backoff.expo, openai.RateLimitError)
    def embeddings_with_backoff(self, **kwargs): # -> CreateEmbeddingResponse:
        ...
    
    def get_embeddings(self, content_strs: list[str]) -> list[list[float]]:
        ...
    
    def get_embedding(self, content_str: str) -> list[float]:
        ...
    
    @backoff.on_exception(backoff.expo, anthropic.RateLimitError)
    def completions_with_backoff(self, **kwargs):
        ...
    
    @backoff.on_exception(backoff.expo, anthropic.RateLimitError)
    def messages_with_backoff(self, **kwargs):
        ...
    
    @retry(wait=wait_random_exponential(min=70, max=600), stop=stop_after_attempt(10))
    def llm_query_with_retry(self, **kwargs): # -> ChatCompletion:
        ...
    
    def llm_query_no_retry(self, messages: list = ..., model: str = ..., max_tokens: int | None = ..., **kwargs): # -> ChatCompletion:
        ...
    
    @retry(wait=wait_random_exponential(min=70, max=600), stop=stop_after_attempt(10))
    def llm_query_functions_with_retry(self, **kwargs): # -> ChatCompletion:
        ...
    
    def llm_query_functions(self, model: str, messages: list, functions: list[dict], max_tokens: int | None = ..., **kwargs): # -> ChatCompletion:
        ...
    
    @staticmethod
    def llm_response_to_json(response) -> str:
        ...
    


