Update embedding model

pagopa · Dec 10, 2024 · 695dfa8 · 695dfa8
1 parent 979f9a8
commit 695dfa8
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/apps/chatbot/.env.example b/apps/chatbot/.env.example
@@ -5,7 +5,8 @@ AUTH_COGNITO_ISSUER=https://cognito-idp.eu-south-1.amazonaws.com/eu-south-1_xxxx
 AUTH_DISABLE_SIGNUP=false
 AUTH_DISABLE_USERNAME_PASSWORD=true
 CHB_AWS_ACCESS_KEY_ID=...
-CHB_AWS_BEDROCK_REGION=eu-west-3
+CHB_AWS_BEDROCK_EMBED_REGION=eu-central-1
+CHB_AWS_BEDROCK_LLM_REGION=eu-west-3
 CHB_AWS_DEFAULT_REGION=eu-south-1
 CHB_AWS_GUARDRAIL_ID=...
 CHB_AWS_GUARDRAIL_VERSION=...

diff --git a/apps/chatbot/src/modules/models.py b/apps/chatbot/src/modules/models.py
@@ -21,7 +21,8 @@
 GOOGLE_API_KEY = get_ssm_parameter(name=os.getenv("CHB_GOOGLE_API_KEY"))
 AWS_ACCESS_KEY_ID = os.getenv("CHB_AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("CHB_AWS_SECRET_ACCESS_KEY")
-AWS_BEDROCK_REGION = os.getenv("CHB_AWS_BEDROCK_REGION")
+AWS_BEDROCK_LLM_REGION = os.getenv("CHB_AWS_BEDROCK_LLM_REGION")
+AWS_BEDROCK_EMBED_REGION = os.getenv("CHB_AWS_BEDROCK_EMBED_REGION")
 AWS_GUARDRAIL_ID = os.getenv("CHB_AWS_GUARDRAIL_ID")
 AWS_GUARDRAIL_VERSION = os.getenv("CHB_AWS_GUARDRAIL_VERSION")
 
@@ -41,7 +42,7 @@ def get_llm():
             max_tokens=int(MODEL_MAXTOKENS),
             aws_access_key_id=AWS_ACCESS_KEY_ID,
             aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
-            region_name=AWS_BEDROCK_REGION
+            region_name=AWS_BEDROCK_LLM_REGION
         )
 
     else:
@@ -71,7 +72,7 @@ def get_embed_model():
             model_name = EMBED_MODEL_ID,
             aws_access_key_id=AWS_ACCESS_KEY_ID,
             aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
-            region_name=AWS_BEDROCK_REGION
+            region_name=AWS_BEDROCK_EMBED_REGION
         )
     else:
         embed_model = GeminiEmbedding(

diff --git a/apps/chatbot/src/modules/vector_database.py b/apps/chatbot/src/modules/vector_database.py
@@ -55,7 +55,8 @@
 EMBED_MODEL_ID = os.getenv("CHB_EMBED_MODEL_ID")
 EMBEDDING_DIMS = {
     "models/text-embedding-004": 768,
-    "cohere.embed-multilingual-v3": 1024
+    "cohere.embed-multilingual-v3": 1024,
+    "amazon.titan-embed-text-v2:0": 1024
 }
 REDIS_SCHEMA = IndexSchema.from_dict({
     "index": {"name": f"{INDEX_ID}", "prefix": f"{INDEX_ID}/vector"},