chatlas 0.9.2__py3-none-any.whl → 0.11.0__py3-none-any.whl

chatlas/_provider_deepseek.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Optional, cast
+
+from ._chat import Chat
+from ._logging import log_model_default
+from ._provider_openai import OpenAIProvider
+from ._turn import Turn
+from ._utils import MISSING, MISSING_TYPE, is_testing
+
+if TYPE_CHECKING:
+    from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
+
+    from .types.openai import ChatClientArgs, SubmitInputArgs
+
+
+def ChatDeepSeek(
+    *,
+    system_prompt: Optional[str] = None,
+    model: Optional[str] = None,
+    api_key: Optional[str] = None,
+    base_url: str = "https://api.deepseek.com",
+    seed: Optional[int] | MISSING_TYPE = MISSING,
+    kwargs: Optional["ChatClientArgs"] = None,
+) -> Chat["SubmitInputArgs", ChatCompletion]:
+    """
+    Chat with a model hosted on DeepSeek.
+
+    DeepSeek is a platform for AI inference with competitive pricing
+    and performance.
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## API key
+
+    Sign up at <https://platform.deepseek.com> to get an API key.
+    :::
+
+    Examples
+    --------
+
+    ```python
+    import os
+    from chatlas import ChatDeepSeek
+
+    chat = ChatDeepSeek(api_key=os.getenv("DEEPSEEK_API_KEY"))
+    chat.chat("What is the capital of France?")
+    ```
+
+    Known limitations
+    --------------
+
+    * Structured data extraction is not supported.
+    * Images are not supported.
+
+    Parameters
+    ----------
+    system_prompt
+        A system prompt to set the behavior of the assistant.
+    model
+        The model to use for the chat. The default, None, will pick a reasonable
+        default, and warn you about it. We strongly recommend explicitly choosing
+        a model for all but the most casual use.
+    api_key
+        The API key to use for authentication. You generally should not supply
+        this directly, but instead set the `DEEPSEEK_API_KEY` environment variable.
+    base_url
+        The base URL to the endpoint; the default uses DeepSeek's API.
+    seed
+        Optional integer seed that DeepSeek uses to try and make output more
+        reproducible.
+    kwargs
+        Additional arguments to pass to the `openai.OpenAI()` client constructor.
+
+    Returns
+    -------
+    Chat
+        A chat object that retains the state of the conversation.
+
+    Note
+    ----
+    This function is a lightweight wrapper around [](`~chatlas.ChatOpenAI`) with
+    the defaults tweaked for DeepSeek.
+
+    Note
+    ----
+    Pasting an API key into a chat constructor (e.g., `ChatDeepSeek(api_key="...")`)
+    is the simplest way to get started, and is fine for interactive use, but is
+    problematic for code that may be shared with others.
+
+    Instead, consider using environment variables or a configuration file to manage
+    your credentials. One popular way to manage credentials is to use a `.env` file
+    to store your credentials, and then use the `python-dotenv` package to load them
+    into your environment.
+
+    ```shell
+    pip install python-dotenv
+    ```
+
+    ```shell
+    # .env
+    DEEPSEEK_API_KEY=...
+    ```
+
+    ```python
+    from chatlas import ChatDeepSeek
+    from dotenv import load_dotenv
+
+    load_dotenv()
+    chat = ChatDeepSeek()
+    chat.console()
+    ```
+
+    Another, more general, solution is to load your environment variables into the shell
+    before starting Python (maybe in a `.bashrc`, `.zshrc`, etc. file):
+
+    ```shell
+    export DEEPSEEK_API_KEY=...
+    ```
+    """
+    if model is None:
+        model = log_model_default("deepseek-chat")
+
+    if api_key is None:
+        api_key = os.getenv("DEEPSEEK_API_KEY")
+
+    if isinstance(seed, MISSING_TYPE):
+        seed = 1014 if is_testing() else None
+
+    return Chat(
+        provider=DeepSeekProvider(
+            api_key=api_key,
+            model=model,
+            base_url=base_url,
+            seed=seed,
+            name="DeepSeek",
+            kwargs=kwargs,
+        ),
+        system_prompt=system_prompt,
+    )
+
+
+class DeepSeekProvider(OpenAIProvider):
+    @staticmethod
+    def _as_message_param(turns: list[Turn]) -> list["ChatCompletionMessageParam"]:
+        from openai.types.chat import (
+            ChatCompletionAssistantMessageParam,
+            ChatCompletionUserMessageParam,
+        )
+
+        params = OpenAIProvider._as_message_param(turns)
+
+        # Content must be a string
+        for i, param in enumerate(params):
+            if param["role"] in ["assistant", "user"]:
+                param = cast(
+                    ChatCompletionAssistantMessageParam
+                    | ChatCompletionUserMessageParam,
+                    param,
+                )
+                contents = param.get("content", None)
+                if not isinstance(contents, list):
+                    continue
+                params[i]["content"] = "".join(
+                    content.get("text", "") for content in contents
+                )
+
+        return params

chatlas/_provider_github.py

@@ -3,9 +3,11 @@ from __future__ import annotations
 import os
 from typing import TYPE_CHECKING, Optional
 
+import requests
+
 from ._chat import Chat
 from ._logging import log_model_default
-from ._provider_openai import OpenAIProvider
+from ._provider_openai import ModelInfo, OpenAIProvider
 from ._utils import MISSING, MISSING_TYPE, is_testing
 
 if TYPE_CHECKING:
@@ -18,7 +20,7 @@ def ChatGithub(
     system_prompt: Optional[str] = None,
     model: Optional[str] = None,
     api_key: Optional[str] = None,
-    base_url: str = "https://models.inference.ai.azure.com/",
+    base_url: str = "https://models.github.ai/inference/",
     seed: Optional[int] | MISSING_TYPE = MISSING,
     kwargs: Optional["ChatClientArgs"] = None,
 ) -> Chat["SubmitInputArgs", ChatCompletion]:
@@ -125,7 +127,7 @@ def ChatGithub(
         seed = 1014 if is_testing() else None
 
     return Chat(
-        provider=OpenAIProvider(
+        provider=GitHubProvider(
             api_key=api_key,
             model=model,
             base_url=base_url,
@@ -135,3 +137,61 @@ def ChatGithub(
         ),
         system_prompt=system_prompt,
     )
+
+
+class GitHubProvider(OpenAIProvider):
+    def __init__(self, base_url: str, **kwargs):
+        super().__init__(**kwargs)
+        self._base_url = base_url
+
+    def list_models(self) -> list[ModelInfo]:
+        # For some reason the OpenAI SDK API fails here? So perform request manually
+        # models = self._client.models.list()
+
+        base_url = self._base_url
+        if not base_url.endswith("/"):
+            base_url += "/"
+
+        if "azure" in base_url:
+            # i.e., https://models.inference.ai.azure.com
+            return list_models_gh_azure(base_url)
+        else:
+            # i.e., https://models.github.ai/inference/
+            return list_models_gh(base_url)
+
+
+def list_models_gh(base_url: str = "https://models.github.ai/inference/"):
+    # replace /inference endpoint with /catalog
+    base_url = base_url.replace("/inference", "/catalog")
+    response = requests.get(f"{base_url}models")
+    response.raise_for_status()
+    models = response.json()
+
+    res: list[ModelInfo] = []
+    for m in models:
+        _id = m["id"].split("/")[-1]
+        info: ModelInfo = {
+            "id": _id,
+            "name": m["name"],
+            "provider": m["publisher"],
+            "url": m["html_url"],
+        }
+        res.append(info)
+
+    return res
+
+
+def list_models_gh_azure(base_url: str = "https://models.inference.ai.azure.com"):
+    response = requests.get(f"{base_url}models")
+    response.raise_for_status()
+    models = response.json()
+
+    res: list[ModelInfo] = []
+    for m in models:
+        info: ModelInfo = {
+            "id": m["name"],
+            "provider": m["publisher"]
+        }
+        res.append(info)
+
+    return res

chatlas/_provider_google.py

@@ -21,8 +21,8 @@ from ._content import (
 )
 from ._logging import log_model_default
 from ._merge import merge_dicts
-from ._provider import Provider, StandardModelParamNames, StandardModelParams
-from ._tokens import tokens_log
+from ._provider import ModelInfo, Provider, StandardModelParamNames, StandardModelParams
+from ._tokens import get_token_pricing, tokens_log
 from ._tools import Tool
 from ._turn import Turn, user_turn
 
@@ -180,6 +180,30 @@ class GoogleProvider(
 
         self._client = genai.Client(**kwargs_full)
 
+    def list_models(self):
+        models = self._client.models.list()
+
+        res: list[ModelInfo] = []
+        for m in models:
+            name = m.name or "[unknown]"
+            pricing = get_token_pricing(self.name, name) or {}
+            info: ModelInfo = {
+                "id": name,
+                "name": m.display_name or "[unknown]",
+                "input": pricing.get("input"),
+                "output": pricing.get("output"),
+                "cached_input": pricing.get("cached_input"),
+            }
+            res.append(info)
+
+        # Sort list by created_by field (more recent first)
+        res.sort(
+            key=lambda x: x.get("created", 0),
+            reverse=True,
+        )
+
+        return res
+
     @overload
     def chat_perform(
         self,

chatlas/_provider_huggingface.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Optional
+
+from ._chat import Chat
+from ._logging import log_model_default
+from ._provider_openai import OpenAIProvider
+
+if TYPE_CHECKING:
+    from openai.types.chat import ChatCompletion
+
+    from .types.openai import ChatClientArgs, SubmitInputArgs
+
+
+def ChatHuggingFace(
+    *,
+    system_prompt: Optional[str] = None,
+    model: Optional[str] = None,
+    api_key: Optional[str] = None,
+    kwargs: Optional["ChatClientArgs"] = None,
+) -> Chat["SubmitInputArgs", ChatCompletion]:
+    """
+    Chat with a model hosted on Hugging Face Inference API.
+
+    [Hugging Face](https://huggingface.co/) hosts a variety of open-source
+    and proprietary AI models available via their Inference API.
+    To use the Hugging Face API, you must have an Access Token, which you can obtain
+    from your [Hugging Face account](https://huggingface.co/settings/tokens).
+    Ensure that at least "Make calls to Inference Providers" and
+    "Make calls to your Inference Endpoints" is checked.
+
+    Prerequisites
+    --------------
+
+    ::: {.callout-note}
+    ## API key
+
+    You will need to create a Hugging Face account and generate an API token
+    from your [account settings](https://huggingface.co/settings/tokens).
+    Make sure to enable "Make calls to Inference Providers" permission.
+    :::
+
+    Examples
+    --------
+    ```python
+    import os
+    from chatlas import ChatHuggingFace
+
+    chat = ChatHuggingFace(api_key=os.getenv("HUGGINGFACE_API_KEY"))
+    chat.chat("What is the capital of France?")
+    ```
+
+    Parameters
+    ----------
+    system_prompt
+        A system prompt to set the behavior of the assistant.
+    model
+        The model to use for the chat. The default, None, will pick a reasonable
+        default, and warn you about it. We strongly recommend explicitly
+        choosing a model for all but the most casual use.
+    api_key
+        The API key to use for authentication. You generally should not supply
+        this directly, but instead set the `HUGGINGFACE_API_KEY` environment
+        variable.
+    kwargs
+        Additional arguments to pass to the underlying OpenAI client
+        constructor.
+
+    Returns
+    -------
+    Chat
+        A chat object that retains the state of the conversation.
+
+    Known limitations
+    -----------------
+
+    * Some models do not support the chat interface or parts of it, for example
+      `google/gemma-2-2b-it` does not support a system prompt. You will need to
+      carefully choose the model.
+    * Tool calling support varies by model - many models do not support it.
+
+    Note
+    ----
+    This function is a lightweight wrapper around [](`~chatlas.ChatOpenAI`), with
+    the defaults tweaked for Hugging Face.
+
+    Note
+    ----
+    Pasting an API key into a chat constructor (e.g., `ChatHuggingFace(api_key="...")`)
+    is the simplest way to get started, and is fine for interactive use, but is
+    problematic for code that may be shared with others.
+
+    Instead, consider using environment variables or a configuration file to manage
+    your credentials. One popular way to manage credentials is to use a `.env` file
+    to store your credentials, and then use the `python-dotenv` package to load them
+    into your environment.
+
+    ```shell
+    pip install python-dotenv
+    ```
+
+    ```shell
+    # .env
+    HUGGINGFACE_API_KEY=...
+    ```
+
+    ```python
+    from chatlas import ChatHuggingFace
+    from dotenv import load_dotenv
+
+    load_dotenv()
+    chat = ChatHuggingFace()
+    chat.console()
+    ```
+
+    Another, more general, solution is to load your environment variables into the shell
+    before starting Python (maybe in a `.bashrc`, `.zshrc`, etc. file):
+
+    ```shell
+    export HUGGINGFACE_API_KEY=...
+    ```
+    """
+    if api_key is None:
+        api_key = os.getenv("HUGGINGFACE_API_KEY")
+
+    if model is None:
+        model = log_model_default("meta-llama/Llama-3.1-8B-Instruct")
+
+    return Chat(
+        provider=HuggingFaceProvider(
+            api_key=api_key,
+            model=model,
+            kwargs=kwargs,
+        ),
+        system_prompt=system_prompt,
+    )
+
+
+class HuggingFaceProvider(OpenAIProvider):
+    def __init__(
+        self,
+        *,
+        api_key: Optional[str] = None,
+        model: str,
+        kwargs: Optional["ChatClientArgs"] = None,
+    ):
+        # https://huggingface.co/docs/inference-providers/en/index?python-clients=requests#http--curl
+        super().__init__(
+            name="HuggingFace",
+            model=model,
+            api_key=api_key,
+            base_url="https://router.huggingface.co/v1",
+            kwargs=kwargs,
+        )

chatlas/_provider_mistral.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Optional
+
+from ._chat import Chat
+from ._logging import log_model_default
+from ._provider_openai import OpenAIProvider
+from ._utils import MISSING, MISSING_TYPE, is_testing
+
+if TYPE_CHECKING:
+    from openai.types.chat import ChatCompletion
+
+    from .types.openai import ChatClientArgs, SubmitInputArgs
+
+
+def ChatMistral(
+    *,
+    system_prompt: Optional[str] = None,
+    model: Optional[str] = None,
+    api_key: Optional[str] = None,
+    base_url: str = "https://api.mistral.ai/v1/",
+    seed: int | None | MISSING_TYPE = MISSING,
+    kwargs: Optional["ChatClientArgs"] = None,
+) -> Chat["SubmitInputArgs", ChatCompletion]:
+    """
+    Chat with a model hosted on Mistral's La Plateforme.
+
+    Mistral AI provides high-performance language models through their API platform.
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## API credentials
+
+    Get your API key from https://console.mistral.ai/api-keys.
+    :::
+
+    Examples
+    --------
+    ```python
+    import os
+    from chatlas import ChatMistral
+
+    chat = ChatMistral(api_key=os.getenv("MISTRAL_API_KEY"))
+    chat.chat("Tell me three jokes about statisticians")
+    ```
+
+    Known limitations
+    -----------------
+
+    * Tool calling may be unstable.
+    * Images require a model that supports vision.
+
+    Parameters
+    ----------
+    system_prompt
+        A system prompt to set the behavior of the assistant.
+    model
+        The model to use for the chat. The default, None, will pick a reasonable
+        default, and warn you about it. We strongly recommend explicitly
+        choosing a model for all but the most casual use.
+    api_key
+        The API key to use for authentication. You generally should not supply
+        this directly, but instead set the `MISTRAL_API_KEY` environment
+        variable.
+    base_url
+        The base URL to the endpoint; the default uses Mistral AI.
+    seed
+        Optional integer seed that Mistral uses to try and make output more
+        reproducible.
+    kwargs
+        Additional arguments to pass to the `openai.OpenAI()` client
+        constructor (Mistral uses OpenAI-compatible API).
+
+    Returns
+    -------
+    Chat
+        A chat object that retains the state of the conversation.
+
+    Note
+    ----
+    Pasting an API key into a chat constructor (e.g., `ChatMistral(api_key="...")`)
+    is the simplest way to get started, and is fine for interactive use, but is
+    problematic for code that may be shared with others.
+
+    Instead, consider using environment variables or a configuration file to manage
+    your credentials. One popular way to manage credentials is to use a `.env` file
+    to store your credentials, and then use the `python-dotenv` package to load them
+    into your environment.
+
+    ```shell
+    pip install python-dotenv
+    ```
+
+    ```shell
+    # .env
+    MISTRAL_API_KEY=...
+    ```
+
+    ```python
+    from chatlas import ChatMistral
+    from dotenv import load_dotenv
+
+    load_dotenv()
+    chat = ChatMistral()
+    chat.console()
+    ```
+
+    Another, more general, solution is to load your environment variables into the shell
+    before starting Python (maybe in a `.bashrc`, `.zshrc`, etc. file):
+
+    ```shell
+    export MISTRAL_API_KEY=...
+    ```
+    """
+    if isinstance(seed, MISSING_TYPE):
+        seed = 1014 if is_testing() else None
+
+    if model is None:
+        model = log_model_default("mistral-large-latest")
+
+    if api_key is None:
+        api_key = os.getenv("MISTRAL_API_KEY")
+
+    return Chat(
+        provider=MistralProvider(
+            api_key=api_key,
+            model=model,
+            base_url=base_url,
+            seed=seed,
+            kwargs=kwargs,
+        ),
+        system_prompt=system_prompt,
+    )
+
+
+class MistralProvider(OpenAIProvider):
+    def __init__(
+        self,
+        *,
+        api_key: Optional[str] = None,
+        model: str,
+        base_url: str = "https://api.mistral.ai/v1/",
+        seed: Optional[int] = None,
+        name: str = "Mistral",
+        kwargs: Optional["ChatClientArgs"] = None,
+    ):
+        super().__init__(
+            api_key=api_key,
+            model=model,
+            base_url=base_url,
+            seed=seed,
+            name=name,
+            kwargs=kwargs,
+        )
+
+    # Mistral is essentially OpenAI-compatible, with a couple small differences.
+    # We _could_ bring in the Mistral SDK and use it directly for more precise typing,
+    # etc., but for now that doesn't seem worth it.
+    def _chat_perform_args(
+        self, stream, turns, tools, data_model=None, kwargs=None
+    ) -> "SubmitInputArgs":
+        # Get the base arguments from OpenAI provider
+        kwargs2 = super()._chat_perform_args(stream, turns, tools, data_model, kwargs)
+
+        # Mistral doesn't support stream_options
+        if "stream_options" in kwargs2:
+            del kwargs2["stream_options"]
+
+        # Mistral wants random_seed, not seed
+        if seed := kwargs2.pop("seed", None):
+            if isinstance(seed, int):
+                kwargs2["extra_body"] = {"random_seed": seed}
+            elif seed is not None:
+                raise ValueError(
+                    "MistralProvider only accepts an integer seed, or None."
+                )
+
+        return kwargs2