extending to bedrock

example.env

@@ -11,6 +11,11 @@ API_KEY_OPENAI_AZURE=
 OPENAI_AZURE_ENDPOINT=
 OPENAI_API_VERSION=
 
+REGION_NAME=
+AWS_ACCESS_KEY_ID=
+AWS_SECRET_ACCESS_KEY=
+AWS_SESSION_TOKEN=
+
 HF_TOKEN=
 
 

initialize.py

@@ -5,7 +5,7 @@
 def initialize():
 
     # main chat model used by agents (smarter, more accurate)
-    chat_llm = models.get_openai_chat(model_name="gpt-4o-mini", temperature=0)
+    # chat_llm = models.get_openai_chat(model_name="gpt-4o-mini", temperature=0)
     # chat_llm = models.get_ollama_chat(model_name="llama3.2:3b-instruct-fp16", temperature=0)
     # chat_llm = models.get_lmstudio_chat(model_name="lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF", temperature=0)
     # chat_llm = models.get_openrouter_chat(model_name="openai/o1-mini-2024-09-12")
@@ -15,13 +15,14 @@ def initialize():
     # chat_llm = models.get_mistral_chat(model_name="mistral-small-latest", temperature=0)
     # chat_llm = models.get_groq_chat(model_name="llama-3.2-90b-text-preview", temperature=0)
     # chat_llm = models.get_sambanova_chat(model_name="Meta-Llama-3.1-70B-Instruct-8k", temperature=0)
+     chat_llm = models.get_bedrock_chat(model_id = "anthropic.claude-3-sonnet-20240229-v1:0")
 
     # utility model used for helper functions (cheaper, faster)
     utility_llm = chat_llm
 
     # embedding model used for memory
-    embedding_llm = models.get_openai_embedding(model_name="text-embedding-3-small")
-    # embedding_llm = models.get_ollama_embedding(model_name="nomic-embed-text")
+    # embedding_llm = models.get_openai_embedding(model_name="text-embedding-3-small")
+     embedding_llm = models.get_ollama_embedding(model_name="nomic-embed-text")
     # embedding_llm = models.get_huggingface_embedding(model_name="sentence-transformers/all-MiniLM-L6-v2")
     # embedding_llm = models.get_lmstudio_embedding(model_name="nomic-ai/nomic-embed-text-v1.5-GGUF")
 

models.py

@@ -8,10 +8,15 @@
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_google_genai import GoogleGenerativeAI, HarmBlockThreshold, HarmCategory
 from langchain_mistralai import ChatMistralAI
+import boto3
+from langchain_aws import ChatBedrock
 from pydantic.v1.types import SecretStr
 from python.helpers.dotenv import load_dotenv
 
-# environment variables
+# environment variables, clean out old AWS creds and use ones in .env
+for key in ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN"]:
+    if key in os.environ:
+        del os.environ[key]
 load_dotenv()
 
 # Configuration
@@ -88,3 +93,14 @@ def get_openrouter_embedding(model_name: str, api_key=get_api_key("openrouter"),
 def get_sambanova_chat(model_name: str, api_key=get_api_key("sambanova"), temperature=DEFAULT_TEMPERATURE, base_url=os.getenv("SAMBANOVA_BASE_URL") or "https://fast-api.snova.ai/v1", max_tokens=1024):
     return ChatOpenAI(api_key=api_key, model=model_name, temperature=temperature, base_url=base_url, max_tokens=max_tokens) # type: ignore
 
+# Bedrock models
+def get_bedrock_chat(model_id: str):
+    session = boto3.Session(
+        region_name = os.getenv("REGION_NAME")
+        aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID"), 
+        aws_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY"),
+        aws_session_token = os.getenv("AWS_SESSION_TOKEN")
+    )
+    bedrock_client = session.client(service_name="bedrock-runtime")
+
+    return ChatBedrock(client=bedrock_client, model_id=model_id)

requirements.txt

@@ -23,4 +23,6 @@ python-dotenv==1.0.1
 sentence-transformers==3.0.1
 unstructured==0.15.13
 unstructured-client==0.25.9
-webcolors==24.6.0
+webcolors==24.6.0
+langchain_aws==0.1.7
+boto3