chatlas 0.2.0__py3-none-any.whl

chatlas/_perplexity.py

@@ -0,0 +1,148 @@
+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 ChatPerplexity(
+    *,
+    system_prompt: Optional[str] = None,
+    turns: Optional[list[Turn]] = None,
+    model: Optional[str] = None,
+    api_key: Optional[str] = None,
+    base_url: str = "https://api.perplexity.ai/",
+    seed: Optional[int] | MISSING_TYPE = MISSING,
+    kwargs: Optional["ChatClientArgs"] = None,
+) -> Chat["SubmitInputArgs", ChatCompletion]:
+    """
+    Chat with a model hosted on perplexity.ai.
+
+    Perplexity AI is a platform for running LLMs that are capable of
+    searching the web in real-time to help them answer questions with
+    information that may not have been available when the model was trained.
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## API key
+
+    Sign up at <https://www.perplexity.ai> to get an API key.
+    :::
+
+    ::: {.callout-note}
+    ## Python requirements
+
+    `ChatPerplexity` requires the `openai` package (e.g., `pip install openai`).
+    :::
+
+
+    Examples
+    --------
+
+    ```python
+    import os
+    from chatlas import ChatPerplexity
+
+    chat = ChatPerplexity(api_key=os.getenv("PERPLEXITY_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 `PERPLEXITY_API_KEY` environment
+        variable.
+    base_url
+        The base URL to the endpoint; the default uses Perplexity'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 perplexity.ai.
+
+    Note
+    ----
+    Pasting an API key into a chat constructor (e.g., `ChatPerplexity(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
+    PERPLEXITY_API_KEY=...
+    ```
+
+    ```python
+    from chatlas import ChatPerplexity
+    from dotenv import load_dotenv
+
+    load_dotenv()
+    chat = ChatPerplexity()
+    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 PERPLEXITY_API_KEY=...
+    ```
+    """
+    if model is None:
+        model = log_model_default("llama-3.1-sonar-small-128k-online")
+    if api_key is None:
+        api_key = os.getenv("PERPLEXITY_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/_provider.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import (
+    Any,
+    AsyncIterable,
+    Generic,
+    Iterable,
+    Literal,
+    Optional,
+    TypeVar,
+    overload,
+)
+
+from pydantic import BaseModel
+
+from ._tools import Tool
+from ._turn import Turn
+
+ChatCompletionT = TypeVar("ChatCompletionT")
+ChatCompletionChunkT = TypeVar("ChatCompletionChunkT")
+# A dictionary representation of a chat completion
+ChatCompletionDictT = TypeVar("ChatCompletionDictT")
+
+
+class Provider(
+    ABC, Generic[ChatCompletionT, ChatCompletionChunkT, ChatCompletionDictT]
+):
+    """
+    A model provider interface for a [](`~chatlas.Chat`).
+
+    This abstract class defines the interface a model provider must implement in
+    order to be used with a [](`~chatlas.Chat`) instance. The provider is
+    responsible for performing the actual chat completion, and for handling the
+    streaming of the completion results.
+
+    Note that this class is exposed for developers who wish to implement their
+    own provider. In general, you should not need to interact with this class
+    directly.
+    """
+
+    @overload
+    @abstractmethod
+    def chat_perform(
+        self,
+        *,
+        stream: Literal[False],
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+        kwargs: Any,
+    ) -> ChatCompletionT: ...
+
+    @overload
+    @abstractmethod
+    def chat_perform(
+        self,
+        *,
+        stream: Literal[True],
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+        kwargs: Any,
+    ) -> Iterable[ChatCompletionChunkT]: ...
+
+    @abstractmethod
+    def chat_perform(
+        self,
+        *,
+        stream: bool,
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+        kwargs: Any,
+    ) -> Iterable[ChatCompletionChunkT] | ChatCompletionT: ...
+
+    @overload
+    @abstractmethod
+    async def chat_perform_async(
+        self,
+        *,
+        stream: Literal[False],
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+        kwargs: Any,
+    ) -> ChatCompletionT: ...
+
+    @overload
+    @abstractmethod
+    async def chat_perform_async(
+        self,
+        *,
+        stream: Literal[True],
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+        kwargs: Any,
+    ) -> AsyncIterable[ChatCompletionChunkT]: ...
+
+    @abstractmethod
+    async def chat_perform_async(
+        self,
+        *,
+        stream: bool,
+        turns: list[Turn],
+        tools: dict[str, Tool],
+        data_model: Optional[type[BaseModel]],
+        kwargs: Any,
+    ) -> AsyncIterable[ChatCompletionChunkT] | ChatCompletionT: ...
+
+    @abstractmethod
+    def stream_text(self, chunk: ChatCompletionChunkT) -> Optional[str]: ...
+
+    @abstractmethod
+    def stream_merge_chunks(
+        self,
+        completion: Optional[ChatCompletionDictT],
+        chunk: ChatCompletionChunkT,
+    ) -> ChatCompletionDictT: ...
+
+    @abstractmethod
+    def stream_turn(
+        self,
+        completion: ChatCompletionDictT,
+        has_data_model: bool,
+        stream: Any,
+    ) -> Turn: ...
+
+    @abstractmethod
+    async def stream_turn_async(
+        self,
+        completion: ChatCompletionDictT,
+        has_data_model: bool,
+        stream: Any,
+    ) -> Turn: ...
+
+    @abstractmethod
+    def value_turn(
+        self,
+        completion: ChatCompletionT,
+        has_data_model: bool,
+    ) -> Turn: ...

chatlas/_tokens.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import copy
+from threading import Lock
+from typing import TYPE_CHECKING
+
+from ._logging import logger
+from ._typing_extensions import TypedDict
+
+if TYPE_CHECKING:
+    from ._provider import Provider
+
+
+class TokenUsage(TypedDict):
+    """
+    Token usage for a given provider (name).
+    """
+
+    name: str
+    input: int
+    output: int
+
+
+class ThreadSafeTokenCounter:
+    def __init__(self):
+        self._lock = Lock()
+        self._tokens: dict[str, TokenUsage] = {}
+
+    def log_tokens(self, name: str, input_tokens: int, output_tokens: int) -> None:
+        logger.info(
+            f"Provider '{name}' generated a response of {output_tokens} tokens "
+            f"from an input of {input_tokens} tokens."
+        )
+
+        with self._lock:
+            if name not in self._tokens:
+                self._tokens[name] = {
+                    "name": name,
+                    "input": input_tokens,
+                    "output": output_tokens,
+                }
+            else:
+                self._tokens[name]["input"] += input_tokens
+                self._tokens[name]["output"] += output_tokens
+
+    def get_usage(self) -> list[TokenUsage] | None:
+        with self._lock:
+            if not self._tokens:
+                return None
+            # Create a deep copy to avoid external modifications
+            return copy.deepcopy(list(self._tokens.values()))
+
+
+# Global instance
+_token_counter = ThreadSafeTokenCounter()
+
+
+def tokens_log(provider: "Provider", tokens: tuple[int, int]) -> None:
+    """
+    Log token usage for a provider in a thread-safe manner.
+    """
+    name = provider.__class__.__name__.replace("Provider", "")
+    _token_counter.log_tokens(name, tokens[0], tokens[1])
+
+
+def tokens_reset() -> None:
+    """
+    Reset the token usage counter
+    """
+    global _token_counter  # noqa: PLW0603
+    _token_counter = ThreadSafeTokenCounter()
+
+
+def token_usage() -> list[TokenUsage] | None:
+    """
+    Report on token usage in the current session
+
+    Call this function to find out the cumulative number of tokens that you
+    have sent and received in the current session.
+
+    Returns
+    -------
+    list[TokenUsage] | None
+        A list of dictionaries with the following keys: "name", "input", and "output".
+        If no tokens have been logged, then None is returned.
+    """
+    return _token_counter.get_usage()

chatlas/_tokens_old.py

@@ -0,0 +1,148 @@
+# # A duck type for tiktoken.Encoding
+# class TiktokenEncoding(Protocol):
+#     name: str
+#
+#     def encode(
+#         self,
+#         text: str,
+#         *,
+#         allowed_special: Union[Literal["all"], AbstractSet[str]] = set(),  # noqa: B006
+#         disallowed_special: Union[Literal["all"], Collection[str]] = "all",
+#     ) -> list[int]: ...
+#
+#
+# # A duck type for tokenizers.Encoding
+# @runtime_checkable
+# class TokenizersEncoding(Protocol):
+#     @property
+#     def ids(self) -> list[int]: ...
+#
+#
+# # A duck type for tokenizers.Tokenizer
+# class TokenizersTokenizer(Protocol):
+#     def encode(
+#         self,
+#         sequence: Any,
+#         pair: Any = None,
+#         is_pretokenized: bool = False,
+#         add_special_tokens: bool = True,
+#     ) -> TokenizersEncoding: ...
+#
+#
+# TokenEncoding = Union[TiktokenEncoding, TokenizersTokenizer]
+#
+#
+# def get_default_tokenizer() -> TokenizersTokenizer | None:
+#     try:
+#         from tokenizers import Tokenizer
+#
+#         return Tokenizer.from_pretrained("bert-base-cased")  # type: ignore
+#     except Exception:
+#         pass
+#
+#     return None
+
+
+# def _get_token_count(
+#     self,
+#     content: str,
+# ) -> int:
+#     if self._tokenizer is None:
+#         self._tokenizer = get_default_tokenizer()
+#
+#     if self._tokenizer is None:
+#         raise ValueError(
+#             "A tokenizer is required to impose `token_limits` on messages. "
+#             "To get a generic default tokenizer, install the `tokenizers` "
+#             "package (`pip install tokenizers`). "
+#             "To get a more precise token count, provide a specific tokenizer "
+#             "to the `Chat` constructor."
+#         )
+#
+#     encoded = self._tokenizer.encode(content)
+#     if isinstance(encoded, TokenizersEncoding):
+#         return len(encoded.ids)
+#     else:
+#         return len(encoded)
+
+
+# def _trim_messages(
+#         self,
+#         messages: tuple[TransformedMessage, ...],
+#         token_limits: tuple[int, int],
+#     ) -> tuple[TransformedMessage, ...]:
+
+#         n_total, n_reserve = token_limits
+#         if n_total <= n_reserve:
+#             raise ValueError(
+#                 f"Invalid token limits: {token_limits}. The 1st value must be greater "
+#                 "than the 2nd value."
+#             )
+
+#         # Since don't trim system messages, 1st obtain their total token count
+#         # (so we can determine how many non-system messages can fit)
+#         n_system_tokens: int = 0
+#         n_system_messages: int = 0
+#         n_other_messages: int = 0
+#         token_counts: list[int] = []
+#         for m in messages:
+#             content = (
+#                 m.content_server if isinstance(m, TransformedMessage) else m.content
+#             )
+#             count = self._get_token_count(content)
+#             token_counts.append(count)
+#             if m.role == "system":
+#                 n_system_tokens += count
+#                 n_system_messages += 1
+#             else:
+#                 n_other_messages += 1
+
+#         remaining_non_system_tokens = n_total - n_reserve - n_system_tokens
+
+#         if remaining_non_system_tokens <= 0:
+#             raise ValueError(
+#                 f"System messages exceed `.messages(token_limits={token_limits})`. "
+#                 "Consider increasing the 1st value of `token_limit` or setting it to "
+#                 "`token_limit=None` to disable token limits."
+#             )
+
+#         # Now, iterate through the messages in reverse order and appending
+#         # until we run out of tokens
+#         messages2: list[TransformedMessage] = []
+#         n_other_messages2: int = 0
+#         token_counts.reverse()
+#         for i, m in enumerate(reversed(messages)):
+#             if m.role == "system":
+#                 messages2.append(m)
+#                 continue
+#             remaining_non_system_tokens -= token_counts[i]
+#             if remaining_non_system_tokens >= 0:
+#                 messages2.append(m)
+#                 n_other_messages2 += 1
+
+#         messages2.reverse()
+
+#         if len(messages2) == n_system_messages and n_other_messages2 > 0:
+#             raise ValueError(
+#                 f"Only system messages fit within `.messages(token_limits={token_limits})`. "
+#                 "Consider increasing the 1st value of `token_limit` or setting it to "
+#                 "`token_limit=None` to disable token limits."
+#             )
+
+#         return tuple(messages2)
+
+#   def _trim_anthropic_messages(
+#       self,
+#       messages: tuple[TransformedMessage, ...],
+#   ) -> tuple[TransformedMessage, ...]:
+
+#       if any(m.role == "system" for m in messages):
+#           raise ValueError(
+#               "Anthropic requires a system prompt to be specified in it's `.create()` method "
+#               "(not in the chat messages with `role: system`)."
+#           )
+#       for i, m in enumerate(messages):
+#           if m.role == "user":
+#               return messages[i:]
+
+#       return ()

chatlas/_tools.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import inspect
+import warnings
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional
+
+from pydantic import BaseModel, Field, create_model
+
+from . import _utils
+
+__all__ = ("Tool",)
+
+if TYPE_CHECKING:
+    from openai.types.chat import ChatCompletionToolParam
+
+
+class Tool:
+    """
+    Define a tool
+
+    Define a Python function for use by a chatbot. The function will always be
+    invoked in the current Python process.
+
+    Parameters
+    ----------
+    func
+        The function to be invoked when the tool is called.
+    model
+        A Pydantic model that describes the input parameters for the function.
+        If not provided, the model will be inferred from the function's type hints.
+        The primary reason why you might want to provide a model in
+        Note that the name and docstring of the model takes precedence over the
+        name and docstring of the function.
+    """
+
+    func: Callable[..., Any] | Callable[..., Awaitable[Any]]
+
+    def __init__(
+        self,
+        func: Callable[..., Any] | Callable[..., Awaitable[Any]],
+        *,
+        model: Optional[type[BaseModel]] = None,
+    ):
+        self.func = func
+        self._is_async = _utils.is_async_callable(func)
+        self.schema = func_to_schema(func, model)
+        self.name = self.schema["function"]["name"]
+
+
+def func_to_schema(
+    func: Callable[..., Any] | Callable[..., Awaitable[Any]],
+    model: Optional[type[BaseModel]] = None,
+) -> "ChatCompletionToolParam":
+    if model is None:
+        model = func_to_basemodel(func)
+
+    # Throw if there is a mismatch between the model and the function parameters
+    params = inspect.signature(func).parameters
+    fields = model.model_fields
+    diff = set(params) ^ set(fields)
+    if diff:
+        raise ValueError(
+            f"`model` fields must match tool function parameters exactly. "
+            f"Fields found in one but not the other: {diff}"
+        )
+
+    params = basemodel_to_param_schema(model)
+
+    return {
+        "type": "function",
+        "function": {
+            "name": model.__name__ or func.__name__,
+            "description": model.__doc__ or func.__doc__ or "",
+            "parameters": params,
+        },
+    }
+
+
+def func_to_basemodel(func: Callable) -> type[BaseModel]:
+    params = inspect.signature(func).parameters
+    fields = {}
+
+    for name, param in params.items():
+        annotation = param.annotation
+
+        if annotation == inspect.Parameter.empty:
+            warnings.warn(
+                f"Parameter `{name}` of function `{name}` has no type hint. "
+                "Using `Any` as a fallback."
+            )
+            annotation = Any
+
+        if param.default != inspect.Parameter.empty:
+            field = Field(default=param.default)
+        else:
+            field = Field()
+
+        # Add the field to our fields dict
+        fields[name] = (annotation, field)
+
+    return create_model(func.__name__, **fields)
+
+
+def basemodel_to_param_schema(model: type[BaseModel]) -> dict[str, object]:
+    try:
+        import openai
+    except ImportError:
+        raise ImportError(
+            "The openai package is required for this functionality. "
+            "Please install it with `pip install openai`."
+        )
+
+    # Lean on openai's ability to translate BaseModel.model_json_schema()
+    # to a valid tool schema (this wouldn't be impossible to do ourselves,
+    # but it's fair amount of logic to substitute `$refs`, etc.)
+    tool = openai.pydantic_function_tool(model)
+
+    fn = tool["function"]
+    if "parameters" not in fn:
+        raise ValueError("Expected `parameters` in function definition.")
+
+    params = fn["parameters"]
+
+    # For some reason, openai (or pydantic?) wants to include a title
+    # at the model and field level. I don't think we actually need or want this.
+    if "title" in params:
+        del params["title"]
+
+    if "properties" in params and isinstance(params["properties"], dict):
+        for prop in params["properties"].values():
+            if "title" in prop:
+                del prop["title"]
+
+    return params