chatlas 0.2.0__py3-none-any.whl

chatlas/_groq.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Optional
+
+from ._chat import Chat
+from ._logging import log_model_default
+from ._openai import ChatOpenAI
+from ._turn import Turn
+from ._utils import MISSING, MISSING_TYPE
+
+if TYPE_CHECKING:
+    from ._openai import ChatCompletion
+    from .types.openai import ChatClientArgs, SubmitInputArgs
+
+
+def ChatGroq(
+    *,
+    system_prompt: Optional[str] = None,
+    turns: Optional[list[Turn]] = None,
+    model: Optional[str] = None,
+    api_key: Optional[str] = None,
+    base_url: str = "https://api.groq.com/openai/v1",
+    seed: Optional[int] | MISSING_TYPE = MISSING,
+    kwargs: Optional["ChatClientArgs"] = None,
+) -> Chat["SubmitInputArgs", ChatCompletion]:
+    """
+    Chat with a model hosted on Groq.
+
+    Groq provides a platform for highly efficient AI inference.
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## API key
+
+    Sign up at <https://groq.com> to get an API key.
+    :::
+
+    ::: {.callout-note}
+    ## Python requirements
+
+    `ChatGroq` requires the `openai` package (e.g., `pip install openai`).
+    :::
+
+    Examples
+    --------
+
+    ```python
+    import os
+    from chatlas import ChatGroq
+
+    chat = ChatGroq(api_key=os.getenv("GROQ_API_KEY"))
+    chat.chat("What is the capital of France?")
+    ```
+
+    Parameters
+    ----------
+    system_prompt
+        A system prompt to set the behavior of the assistant.
+    turns
+        A list of turns to start the chat with (i.e., continuing a previous
+        conversation). If not provided, the conversation begins from scratch.
+        Do not provide non-`None` values for both `turns` and `system_prompt`.
+        Each message in the list should be a dictionary with at least `role`
+        (usually `system`, `user`, or `assistant`, but `tool` is also possible).
+        Normally there is also a `content` field, which is a string.
+    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 `GROQ_API_KEY` environment variable.
+    base_url
+        The base URL to the endpoint; the default uses Groq's API.
+    seed
+        Optional integer seed that ChatGPT 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 groq.
+
+    Note
+    ----
+    Pasting an API key into a chat constructor (e.g., `ChatGroq(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
+    GROQ_API_KEY=...
+    ```
+
+    ```python
+    from chatlas import ChatGroq
+    from dotenv import load_dotenv
+
+    load_dotenv()
+    chat = ChatGroq()
+    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 GROQ_API_KEY=...
+    ```
+    """
+    if model is None:
+        model = log_model_default("llama3-8b-8192")
+    if api_key is None:
+        api_key = os.getenv("GROQ_API_KEY")
+
+    return ChatOpenAI(
+        system_prompt=system_prompt,
+        turns=turns,
+        model=model,
+        api_key=api_key,
+        base_url=base_url,
+        seed=seed,
+        kwargs=kwargs,
+    )

chatlas/_interpolate.py

@@ -0,0 +1,133 @@
+import inspect
+from pathlib import Path
+from typing import Any, Optional, Union
+
+from jinja2 import Environment
+
+__all__ = (
+    "interpolate",
+    "interpolate_file",
+)
+
+
+def interpolate(
+    prompt: str,
+    *,
+    variables: Optional[dict[str, Any]] = None,
+    variable_start: str = "{{",
+    variable_end: str = "}}",
+) -> str:
+    """
+    Interpolate variables into a prompt
+
+    This is a light-weight wrapper around the Jinja2 templating engine, making
+    it easier to interpolate dynamic data into a prompt template. Compared to
+    f-strings, which expects you to wrap dynamic values in `{ }`, this function
+    expects `{{ }}` instead, making it easier to include Python code and JSON in
+    your prompt.
+
+    Parameters
+    ----------
+    prompt
+        The prompt to interpolate (as a string).
+    variables
+        A dictionary of variables to interpolate into the prompt. If not
+        provided, the caller's global and local variables are used.
+    variable_start
+        The string that marks the beginning of a variable.
+    variable_end
+        The string that marks the end of a variable.
+
+    Returns
+    -------
+    str
+        The prompt with variables interpolated.
+
+    Examples
+    --------
+
+    ```python
+    from chatlas import interpolate
+
+    x = 1
+    interpolate("The value of `x` is: {{ x }}")
+    ```
+    """
+    if variables is None:
+        frame = inspect.currentframe()
+        variables = _infer_variables(frame)
+        del frame
+
+    env = Environment(
+        variable_start_string=variable_start,
+        variable_end_string=variable_end,
+    )
+
+    template = env.from_string(prompt)
+    return template.render(variables)
+
+
+def interpolate_file(
+    path: Union[str, Path],
+    *,
+    variables: Optional[dict[str, Any]] = None,
+    variable_start: str = "{{",
+    variable_end: str = "}}",
+) -> str:
+    """
+    Interpolate variables into a prompt from a file
+
+    This is a light-weight wrapper around the Jinja2 templating engine, making
+    it easier to interpolate dynamic data into a static prompt. Compared to
+    f-strings, which expects you to wrap dynamic values in `{ }`, this function
+    expects `{{ }}` instead, making it easier to include Python code and JSON in
+    your prompt.
+
+    Parameters
+    ----------
+    path
+        The path to the file containing the prompt to interpolate.
+    variables
+        A dictionary of variables to interpolate into the prompt. If not
+        provided, the caller's global and local variables are used.
+    variable_start
+        The string that marks the beginning of a variable.
+    variable_end
+        The string that marks the end of a variable.
+
+    Returns
+    -------
+    str
+        The prompt with variables interpolated.
+
+    See Also
+    --------
+    interpolate
+        Interpolating data into a system prompt
+    """
+    if variables is None:
+        frame = inspect.currentframe()
+        variables = _infer_variables(frame)
+        del frame
+
+    with open(path, "r") as file:
+        return interpolate(
+            file.read(),
+            variables=variables,
+            variable_start=variable_start,
+            variable_end=variable_end,
+        )
+
+
+def _infer_variables(frame) -> dict[str, Any]:
+    if not inspect.isframe(frame) or frame.f_back is None:
+        raise RuntimeError(
+            "`interpolate()` was unable to infer the caller's global and local "
+            "variables (because the caller's frame is not available). Consider "
+            "passing `variables` explicitly to `interpolate()`."
+        )
+
+    return {
+        **frame.f_back.f_globals,
+        **frame.f_back.f_locals,
+    }

chatlas/_logging.py

@@ -0,0 +1,61 @@
+import logging
+import os
+import warnings
+
+from rich.logging import RichHandler
+
+
+def _rich_handler() -> RichHandler:
+    formatter = logging.Formatter("%(name)s - %(message)s")
+    handler = RichHandler()
+    handler.setFormatter(formatter)
+    return handler
+
+
+logger = logging.getLogger("chatlas")
+
+if os.environ.get("CHATLAS_LOG") == "info":
+    # By adding a RichHandler to chatlas' logger, we can guarantee that they
+    # never get dropped, even if the root logger's handlers are not
+    # RichHandlers.
+    logger.setLevel(logging.INFO)
+    logger.addHandler(_rich_handler())
+    logger.propagate = False
+
+    # Add a RichHandler to the root logger if there are no handlers. Note that
+    # if chatlas is imported before other libraries that set up logging, (like
+    # openai, anthropic, or httpx), this will ensure that logs from those
+    # libraries are also displayed in the rich console.
+    root = logging.getLogger()
+    if not root.handlers:
+        root.addHandler(_rich_handler())
+
+    # Warn if there are non-RichHandler handlers on the root logger.
+    # TODO: we could consider something a bit more abusive here, like removing
+    # non-RichHandler handlers from the root logger, but that could be
+    # surprising to users.
+    bad_handlers = [
+        h.get_name() for h in root.handlers if not isinstance(h, RichHandler)
+    ]
+    if len(bad_handlers) > 0:
+        warnings.warn(
+            "When setting up logging handlers for CHATLAS_LOG, chatlas detected "
+            f"non-rich handler(s) on the root logger named {bad_handlers}. "
+            "As a result, logs handled those handlers may be dropped when the "
+            "`echo` argument of `.chat()`, `.stream()`, etc., is something "
+            "other than 'none'. This problem can likely be fixed by importing "
+            "`chatlas` before other libraries that set up logging, or adding a "
+            "RichHandler to the root logger before loading other libraries.",
+        )
+
+
+def log_model_default(model: str) -> str:
+    logger.info(f"Defaulting to `model = '{model}'`.")
+    return model
+
+
+def log_tool_error(name: str, arguments: str, e: Exception):
+    logger.info(
+        f"Error invoking tool function '{name}' with arguments: {arguments}. "
+        f"The error message is: '{e}'",
+    )

chatlas/_merge.py

@@ -0,0 +1,103 @@
+# Adapted from https://github.com/langchain-ai/langchain/blob/master/libs/core/langchain_core/utils/_merge.py
+# Also tweaked to more closely match https://github.com/hadley/elmer/blob/main/R/utils-merge.R
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+
+def merge_dicts(left: dict[str, Any], *others: dict[str, Any]) -> dict[str, Any]:
+    """Merge many dicts, handling specific scenarios where a key exists in both
+    dictionaries but has a value of None in 'left'. In such cases, the method uses the
+    value from 'right' for that key in the merged dictionary.
+
+    Args:
+        left: The first dictionary to merge.
+        others: The other dictionaries to merge.
+
+    Returns:
+        The merged dictionary.
+
+    Raises:
+        TypeError: If the key exists in both dictionaries but has a different type.
+        TypeError: If the value has an unsupported type.
+
+    Example:
+        If left = {"function_call": {"arguments": None}} and
+        right = {"function_call": {"arguments": "{\n"}}
+        then, after merging, for the key "function_call",
+        the value from 'right' is used,
+        resulting in merged = {"function_call": {"arguments": "{\n"}}.
+    """
+    merged = left.copy()
+    for right in others:
+        for right_k, right_v in right.items():
+            left_v = merged.get(right_k, None)
+
+            if right_v is None:
+                if right_k not in merged:
+                    merged[right_k] = None
+            elif left_v is None:
+                merged[right_k] = right_v
+            elif left_v == right_v:
+                continue
+            elif isinstance(left_v, str):
+                merged[right_k] += right_v
+            elif isinstance(left_v, (int, float)):
+                merged[right_k] = right_v
+            elif isinstance(merged[right_k], dict):
+                merged[right_k] = merge_dicts(merged[right_k], right_v)
+            elif isinstance(merged[right_k], list):
+                merged[right_k] = merge_lists(merged[right_k], right_v)
+            elif type(merged[right_k]) is not type(right_v):
+                raise TypeError(
+                    f'additional_kwargs["{right_k}"] already exists in this message,'
+                    " but with a different type."
+                )
+            else:
+                raise TypeError(
+                    f"Additional kwargs key {right_k} already exists in left dict and "
+                    f"value has unsupported type {type(merged[right_k])}."
+                )
+    return merged
+
+
+def merge_lists(
+    left: Optional[list[Any]], *others: Optional[list[Any]]
+) -> Optional[list[Any]]:
+    """Add many lists, handling None.
+
+    Args:
+        left: The first list to merge.
+        others: The other lists to merge.
+
+    Returns:
+        The merged list.
+    """
+    merged = left.copy() if left is not None else None
+    for other in others:
+        if other is None:
+            continue
+        elif merged is None:
+            merged = other.copy()
+        else:
+            for e in other:
+                if isinstance(e, dict) and "index" in e and isinstance(e["index"], int):
+                    to_merge = [
+                        i
+                        for i, e_left in enumerate(merged)
+                        if e_left["index"] == e["index"]
+                    ]
+                    if to_merge:
+                        # TODO: Remove this once merge_dict is updated with special
+                        # handling for 'type'.
+                        if "type" in e:
+                            e: dict[str, Any] = {  # noqa: PLW2901
+                                k: v for k, v in e.items() if k != "type"
+                            }
+                        merged[to_merge[0]] = merge_dicts(merged[to_merge[0]], e)
+                    else:
+                        merged.append(e)
+                else:
+                    merged.append(e)
+    return merged

chatlas/_ollama.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import json
+import re
+import urllib.request
+from typing import TYPE_CHECKING, Optional
+
+from ._chat import Chat
+from ._openai import ChatOpenAI
+from ._turn import Turn
+
+if TYPE_CHECKING:
+    from ._openai import ChatCompletion
+    from .types.openai import ChatClientArgs, SubmitInputArgs
+
+
+def ChatOllama(
+    model: Optional[str] = None,
+    *,
+    system_prompt: Optional[str] = None,
+    turns: Optional[list[Turn]] = None,
+    base_url: str = "http://localhost:11434",
+    seed: Optional[int] = None,
+    kwargs: Optional["ChatClientArgs"] = None,
+) -> Chat["SubmitInputArgs", ChatCompletion]:
+    """
+    Chat with a local Ollama model.
+
+    [Ollama](https://ollama.com) makes it easy to run a wide-variety of
+    open-source models locally, making it a great choice for privacy
+    and security.
+
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## Ollama runtime
+
+    `ChatOllama` requires the [ollama](https://ollama.com/download) executable
+    to be installed and running on your machine.
+    :::
+
+    ::: {.callout-note}
+    ## Pull model(s)
+
+    Once ollama is running locally, download a model from the command line
+    (e.g. `ollama pull llama3.2`).
+    :::
+
+    Examples
+    --------
+
+    ```python
+    from chatlas import ChatOllama
+
+    chat = ChatOllama(model="llama3.2")
+    chat.chat("What is the capital of France?")
+    ```
+
+    Parameters
+    ----------
+    model
+        The model to use for the chat. If `None`, a list of locally installed
+        models will be printed.
+    system_prompt
+        A system prompt to set the behavior of the assistant.
+    turns
+        A list of turns to start the chat with (i.e., continuing a previous
+        conversation). If not provided, the conversation begins from scratch. Do
+        not provide non-`None` values for both `turns` and `system_prompt`. Each
+        message in the list should be a dictionary with at least `role` (usually
+        `system`, `user`, or `assistant`, but `tool` is also possible). Normally
+        there is also a `content` field, which is a string.
+    base_url
+        The base URL to the endpoint; the default uses ollama's API.
+    seed
+        Optional integer seed that helps to make output more reproducible.
+    kwargs
+        Additional arguments to pass to the `openai.OpenAI()` client constructor.
+
+    Note
+    ----
+    This function is a lightweight wrapper around [](`~chatlas.ChatOpenAI`) with
+    the defaults tweaked for ollama.
+
+    Limitations
+    -----------
+    `ChatOllama` currently doesn't work with streaming tools, and tool calling more
+    generally doesn't seem to work very well with currently available models.
+    """
+
+    base_url = re.sub("/+$", "", base_url)
+
+    if not has_ollama(base_url):
+        raise RuntimeError("Can't find locally running ollama.")
+
+    if model is None:
+        models = ollama_models(base_url)
+        raise ValueError(
+            f"Must specify model. Locally installed models: {', '.join(models)}"
+        )
+
+    return ChatOpenAI(
+        system_prompt=system_prompt,
+        turns=turns,
+        base_url=f"{base_url}/v1",
+        model=model,
+        seed=seed,
+        kwargs=kwargs,
+    )
+
+
+def ollama_models(base_url: str) -> list[str]:
+    res = urllib.request.urlopen(url=f"{base_url}/api/tags")
+    data = json.loads(res.read())
+    return [re.sub(":latest$", "", x["name"]) for x in data["models"]]
+
+
+def has_ollama(base_url):
+    try:
+        urllib.request.urlopen(url=f"{base_url}/api/tags")
+        return True
+    except Exception:
+        return False