chore: use primitives instead of typing imports and fixes completion … (#149)

Signed-off-by: Adrian Cole <adrian.cole@elastic.co>

packages/exchange/src/exchange/checkpoint.py

@@ -1,5 +1,4 @@
 from copy import deepcopy
-from typing import List
 from attrs import define, field
 
 
@@ -31,7 +30,7 @@ class CheckpointData:
     total_token_count: int = field(default=0)
 
     # in order list of individual checkpoints in the exchange
-    checkpoints: List[Checkpoint] = field(factory=list)
+    checkpoints: list[Checkpoint] = field(factory=list)
 
     # the offset to apply to the message index when calculating the last message index
     # this is useful because messages on the exchange behave like a queue, where you can only

packages/exchange/src/exchange/content.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Optional
 
 from attrs import define, asdict
 
@@ -7,11 +7,11 @@ CONTENT_TYPES = {}
 
 
 class Content:
-    def __init_subclass__(cls, **kwargs: Dict[str, Any]) -> None:
+    def __init_subclass__(cls, **kwargs: dict[str, any]) -> None:
         super().__init_subclass__(**kwargs)
         CONTENT_TYPES[cls.__name__] = cls
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, any]:
         data = asdict(self, recurse=True)
         data["type"] = self.__class__.__name__
         return data
@@ -26,7 +26,7 @@ class Text(Content):
 class ToolUse(Content):
     id: str
     name: str
-    parameters: Any
+    parameters: any
     is_error: bool = False
     error_message: Optional[str] = None
 

packages/exchange/src/exchange/exchange.py

@@ -1,8 +1,7 @@
 import json
 import traceback
 from copy import deepcopy
-from typing import Any, Dict, List, Mapping, Tuple
-
+from typing import Mapping
 from attrs import define, evolve, field, Factory
 from tiktoken import get_encoding
 
@@ -41,8 +40,8 @@ class Exchange:
     model: str
     system: str
     moderator: Moderator = field(default=ContextTruncate())
-    tools: Tuple[Tool] = field(factory=tuple, converter=tuple)
-    messages: List[Message] = field(factory=list)
+    tools: tuple[Tool, ...] = field(factory=tuple, converter=tuple)
+    messages: list[Message] = field(factory=list)
     checkpoint_data: CheckpointData = field(factory=CheckpointData)
     generation_args: dict = field(default=Factory(dict))
 
@@ -50,7 +49,7 @@ class Exchange:
     def _toolmap(self) -> Mapping[str, Tool]:
         return {tool.name: tool for tool in self.tools}
 
-    def replace(self, **kwargs: Dict[str, Any]) -> "Exchange":
+    def replace(self, **kwargs: dict[str, any]) -> "Exchange":
         """Make a copy of the exchange, replacing any passed arguments"""
         # TODO: ensure that the checkpoint data is updated correctly. aka,
         # if we replace the messages, we need to update the checkpoint data
@@ -264,7 +263,7 @@ class Exchange:
             # we've removed all the checkpoints, so we need to reset the message index offset
             self.checkpoint_data.message_index_offset = 0
 
-    def pop_last_checkpoint(self) -> Tuple[Checkpoint, List[Message]]:
+    def pop_last_checkpoint(self) -> tuple[Checkpoint, list[Message]]:
         """
         Reverts the exchange back to the last checkpoint, removing associated messages
         """
@@ -275,7 +274,7 @@ class Exchange:
             messages.append(self.messages.pop())
         return removed_checkpoint, messages
 
-    def pop_first_checkpoint(self) -> Tuple[Checkpoint, List[Message]]:
+    def pop_first_checkpoint(self) -> tuple[Checkpoint, list[Message]]:
         """
         Pop the first checkpoint from the exchange, removing associated messages
         """
@@ -332,5 +331,6 @@ class Exchange:
         # this to be a required method of the provider instead.
         return len(self.messages) > 0 and self.messages[-1].role == "user"
 
-    def get_token_usage(self) -> Dict[str, Usage]:
+    @staticmethod
+    def get_token_usage() -> dict[str, Usage]:
         return _token_usage_collector.get_token_usage_group_by_model()

packages/exchange/src/exchange/invalid_choice_error.py

@@ -1,8 +1,5 @@
-from typing import List
-
-
 class InvalidChoiceError(Exception):
-    def __init__(self, attribute_name: str, attribute_value: str, available_values: List[str]) -> None:
+    def __init__(self, attribute_name: str, attribute_value: str, available_values: list[str]) -> None:
         self.attribute_name = attribute_name
         self.attribute_value = attribute_value
         self.available_values = available_values

packages/exchange/src/exchange/message.py

@@ -1,7 +1,7 @@
 import inspect
 import time
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Type
+from typing import Literal
 
 from attrs import define, field
 from jinja2 import Environment, FileSystemLoader
@@ -12,7 +12,7 @@ from exchange.utils import create_object_id
 Role = Literal["user", "assistant"]
 
 
-def validate_role_and_content(instance: "Message", *_: Any) -> None:  # noqa: ANN401
+def validate_role_and_content(instance: "Message", *_: any) -> None:  # noqa: ANN401
     if instance.role == "user":
         if not (instance.text or instance.tool_result):
             raise ValueError("User message must include a Text or ToolResult")
@@ -25,7 +25,7 @@ def validate_role_and_content(instance: "Message", *_: Any) -> None:  # noqa: AN
             raise ValueError("Assistant message does not support ToolResult")
 
 
-def content_converter(contents: List[Dict[str, Any]]) -> List[Content]:
+def content_converter(contents: list[dict[str, any]]) -> list[Content]:
     return [(CONTENT_TYPES[c.pop("type")](**c) if c.__class__ not in CONTENT_TYPES.values() else c) for c in contents]
 
 
@@ -48,9 +48,9 @@ class Message:
     role: Role = field(default="user")
     id: str = field(factory=lambda: str(create_object_id(prefix="msg")))
     created: int = field(factory=lambda: int(time.time()))
-    content: List[Content] = field(factory=list, validator=validate_role_and_content, converter=content_converter)
+    content: list[Content] = field(factory=list, validator=validate_role_and_content, converter=content_converter)
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, any]:
         return {
             "role": self.role,
             "id": self.id,
@@ -68,7 +68,7 @@ class Message:
         return "\n".join(result)
 
     @property
-    def tool_use(self) -> List[ToolUse]:
+    def tool_use(self) -> list[ToolUse]:
         """All tool use content of this message."""
         result = []
         for content in self.content:
@@ -77,7 +77,7 @@ class Message:
         return result
 
     @property
-    def tool_result(self) -> List[ToolResult]:
+    def tool_result(self) -> list[ToolResult]:
         """All tool result content of this message."""
         result = []
         for content in self.content:
@@ -87,10 +87,10 @@ class Message:
 
     @classmethod
     def load(
-        cls: Type["Message"],
+        cls: type["Message"],
         filename: str,
         role: Role = "user",
-        **kwargs: Dict[str, Any],
+        **kwargs: dict[str, any],
     ) -> "Message":
         """Load the message from filename relative to where the load is called.
 
@@ -113,9 +113,9 @@ class Message:
         return cls(role=role, content=[Text(text=rendered_content)])
 
     @classmethod
-    def user(cls: Type["Message"], text: str) -> "Message":
+    def user(cls: type["Message"], text: str) -> "Message":
         return cls(role="user", content=[Text(text)])
 
     @classmethod
-    def assistant(cls: Type["Message"], text: str) -> "Message":
+    def assistant(cls: type["Message"], text: str) -> "Message":
         return cls(role="assistant", content=[Text(text)])

packages/exchange/src/exchange/moderators/__init__.py

@@ -1,5 +1,4 @@
 from functools import cache
-from typing import Type
 
 from exchange.invalid_choice_error import InvalidChoiceError
 from exchange.moderators.base import Moderator
@@ -10,7 +9,7 @@ from exchange.moderators.summarizer import ContextSummarizer  # noqa
 
 
 @cache
-def get_moderator(name: str) -> Type[Moderator]:
+def get_moderator(name: str) -> type[Moderator]:
     moderators = load_plugins(group="exchange.moderator")
     if name not in moderators:
         raise InvalidChoiceError("moderator", name, moderators.keys())

packages/exchange/src/exchange/moderators/base.py

@@ -1,8 +1,7 @@
 from abc import ABC, abstractmethod
-from typing import Type
 
 
 class Moderator(ABC):
     @abstractmethod
-    def rewrite(self, exchange: Type["exchange.exchange.Exchange"]) -> None:  # noqa: F821
+    def rewrite(self, exchange: type["exchange.exchange.Exchange"]) -> None:  # noqa: F821
         pass

packages/exchange/src/exchange/moderators/passive.py

@@ -1,7 +1,6 @@
-from typing import Type
 from exchange.moderators.base import Moderator
 
 
 class PassiveModerator(Moderator):
-    def rewrite(self, _: Type["exchange.exchange.Exchange"]) -> None:  # noqa: F821
+    def rewrite(self, _: type["exchange.exchange.Exchange"]) -> None:  # noqa: F821
         pass

packages/exchange/src/exchange/moderators/summarizer.py

@@ -1,12 +1,10 @@
-from typing import Type
-
 from exchange import Message
 from exchange.checkpoint import CheckpointData
 from exchange.moderators import ContextTruncate, PassiveModerator
 
 
 class ContextSummarizer(ContextTruncate):
-    def rewrite(self, exchange: Type["exchange.exchange.Exchange"]) -> None:  # noqa: F821
+    def rewrite(self, exchange: type["exchange.exchange.Exchange"]) -> None:  # noqa: F821
         """Summarize the context history up to the last few messages in the exchange"""
 
         self._update_system_prompt_token_count(exchange)

packages/exchange/src/exchange/moderators/truncate.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Optional
 
 from exchange.checkpoint import CheckpointData
 from exchange.message import Message
@@ -62,7 +62,7 @@ class ContextTruncate(Moderator):
             exchange.checkpoint_data.total_token_count -= last_system_prompt_token_count
             exchange.checkpoint_data.total_token_count += self.system_prompt_token_count
 
-    def _get_messages_to_remove(self, exchange: Exchange) -> List[Message]:
+    def _get_messages_to_remove(self, exchange: Exchange) -> list[Message]:
         # this keeps all the messages/checkpoints
         throwaway_exchange = exchange.replace(
             moderator=PassiveModerator(),

packages/exchange/src/exchange/providers/__init__.py

@@ -1,5 +1,4 @@
 from functools import cache
-from typing import Type
 
 from exchange.invalid_choice_error import InvalidChoiceError
 from exchange.providers.anthropic import AnthropicProvider  # noqa
@@ -15,7 +14,7 @@ from exchange.utils import load_plugins
 
 
 @cache
-def get_provider(name: str) -> Type[Provider]:
+def get_provider(name: str) -> type[Provider]:
     providers = load_plugins(group="exchange.provider")
     if name not in providers:
         raise InvalidChoiceError("provider", name, providers.keys())

packages/exchange/src/exchange/providers/anthropic.py

@@ -1,5 +1,4 @@
 import os
-from typing import Any, Dict, List, Tuple, Type
 
 import httpx
 
@@ -29,7 +28,7 @@ class AnthropicProvider(Provider):
         self.client = client
 
     @classmethod
-    def from_env(cls: Type["AnthropicProvider"]) -> "AnthropicProvider":
+    def from_env(cls: type["AnthropicProvider"]) -> "AnthropicProvider":
         cls.check_env_vars()
         url = os.environ.get("ANTHROPIC_HOST", ANTHROPIC_HOST)
         key = os.environ.get("ANTHROPIC_API_KEY")
@@ -45,7 +44,7 @@ class AnthropicProvider(Provider):
         return cls(client)
 
     @staticmethod
-    def get_usage(data: Dict) -> Usage:  # noqa: ANN401
+    def get_usage(data: dict) -> Usage:  # noqa: ANN401
         usage = data.get("usage")
         input_tokens = usage.get("input_tokens")
         output_tokens = usage.get("output_tokens")
@@ -61,7 +60,7 @@ class AnthropicProvider(Provider):
         )
 
     @staticmethod
-    def anthropic_response_to_message(response: Dict) -> Message:
+    def anthropic_response_to_message(response: dict) -> Message:
         content_blocks = response.get("content", [])
         content = []
         for block in content_blocks:
@@ -78,7 +77,7 @@ class AnthropicProvider(Provider):
         return Message(role="assistant", content=content)
 
     @staticmethod
-    def tools_to_anthropic_spec(tools: Tuple[Tool]) -> List[Dict[str, Any]]:
+    def tools_to_anthropic_spec(tools: tuple[Tool, ...]) -> list[dict[str, any]]:
         return [
             {
                 "name": tool.name,
@@ -89,7 +88,7 @@ class AnthropicProvider(Provider):
         ]
 
     @staticmethod
-    def messages_to_anthropic_spec(messages: List[Message]) -> List[Dict[str, Any]]:
+    def messages_to_anthropic_spec(messages: list[Message]) -> list[dict[str, any]]:
         messages_spec = []
         # if messages is empty - just make a default
         for message in messages:
@@ -127,10 +126,12 @@ class AnthropicProvider(Provider):
         self,
         model: str,
         system: str,
-        messages: List[Message],
-        tools: List[Tool] = [],
-        **kwargs: Dict[str, Any],
-    ) -> Tuple[Message, Usage]:
+        messages: list[Message],
+        tools: list[Tool] = None,
+        **kwargs: dict[str, any],
+    ) -> tuple[Message, Usage]:
+        if tools is None:
+            tools = []
         tools_set = set()
         unique_tools = []
         for tool in tools:

packages/exchange/src/exchange/providers/azure.py

@@ -1,5 +1,3 @@
-from typing import Type
-
 import httpx
 import os
 
@@ -21,7 +19,7 @@ class AzureProvider(OpenAiProvider):
         super().__init__(client)
 
     @classmethod
-    def from_env(cls: Type["AzureProvider"]) -> "AzureProvider":
+    def from_env(cls: type["AzureProvider"]) -> "AzureProvider":
         cls.check_env_vars()
         url = os.environ.get("AZURE_CHAT_COMPLETIONS_HOST_NAME")
         deployment_name = os.environ.get("AZURE_CHAT_COMPLETIONS_DEPLOYMENT_NAME")

packages/exchange/src/exchange/providers/base.py

@@ -1,7 +1,7 @@
 import os
 from abc import ABC, abstractmethod
 from attrs import define, field
-from typing import List, Optional, Tuple, Type
+from typing import Optional
 
 from exchange.message import Message
 from exchange.tool import Tool
@@ -19,11 +19,11 @@ class Provider(ABC):
     REQUIRED_ENV_VARS: list[str] = []
 
     @classmethod
-    def from_env(cls: Type["Provider"]) -> "Provider":
+    def from_env(cls: type["Provider"]) -> "Provider":
         return cls()
 
     @classmethod
-    def check_env_vars(cls: Type["Provider"], instructions_url: Optional[str] = None) -> None:
+    def check_env_vars(cls: type["Provider"], instructions_url: Optional[str] = None) -> None:
         for env_var in cls.REQUIRED_ENV_VARS:
             if env_var not in os.environ:
                 raise MissingProviderEnvVariableError(env_var, cls.PROVIDER_NAME, instructions_url)
@@ -33,9 +33,10 @@ class Provider(ABC):
         self,
         model: str,
         system: str,
-        messages: List[Message],
-        tools: Tuple[Tool],
-    ) -> Tuple[Message, Usage]:
+        messages: list[Message],
+        tools: tuple[Tool, ...],
+        **kwargs: dict[str, any],
+    ) -> tuple[Message, Usage]:
         """Generate the next message using the specified model"""
         pass
 

packages/exchange/src/exchange/providers/bedrock.py

@@ -4,7 +4,7 @@ import json
 import logging
 import os
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Tuple, Type
+from typing import Optional
 from urllib.parse import quote, urlparse
 
 import httpx
@@ -36,7 +36,7 @@ class AwsClient(httpx.Client):
         aws_access_key: str,
         aws_secret_key: str,
         aws_session_token: Optional[str] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: dict[str, any],
     ) -> None:
         self.region = aws_region
         self.host = f"https://{SERVICE}.{aws_region}.amazonaws.com/"
@@ -45,7 +45,7 @@ class AwsClient(httpx.Client):
         self.session_token = aws_session_token
         super().__init__(base_url=self.host, timeout=600, **kwargs)
 
-    def post(self, path: str, json: Dict, **kwargs: Dict[str, Any]) -> httpx.Response:
+    def post(self, path: str, json: dict, **kwargs: dict[str, any]) -> httpx.Response:
         signed_headers = self.sign_and_get_headers(
             method="POST",
             url=path,
@@ -60,7 +60,7 @@ class AwsClient(httpx.Client):
         url: str,
         payload: dict,
         service: str,
-    ) -> Dict[str, str]:
+    ) -> dict[str, str]:
         """
         Sign the request and generate the necessary headers for AWS authentication.
 
@@ -72,10 +72,10 @@ class AwsClient(httpx.Client):
             region (str): The AWS region.
             access_key (str): The AWS access key.
             secret_key (str): The AWS secret key.
-            session_token (Optional[str]): The AWS session token, if any.
+            session_token (optional[str]): The AWS session token, if any.
 
         Returns:
-            Dict[str, str]: The headers required for the request.
+            dict[str, str]: The headers required for the request.
         """
 
         def sign(key: bytes, msg: str) -> bytes:
@@ -160,7 +160,7 @@ class BedrockProvider(Provider):
         self.client = client
 
     @classmethod
-    def from_env(cls: Type["BedrockProvider"]) -> "BedrockProvider":
+    def from_env(cls: type["BedrockProvider"]) -> "BedrockProvider":
         cls.check_env_vars()
         aws_region = os.environ.get("AWS_REGION", "us-east-1")
         aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID")
@@ -179,22 +179,22 @@ class BedrockProvider(Provider):
         self,
         model: str,
         system: str,
-        messages: List[Message],
-        tools: Tuple[Tool],
-        **kwargs: Dict[str, Any],
-    ) -> Tuple[Message, Usage]:
+        messages: list[Message],
+        tools: tuple[Tool, ...],
+        **kwargs: dict[str, any],
+    ) -> tuple[Message, Usage]:
         """
         Generate a completion response from the Bedrock gateway.
 
         Args:
             model (str): The model identifier.
             system (str): The system prompt or configuration.
-            messages (List[Message]): A list of messages to be processed by the model.
-            tools (Tuple[Tool]): A tuple of tools to be used in the completion process.
+            messages (list[Message]): A list of messages to be processed by the model.
+            tools (tuple[Tool]): A tuple of tools to be used in the completion process.
             **kwargs: Additional keyword arguments for inference configuration.
 
         Returns:
-            Tuple[Message, Usage]: A tuple containing the response message and usage data.
+            tuple[Message, Usage]: A tuple containing the response message and usage data.
         """
 
         inference_config = dict(
@@ -231,7 +231,7 @@ class BedrockProvider(Provider):
         return self.response_to_message(response_message), usage
 
     @retry_procedure
-    def _post(self, payload: Any, path: str) -> dict:  # noqa: ANN401
+    def _post(self, payload: any, path: str) -> dict:  # noqa: ANN401
         response = self.client.post(path, json=payload)
         return raise_for_status(response).json()
 
@@ -311,7 +311,7 @@ class BedrockProvider(Provider):
         raise Exception("Invalid response")
 
     @staticmethod
-    def tools_to_bedrock_spec(tools: Tuple[Tool]) -> Optional[dict]:
+    def tools_to_bedrock_spec(tools: tuple[Tool, ...]) -> Optional[dict]:
         if len(tools) == 0:
             return None  # API requires a non-empty tool config or None
         tools_added = set()

packages/exchange/src/exchange/providers/databricks.py

@@ -1,5 +1,3 @@
-from typing import Any, Dict, List, Tuple, Type
-
 import httpx
 import os
 
@@ -43,7 +41,7 @@ class DatabricksProvider(Provider):
         self.client = client
 
     @classmethod
-    def from_env(cls: Type["DatabricksProvider"]) -> "DatabricksProvider":
+    def from_env(cls: type["DatabricksProvider"]) -> "DatabricksProvider":
         cls.check_env_vars(cls.instructions_url)
         url = os.environ.get("DATABRICKS_HOST")
         key = os.environ.get("DATABRICKS_TOKEN")
@@ -73,10 +71,10 @@ class DatabricksProvider(Provider):
         self,
         model: str,
         system: str,
-        messages: List[Message],
-        tools: Tuple[Tool],
-        **kwargs: Dict[str, Any],
-    ) -> Tuple[Message, Usage]:
+        messages: list[Message],
+        tools: tuple[Tool, ...],
+        **kwargs: dict[str, any],
+    ) -> tuple[Message, Usage]:
         payload = dict(
             messages=[
                 {"role": "system", "content": system},

packages/exchange/src/exchange/providers/google.py

@@ -1,5 +1,4 @@
 import os
-from typing import Any, Dict, List, Tuple, Type
 
 import httpx
 
@@ -30,7 +29,7 @@ class GoogleProvider(Provider):
         self.client = client
 
     @classmethod
-    def from_env(cls: Type["GoogleProvider"]) -> "GoogleProvider":
+    def from_env(cls: type["GoogleProvider"]) -> "GoogleProvider":
         cls.check_env_vars(cls.instructions_url)
         url = os.environ.get("GOOGLE_HOST", GOOGLE_HOST)
         key = os.environ.get("GOOGLE_API_KEY")
@@ -45,7 +44,7 @@ class GoogleProvider(Provider):
         return cls(client)
 
     @staticmethod
-    def get_usage(data: Dict) -> Usage:  # noqa: ANN401
+    def get_usage(data: dict) -> Usage:  # noqa: ANN401
         usage = data.get("usageMetadata")
         input_tokens = usage.get("promptTokenCount")
         output_tokens = usage.get("candidatesTokenCount")
@@ -61,7 +60,7 @@ class GoogleProvider(Provider):
         )
 
     @staticmethod
-    def google_response_to_message(response: Dict) -> Message:
+    def google_response_to_message(response: dict) -> Message:
         candidates = response.get("candidates", [])
         if candidates:
             # Only use first candidate for now
@@ -85,12 +84,12 @@ class GoogleProvider(Provider):
         return Message(role="assistant", content=[])
 
     @staticmethod
-    def tools_to_google_spec(tools: Tuple[Tool]) -> Dict[str, List[Dict[str, Any]]]:
+    def tools_to_google_spec(tools: tuple[Tool, ...]) -> dict[str, list[dict[str, any]]]:
         if not tools:
             return {}
         converted_tools = []
         for tool in tools:
-            converted_tool: Dict[str, Any] = {
+            converted_tool: dict[str, any] = {
                 "name": tool.name,
                 "description": tool.description or "",
             }
@@ -100,7 +99,7 @@ class GoogleProvider(Provider):
         return {"functionDeclarations": converted_tools}
 
     @staticmethod
-    def messages_to_google_spec(messages: List[Message]) -> List[Dict[str, Any]]:
+    def messages_to_google_spec(messages: list[Message]) -> list[dict[str, any]]:
         messages_spec = []
         for message in messages:
             role = "user" if message.role == "user" else "model"
@@ -136,10 +135,10 @@ class GoogleProvider(Provider):
         self,
         model: str,
         system: str,
-        messages: List[Message],
-        tools: List[Tool] = [],
-        **kwargs: Dict[str, Any],
-    ) -> Tuple[Message, Usage]:
+        messages: list[Message],
+        tools: list[Tool] = None,
+        **kwargs: dict[str, any],
+    ) -> tuple[Message, Usage]:
         tools_set = set()
         unique_tools = []
         for tool in tools:

packages/exchange/src/exchange/providers/groq.py

@@ -1,5 +1,4 @@
 import os
-from typing import Any, Dict, List, Tuple, Type
 
 import httpx
 
@@ -37,7 +36,7 @@ class GroqProvider(Provider):
         self.client = client
 
     @classmethod
-    def from_env(cls: Type["GroqProvider"]) -> "GroqProvider":
+    def from_env(cls: type["GroqProvider"]) -> "GroqProvider":
         cls.check_env_vars(cls.instructions_url)
         url = os.environ.get("GROQ_HOST", GROQ_HOST)
         key = os.environ.get("GROQ_API_KEY")
@@ -69,10 +68,10 @@ class GroqProvider(Provider):
         self,
         model: str,
         system: str,
-        messages: List[Message],
-        tools: Tuple[Tool],
-        **kwargs: Dict[str, Any],
-    ) -> Tuple[Message, Usage]:
+        messages: list[Message],
+        tools: tuple[Tool, ...],
+        **kwargs: dict[str, any],
+    ) -> tuple[Message, Usage]:
         system_message = [{"role": "system", "content": system}]
         payload = dict(
             messages=system_message + messages_to_openai_spec(messages),

packages/exchange/src/exchange/providers/ollama.py

@@ -1,5 +1,4 @@
 import os
-from typing import Type
 
 import httpx
 
@@ -31,7 +30,7 @@ ollama:
         super().__init__(client)
 
     @classmethod
-    def from_env(cls: Type["OllamaProvider"]) -> "OllamaProvider":
+    def from_env(cls: type["OllamaProvider"]) -> "OllamaProvider":
         ollama_url = os.environ.get("OLLAMA_HOST", OLLAMA_HOST)
         timeout = httpx.Timeout(60 * 10)
 

packages/exchange/src/exchange/providers/openai.py

@@ -1,5 +1,4 @@
 import os
-from typing import Any, Dict, List, Tuple, Type
 
 import httpx
 
@@ -37,7 +36,7 @@ class OpenAiProvider(Provider):
         self.client = client
 
     @classmethod
-    def from_env(cls: Type["OpenAiProvider"]) -> "OpenAiProvider":
+    def from_env(cls: type["OpenAiProvider"]) -> "OpenAiProvider":
         cls.check_env_vars(cls.instructions_url)
         url = os.environ.get("OPENAI_HOST", OPENAI_HOST)
         key = os.environ.get("OPENAI_API_KEY")
@@ -69,10 +68,10 @@ class OpenAiProvider(Provider):
         self,
         model: str,
         system: str,
-        messages: List[Message],
-        tools: Tuple[Tool],
-        **kwargs: Dict[str, Any],
-    ) -> Tuple[Message, Usage]:
+        messages: list[Message],
+        tools: tuple[Tool, ...],
+        **kwargs: dict[str, any],
+    ) -> tuple[Message, Usage]:
         system_message = [] if model.startswith("o1") else [{"role": "system", "content": system}]
         payload = dict(
             messages=system_message + messages_to_openai_spec(messages),

packages/exchange/src/exchange/providers/utils.py

@@ -1,7 +1,7 @@
 import base64
 import json
 import re
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Optional
 
 import httpx
 from exchange.content import Text, ToolResult, ToolUse
@@ -10,10 +10,10 @@ from exchange.tool import Tool
 from tenacity import retry_if_exception
 
 
-def retry_if_status(codes: Optional[List[int]] = None, above: Optional[int] = None) -> Callable:
+def retry_if_status(codes: Optional[list[int]] = None, above: Optional[int] = None) -> callable:
     codes = codes or []
 
-    def predicate(exc: Exception) -> bool:
+    def predicate(exc: BaseException) -> bool:
         if isinstance(exc, httpx.HTTPStatusError):
             if exc.response.status_code in codes:
                 return True
@@ -42,7 +42,7 @@ def encode_image(image_path: str) -> str:
         return base64.b64encode(image_file.read()).decode("utf-8")
 
 
-def messages_to_openai_spec(messages: List[Message]) -> List[Dict[str, Any]]:
+def messages_to_openai_spec(messages: list[Message]) -> list[dict[str, any]]:
     messages_spec = []
     for message in messages:
         converted = {"role": message.role}
@@ -106,7 +106,7 @@ def messages_to_openai_spec(messages: List[Message]) -> List[Dict[str, Any]]:
     return messages_spec
 
 
-def tools_to_openai_spec(tools: Tuple[Tool]) -> Dict[str, Any]:
+def tools_to_openai_spec(tools: tuple[Tool, ...]) -> dict[str, any]:
     tools_names = set()
     result = []
     for tool in tools:

packages/exchange/src/exchange/token_usage_collector.py

@@ -1,5 +1,4 @@
 from collections import defaultdict
-from typing import Dict
 
 from exchange.providers.base import Usage
 
@@ -11,7 +10,7 @@ class _TokenUsageCollector:
     def collect(self, model: str, usage: Usage) -> None:
         self.usage_data.append((model, usage))
 
-    def get_token_usage_group_by_model(self) -> Dict[str, Usage]:
+    def get_token_usage_group_by_model(self) -> dict[str, Usage]:
         usage_group_by_model = defaultdict(lambda: Usage(0, 0, 0))
         for model, usage in self.usage_data:
             usage_by_model = usage_group_by_model[model]

packages/exchange/src/exchange/tool.py

@@ -1,5 +1,4 @@
 import inspect
-from typing import Any, Callable, Type
 
 from attrs import define
 
@@ -13,17 +12,17 @@ class Tool:
     Attributes:
         name (str): The name of the tool
         description (str): A description of what the tool does
-        parameters dict[str, Any]: A json schema of the function signature
+        parameters dict[str, any]: A json schema of the function signature
         function (Callable): The python function that powers the tool
     """
 
     name: str
     description: str
-    parameters: dict[str, Any]
-    function: Callable
+    parameters: dict[str, any]
+    function: callable
 
     @classmethod
-    def from_function(cls: Type["Tool"], func: Any) -> "Tool":  # noqa: ANN401
+    def from_function(cls: type["Tool"], func: any) -> "Tool":  # noqa: ANN401
         """Create a tool instance from a function and its docstring
 
         The function must have a docstring - we require it to load the description

packages/exchange/src/exchange/utils.py

@@ -1,7 +1,7 @@
 import inspect
 import uuid
 from importlib.metadata import entry_points
-from typing import Any, Callable, Dict, List, Type, get_args, get_origin
+from typing import get_args, get_origin
 
 from griffe import (
     Docstring,
@@ -20,7 +20,7 @@ def compact(content: str) -> str:
     return " ".join(content.split())
 
 
-def parse_docstring(func: Callable) -> tuple[str, List[Dict]]:
+def parse_docstring(func: callable) -> tuple[str, list[dict]]:
     """Get description and parameters from function docstring"""
     function_args = list(inspect.signature(func).parameters.keys())
     text = str(func.__doc__)
@@ -71,7 +71,7 @@ def parse_docstring(func: Callable) -> tuple[str, List[Dict]]:
 
 
 def _check_section_is_present(
-    parsed_docstring: List[DocstringSection], section_type: Type[DocstringSectionText]
+    parsed_docstring: list[DocstringSection], section_type: type[DocstringSectionText]
 ) -> bool:
     for section in parsed_docstring:
         if isinstance(section, section_type):
@@ -79,7 +79,7 @@ def _check_section_is_present(
     return False
 
 
-def json_schema(func: Any) -> dict[str, Any]:  # noqa: ANN401
+def json_schema(func: any) -> dict[str, any]:  # noqa: ANN401
     """Get the json schema for a function"""
     signature = inspect.signature(func)
     parameters = signature.parameters
@@ -107,16 +107,16 @@ def json_schema(func: Any) -> dict[str, Any]:  # noqa: ANN401
     return schema
 
 
-def _map_type_to_schema(py_type: Type) -> Dict[str, Any]:  # noqa: ANN401
+def _map_type_to_schema(py_type: type) -> dict[str, any]:  # noqa: ANN401
     origin = get_origin(py_type)
     args = get_args(py_type)
 
     if origin is list or origin is tuple:
-        return {"type": "array", "items": _map_type_to_schema(args[0] if args else Any)}
+        return {"type": "array", "items": _map_type_to_schema(args[0] if args else any)}
     elif origin is dict:
         return {
             "type": "object",
-            "additionalProperties": _map_type_to_schema(args[1] if len(args) > 1 else Any),
+            "additionalProperties": _map_type_to_schema(args[1] if len(args) > 1 else any),
         }
     elif py_type is int:
         return {"type": "integer"}

packages/exchange/tests/providers/conftest.py

@@ -1,7 +1,6 @@
 import json
 import os
 import re
-from typing import Type, Tuple
 
 import pytest
 import yaml
@@ -189,14 +188,14 @@ def scrub_response_headers(response):
     return response
 
 
-def complete(provider_cls: Type[Provider], model: str, **kwargs) -> Tuple[Message, Usage]:
+def complete(provider_cls: type[Provider], model: str, **kwargs) -> tuple[Message, Usage]:
     provider = provider_cls.from_env()
     system = "You are a helpful assistant."
     messages = [Message.user("Hello")]
     return provider.complete(model=model, system=system, messages=messages, tools=(), **kwargs)
 
 
-def tools(provider_cls: Type[Provider], model: str, **kwargs) -> Tuple[Message, Usage]:
+def tools(provider_cls: type[Provider], model: str, **kwargs) -> tuple[Message, Usage]:
     provider = provider_cls.from_env()
     system = "You are a helpful assistant. Expect to need to read a file using read_file."
     messages = [Message.user("What are the contents of this file? test.txt")]
@@ -205,7 +204,7 @@ def tools(provider_cls: Type[Provider], model: str, **kwargs) -> Tuple[Message,
     )
 
 
-def vision(provider_cls: Type[Provider], model: str, **kwargs) -> Tuple[Message, Usage]:
+def vision(provider_cls: type[Provider], model: str, **kwargs) -> tuple[Message, Usage]:
     provider = provider_cls.from_env()
     system = "You are a helpful assistant."
     messages = [

packages/exchange/tests/providers/test_azure.py

@@ -14,10 +14,10 @@ AZURE_MODEL = os.getenv("AZURE_MODEL", "gpt-4o-mini")
 @pytest.mark.parametrize(
     "env_var_name",
     [
-        ("AZURE_CHAT_COMPLETIONS_HOST_NAME"),
-        ("AZURE_CHAT_COMPLETIONS_DEPLOYMENT_NAME"),
-        ("AZURE_CHAT_COMPLETIONS_DEPLOYMENT_API_VERSION"),
-        ("AZURE_CHAT_COMPLETIONS_KEY"),
+        "AZURE_CHAT_COMPLETIONS_HOST_NAME",
+        "AZURE_CHAT_COMPLETIONS_DEPLOYMENT_NAME",
+        "AZURE_CHAT_COMPLETIONS_DEPLOYMENT_API_VERSION",
+        "AZURE_CHAT_COMPLETIONS_KEY",
     ],
 )
 def test_from_env_throw_error_when_missing_env_var(env_var_name):

packages/exchange/tests/providers/test_bedrock.py

@@ -15,9 +15,9 @@ logger = logging.getLogger(__name__)
 @pytest.mark.parametrize(
     "env_var_name",
     [
-        ("AWS_ACCESS_KEY_ID"),
-        ("AWS_SECRET_ACCESS_KEY"),
-        ("AWS_SESSION_TOKEN"),
+        "AWS_ACCESS_KEY_ID",
+        "AWS_SECRET_ACCESS_KEY",
+        "AWS_SESSION_TOKEN",
     ],
 )
 def test_from_env_throw_error_when_missing_env_var(env_var_name):

packages/exchange/tests/providers/test_databricks.py

@@ -10,8 +10,8 @@ from exchange.providers.databricks import DatabricksProvider
 @pytest.mark.parametrize(
     "env_var_name",
     [
-        ("DATABRICKS_HOST"),
-        ("DATABRICKS_TOKEN"),
+        "DATABRICKS_HOST",
+        "DATABRICKS_TOKEN",
     ],
 )
 def test_from_env_throw_error_when_missing_env_var(env_var_name):

packages/exchange/tests/providers/test_provider_utils.py

@@ -107,9 +107,9 @@ def test_messages_to_openai_spec() -> None:
         Message(role="user", content=[Text("How are you?")]),
         Message(
             role="assistant",
-            content=[ToolUse(id=1, name="tool1", parameters={"param1": "value1"})],
+            content=[ToolUse(id="1", name="tool1", parameters={"param1": "value1"})],
         ),
-        Message(role="user", content=[ToolResult(tool_use_id=1, output="Result")]),
+        Message(role="user", content=[ToolResult(tool_use_id="1", output="Result")]),
     ]
 
     spec = messages_to_openai_spec(messages)
@@ -121,7 +121,7 @@ def test_messages_to_openai_spec() -> None:
             "role": "assistant",
             "tool_calls": [
                 {
-                    "id": 1,
+                    "id": "1",
                     "type": "function",
                     "function": {
                         "name": "tool1",
@@ -133,7 +133,7 @@ def test_messages_to_openai_spec() -> None:
         {
             "role": "tool",
             "content": "Result",
-            "tool_call_id": 1,
+            "tool_call_id": "1",
         },
     ]
 
@@ -216,7 +216,7 @@ def test_openai_response_to_message_valid_tooluse() -> None:
     expect = asdict(
         Message(
             role="assistant",
-            content=[ToolUse(id=1, name="example_fn", parameters={"param": "value"})],
+            content=[ToolUse(id="1", name="example_fn", parameters={"param": "value"})],
         )
     )
     actual.pop("id")

packages/exchange/tests/test_exchange.py

@@ -1,5 +1,3 @@
-from typing import List, Tuple
-
 import pytest
 
 from exchange.checkpoint import Checkpoint, CheckpointData
@@ -29,12 +27,12 @@ def no_overlapping_checkpoints(exchange: Exchange) -> bool:
     return True
 
 
-def checkpoint_to_index_pairs(checkpoints: List[Checkpoint]) -> List[Tuple[int, int]]:
+def checkpoint_to_index_pairs(checkpoints: list[Checkpoint]) -> list[tuple[int, int]]:
     return [(checkpoint.start_index, checkpoint.end_index) for checkpoint in checkpoints]
 
 
 class MockProvider(Provider):
-    def __init__(self, sequence: List[Message], usage_dicts: List[dict]):
+    def __init__(self, sequence: list[Message], usage_dicts: list[dict]):
         # We'll use init to provide a preplanned reply sequence
         self.sequence = sequence
         self.call_count = 0
@@ -56,11 +54,18 @@ class MockProvider(Provider):
             total_tokens=total_tokens,
         )
 
-    def complete(self, model: str, system: str, messages: List[Message], tools: List[Tool]) -> Message:
+    def complete(
+        self,
+        model: str,
+        system: str,
+        messages: list[Message],
+        tools: tuple[Tool, ...],
+        **kwargs: dict[str, any],
+    ) -> tuple[Message, Usage]:
         output = self.sequence[self.call_count]
         usage = self.get_usage(self.usage_dicts[self.call_count])
         self.call_count += 1
-        return (output, usage)
+        return output, usage
 
 
 def test_reply_with_unsupported_tool():
@@ -116,7 +121,7 @@ def test_invalid_tool_parameters():
         ),
         model="gpt-4o-2024-05-13",
         system="You are a helpful assistant.",
-        tools=[Tool.from_function(dummy_tool)],
+        tools=(Tool.from_function(dummy_tool),),
         moderator=PassiveModerator(),
     )
 
@@ -154,7 +159,7 @@ def test_max_tool_use_when_limit_reached():
         ),
         model="gpt-4o-2024-05-13",
         system="You are a helpful assistant.",
-        tools=[Tool.from_function(dummy_tool)],
+        tools=(Tool.from_function(dummy_tool),),
         moderator=PassiveModerator(),
     )
 
@@ -195,7 +200,7 @@ def test_tool_output_too_long_character_error():
         ),
         model="gpt-4o-2024-05-13",
         system="You are a helpful assistant.",
-        tools=[Tool.from_function(long_output_tool_char)],
+        tools=(Tool.from_function(long_output_tool_char),),
         moderator=PassiveModerator(),
     )
 
@@ -236,7 +241,7 @@ def test_tool_output_too_long_token_error():
         ),
         model="gpt-4o-2024-05-13",
         system="You are a helpful assistant.",
-        tools=[Tool.from_function(long_output_tool_token)],
+        tools=(Tool.from_function(long_output_tool_token),),
         moderator=PassiveModerator(),
     )
 
@@ -301,7 +306,7 @@ def resumed_exchange() -> Exchange:
     ex = Exchange(
         provider=provider,
         messages=messages,
-        tools=[],
+        tools=(),
         model="gpt-4o-2024-05-13",
         system="You are a helpful assistant.",
         checkpoint_data=CheckpointData(),
@@ -399,7 +404,7 @@ def test_pop_first_message_no_messages():
         provider=MockProvider(sequence=[], usage_dicts=[]),
         model="gpt-4o-2024-05-13",
         system="You are a helpful assistant.",
-        tools=[Tool.from_function(dummy_tool)],
+        tools=(Tool.from_function(dummy_tool),),
         moderator=PassiveModerator(),
     )
 
@@ -741,7 +746,7 @@ def test_rewind_with_tool_usage():
         ),
         model="gpt-4o-2024-05-13",
         system="You are a helpful assistant.",
-        tools=[Tool.from_function(dummy_tool)],
+        tools=(Tool.from_function(dummy_tool),),
         moderator=PassiveModerator(),
     )
     ex.add(Message(role="user", content=[Text(text="test")]))

packages/exchange/tests/test_exchange_frozen.py

@@ -9,7 +9,7 @@ from exchange.tool import Tool
 
 
 class MockProvider(Provider):
-    def complete(self, model, system, messages, tools=None):
+    def complete(self, model, system, messages, tools, **kwargs):
         return Message(role="assistant", content=[Text(text="This is a mock response.")]), Usage.from_dict(
             {"total_tokens": 35}
         )

packages/exchange/tests/test_summarizer.py

@@ -3,11 +3,11 @@ from exchange import Exchange, Message
 from exchange.content import ToolResult, ToolUse
 from exchange.moderators.passive import PassiveModerator
 from exchange.moderators.summarizer import ContextSummarizer
-from exchange.providers import Usage
+from exchange.providers import Usage, Provider
 
 
-class MockProvider:
-    def complete(self, model, system, messages, tools):
+class MockProvider(Provider):
+    def complete(self, model, system, messages, tools, **kwargs):
         assistant_message_text = "Summarized content here."
         output_tokens = len(assistant_message_text)
         total_input_tokens = sum(len(msg.text) for msg in messages)
@@ -138,14 +138,14 @@ MESSAGE_SEQUENCE = [
 ]
 
 
-class AnotherMockProvider:
+class AnotherMockProvider(Provider):
     def __init__(self):
         self.sequence = MESSAGE_SEQUENCE
         self.current_index = 1
         self.summarize_next = False
         self.summarized_count = 0
 
-    def complete(self, model, system, messages, tools):
+    def complete(self, model, system, messages, tools, **kwargs):
         system_prompt_tokens = 100
         input_token_count = system_prompt_tokens
 

packages/exchange/tests/test_truncate.py

@@ -73,7 +73,7 @@ class TruncateLinearProvider(Provider):
         self.summarize_next = False
         self.summarized_count = 0
 
-    def complete(self, model, system, messages, tools):
+    def complete(self, model, system, messages, tools, **kwargs):
         input_token_count = SYSTEM_PROMPT_TOKENS
 
         message = self.sequence[self.current_index]