Merge pull request #221 from cloudera/liu/refactor

Move model get logic to model provider classes

Author: ewilliams-cloudera

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

@@ -39,9 +39,6 @@
 
 from fastapi import HTTPException
 from llama_index.core.base.embeddings.base import BaseEmbedding
-from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
-from llama_index.embeddings.bedrock import BedrockEmbedding
-from llama_index.embeddings.openai import OpenAIEmbedding
 
 from . import _model_type, _noop
 from .providers import (
@@ -50,9 +47,7 @@
     CAIIModelProvider,
 )
 from .providers.openai import OpenAiModelProvider
-from ..caii.caii import get_embedding_model as caii_embedding
 from ..caii.types import ModelResponse
-from ...config import settings
 
 
 class Embedding(_model_type.ModelType[BaseEmbedding]):
@@ -62,24 +57,12 @@ def get(cls, model_name: Optional[str] = None) -> BaseEmbedding:
             model_name = cls.list_available()[0].model_id
 
         if AzureModelProvider.is_enabled():
-            return AzureOpenAIEmbedding(
-                model_name=model_name,
-                deployment_name=model_name,
-                # must be passed manually otherwise AzureOpenAIEmbedding checks OPENAI_API_KEY
-                api_key=settings.azure_openai_api_key,
-            )
-
+            return AzureModelProvider.get_embedding_model(model_name)
         if CAIIModelProvider.is_enabled():
-            return caii_embedding(model_name=model_name)
-
+            return CAIIModelProvider.get_embedding_model(model_name)
         if OpenAiModelProvider.is_enabled():
-            return OpenAIEmbedding(
-                model_name=model_name,
-                api_key=settings.openai_api_key,
-                api_base=settings.openai_api_base,
-            )
-
-        return BedrockEmbedding(model_name=model_name)
+            return OpenAiModelProvider.get_embedding_model(model_name)
+        return BedrockModelProvider.get_embedding_model(model_name)
 
     @staticmethod
     def get_noop() -> BaseEmbedding:
@@ -88,15 +71,12 @@ def get_noop() -> BaseEmbedding:
     @staticmethod
     def list_available() -> list[ModelResponse]:
         if AzureModelProvider.is_enabled():
-            return AzureModelProvider.get_embedding_models()
-
+            return AzureModelProvider.list_embedding_models()
         if CAIIModelProvider.is_enabled():
-            return CAIIModelProvider.get_embedding_models()
-
+            return CAIIModelProvider.list_embedding_models()
         if OpenAiModelProvider.is_enabled():
-            return OpenAiModelProvider.get_embedding_models()
-
-        return BedrockModelProvider.get_embedding_models()
+            return OpenAiModelProvider.list_embedding_models()
+        return BedrockModelProvider.list_embedding_models()
 
     @classmethod
     def test(cls, model_name: str) -> str:

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

@@ -40,9 +40,6 @@
 from fastapi import HTTPException
 from llama_index.core import llms
 from llama_index.core.base.llms.types import ChatMessage, MessageRole
-from llama_index.llms.azure_openai import AzureOpenAI
-from llama_index.llms.bedrock_converse import BedrockConverse
-from llama_index.llms.openai import OpenAI
 
 from . import _model_type, _noop
 from .providers import (
@@ -51,10 +48,7 @@
     CAIIModelProvider,
 )
 from .providers.openai import OpenAiModelProvider
-from ..caii.caii import get_llm as caii_llm
 from ..caii.types import ModelResponse
-from ..llama_utils import completion_to_prompt, messages_to_prompt
-from ...config import settings
 
 
 class LLM(_model_type.ModelType[llms.LLM]):
@@ -64,37 +58,12 @@ def get(cls, model_name: Optional[str] = None) -> llms.LLM:
             model_name = cls.list_available()[0].model_id
 
         if AzureModelProvider.is_enabled():
-            return AzureOpenAI(
-                model=model_name,
-                engine=model_name,
-                messages_to_prompt=messages_to_prompt,
-                completion_to_prompt=completion_to_prompt,
-                max_tokens=2048,
-            )
-
-        if OpenAiModelProvider.is_enabled():
-            return OpenAI(
-                model=model_name,
-                messages_to_prompt=messages_to_prompt,
-                completion_to_prompt=completion_to_prompt,
-                max_tokens=2048,
-                api_base=settings.openai_api_base,
-                api_key=settings.openai_api_key,
-            )
-
+            return AzureModelProvider.get_llm_model(model_name)
         if CAIIModelProvider.is_enabled():
-            return caii_llm(
-                endpoint_name=model_name,
-                messages_to_prompt=messages_to_prompt,
-                completion_to_prompt=completion_to_prompt,
-            )
-
-        return BedrockConverse(
-            model=model_name,
-            messages_to_prompt=messages_to_prompt,
-            completion_to_prompt=completion_to_prompt,
-            max_tokens=2048,
-        )
+            return CAIIModelProvider.get_llm_model(model_name)
+        if OpenAiModelProvider.is_enabled():
+            return OpenAiModelProvider.get_llm_model(model_name)
+        return BedrockModelProvider.get_llm_model(model_name)
 
     @staticmethod
     def get_noop() -> llms.LLM:
@@ -103,15 +72,12 @@ def get_noop() -> llms.LLM:
     @staticmethod
     def list_available() -> list[ModelResponse]:
         if AzureModelProvider.is_enabled():
-            return AzureModelProvider.get_llm_models()
-
+            return AzureModelProvider.list_llm_models()
         if CAIIModelProvider.is_enabled():
-            return CAIIModelProvider.get_llm_models()
-
+            return CAIIModelProvider.list_llm_models()
         if OpenAiModelProvider.is_enabled():
-            return OpenAiModelProvider.get_llm_models()
-
-        return BedrockModelProvider.get_llm_models()
+            return OpenAiModelProvider.list_llm_models()
+        return BedrockModelProvider.list_llm_models()
 
     @classmethod
     def test(cls, model_name: str) -> Literal["ok"]:

llm-service/app/services/models/providers/_model_provider.py

@@ -37,7 +37,10 @@
 #
 import abc
 import os
-from typing import List
+
+from llama_index.core.base.embeddings.base import BaseEmbedding
+from llama_index.core.llms import LLM
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
 
 from ...caii.types import ModelResponse
 
@@ -56,18 +59,36 @@ def get_env_var_names() -> set[str]:
 
     @staticmethod
     @abc.abstractmethod
-    def get_llm_models() -> List[ModelResponse]:
-        """Return available LLM models."""
+    def list_llm_models() -> list[ModelResponse]:
+        """Return names and IDs of available LLM models."""
+        raise NotImplementedError
+
+    @staticmethod
+    @abc.abstractmethod
+    def list_embedding_models() -> list[ModelResponse]:
+        """Return names and IDs of available embedding models."""
+        raise NotImplementedError
+
+    @staticmethod
+    @abc.abstractmethod
+    def list_reranking_models() -> list[ModelResponse]:
+        """Return names and IDs of available reranking models."""
+        raise NotImplementedError
+
+    @staticmethod
+    @abc.abstractmethod
+    def get_llm_model(name: str) -> LLM:
+        """Return LLM model with `name`."""
         raise NotImplementedError
 
     @staticmethod
     @abc.abstractmethod
-    def get_embedding_models() -> List[ModelResponse]:
-        """Return available embedding models."""
+    def get_embedding_model(name: str) -> BaseEmbedding:
+        """Return embedding model with `name`."""
         raise NotImplementedError
 
     @staticmethod
     @abc.abstractmethod
-    def get_reranking_models() -> List[ModelResponse]:
-        """Return available reranking models."""
+    def get_reranking_model(name: str, top_n: int) -> BaseNodePostprocessor:
+        """Return reranking model with `name`."""
         raise NotImplementedError

llm-service/app/services/models/providers/azure.py

@@ -35,11 +35,14 @@
 #  BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
 #  DATA.
 #
+from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
+from llama_index.llms.azure_openai import AzureOpenAI
 
-from typing import List
-
-from ...caii.types import ModelResponse
 from ._model_provider import ModelProvider
+from ...caii.types import ModelResponse
+from ...llama_utils import completion_to_prompt, messages_to_prompt
+from ...query.simple_reranker import SimpleReranker
+from ....config import settings
 
 
 class AzureModelProvider(ModelProvider):
@@ -48,7 +51,7 @@ def get_env_var_names() -> set[str]:
         return {"AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "OPENAI_API_VERSION"}
 
     @staticmethod
-    def get_llm_models() -> List[ModelResponse]:
+    def list_llm_models() -> list[ModelResponse]:
         return [
             ModelResponse(
                 model_id="gpt-4o",
@@ -61,7 +64,7 @@ def get_llm_models() -> List[ModelResponse]:
         ]
 
     @staticmethod
-    def get_embedding_models() -> List[ModelResponse]:
+    def list_embedding_models() -> list[ModelResponse]:
         return [
             ModelResponse(
                 model_id="text-embedding-ada-002",
@@ -74,9 +77,32 @@ def get_embedding_models() -> List[ModelResponse]:
         ]
 
     @staticmethod
-    def get_reranking_models() -> List[ModelResponse]:
+    def list_reranking_models() -> list[ModelResponse]:
         return []
 
+    @staticmethod
+    def get_llm_model(name: str) -> AzureOpenAI:
+        return AzureOpenAI(
+            model=name,
+            engine=name,
+            messages_to_prompt=messages_to_prompt,
+            completion_to_prompt=completion_to_prompt,
+            max_tokens=2048,
+        )
+
+    @staticmethod
+    def get_embedding_model(name: str) -> AzureOpenAIEmbedding:
+        return AzureOpenAIEmbedding(
+            model_name=name,
+            deployment_name=name,
+            # must be passed manually otherwise AzureOpenAIEmbedding checks OPENAI_API_KEY
+            api_key=settings.azure_openai_api_key,
+        )
+
+    @staticmethod
+    def get_reranking_model(name: str, top_n: int) -> SimpleReranker:
+        return SimpleReranker(top_n=top_n)
+
 
 # ensure interface is implemented
 _ = AzureModelProvider()

llm-service/app/services/models/providers/bedrock.py

@@ -35,13 +35,17 @@
 #  BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
 #  DATA.
 #
-from typing import List, Optional, cast
+from typing import Optional, cast
 
 import boto3
+from llama_index.embeddings.bedrock import BedrockEmbedding
+from llama_index.llms.bedrock_converse import BedrockConverse
+from llama_index.postprocessor.bedrock_rerank import AWSBedrockRerank
 
 from app.config import settings
-from ...caii.types import ModelResponse
 from ._model_provider import ModelProvider
+from ...caii.types import ModelResponse
+from ...llama_utils import completion_to_prompt, messages_to_prompt
 
 DEFAULT_BEDROCK_LLM_MODEL = "meta.llama3-1-8b-instruct-v1:0"
 DEFAULT_BEDROCK_RERANK_MODEL = "cohere.rerank-v3-5:0"
@@ -53,7 +57,7 @@ def get_env_var_names() -> set[str]:
         return {"AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_DEFAULT_REGION"}
 
     @staticmethod
-    def get_llm_models() -> List[ModelResponse]:
+    def list_llm_models() -> list[ModelResponse]:
         models = [
             ModelResponse(
                 model_id=DEFAULT_BEDROCK_LLM_MODEL, name="Llama3.1 8B Instruct v1"
@@ -91,7 +95,8 @@ def get_llm_models() -> List[ModelResponse]:
 
     @staticmethod
     def _get_model_arn_by_profiles(
-        suffix: str, profiles: List[dict[str, str]]
+        suffix: str,
+        profiles: list[dict[str, str]],
     ) -> Optional[ModelResponse]:
         for profile in profiles:
             if profile["inferenceProfileId"].endswith(suffix):
@@ -102,13 +107,16 @@ def _get_model_arn_by_profiles(
         return None
 
     @staticmethod
-    def _get_model_arns() -> List[dict[str, str]]:
-        bedrock_client = boto3.client("bedrock", region_name=settings.aws_default_region)
+    def _get_model_arns() -> list[dict[str, str]]:
+        bedrock_client = boto3.client(
+            "bedrock",
+            region_name=settings.aws_default_region,
+        )
         profiles = bedrock_client.list_inference_profiles()["inferenceProfileSummaries"]
-        return cast(List[dict[str, str]], profiles)
+        return cast(list[dict[str, str]], profiles)
 
     @staticmethod
-    def get_embedding_models() -> List[ModelResponse]:
+    def list_embedding_models() -> list[ModelResponse]:
         return [
             ModelResponse(
                 model_id="cohere.embed-english-v3",
@@ -121,7 +129,7 @@ def get_embedding_models() -> List[ModelResponse]:
         ]
 
     @staticmethod
-    def get_reranking_models() -> List[ModelResponse]:
+    def list_reranking_models() -> list[ModelResponse]:
         return [
             ModelResponse(
                 model_id=DEFAULT_BEDROCK_RERANK_MODEL,
@@ -133,6 +141,23 @@ def get_reranking_models() -> List[ModelResponse]:
             ),
         ]
 
+    @staticmethod
+    def get_llm_model(name: str) -> BedrockConverse:
+        return BedrockConverse(
+            model=name,
+            messages_to_prompt=messages_to_prompt,
+            completion_to_prompt=completion_to_prompt,
+            max_tokens=2048,
+        )
+
+    @staticmethod
+    def get_embedding_model(name: str) -> BedrockEmbedding:
+        return BedrockEmbedding(model_name=name)
+
+    @staticmethod
+    def get_reranking_model(name: str, top_n: int) -> AWSBedrockRerank:
+        return AWSBedrockRerank(rerank_model_name=name, top_n=top_n)
+
 
 # ensure interface is implemented
 _ = BedrockModelProvider()