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

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

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]