chatlas 0.9.1__py3-none-any.whl → 0.10.0__py3-none-any.whl

chatlas/_provider_portkey.py

@@ -0,0 +1,123 @@
+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 drop_none
+
+if TYPE_CHECKING:
+    from ._provider_openai import ChatCompletion
+    from .types.openai import ChatClientArgs, SubmitInputArgs
+
+
+def ChatPortkey(
+    *,
+    system_prompt: Optional[str] = None,
+    model: Optional[str] = None,
+    api_key: Optional[str] = None,
+    virtual_key: Optional[str] = None,
+    base_url: str = "https://api.portkey.ai/v1",
+    kwargs: Optional["ChatClientArgs"] = None,
+) -> Chat["SubmitInputArgs", ChatCompletion]:
+    """
+    Chat with a model hosted on PortkeyAI
+
+    [PortkeyAI](https://portkey.ai/docs/product/ai-gateway/universal-api)
+    provides an interface (AI Gateway) to connect through its Universal API to a
+    variety of LLMs providers with a single endpoint.
+
+    Prerequisites
+    -------------
+
+    ::: {.callout-note}
+    ## Portkey credentials
+
+    Follow the instructions at <https://portkey.ai/docs/introduction/make-your-first-request>
+    to get started making requests to PortkeyAI. You will need to set the
+    `PORTKEY_API_KEY` environment variable to your Portkey API key, and optionally
+    the `PORTKEY_VIRTUAL_KEY` environment variable to your virtual key.
+    :::
+
+    Examples
+    --------
+    ```python
+    import os
+    from chatlas import ChatPortkey
+
+    chat = ChatPortkey(api_key=os.getenv("PORTKEY_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 `PORTKEY_API_KEY` environment variable.
+    virtual_key
+        An (optional) virtual identifier, storing the LLM provider's API key. See
+        [documentation](https://portkey.ai/docs/product/ai-gateway/virtual-keys).
+        You generally should not supply this directly, but instead set the
+        `PORTKEY_VIRTUAL_KEY` environment variable.
+    base_url
+        The base URL for the Portkey API. The default is suitable for most users.
+    kwargs
+        Additional arguments to pass to the OpenAIProvider, such as headers or
+        other client configuration options.
+
+    Returns
+    -------
+    Chat
+        A chat object that retains the state of the conversation.
+
+    Notes
+    -----
+    This function is a lightweight wrapper around [](`~chatlas.ChatOpenAI`) with
+    the defaults tweaked for PortkeyAI.
+
+    """
+    if model is None:
+        model = log_model_default("gpt-4.1")
+    if api_key is None:
+        api_key = os.getenv("PORTKEY_API_KEY")
+
+    kwargs2 = add_default_headers(
+        kwargs or {},
+        api_key=api_key,
+        virtual_key=virtual_key,
+    )
+
+    return Chat(
+        provider=OpenAIProvider(
+            api_key=api_key,
+            model=model,
+            base_url=base_url,
+            name="Portkey",
+            kwargs=kwargs2,
+        ),
+        system_prompt=system_prompt,
+    )
+
+
+def add_default_headers(
+    kwargs: "ChatClientArgs",
+    api_key: Optional[str] = None,
+    virtual_key: Optional[str] = None,
+) -> "ChatClientArgs":
+    headers = kwargs.get("default_headers", None)
+    default_headers = drop_none(
+        {
+            "x-portkey-api-key": api_key,
+            "x-portkey-virtual-key": virtual_key,
+            **(headers or {}),
+        }
+    )
+    return {"default_headers": default_headers, **kwargs}

chatlas/{_snowflake.py → _provider_snowflake.py}

@@ -537,12 +537,12 @@ class SnowflakeProvider(
                         arguments=params,
                     )
                 )
-
+        # Snowflake does not currently appear to support caching, so we set cached tokens to 0
         usage = completion.usage
         if usage is None:
-            tokens = (0, 0)
+            tokens = (0, 0, 0)
         else:
-            tokens = (usage.prompt_tokens or 0, usage.completion_tokens or 0)
+            tokens = (usage.prompt_tokens or 0, usage.completion_tokens or 0, 0)
 
         tokens_log(self, tokens)
 

chatlas/_tokens.py

@@ -8,7 +8,7 @@ from typing import TYPE_CHECKING
 import orjson
 
 from ._logging import logger
-from ._typing_extensions import TypedDict
+from ._typing_extensions import NotRequired, TypedDict
 
 if TYPE_CHECKING:
     from ._provider import Provider
@@ -23,6 +23,7 @@ class TokenUsage(TypedDict):
     model: str
     input: int
     output: int
+    cached_input: int
     cost: float | None
 
 
@@ -32,11 +33,16 @@ class ThreadSafeTokenCounter:
         self._tokens: dict[str, TokenUsage] = {}
 
     def log_tokens(
-        self, name: str, model: str, input_tokens: int, output_tokens: int
+        self,
+        name: str,
+        model: str,
+        input_tokens: int,
+        output_tokens: int,
+        cached_tokens: int,
     ) -> None:
         logger.info(
             f"Provider '{name}' generated a response of {output_tokens} tokens "
-            f"from an input of {input_tokens} tokens."
+            f"from an input of {input_tokens} tokens and {cached_tokens} cached input tokens."
         )
 
         with self._lock:
@@ -46,12 +52,18 @@ class ThreadSafeTokenCounter:
                     "model": model,
                     "input": input_tokens,
                     "output": output_tokens,
-                    "cost": compute_cost(name, model, input_tokens, output_tokens),
+                    "cached_input": cached_tokens,
+                    "cost": compute_cost(
+                        name, model, input_tokens, output_tokens, cached_tokens
+                    ),
                 }
             else:
                 self._tokens[name]["input"] += input_tokens
                 self._tokens[name]["output"] += output_tokens
-                price = compute_cost(name, model, input_tokens, output_tokens)
+                self._tokens[name]["cached_input"] += cached_tokens
+                price = compute_cost(
+                    name, model, input_tokens, output_tokens, cached_tokens
+                )
                 if price is not None:
                     cost = self._tokens[name]["cost"]
                     if cost is None:
@@ -71,11 +83,13 @@ class ThreadSafeTokenCounter:
 _token_counter = ThreadSafeTokenCounter()
 
 
-def tokens_log(provider: "Provider", tokens: tuple[int, int]) -> None:
+def tokens_log(provider: "Provider", tokens: tuple[int, int, int]) -> None:
     """
     Log token usage for a provider in a thread-safe manner.
     """
-    _token_counter.log_tokens(provider.name, provider.model, tokens[0], tokens[1])
+    _token_counter.log_tokens(
+        provider.name, provider.model, tokens[0], tokens[1], tokens[2]
+    )
 
 
 def tokens_reset() -> None:
@@ -95,11 +109,11 @@ class TokenPrice(TypedDict):
     """The provider name (e.g., "OpenAI", "Anthropic", etc.)"""
     model: str
     """The model name (e.g., "gpt-3.5-turbo", "claude-2", etc.)"""
-    cached_input: float
+    cached_input: NotRequired[float]
     """The cost per user token in USD per million tokens for cached input"""
     input: float
     """The cost per user token in USD per million tokens"""
-    output: float
+    output: NotRequired[float]
     """The cost per assistant token in USD per million tokens"""
 
 
@@ -132,7 +146,7 @@ def get_token_pricing(name: str, model: str) -> TokenPrice | None:
 
 
 def compute_cost(
-    name: str, model: str, input_tokens: int, output_tokens: int
+    name: str, model: str, input_tokens: int, output_tokens: int, cached_tokens: int = 0
 ) -> float | None:
     """
     Compute the cost of a turn.
@@ -146,8 +160,9 @@ def compute_cost(
     if price is None:
         return None
     input_price = input_tokens * (price["input"] / 1e6)
-    output_price = output_tokens * (price["output"] / 1e6)
-    return input_price + output_price
+    output_price = output_tokens * (price.get("output", 0) / 1e6)
+    cached_price = cached_tokens * (price.get("cached_input", 0) / 1e6)
+    return input_price + output_price + cached_price
 
 
 def token_usage() -> list[TokenUsage] | None:

chatlas/_turn.py

@@ -55,7 +55,7 @@ class Turn(BaseModel, Generic[CompletionT]):
     contents
         A list of [](`~chatlas.types.Content`) objects.
     tokens
-        A numeric vector of length 2 representing the number of input and output
+        A numeric vector of length 3 representing the number of input, output, and cached
         tokens (respectively) used in this turn. Currently only recorded for
         assistant turns.
     finish_reason
@@ -69,7 +69,7 @@ class Turn(BaseModel, Generic[CompletionT]):
 
     role: Literal["user", "assistant", "system"]
     contents: list[ContentUnion] = Field(default_factory=list)
-    tokens: Optional[tuple[int, int]] = None
+    tokens: Optional[tuple[int, int, int]] = None
     finish_reason: Optional[str] = None
     completion: Optional[CompletionT] = Field(default=None, exclude=True)
 
@@ -80,7 +80,7 @@ class Turn(BaseModel, Generic[CompletionT]):
         role: Literal["user", "assistant", "system"],
         contents: str | Sequence[Content | str],
         *,
-        tokens: Optional[tuple[int, int]] = None,
+        tokens: Optional[tuple[int, int, int]] = None,
         finish_reason: Optional[str] = None,
         completion: Optional[CompletionT] = None,
         **kwargs,
@@ -134,4 +134,3 @@ def user_turn(*args: Content | str) -> Turn:
         raise ValueError("Must supply at least one input.")
 
     return Turn("user", args)
-

chatlas/_typing_extensions.py

@@ -14,13 +14,13 @@ else:
 # they should both come from the same typing module.
 # https://peps.python.org/pep-0655/#usage-in-python-3-11
 if sys.version_info >= (3, 11):
-    from typing import Required, TypedDict
+    from typing import NotRequired, Required, TypedDict
 else:
-    from typing_extensions import Required, TypedDict
+    from typing_extensions import NotRequired, Required, TypedDict
 
 
 # The only purpose of the following line is so that pyright will put all of the
 # conditional imports into the .pyi file when generating type stubs. Without this line,
 # pyright will not include the above imports in the generated .pyi file, and it will
 # result in a lot of red squiggles in user code.
-_: "ParamSpec | TypeGuard | is_typeddict | Required | TypedDict"  # type: ignore
+_: "ParamSpec | TypeGuard | is_typeddict | NotRequired | Required | TypedDict"  # type: ignore

chatlas/_version.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     from typing import Union
 
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.1'
-__version_tuple__ = version_tuple = (0, 9, 1)
+__version__ = version = '0.10.0'
+__version_tuple__ = version_tuple = (0, 10, 0)
+
+__commit_id__ = commit_id = None