From 8d83050a64c35a2cb163ffe3799abf9904e2236a Mon Sep 17 00:00:00 2001
From: Stan Zubarev <szubarev@gmail.com>
Date: Thu, 31 Oct 2024 19:34:50 -0400
Subject: [PATCH] add Amazon Bedrock provider

---
 simplemind/providers/__init__.py |  3 +-
 simplemind/providers/amazon.py   | 85 ++++++++++++++++++++++++++++++++
 simplemind/settings.py           |  3 ++
 3 files changed, 90 insertions(+), 1 deletion(-)
 create mode 100644 simplemind/providers/amazon.py

diff --git a/simplemind/providers/__init__.py b/simplemind/providers/__init__.py
index 201f851..cc6c465 100644
--- a/simplemind/providers/__init__.py
+++ b/simplemind/providers/__init__.py
@@ -7,5 +7,6 @@ from .groq import Groq
 from .ollama import Ollama
 from .openai import OpenAI
 from .xai import XAI
+from .amazon import Amazon
 
-providers: List[Type[BaseProvider]] = [Anthropic, Gemini, Groq, OpenAI, Ollama, XAI]
+providers: List[Type[BaseProvider]] = [Anthropic, Gemini, Groq, OpenAI, Ollama, XAI, Amazon]
diff --git a/simplemind/providers/amazon.py b/simplemind/providers/amazon.py
new file mode 100644
index 0000000..f4a9ede
--- /dev/null
+++ b/simplemind/providers/amazon.py
@@ -0,0 +1,85 @@
+from typing import Union
+import json
+
+import instructor
+import anthropic
+
+from ._base import BaseProvider
+from ..settings import settings
+
+PROVIDER_NAME = "amazon"
+DEFAULT_MODEL = "anthropic.claude-3-sonnet-20240229-v1:0"
+DEFAULT_MAX_TOKENS = 5000
+
+class Amazon(BaseProvider):
+    NAME = PROVIDER_NAME
+    DEFAULT_MODEL = DEFAULT_MODEL
+
+    def __init__(self, api_key: Union[str, None] = None):
+        self.api_key = api_key or settings.get_api_key(PROVIDER_NAME)
+
+    @property
+    def client(self):
+        """The AnthropicBedrock client."""
+        if not self.api_key:
+            raise ValueError("Profile name is not provided")
+
+        return anthropic.AnthropicBedrock(aws_profile=self.api_key)
+
+    @property
+    def structured_client(self):
+        """A client patched with Instructor."""
+        return instructor.from_anthropic(self.client)
+
+    def send_conversation(self, conversation: "Conversation", **kwargs):
+        """Send a conversation to the OpenAI API."""
+        from ..models import Message
+
+        messages = [
+            {"role": msg.role, "content": msg.text} for msg in conversation.messages
+        ]
+
+        response = self.client.chat.completions.create(
+            model=conversation.llm_model or DEFAULT_MODEL, messages=messages, **kwargs
+        )
+
+        # Get the response content from the OpenAI response
+        assistant_message = response.choices[0].message
+
+        # Create and return a properly formatted Message instance
+        return Message(
+            role="assistant",
+            text=assistant_message.content or "",
+            raw=response,
+            llm_model=conversation.llm_model or DEFAULT_MODEL,
+            llm_provider=PROVIDER_NAME,
+        )
+
+    def structured_response(self, prompt, response_model, *, llm_model: str, **kwargs):
+        # Ensure messages are provided in kwargs
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.structured_client.chat.completions.create(
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            response_model=response_model,
+            max_tokens=DEFAULT_MAX_TOKENS,
+            **kwargs,
+        )
+        return response
+
+    def generate_text(self, prompt, *, llm_model, **kwargs):
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.client.messages.create(
+            model=llm_model or self.DEFAULT_MODEL,
+            messages=messages,
+            max_tokens=DEFAULT_MAX_TOKENS,
+            **kwargs,
+        )
+
+        return response.content[0].text
\ No newline at end of file
diff --git a/simplemind/settings.py b/simplemind/settings.py
index d561228..d840ae7 100644
--- a/simplemind/settings.py
+++ b/simplemind/settings.py
@@ -18,6 +18,9 @@ class LoggingConfig(BaseSettings):
 class Settings(BaseSettings):
     """The class that holds all the API keys for the application."""
 
+    AMAZON_API_KEY: Optional[SecretStr] = Field(
+        "default", description="AWS Named Profile"
+    )
     ANTHROPIC_API_KEY: Optional[SecretStr] = Field(
         None, description="API key for Anthropic"
     )