Switch from Azure OpenAI to Bedrock

chat/src/agent/agent_handler.py

@@ -52,7 +52,8 @@ def on_tool_end(self, output: ToolMessage, **kwargs: Dict[str, Any]):
                 pass
             case "search":
                 try:
-                    docs: List[Dict[str, Any]] = [doc.metadata for doc in output.artifact]
+                    result_fields = ("id", "title", "visibility", "work_type", "thumbnail")
+                    docs: List[Dict[str, Any]] = [{k: doc.metadata.get(k) for k in result_fields} for doc in output.artifact]
                     self.socket.send({"type": "search_result", "ref": self.ref, "message": docs})
                 except json.decoder.JSONDecodeError as e:
                     print(f"Invalid json ({e}) returned from {output.name} tool: {output.content}")

chat/src/agent/search_agent.py

@@ -8,21 +8,21 @@
 from langchain_core.messages.system import SystemMessage
 from langgraph.graph import END, START, StateGraph, MessagesState
 from langgraph.prebuilt import ToolNode
-from setup import openai_chat_client
+from setup import chat_client
 
 # Keep your answer concise and keep reading time under 45 seconds.
 
 system_message = """
 Please provide a brief answer to the question using the tools provided. Include specific details from multiple documents that 
 support your answer. Answer in raw markdown, but not within a code block. When citing source documents, construct Markdown 
-links using the document's canonical_link field.
+links using the document's canonical_link field. Do not include intermediate messages explaining your process.
 """
 
 tools = [discover_fields, search, aggregate]
 
 tool_node = ToolNode(tools)
 
-model = openai_chat_client(streaming=True).bind_tools(tools)
+model = chat_client(streaming=True).bind_tools(tools)
 
 # Define the function that determines whether to continue or not
 def should_continue(state: MessagesState) -> Literal["tools", END]:
@@ -38,7 +38,7 @@ def should_continue(state: MessagesState) -> Literal["tools", END]:
 # Define the function that calls the model
 def call_model(state: MessagesState):
     messages = [SystemMessage(content=system_message)] + state["messages"]
-    response: BaseMessage = model.invoke(messages, model=os.getenv("AZURE_OPENAI_LLM_DEPLOYMENT_ID"))
+    response: BaseMessage = model.invoke(messages) # , model=os.getenv("AZURE_OPENAI_LLM_DEPLOYMENT_ID")
     # We return a list, because this will get added to the existing list
     # if socket is not none and the response content is not an empty string
     return {"messages": [response]}

chat/src/handlers/chat.py

@@ -19,11 +19,9 @@
     "debug": [
         "answer",
         "attributes",
-        "azure_endpoint",
         "deployment_name",
         "is_superuser",
         "k",
-        "openai_api_version",
         "prompt",
         "question",
         "ref",
@@ -36,7 +34,6 @@
         "deployment_name",
         "is_superuser",
         "k",
-        "openai_api_version",
         "prompt",
         "question",
         "ref",
@@ -79,8 +76,8 @@ def handler(event, context):
     try:
         search_agent.invoke(
             {"messages": [HumanMessage(content=config.question)]},
-            config={"configurable": {"thread_id": config.ref}, "callbacks": callbacks, "metadata": {"model_deployment": os.getenv("AZURE_OPENAI_LLM_DEPLOYMENT_ID")}},
-            debug=False
+            config={"configurable": {"thread_id": config.ref}, "callbacks": callbacks},
+            debug=True
         )
     except Exception as e:
         print(f"Error: {e}")

chat/src/helpers/metrics.py

@@ -6,12 +6,10 @@ def debug_response(config, response, original_question):
     return {
         "answer": response,
         "attributes": config.attributes,
-        "azure_endpoint": config.azure_endpoint,
         "deployment_name": config.deployment_name,
         "is_dev_team": config.api_token.is_dev_team(),
         "is_superuser": config.api_token.is_superuser(),
         "k": config.k,
-        "openai_api_version": config.openai_api_version,
         "prompt": config.prompt_text,
         "question": config.question,
         "ref": config.ref,

chat/src/secrets.py

@@ -7,16 +7,12 @@ def load_secrets():
   EnvironmentMap = [
     ['API_TOKEN_SECRET', 'dcapi', 'api_token_secret'],
     ['OPENSEARCH_ENDPOINT', 'index', 'endpoint'],
-    ['OPENSEARCH_MODEL_ID', 'index', 'embedding_model'],
-    ['AZURE_OPENAI_API_KEY', 'azure_openai', 'api_key'],
-    ['AZURE_OPENAI_LLM_DEPLOYMENT_ID', 'azure_openai', 'llm_deployment_id'],
-    ['AZURE_OPENAI_RESOURCE_NAME', 'azure_openai', 'resource_name']
+    ['OPENSEARCH_MODEL_ID', 'index', 'embedding_model']
   ]
 
   client = boto3.client("secretsmanager")
   response = client.batch_get_secret_value(SecretIdList=[
     f'{SecretsPath}/infrastructure/index',
-    f'{SecretsPath}/infrastructure/azure_openai',
     f'{SecretsPath}/config/dcapi'
   ])
 

chat/src/setup.py

@@ -1,4 +1,4 @@
-from langchain_openai import AzureChatOpenAI
+from langchain_aws import ChatBedrock
 from handlers.opensearch_neural_search import OpenSearchNeuralSearch
 from opensearchpy import OpenSearch, RequestsHttpConnection
 from requests_aws4auth import AWS4Auth
@@ -13,11 +13,9 @@ def prefix(value):
     return "-".join(filter(None, [env_prefix, value]))
 
 
-def openai_chat_client(**kwargs):
-    return AzureChatOpenAI(
-        openai_api_key=os.getenv("AZURE_OPENAI_API_KEY"),
-        api_version=os.getenv("AZURE_API_VERSION", "2024-08-01-preview"),
-        azure_endpoint=f"https://{os.getenv("AZURE_OPENAI_RESOURCE_NAME")}.openai.azure.com",
+def chat_client(**kwargs):
+    return ChatBedrock(
+        model="us.anthropic.claude-3-5-sonnet-20241022-v2:0",
         **kwargs,
     )
 

chat/template.yaml

@@ -251,6 +251,12 @@ Resources:
           Resource:
           - !Sub "arn:aws:s3:::${CheckpointBucket}"
           - !Sub "arn:aws:s3:::${CheckpointBucket}/*"
+      - Statement:
+        - Effect: Allow
+          Action:
+          - bedrock:InvokeModel
+          - bedrock:InvokeModelWithResponseStream
+          Resource: "*"
     #* Metadata:
     #*   BuildMethod: nodejs20.x
   # ChatSyncFunction: