From 3aba9019ac27481361c3fa87956907f5c41dd4ca Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 2 Mar 2024 10:06:33 -0800
Subject: [PATCH] (feat) track used api_base in response

---
 litellm/router.py | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/litellm/router.py b/litellm/router.py
index 6f33d0b0d5c..59cbf43d56c 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -357,6 +357,11 @@ class Router:
             verbose_router_logger.info(
                 f"litellm.completion(model={model_name})\033[32m 200 OK\033[0m"
             )
+            # set used api_base in response
+            if hasattr(response, "_hidden_params"):
+                response._hidden_params["api_base"] = deployment.get(
+                    "litellm_params", {}
+                ).get("api_base", None)
             return response
         except Exception as e:
             verbose_router_logger.info(
@@ -446,6 +451,12 @@ class Router:
             verbose_router_logger.info(
                 f"litellm.acompletion(model={model_name})\033[32m 200 OK\033[0m"
             )
+
+            # set used api_base in response
+            if hasattr(response, "_hidden_params"):
+                response._hidden_params["api_base"] = deployment.get(
+                    "litellm_params", {}
+                ).get("api_base", None)
             return response
         except Exception as e:
             verbose_router_logger.info(
@@ -602,6 +613,12 @@ class Router:
             verbose_router_logger.info(
                 f"litellm.aimage_generation(model={model_name})\033[32m 200 OK\033[0m"
             )
+
+            # set used api_base in response
+            if hasattr(response, "_hidden_params"):
+                response._hidden_params["api_base"] = deployment.get(
+                    "litellm_params", {}
+                ).get("api_base", None)
             return response
         except Exception as e:
             verbose_router_logger.info(
@@ -694,6 +711,11 @@ class Router:
             verbose_router_logger.info(
                 f"litellm.amoderation(model={model_name})\033[32m 200 OK\033[0m"
             )
+            # set used api_base in response
+            if hasattr(response, "_hidden_params"):
+                response._hidden_params["api_base"] = deployment.get(
+                    "litellm_params", {}
+                ).get("api_base", None)
             return response
         except Exception as e:
             verbose_router_logger.info(
@@ -825,6 +847,12 @@ class Router:
             verbose_router_logger.info(
                 f"litellm.atext_completion(model={model_name})\033[32m 200 OK\033[0m"
             )
+
+            # set used api_base in response
+            if hasattr(response, "_hidden_params"):
+                response._hidden_params["api_base"] = deployment.get(
+                    "litellm_params", {}
+                ).get("api_base", None)
             return response
         except Exception as e:
             verbose_router_logger.info(
@@ -954,6 +982,12 @@ class Router:
             verbose_router_logger.info(
                 f"litellm.aembedding(model={model_name})\033[32m 200 OK\033[0m"
             )
+
+            # set used api_base in response
+            if hasattr(response, "_hidden_params"):
+                response._hidden_params["api_base"] = deployment.get(
+                    "litellm_params", {}
+                ).get("api_base", None)
             return response
         except Exception as e:
             verbose_router_logger.info(