Add support for Bedrock ARNs for regional support (#149)

* fix issue with async call deep within the mistral library

* simplify the async stuff to encapsulate in the evaluators

* get model arn by suffix

* merge main and resolve conflicts

* fix pytest

* fix mypy

---------

Co-authored-by: jwatson <[email protected]>

Author: ewilliams-cloudera

llm-service/app/services/evaluators.py

@@ -35,36 +35,53 @@
 #  BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
 #  DATA.
 # ##############################################################################
+import asyncio
 
 from llama_index.core.base.response.schema import Response
 from llama_index.core.chat_engine.types import AgentChatResponse
-from llama_index.core.evaluation import FaithfulnessEvaluator, RelevancyEvaluator
+from llama_index.core.evaluation import (
+    FaithfulnessEvaluator,
+    RelevancyEvaluator,
+    EvaluationResult,
+)
+from llama_index.core.llms import LLM
 
 from ..services import models
 
 
 def evaluate_response(
-    query: str, chat_response: AgentChatResponse, model_name: str
+        query: str, chat_response: AgentChatResponse, model_name: str
 ) -> tuple[float, float]:
     # todo: pass in the correct llm model and use it, rather than requiring querying for it like this.
     evaluator_llm = models.LLM.get(model_name)
+    return asyncio.run(_async_evaluate_response(query, chat_response, evaluator_llm))
 
-    relevancy_evaluator = RelevancyEvaluator(llm=evaluator_llm)
-    relevance = relevancy_evaluator.evaluate_response(
+
+async def _async_evaluate_response(query: str, chat_response: AgentChatResponse, evaluator_llm: LLM) -> tuple[float, float] :
+    relevance = await _evaluate_relevancy(chat_response, evaluator_llm, query)
+    faithfulness = await _evaluate_faithfulness(chat_response, evaluator_llm, query)
+    return relevance.score or 0, faithfulness.score or 0
+
+
+async def _evaluate_faithfulness(chat_response: AgentChatResponse, evaluator_llm: LLM, query: str) -> EvaluationResult:
+    faithfulness_evaluator = FaithfulnessEvaluator(llm=evaluator_llm)
+    return await faithfulness_evaluator.aevaluate_response(
         query=query,
         response=Response(
             response=chat_response.response,
             source_nodes=chat_response.source_nodes,
             metadata=chat_response.metadata,
         ),
     )
-    faithfulness_evaluator = FaithfulnessEvaluator(llm=evaluator_llm)
-    faithfulness = faithfulness_evaluator.evaluate_response(
+
+
+async def _evaluate_relevancy(chat_response: AgentChatResponse, evaluator_llm: LLM, query: str) -> EvaluationResult:
+    relevancy_evaluator = RelevancyEvaluator(llm=evaluator_llm)
+    return await relevancy_evaluator.aevaluate_response(
         query=query,
         response=Response(
             response=chat_response.response,
             source_nodes=chat_response.source_nodes,
             metadata=chat_response.metadata,
         ),
     )
-    return relevance.score or 0, faithfulness.score or 0

llm-service/app/services/models/_bedrock.py

@@ -35,7 +35,10 @@
 #  BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
 #  DATA.
 #
-from typing import List
+import os
+from typing import List, Optional
+
+import boto3
 
 from app.services.caii.types import ModelResponse
 from ._model_provider import ModelProvider
@@ -51,20 +54,43 @@ def get_env_var_names() -> set[str]:
 
     @staticmethod
     def get_llm_models() -> List[ModelResponse]:
-        return [
+        models = [
             ModelResponse(
-                model_id=DEFAULT_BEDROCK_LLM_MODEL,
-                name="Llama3.1 8B Instruct v1",
+                model_id=DEFAULT_BEDROCK_LLM_MODEL, name="Llama3.1 8B Instruct v1"
             ),
             ModelResponse(
                 model_id="meta.llama3-1-70b-instruct-v1:0",
                 name="Llama3.1 70B Instruct v1",
             ),
             ModelResponse(
-                model_id="cohere.command-r-plus-v1:0",
-                name="Cohere Command R Plus v1",
+                model_id="cohere.command-r-plus-v1:0", name="Cohere Command R Plus v1"
             ),
         ]
+        llama323b = BedrockModelProvider._get_model_arn_by_suffix(
+            "meta.llama3-2-3b-instruct-v1:0"
+        )
+        if llama323b:
+            models.append(llama323b)
+        llama321b = BedrockModelProvider._get_model_arn_by_suffix(
+            "meta.llama3-2-1b-instruct-v1:0"
+        )
+        if llama321b:
+            models.append(llama321b)
+
+        return models
+
+    @staticmethod
+    def _get_model_arn_by_suffix(suffix: str) -> Optional[ModelResponse]:
+        default_region = os.environ.get("AWS_DEFAULT_REGION") or None
+        bedrock_client = boto3.client("bedrock", region_name=default_region)
+        profiles = bedrock_client.list_inference_profiles()["inferenceProfileSummaries"]
+        for profile in profiles:
+            if profile["inferenceProfileId"].endswith(suffix):
+                return ModelResponse(
+                    model_id=profile["inferenceProfileId"],
+                    name=profile["inferenceProfileName"],
+                )
+        return None
 
     @staticmethod
     def get_embedding_models() -> List[ModelResponse]:

llm-service/app/tests/conftest.py

@@ -52,9 +52,11 @@
 
 from app.ai.vector_stores.qdrant import QdrantVectorStore
 from app.main import app
+from app.services.caii.types import ModelResponse
 from app.services.metadata_apis import data_sources_metadata_api
 from app.services import models
 from app.services.metadata_apis.data_sources_metadata_api import RagDataSource
+from app.services.models._bedrock import BedrockModelProvider
 
 
 @dataclass
@@ -215,3 +217,12 @@ def client() -> Iterator[TestClient]:
     """
     with TestClient(app) as test_client:
         yield test_client
+
+
+@pytest.fixture(autouse=True)
+def _get_model_arn_by_suffix(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        BedrockModelProvider,
+        "_get_model_arn_by_suffix",
+        lambda name: ModelResponse(model_id=f"us.{name}", name=name),
+    )