Add CAII endpoint description retrieval and update related models

Author: baasitsharief

llm-service/app/routers/index/models/__init__.py

@@ -41,7 +41,8 @@
 
 from .... import exceptions
 from ....services import models
-from ....services.caii.types import ModelResponse
+from ....services.caii.caii import describe_endpoint, build_model_response
+from ....services.caii.types import ModelResponse, Endpoint
 
 router = APIRouter(prefix="/models", tags=["Models"])
 
@@ -72,6 +73,17 @@ def get_model() -> models.ModelSource:
     return models.get_model_source()
 
 
+@router.get(path="/caii/endpoint/{endpoint_name}", summary="Get CAII endpoint details.")
+@exceptions.propagates
+def get_endpoint_description(endpoint_name: str) -> ModelResponse:
+    """
+    Get the details of a specific CAII endpoint by its name.
+    """
+    endpoint = describe_endpoint(endpoint_name)
+    model_response = build_model_response(endpoint)
+    return model_response
+
+
 @router.get("/llm/{model_name}/test", summary="Test LLM Inference model.")
 @exceptions.propagates
 def llm_model_test(model_name: str) -> Literal["ok"]:

llm-service/app/services/caii/CaiiEmbeddingModel.py

@@ -42,17 +42,23 @@
 from llama_index.core.base.embeddings.base import BaseEmbedding, Embedding
 from pydantic import Field
 
-from .types import Endpoint
+from .types import Endpoint, ListEndpointEntry
 from .utils import build_auth_headers
 
 
 class CaiiEmbeddingModel(BaseEmbedding):
-    endpoint: Endpoint = Field(
-        Endpoint, description="The endpoint to use for embeddings"
+    endpoint: ListEndpointEntry = Field(
+        ListEndpointEntry, description="The endpoint to use for embeddings"
+    )
+    http_client: httpx.Client = Field(
+        httpx.Client, description="The http client to use for requests"
     )
-    http_client: httpx.Client = Field(httpx.Client, description="The http client to use for requests")
 
-    def __init__(self, endpoint: Endpoint, http_client: httpx.Client | None = httpx.Client()):
+    def __init__(
+        self,
+        endpoint: ListEndpointEntry,
+        http_client: httpx.Client | None = httpx.Client(),
+    ):
         super().__init__()
         self.endpoint = endpoint
         self.http_client = http_client or httpx.Client()
@@ -86,7 +92,9 @@ def _get_embedding(self, query: str, input_type: str) -> Embedding:
     def make_embedding_request(self, body: str) -> Any:
         headers = build_auth_headers()
         headers["Content-Type"] = "application/json"
-        response = self.http_client.post(url=self.endpoint.url, content=body, headers=headers)
+        response = self.http_client.post(
+            url=self.endpoint.url, content=body, headers=headers
+        )
         res = response.content
         json_response = res.decode("utf-8")
         structured_response = json.loads(json_response)

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

@@ -62,30 +62,37 @@
 logger = logging.getLogger(__name__)
 
 
-def describe_endpoint(
+def describe_endpoint_entry(
     endpoint_name: str, endpoints: Optional[List[ListEndpointEntry]] = None
-) -> Endpoint:
+) -> ListEndpointEntry:
     if not endpoints:
-        logger.info("Fetching endpoint details from CAII REST API")
-        domain = settings.caii_domain
-        headers = build_auth_headers()
-        describe_url = f"https://{domain}/api/v1alpha1/describeEndpoint"
-        desc_json = {"name": endpoint_name, "namespace": DEFAULT_NAMESPACE}
-
-        response = requests.post(describe_url, headers=headers, json=desc_json)
-        raise_for_http_error(response)
-        return Endpoint(**body_to_json(response))
+        endpoint = describe_endpoint(endpoint_name)
+        return ListEndpointEntry(**endpoint.model_dump())
 
     logger.info("Fetching endpoint details from cached list of endpoints")
     for endpoint in endpoints:
         if endpoint.name == endpoint_name:
-            return Endpoint(**endpoint.model_dump())
+            return ListEndpointEntry(**endpoint.model_dump())
 
     raise HTTPException(
         status_code=404, detail=f"Endpoint '{endpoint_name}' not found."
     )
 
 
+def describe_endpoint(
+    endpoint_name: str,
+) -> Endpoint:
+    logger.info("Fetching endpoint details from CAII REST API")
+    domain = settings.caii_domain
+    headers = build_auth_headers()
+    describe_url = f"https://{domain}/api/v1alpha1/describeEndpoint"
+    desc_json = {"name": endpoint_name, "namespace": DEFAULT_NAMESPACE}
+
+    response = requests.post(describe_url, headers=headers, json=desc_json)
+    raise_for_http_error(response)
+    return Endpoint(**body_to_json(response))
+
+
 def list_endpoints() -> list[ListEndpointEntry]:
     try:
         import cmlapi
@@ -123,7 +130,7 @@ def list_endpoints() -> list[ListEndpointEntry]:
 
 
 def get_reranking_model(model_name: str, top_n: int) -> BaseNodePostprocessor:
-    endpoint = describe_endpoint(endpoint_name=model_name)
+    endpoint = describe_endpoint_entry(endpoint_name=model_name)
     token = get_caii_access_token()
     return CaiiRerankingModel(
         model=endpoint.model_name,
@@ -167,7 +174,7 @@ def get_llm(
 
 def get_embedding_model(model_name: str) -> BaseEmbedding:
     endpoint_name = model_name
-    endpoint = describe_endpoint(endpoint_name=endpoint_name)
+    endpoint = describe_endpoint_entry(endpoint_name=endpoint_name)
 
     if os.path.exists("/etc/ssl/certs/ca-certificates.crt"):
         http_client = httpx.Client(verify="/etc/ssl/certs/ca-certificates.crt")
@@ -223,7 +230,10 @@ def get_caii_embedding_models() -> List[ModelResponse]:
 def get_models_with_task(task_type: str) -> List[Endpoint]:
     endpoints = list_endpoints()
     endpoint_details = list(
-        map(lambda endpoint: describe_endpoint(endpoint.name, endpoints), endpoints)
+        map(
+            lambda endpoint: describe_endpoint_entry(endpoint.name, endpoints),
+            endpoints,
+        )
     )
     llm_endpoints = list(
         filter(

ui/src/api/modelsApi.ts

@@ -147,6 +147,21 @@ const getModelSource = async (): Promise<ModelSource> => {
   return await getRequest(`${llmServicePath}/models/model_source`);
 };
 
+export const useGetCAIIModelStatus = (endpoint_name: string) => {
+  return useQuery({
+    queryKey: [QueryKeys.getCAIIModelStatus, { endpoint_name }],
+    queryFn: async () => {
+      return await getCAIIModelStatus(endpoint_name);
+    },
+  });
+};
+
+const getCAIIModelStatus = async (endpoint_name: string): Promise<Model> => {
+  return await getRequest(
+    `${llmServicePath}/models/caii/endpoint/${endpoint_name}`,
+  );
+};
+
 export const useTestLlmModel = ({
   onSuccess,
   onError,

ui/src/api/utils.ts

@@ -114,6 +114,7 @@ export enum QueryKeys {
   "getAmpConfig" = "getAmpConfig",
   "getTools" = "getTools",
   "getPollingAmpConfig" = "getPollingAmpConfig",
+  "getCAIIModelStatus" = "getCAIIModelStatus",
 }
 
 export const commonHeaders = {