Merge pull request #9 from vatsalsaglani/main

AWS Bedrock Support

Readme.md

@@ -11,6 +11,7 @@ With `claudetools` one can now use any model from the **Claude 3** family of mod
 - **Flexible Tool Definition**: Specify function names, descriptions, and parameters using the Pydantic library for type safety and validation.
 - **Multiple Tools Support**: Opt to call multiple tools within a single prompt, enabling more complex interactions.
 - **Customizable System Prompts**: Attach custom system prompts to your conversations for better context and control.
+- **Bedrock Client**: Use Claude 3 via AWS Bedrock for function calling.
 
 
 ## Installation
@@ -167,6 +168,87 @@ if __name__ == "__main__":
  asyncio.run(main())
 ```
 
+## Bedrock Example
+
+> Currently, only the sync client supports AWS Bedrock.
+
+```py
+import asyncio
+from claudetools.tools.tool import Tool
+from pydantic import BaseModel, Field
+from typing import List, Dict
+import json
+
+AWS_ACCESS_KEY=""
+AWS_SECRET_KEY=""
+AWS_REGION=""
+# or
+AWS_SESSION_TOKEN=""
+
+# create a tool instance with your aws access and secret keys
+tool = Tool(aws_access_key=AWS_ACCESS_KEY,
+ aws_secret_key=AWS_SECRET_KEY,
+ aws_region=AWS_REGION)
+# to use session token from AWS STS use the following
+# tool = Tool(aws_session_token=AWS_SESSION_TOKEN, aws_region=AWS_REGION)
+
+
+# define your function parameters
+class AddTodo(BaseModel):
+ text: str = Field(..., description="Text to add for the TODO to remember.")
+
+class MarkCompleted(BaseModel):
+ text: str = Field(..., description="Text of the completed TODO.")
+
+
+class ReOpen(BaseModel):
+ text: str = Field(..., description="Text of the TODO to reopen.")
+
+
+# specify the functions you want to use
+functions = [{
+ "name": "AddTodo",
+ "description": "Add a TODO with text to remember.",
+ "parameters": AddTodo.model_json_schema()
+}, {
+ "name": "MarkCompleted",
+ "description": "Get text of the todo mark it complete",
+ "parameters": MarkCompleted.model_json_schema()
+}, {
+ "name": "ReOpen",
+ "description": "Get text of the todo reopen it.",
+ "parameters": ReOpen.model_json_schema()
+}]
+
+# set up the user messages
+user_messages = [{
+ "role":
+ "user",
+ "content":
+ """I have to pick up my daughter from school. After which I've to do the laundary. And now I need to cook lunch."""
+}]
+
+# dependency prompt to attach to the main system prompt
+DEPENDENCY_PROMPT = """You are a helpful assistant that helps a user with their tasks and todos. The user can add a todos, mark todos as completed, or reopen certain todos.
+The user can provide multiple actions at once so you've to break those down and call the appropriate functions in the correct sequence."""
+
+
+# call the tool with the required parameters
+output = tool(model="claude-3-sonnet-20240229",
+ messages=user_messages,
+ tools=functions,
+ tool_choice=None,
+ multiple_tools=True,
+ attach_system=DEPENDENCY_PROMPT,
+ max_tokens=3000)
+
+if output:
+ print(json.dumps(output, indent=4))
+else:
+ print("Unable to find a function!")
+```
+
+
 ## Requirements
 
 Python 3.7 or higher.

claudetools/completion/complete.py

@@ -1,11 +1,26 @@
 from anthropic import Anthropic
-from typing import List, Dict
+from typing import List, Dict, Union
 
 
 class Complete:
 
- def __init__(self, anthropic_api_key: str):
- self.client = Anthropic(api_key=anthropic_api_key)
+ def __init__(self,
+ 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 = Anthropic(api_key=anthropic_api_key)
+ else:
+
+ if aws_session_token:
+ self.client = Anthropic(aws_session_token=aws_session_token,
+ aws_region=aws_region)
+ else:
+ self.client = Anthropic(aws_access_key=aws_access_key,
+ aws_secret_key=aws_secret_key,
+ aws_region=aws_region)
 
  def __call__(self, model: str, messages: List[Dict], **kwargs):
  output = self.client.messages.create(model=model,

claudetools/tools/tool.py

@@ -75,8 +75,17 @@ async def perform_model_call(self, model, messages, system, **kwargs):
 
 class Tool(BaseTool):
 
- def __init__(self, anthropic_api_key: str):
- self.complete = Complete(anthropic_api_key)
+ def __init__(self,
+ 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 = Complete(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 self.complete(model, messages, system=system, **kwargs)

requirements.txt

@@ -1,5 +1,5 @@
 httpx==0.25.0
-anthropic==0.19.1
+anthropic[bedrock]==0.19.1
 pydantic==2.4.2
 python-dotenv==1.0.0
 

setup.py

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