Add AWS4Auth to opensearch client

chat/src/setup.py

@@ -2,11 +2,13 @@
 from langchain.chat_models import AzureChatOpenAI
 from langchain_community.embeddings import SagemakerEndpointEmbeddings
 from langchain_community.vectorstores import OpenSearchVectorSearch
+from opensearchpy import RequestsHttpConnection
+from requests_aws4auth import AWS4Auth
 import os
 import boto3
 
 def prefix(value):
-    env_prefix = os.getenv('ENV_PREFIX')
+    env_prefix = os.getenv("ENV_PREFIX")
     env_prefix = None if env_prefix == "" else env_prefix
     return '-'.join(filter(None, [env_prefix, value]))
 
@@ -17,20 +19,23 @@ def openai_chat_client(**kwargs):
     )
 
 def opensearch_client():
-    session = boto3.Session()
+    session = boto3.Session(region_name=os.getenv("AWS_REGION"))
+    awsauth = AWS4Auth(region=session.region_name, service="es", refreshable_credentials=session.get_credentials())
 
     sagemaker_client = session.client(service_name="sagemaker-runtime", region_name=session.region_name)
     embeddings = SagemakerEndpointEmbeddings(
       client=sagemaker_client,
       region_name=session.region_name,
-      endpoint_name="nul-embedding",
+      endpoint_name=os.getenv("EMBEDDING_ENDPOINT"),
       content_handler=ContentHandler()
     )
 
     docsearch = OpenSearchVectorSearch(
         index_name=prefix("dc-v2-work"),
         embedding_function=embeddings,
         opensearch_url="https://" + os.getenv("ELASTICSEARCH_ENDPOINT"),
+        connection_class=RequestsHttpConnection,
+        http_auth=awsauth,
     )
     return docsearch