"""Type stubs for langchain_openai package."""

from collections.abc import Sequence
from typing import Any

from langchain_core.embeddings import Embeddings as _Embeddings
from langchain_core.language_models import BaseLanguageModel

class ChatOpenAI(BaseLanguageModel):
    """OpenAI chat model wrapper."""

    def __init__(
        self,
        *,
        model: str = ...,
        temperature: float = ...,
        max_tokens: int | None = None,
        timeout: float | None = None,
        max_retries: int = ...,
        api_key: Any = None,
        base_url: str | None = None,
        organization: str | None = None,
        streaming: bool = ...,
        n: int | None = None,
        openai_api_base: str | None = None,
        openai_api_key: Any = None,
        **kwargs: Any,
    ) -> None: ...

class AzureChatOpenAI(BaseLanguageModel):
    """Azure OpenAI chat model wrapper."""

    def __init__(
        self,
        *,
        deployment_name: str = ...,
        azure_endpoint: str | None = None,
        api_version: str | None = None,
        api_key: Any = None,
        temperature: float = ...,
        max_tokens: int | None = None,
        timeout: float | None = None,
        max_retries: int = ...,
        streaming: bool = ...,
        **kwargs: Any,
    ) -> None: ...

class OpenAIEmbeddings(_Embeddings):
    """OpenAI embeddings wrapper used for vector stores."""

    def __init__(
        self,
        *,
        model: str = ...,
        dimensions: int | None = None,
        api_key: Any = None,
        **kwargs: Any,
    ) -> None: ...
    def embed_documents(self, texts: Sequence[str]) -> list[list[float]]: ...
    def embed_query(self, text: str) -> list[float]: ...

__all__ = ["AzureChatOpenAI", "ChatOpenAI", "OpenAIEmbeddings"]
