Merge pull request #16 from vatsalsaglani/main

Async Client Update and Bedrock Support in `AsyncTool`

claudetools/completion/async_complete.py

@@ -1,33 +1,59 @@
-import httpx
 import os
-from typing import List, Dict
+from typing import List, Dict, Union
+from anthropic import AsyncAnthropic, AsyncAnthropicBedrock
+
+# class AsyncComplete:
+
+#     def __init__(self,
+#                  anthropic_api_key: str,
+#                  anthropic_version: str = "2023-06-01",
+#                  anthropic_base_url: str = "https://api.anthropic.com/v1"):
+#         self.headers = {
+#             "x-api-key": anthropic_api_key,
+#             "anthropic-version": anthropic_version,
+#             "content-type": "application/json"
+#         }
+#         self.anthropic_base_url = anthropic_base_url
+
+#     async def __call__(self, model: str, messages: List[Dict], **kwargs):
+#         payload = {
+#             **kwargs,
+#             "model": model,
+#             "messages": messages,
+#         }
+#         async with httpx.AsyncClient(timeout=600) as client:
+#             response = await client.post(os.path.join(self.anthropic_base_url,
+#                                                       "messages"),
+#                                          json=payload,
+#                                          headers=self.headers)
+#             response.raise_for_status()
+#             output = response.json()
+#             # print("MODEL OUTPUT\n", output)
+#             return output.get("content")[0].get("text")
 
 
 class AsyncComplete:
 
     def __init__(self,
-                 anthropic_api_key: str,
-                 anthropic_version: str = "2023-06-01",
-                 anthropic_base_url: str = "https://api.anthropic.com/v1"):
-        self.headers = {
-            "x-api-key": anthropic_api_key,
-            "anthropic-version": anthropic_version,
-            "content-type": "application/json"
-        }
-        self.anthropic_base_url = anthropic_base_url
+                 anthropic_api_key: Union[str, None] = None,
+                 aws_access_key: Union[str, None] = None,
+                 aws_secret_key: Union[str, None] = None,
+                 aws_region: Union[str, None] = None,
+                 aws_session_token: Union[str, None] = None):
+        if anthropic_api_key:
+            self.client = AsyncAnthropic(api_key=anthropic_api_key)
+        else:
+            if aws_session_token:
+                self.client = AsyncAnthropicBedrock(
+                    aws_session_token=aws_session_token, aws_region=aws_region)
+            else:
+                self.client = AsyncAnthropicBedrock(
+                    aws_access_key=aws_access_key,
+                    aws_secret_key=aws_secret_key,
+                    aws_region=aws_region)
 
     async def __call__(self, model: str, messages: List[Dict], **kwargs):
-        payload = {
-            **kwargs,
-            "model": model,
-            "messages": messages,
-        }
-        async with httpx.AsyncClient(timeout=600) as client:
-            response = await client.post(os.path.join(self.anthropic_base_url,
-                                                      "messages"),
-                                         json=payload,
-                                         headers=self.headers)
-            response.raise_for_status()
-            output = response.json()
-            # print("MODEL OUTPUT\n", output)
-            return output.get("content")[0].get("text")
+        response = await self.client.messages.create(model=model,
+                                                     messages=messages,
+                                                     **kwargs)
+        return response.content[0].text

claudetools/tools/tool.py

@@ -106,11 +106,18 @@ def __call__(self,
 class AsyncTool(BaseTool):
 
     def __init__(self,
-                 anthropic_api_key: str,
-                 anthropic_version: str = "2023-06-01",
-                 anthropic_base_url: str = "https://api.anthropic.com/v1"):
-        self.complete = AsyncComplete(anthropic_api_key, anthropic_version,
-                                      anthropic_base_url)
+                 anthropic_api_key: Union[str, None] = None,
+                 aws_access_key: Union[str, None] = None,
+                 aws_secret_key: Union[str, None] = None,
+                 aws_region: Union[str, None] = None,
+                 aws_session_token: Union[str, None] = None):
+        # self.complete = AsyncComplete(anthropic_api_key, anthropic_version,
+        #                               anthropic_base_url)
+        self.complete = AsyncComplete(anthropic_api_key=anthropic_api_key,
+                                      aws_secret_key=aws_secret_key,
+                                      aws_access_key=aws_access_key,
+                                      aws_region=aws_region,
+                                      aws_session_token=aws_session_token)
 
     async def perform_model_call(self, model, messages, system, **kwargs):
         return await self.complete(model, messages, system=system, **kwargs)

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="claudetools",
-    version="0.6.0",
+    version="0.7.0",
     author="Vatsal J. Saglani",
     author_email="[email protected]",
     description=