Merge pull request #1878 from BerriAI/litellm_improve_semantic_cache_…

…tracing [Feat] Semantic Caching - Track Cost of using embedding, Use Langfuse Trace ID
BerriAI · Feb 8, 2024 · 717dc78 · 717dc78
2 parents 98b0ace + 8197b3d
commit 717dc78
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/litellm/caching.py b/litellm/caching.py
@@ -427,10 +427,16 @@ async def async_set_cache(self, key, value, **kwargs):
             else []
         )
         if llm_router is not None and self.embedding_model in router_model_names:
+            user_api_key = kwargs.get("metadata", {}).get("user_api_key", "")
             embedding_response = await llm_router.aembedding(
                 model=self.embedding_model,
                 input=prompt,
                 cache={"no-store": True, "no-cache": True},
+                metadata={
+                    "user_api_key": user_api_key,
+                    "semantic-cache-embedding": True,
+                    "trace_id": kwargs.get("metadata", {}).get("trace_id", None),
+                },
             )
         else:
             # convert to embedding
@@ -476,10 +482,16 @@ async def async_get_cache(self, key, **kwargs):
             else []
         )
         if llm_router is not None and self.embedding_model in router_model_names:
+            user_api_key = kwargs.get("metadata", {}).get("user_api_key", "")
             embedding_response = await llm_router.aembedding(
                 model=self.embedding_model,
                 input=prompt,
                 cache={"no-store": True, "no-cache": True},
+                metadata={
+                    "user_api_key": user_api_key,
+                    "semantic-cache-embedding": True,
+                    "trace_id": kwargs.get("metadata", {}).get("trace_id", None),
+                },
             )
         else:
             # convert to embedding