feat: Get ai agent from bedrock

Author: kieran-wilkinson-4

packages/slackBotFunction/app/services/bedrock.py

@@ -42,8 +42,10 @@ def query_bedrock(user_query: str, session_id: str = None) -> RetrieveAndGenerat
             "type": "KNOWLEDGE_BASE",
             "knowledgeBaseConfiguration": {
                 "knowledgeBaseId": config.KNOWLEDGEBASE_ID,
-                "modelArn": config.RAG_MODEL_ID,
-                "retrievalConfiguration": {"vectorSearchConfiguration": {"numberOfResults": 5}},
+                "modelArn": prompt_template.get("model_id", config.RAG_MODEL_ID),
+                "retrievalConfiguration": {
+                    "vectorSearchConfiguration": {"numberOfResults": 10, "overrideSearchType": "SEMANTIC"}
+                },
                 "generationConfiguration": {
                     "guardrailConfiguration": {
                         "guardrailId": config.GUARD_RAIL_ID,
@@ -58,16 +60,6 @@ def query_bedrock(user_query: str, session_id: str = None) -> RetrieveAndGenerat
                         }
                     },
                 },
-                "orchestrationConfiguration": {
-                    "inferenceConfig": {
-                        "textInferenceConfig": {
-                            **inference_config,
-                            "stopSequences": [
-                                "Human:",
-                            ],
-                        }
-                    },
-                },
             },
         },
     }

packages/slackBotFunction/app/services/prompt_loader.py

@@ -92,23 +92,28 @@ def load_prompt(prompt_name: str, prompt_version: str = None) -> dict:
 
         logger.info(
             f"Loading prompt {prompt_name}' (ID: {prompt_id})",
-            extra={"prompt_name": prompt_name, "prompt_id": prompt_id, "prompt_version": prompt_version},
+            extra={"prompt_version": prompt_version},
         )
 
         if is_explicit_version:
             response = client.get_prompt(promptIdentifier=prompt_id, promptVersion=selected_version)
         else:
             response = client.get_prompt(promptIdentifier=prompt_id)
 
+        logger.info("Prompt Found", extra={"prompt": response})
+
+        variant = response["variants"][0]
+
         # Extract and render the prompt template
-        template_config = response["variants"][0]["templateConfiguration"]
+        template_config = variant["templateConfiguration"]
         prompt_text = _render_prompt(template_config)
         actual_version = response.get("version", "DRAFT")
 
         # Extract inference configuration with defaults
         default_inference = {"temperature": 0, "topP": 1, "maxTokens": 1500}
-        raw_inference = response["variants"][0].get("inferenceConfiguration", {})
-        raw_text_config = raw_inference.get("textInferenceConfiguration", {})
+        model_id = variant.get("modelId", "")
+        raw_inference = variant.get("inferenceConfiguration", {})
+        raw_text_config = raw_inference.get("text", {})
         inference_config = {**default_inference, **raw_text_config}
 
         logger.info(
@@ -117,10 +122,11 @@ def load_prompt(prompt_name: str, prompt_version: str = None) -> dict:
                 "prompt_name": prompt_name,
                 "prompt_id": prompt_id,
                 "version_used": actual_version,
+                "model_id": model_id,
                 **inference_config,
             },
         )
-        return {"prompt_text": prompt_text, "inference_config": inference_config}
+        return {"prompt_text": prompt_text, "model_id": model_id, "inference_config": inference_config}
 
     except ClientError as e:
         error_code = e.response.get("Error", {}).get("Code", "Unknown")