adding bedrock as model provider (#37)

* adding bedrock as model provider

* bump version

Author: eleonorecharles

docs/docs/models/bedrock/_category_.yaml

@@ -0,0 +1,2 @@
+label: "Bedrock"
+position: 5

docs/docs/models/bedrock/index.mdx

@@ -0,0 +1,25 @@
+# Bedrock
+
+Jupyter AI Agents supports [Amazon Bedrock](https://aws.amazon.com/bedrock/) models.
+
+To use Bedrock models, you'll need to define the following environment variables:
+
+```bash
+AWS_ACCESS_KEY_ID - The access key for your AWS account.
+AWS_SECRET_ACCESS_KEY - The secret key for your AWS account.
+AWS_SESSION_TOKEN - The session key for your AWS account. This is only needed when you are using temporary credentials.
+```
+
+Example usage with direct Bedrock integration:
+
+```bash
+# Prompt agent example with Bedrock
+jupyter-ai-agents prompt \
+  --url http://localhost:8888 \
+  --token MY_TOKEN \
+  --model-provider bedrock \
+  --model-name anthropic.claude-3-sonnet-20240229-v1:0 \
+  --path test.ipynb \
+  --input "Create a matplotlib example"
+```
+

jupyter_ai_agents/__version__.py

@@ -4,4 +4,4 @@
 
 """Jupyter AI Agents."""
 
-__version__ = "0.10.1"
+__version__ = "0.10.2"

jupyter_ai_agents/agents/utils.py

@@ -7,6 +7,7 @@
 from jupyter_ai_agents.providers.anthropic import create_anthropic_agent
 from jupyter_ai_agents.providers.azure_openai import create_azure_openai_agent
 from jupyter_ai_agents.providers.github_copilot import create_github_copilot_agent
+from jupyter_ai_agents.providers.bedrock import create_bedrock_agent
 from jupyter_ai_agents.providers.openai import create_openai_agent
 
 
@@ -22,6 +23,8 @@ def create_ai_agent(
         agent = create_openai_agent(model_name, system_prompt_final, tools)
     elif model_provider == "anthropic":
         agent = create_anthropic_agent(model_name, system_prompt_final, tools)
+    elif model_provider == "bedrock":
+        agent = create_bedrock_agent(model_name, system_prompt_final, tools)
     else:
         raise ValueError(f"Model provider {model_provider} is not supported.")
     return agent

jupyter_ai_agents/providers/anthropic.py

@@ -8,6 +8,8 @@
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.agents.tool_calling_agent.base import create_tool_calling_agent
 from langchain_core.tools import BaseTool
+from dotenv import load_dotenv
+
 
 def create_anthropic_agent(model_name: str, system_prompt: str, tools: List[BaseTool]) -> AgentExecutor:
     """Create an agent from a set of tools using Anthropic's Claude API.
@@ -20,6 +22,8 @@ def create_anthropic_agent(model_name: str, system_prompt: str, tools: List[Base
     Returns:
         An agent executor that can use tools via Claude
     """
+    
+    load_dotenv()
 
     # Create the Anthropic LLM
     llm = ChatAnthropic(model=model_name)

jupyter_ai_agents/providers/azure_openai.py

@@ -5,10 +5,13 @@
 from langchain.agents import AgentExecutor, create_openai_tools_agent
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_openai import AzureChatOpenAI
+from dotenv import load_dotenv
 
 
 def create_azure_openai_agent(model_name: str, system_prompt: str, tools: list) -> AgentExecutor:
     """Create an agent from a set of tools and an Azure deployment"""
+    
+    load_dotenv()
 
     llm = AzureChatOpenAI(azure_deployment=model_name)
 

jupyter_ai_agents/providers/bedrock.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2023-2024 Datalayer, Inc.
+#
+# BSD 3-Clause License
+
+from typing import List
+from langchain.agents.agent import AgentExecutor
+from langchain_aws import ChatBedrockConverse
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain.agents.tool_calling_agent.base import create_tool_calling_agent
+from langchain_core.tools import BaseTool
+from dotenv import load_dotenv
+
+def create_bedrock_agent(model_name: str, system_prompt: str, tools: List[BaseTool]) -> AgentExecutor:
+    """Create an agent from a set of tools using Anthropic's Claude API.
+    
+    Args:
+        model_name: The name of the Claude model to use (e.g., "claude-3-haiku-20240307")
+        system_prompt: The system prompt to use for the agent
+        tools: A list of tools for the agent to use
+        
+    Returns:
+        An agent executor that can use tools via Claude
+    """
+    
+    load_dotenv()
+
+    # Create the Anthropic LLM
+    llm = ChatBedrockConverse(model_id=model_name)
+
+    # Create a prompt template for the agent with enhanced instructions
+    enhanced_system_prompt = f"""
+{system_prompt}
+
+When you use tools, please include the results in your response to the user.
+Be sure to always provide a text response, even if it's just to acknowledge the tool's output.
+After using a tool, explain what the result means in a clear and helpful way.
+"""
+
+    # Create prompt template
+    prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", enhanced_system_prompt),
+        ("user", "{input}"),
+        MessagesPlaceholder(variable_name="agent_scratchpad"),
+    ])
+
+    # Create a tool-calling agent
+    agent = create_tool_calling_agent(llm, tools, prompt)
+
+    # Create an agent executor with output parsing
+    agent_executor = AgentExecutor(
+        name="AnthropicToolAgent", 
+        agent=agent, 
+        tools=tools, 
+        verbose=True, 
+        handle_parsing_errors=True,
+        return_intermediate_steps=True  # Include intermediate steps in the output
+    )
+    
+    return agent_executor 

jupyter_ai_agents/providers/github_copilot.py

@@ -5,10 +5,13 @@
 from langchain.agents import AgentExecutor, create_openai_tools_agent
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_github_copilot import ChatGitHubCopilot
+from dotenv import load_dotenv
 
 
 def create_github_copilot_agent(model_name: str, system_prompt: str, tools: list) -> AgentExecutor:
     """Create an agent from a set of tools and a Github Copilot model"""
+    
+    load_dotenv()
 
     llm = ChatGitHubCopilot(model_name=model_name)
 

jupyter_ai_agents/providers/openai.py

@@ -6,9 +6,12 @@
 from langchain_openai import ChatOpenAI
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.agents import create_openai_tools_agent
+from dotenv import load_dotenv
 
 def create_openai_agent(model_name: str, system_prompt: str, tools: list) -> AgentExecutor:
     """Create an agent from a set of tools using OpenAI's API."""
+    
+    load_dotenv()
 
     llm = ChatOpenAI(model_name=model_name)
 

pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
   "langchain-openai",
   "langchain-github-copilot>=0.4.0",
   "langchain-anthropic",
+  "langchain-aws",
   "fastapi",
   "jupyter_kernel_client==0.6.0",
   "jupyter_nbmodel_client==0.11.3",