From c58781460503bafb07d9a1c2517be4937076ec09 Mon Sep 17 00:00:00 2001
From: Manuel Fittko <manuel.fittko@sofatutor.com>
Date: Wed, 11 Sep 2024 11:28:03 +0200
Subject: [PATCH] add logging and caching for assistants

---
 litellm/proxy/proxy_server.py | 30 ++++++++++++++++++++++++++++++
 litellm/proxy/utils.py        |  1 +
 2 files changed, 31 insertions(+)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 341f2b5d623f..bd75bc800b79 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -4202,6 +4202,19 @@ async def get_assistants(
             proxy_config=proxy_config,
         )
 
+        # Check cache
+        cache_key = f"assistants:{user_api_key_dict.user_id}"
+        cached_response = await proxy_logging_obj.get_cache(cache_key)
+        if cached_response:
+            return cached_response
+
+        # Log the request
+        await proxy_logging_obj.pre_call_hook(
+            user_api_key_dict=user_api_key_dict,
+            data=data,
+            call_type="assistants",
+        )
+
         # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch
         if llm_router is None:
             raise HTTPException(
@@ -4209,6 +4222,17 @@ async def get_assistants(
             )
         response = await llm_router.aget_assistants(**data)
 
+        # Cache the response
+        await proxy_logging_obj.set_cache(cache_key, response, ttl=300)  # Cache for 5 minutes
+
+        ### LOGGING ###
+        await proxy_logging_obj.success_call_hook(
+            user_api_key_dict=user_api_key_dict,
+            data=data,
+            response=response,
+            call_type="assistants",
+        )
+
         ### ALERTING ###
         asyncio.create_task(
             proxy_logging_obj.update_request_status(
@@ -4236,6 +4260,12 @@ async def get_assistants(
 
         return response
     except Exception as e:
+        await proxy_logging_obj.failure_call_hook(
+            user_api_key_dict=user_api_key_dict,
+            data=data,
+            error=e,
+            call_type="assistants",
+        )
         await proxy_logging_obj.post_call_failure_hook(
             user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
         )
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 4031999d815e..44c2db5108a9 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -376,6 +376,7 @@ async def pre_call_hook(
             "audio_transcription",
             "pass_through_endpoint",
             "rerank",
+            "assistants",
         ],
     ) -> dict:
         """