From 12343650e4ffe15ae0aa47f3130c9aac94dd2a34 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 4 Jul 2024 13:13:42 -0700
Subject: [PATCH] feat - allow looking up model_id in model info

---
 .../common_utils/openai_endpoint_utils.py     | 21 +++++++++++
 litellm/proxy/proxy_server.py                 | 35 ++++++++++++++++---
 2 files changed, 51 insertions(+), 5 deletions(-)
 create mode 100644 litellm/proxy/common_utils/openai_endpoint_utils.py

diff --git a/litellm/proxy/common_utils/openai_endpoint_utils.py b/litellm/proxy/common_utils/openai_endpoint_utils.py
new file mode 100644
index 00000000000..4d3f4220b90
--- /dev/null
+++ b/litellm/proxy/common_utils/openai_endpoint_utils.py
@@ -0,0 +1,21 @@
+"""
+Contains utils used by OpenAI compatible endpoints 
+"""
+
+
+def remove_sensitive_info_from_deployment(deployment_dict: dict) -> dict:
+    """
+    Removes sensitive information from a deployment dictionary.
+
+    Args:
+        deployment_dict (dict): The deployment dictionary to remove sensitive information from.
+
+    Returns:
+        dict: The modified deployment dictionary with sensitive information removed.
+    """
+    deployment_dict["litellm_params"].pop("api_key", None)
+    deployment_dict["litellm_params"].pop("vertex_credentials", None)
+    deployment_dict["litellm_params"].pop("aws_access_key_id", None)
+    deployment_dict["litellm_params"].pop("aws_secret_access_key", None)
+
+    return deployment_dict
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 9f745bb54de..7794f2a4cd1 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -143,6 +143,9 @@ from litellm.proxy.caching_routes import router as caching_router
 from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router
 from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
 from litellm.proxy.common_utils.init_callbacks import initialize_callbacks_on_proxy
+from litellm.proxy.common_utils.openai_endpoint_utils import (
+    remove_sensitive_info_from_deployment,
+)
 from litellm.proxy.guardrails.init_guardrails import initialize_guardrails
 from litellm.proxy.health_check import perform_health_check
 from litellm.proxy.health_endpoints._health_endpoints import router as health_router
@@ -6657,14 +6660,39 @@ async def model_metrics_exceptions(
 )
 async def model_info_v1(
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+    litellm_model_id: Optional[str] = None,
 ):
-    global llm_model_list, general_settings, user_config_file_path, proxy_config
+    global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router
 
     if llm_model_list is None:
         raise HTTPException(
             status_code=500, detail={"error": "LLM Model List not loaded in"}
         )
 
+    if llm_router is None:
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "error": "LLM Router is not loaded in. Make sure you passed models in your config.yaml or on the LiteLLM Admin UI."
+            },
+        )
+
+    if litellm_model_id is not None:
+        # user is trying to get specific model from litellm router
+        deployment_info = llm_router.get_deployment(model_id=litellm_model_id)
+        if deployment_info is None:
+            raise HTTPException(
+                status_code=404,
+                detail={
+                    "error": f"Model id = {litellm_model_id} not found on litellm proxy"
+                },
+            )
+        _deployment_info_dict = deployment_info.model_dump()
+        _deployment_info_dict = remove_sensitive_info_from_deployment(
+            deployment_dict=_deployment_info_dict
+        )
+        return {"data": _deployment_info_dict}
+
     all_models: List[dict] = []
     ## CHECK IF MODEL RESTRICTIONS ARE SET AT KEY/TEAM LEVEL ##
     if llm_model_list is None:
@@ -6726,10 +6754,7 @@ async def model_info_v1(
                 model_info[k] = v
         model["model_info"] = model_info
         # don't return the llm credentials
-        model["litellm_params"].pop("api_key", None)
-        model["litellm_params"].pop("vertex_credentials", None)
-        model["litellm_params"].pop("aws_access_key_id", None)
-        model["litellm_params"].pop("aws_secret_access_key", None)
+        model = remove_sensitive_info_from_deployment(deployment_dict=model)
 
     verbose_proxy_logger.debug("all_models: %s", all_models)
     return {"data": all_models}