Few Refactorings.

Author: rohitprasad15

aisuite/providers/anthropic_provider.py

@@ -1,3 +1,7 @@
+# Anthropic provider
+# Links:
+# Tool calling docs - https://docs.anthropic.com/en/docs/build-with-claude/tool-use
+
 import anthropic
 import json
 from aisuite.provider import Provider
@@ -7,36 +11,34 @@
 # Define a constant for the default max_tokens value
 DEFAULT_MAX_TOKENS = 4096
 
-# Links:
-# Tool calling docs - https://docs.anthropic.com/en/docs/build-with-claude/tool-use
-
-
-class AnthropicProvider(Provider):
-    # Add these at the class level, after the class definition
-    FINISH_REASON_MAPPING = {
-        "end_turn": "stop",
-        "max_tokens": "length",
-        "tool_use": "tool_calls",
-        # Add more mappings as needed
-    }
 
+class AnthropicMessageConverter:
     # Role constants
     ROLE_USER = "user"
     ROLE_ASSISTANT = "assistant"
     ROLE_TOOL = "tool"
     ROLE_SYSTEM = "system"
 
-    def __init__(self, **config):
-        """
-        Initialize the Anthropic provider with the given configuration.
-        Pass the entire configuration dictionary to the Anthropic client constructor.
-        """
-
-        self.client = anthropic.Anthropic(**config)
+    # Finish reason mapping
+    FINISH_REASON_MAPPING = {
+        "end_turn": "stop",
+        "max_tokens": "length",
+        "tool_use": "tool_calls",
+    }
 
     def convert_request(self, messages):
         """Convert framework messages to Anthropic format."""
-        return [self._convert_single_message(msg) for msg in messages]
+        system_message = self._extract_system_message(messages)
+        converted_messages = [self._convert_single_message(msg) for msg in messages]
+        return system_message, converted_messages
+
+    def convert_response(self, response):
+        """Normalize the response from the Anthropic API to match OpenAI's response format."""
+        normalized_response = ChatCompletionResponse()
+        normalized_response.choices[0].finish_reason = self._get_finish_reason(response)
+        normalized_response.usage = self._get_usage_stats(response)
+        normalized_response.choices[0].message = self._get_message(response)
+        return normalized_response
 
     def _convert_single_message(self, msg):
         """Convert a single message to Anthropic format."""
@@ -104,38 +106,45 @@ def _create_assistant_tool_message(self, content, tool_calls):
 
         return {"role": self.ROLE_ASSISTANT, "content": message_content}
 
-    def chat_completions_create(self, model, messages, **kwargs):
-        """Create a chat completion using the Anthropic API."""
-        system_message = self._extract_system_message(messages)
-        kwargs = self._prepare_kwargs(kwargs)
-        converted_messages = self.convert_request(messages)
-
-        response = self.client.messages.create(
-            model=model, system=system_message, messages=converted_messages, **kwargs
-        )
-        return self.convert_response(response)
-
     def _extract_system_message(self, messages):
         """Extract system message if present, otherwise return empty list."""
+        # TODO: This is a temporary solution to extract the system message.
+        # User can pass multiple system messages, which can mingled with other messages.
+        # This needs to be fixed to handle this case.
         if messages and messages[0]["role"] == "system":
             system_message = messages[0]["content"]
             messages.pop(0)
             return system_message
         return []
 
-    def _prepare_kwargs(self, kwargs):
-        """Prepare kwargs for the API call."""
-        kwargs = kwargs.copy()  # Create a copy to avoid modifying the original
-        kwargs.setdefault("max_tokens", DEFAULT_MAX_TOKENS)
+    def _get_finish_reason(self, response):
+        """Get the normalized finish reason."""
+        return self.FINISH_REASON_MAPPING.get(response.stop_reason, "stop")
 
-        if "tools" in kwargs:
-            kwargs["tools"] = self._convert_tool_spec(kwargs["tools"])
+    def _get_usage_stats(self, response):
+        """Get the usage statistics."""
+        return {
+            "prompt_tokens": response.usage.input_tokens,
+            "completion_tokens": response.usage.output_tokens,
+            "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
+        }
 
-        return kwargs
+    def _get_message(self, response):
+        """Get the appropriate message based on response type."""
+        if response.stop_reason == "tool_use":
+            tool_message = self.convert_response_with_tool_use(response)
+            if tool_message:
+                return tool_message
+
+        return Message(
+            content=response.content[0].text,
+            role="assistant",
+            tool_calls=None,
+            refusal=None,
+        )
 
     def convert_response_with_tool_use(self, response):
         """Convert Anthropic tool use response to the framework's format."""
-        # Find the tool_use content
         tool_call = next(
             (content for content in response.content if content.type == "tool_use"),
             None,
@@ -148,7 +157,6 @@ def convert_response_with_tool_use(self, response):
             tool_call_obj = ChatCompletionMessageToolCall(
                 id=tool_call.id, function=function, type="function"
             )
-            # Get the text content if any
             text_content = next(
                 (
                     content.text
@@ -166,53 +174,15 @@ def convert_response_with_tool_use(self, response):
             )
         return None
 
-    def convert_response(self, response):
-        """Normalize the response from the Anthropic API to match OpenAI's response format."""
-        normalized_response = ChatCompletionResponse()
-
-        normalized_response.choices[0].finish_reason = self._get_finish_reason(response)
-        normalized_response.usage = self._get_usage_stats(response)
-        normalized_response.choices[0].message = self._get_message(response)
-
-        return normalized_response
-
-    def _get_finish_reason(self, response):
-        """Get the normalized finish reason."""
-        return self.FINISH_REASON_MAPPING.get(response.stop_reason, "stop")
-
-    def _get_usage_stats(self, response):
-        """Get the usage statistics."""
-        return {
-            "prompt_tokens": response.usage.input_tokens,
-            "completion_tokens": response.usage.output_tokens,
-            "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
-        }
-
-    def _get_message(self, response):
-        """Get the appropriate message based on response type."""
-        if response.stop_reason == "tool_use":
-            tool_message = self.convert_response_with_tool_use(response)
-            if tool_message:
-                return tool_message
-
-        return Message(
-            content=response.content[0].text,
-            role="assistant",
-            tool_calls=None,
-            refusal=None,
-        )
-
-    def _convert_tool_spec(self, openai_tools):
+    def convert_tool_spec(self, openai_tools):
         """Convert OpenAI tool specification to Anthropic format."""
         anthropic_tools = []
 
         for tool in openai_tools:
-            # Only handle function-type tools from OpenAI
             if tool.get("type") != "function":
                 continue
 
             function = tool["function"]
-
             anthropic_tool = {
                 "name": function["name"],
                 "description": function["description"],
@@ -222,7 +192,33 @@ def _convert_tool_spec(self, openai_tools):
                     "required": function["parameters"].get("required", []),
                 },
             }
-
             anthropic_tools.append(anthropic_tool)
 
         return anthropic_tools
+
+
+class AnthropicProvider(Provider):
+    def __init__(self, **config):
+        """Initialize the Anthropic provider with the given configuration."""
+        self.client = anthropic.Anthropic(**config)
+        self.converter = AnthropicMessageConverter()
+
+    def chat_completions_create(self, model, messages, **kwargs):
+        """Create a chat completion using the Anthropic API."""
+        kwargs = self._prepare_kwargs(kwargs)
+        system_message, converted_messages = self.converter.convert_request(messages)
+
+        response = self.client.messages.create(
+            model=model, system=system_message, messages=converted_messages, **kwargs
+        )
+        return self.converter.convert_response(response)
+
+    def _prepare_kwargs(self, kwargs):
+        """Prepare kwargs for the API call."""
+        kwargs = kwargs.copy()
+        kwargs.setdefault("max_tokens", DEFAULT_MAX_TOKENS)
+
+        if "tools" in kwargs:
+            kwargs["tools"] = self.converter.convert_tool_spec(kwargs["tools"])
+
+        return kwargs

aisuite/providers/aws_provider.py

@@ -33,6 +33,10 @@ def convert_request(
             for message in messages
         ]
 
+        import pprint
+
+        pprint.pprint(messages)
+
         # Handle system message
         system_message = []
         if messages and messages[0]["role"] == "system":
@@ -61,6 +65,8 @@ def convert_request(
                     }
                 )
 
+        pprint.pprint(formatted_messages)
+
         return system_message, formatted_messages
 
     @staticmethod
@@ -150,21 +156,14 @@ def convert_assistant(message: Dict[str, Any]) -> Optional[Dict[str, Any]]:
 
         return {"role": "assistant", "content": content} if content else None
 
-
-class AwsProvider(Provider):
-    def __init__(self, **config):
-        """Initialize the AWS Bedrock provider with the given configuration."""
-        self.config = BedrockConfig(**config)
-        self.client = self.config.create_client()
-        self.transformer = BedrockMessageConverter()
-
-    def convert_response(self, response: Dict[str, Any]) -> ChatCompletionResponse:
+    @staticmethod
+    def convert_response(response: Dict[str, Any]) -> ChatCompletionResponse:
         """Normalize the response from the Bedrock API to match OpenAI's response format."""
         norm_response = ChatCompletionResponse()
 
         # Check if the model is requesting tool use
         if response.get("stopReason") == "tool_use":
-            tool_message = self.transformer.convert_response_tool_call(response)
+            tool_message = BedrockMessageConverter.convert_response_tool_call(response)
             if tool_message:
                 norm_response.choices[0].message = Message(**tool_message)
                 norm_response.choices[0].finish_reason = "tool_calls"
@@ -176,6 +175,18 @@ def convert_response(self, response: Dict[str, Any]) -> ChatCompletionResponse:
         ][0]["text"]
         return norm_response
 
+
+class AwsProvider(Provider):
+    def __init__(self, **config):
+        """Initialize the AWS Bedrock provider with the given configuration."""
+        self.config = BedrockConfig(**config)
+        self.client = self.config.create_client()
+        self.transformer = BedrockMessageConverter()
+
+    def convert_response(self, response: Dict[str, Any]) -> ChatCompletionResponse:
+        """Normalize the response from the Bedrock API to match OpenAI's response format."""
+        return self.transformer.convert_response(response)
+
     def _convert_tool_spec(self, kwargs: Dict[str, Any]) -> Optional[Dict[str, Any]]:
         """Convert tool specifications to Bedrock format."""
         if "tools" not in kwargs:
@@ -213,12 +224,16 @@ def _prepare_request_config(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
             if key not in BedrockConfig.INFERENCE_PARAMETERS
         }
 
-        return {
+        request_config = {
             "inferenceConfig": inference_config,
             "additionalModelRequestFields": additional_fields,
-            "toolConfig": tool_config,
         }
 
+        if tool_config is not None:
+            request_config["toolConfig"] = tool_config
+
+        return request_config
+
     def chat_completions_create(
         self, model: str, messages: List[Dict[str, Any]], **kwargs
     ) -> ChatCompletionResponse:

examples/client.ipynb

@@ -231,6 +231,18 @@
     "response = client.chat.completions.create(model=togetherai_model, messages=messages, temperature=0.75, top_p=0.7, top_k=50)\n",
     "print(response.choices[0].message.content)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dcf63a11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "gemini_15_flash = \"google:gemini-1.5-flash\"\n",
+    "response = client.chat.completions.create(model=gemini_15_flash, messages=messages, temperature=0.75)\n",
+    "print(response.choices[0].message.content)"
+   ]
   }
  ],
  "metadata": {