From a4e51362b51bc264309456884cb98b0fb2182ac6 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 9 Jan 2024 13:40:37 +0530
Subject: [PATCH] docs(streaming_logging.md): fix tutorial

---
 .../docs/proxy/streaming_logging.md           | 30 +++----------------
 1 file changed, 4 insertions(+), 26 deletions(-)

diff --git a/docs/my-website/docs/proxy/streaming_logging.md b/docs/my-website/docs/proxy/streaming_logging.md
index 48d813b503..6bc5882d1f 100644
--- a/docs/my-website/docs/proxy/streaming_logging.md
+++ b/docs/my-website/docs/proxy/streaming_logging.md
@@ -1,4 +1,4 @@
-# Track Token Usage (Streaming) 
+# Track Token Usage
 
 ### Step 1 - Create your custom `litellm` callback class
 We use `litellm.integrations.custom_logger` for this, **more details about litellm custom callbacks [here](https://docs.litellm.ai/docs/observability/custom_callback)**
@@ -25,34 +25,12 @@ class MyCustomHandler(CustomLogger):
                     datefmt='%Y-%m-%d %H:%M:%S'
             )
 
-            # check if it has collected an entire stream response
-            if "complete_streaming_response" in kwargs:
-                # for tracking streaming cost we pass the "messages" and the output_text to litellm.completion_cost 
-                completion_response=kwargs["complete_streaming_response"]
-                input_text = kwargs["messages"]
-                output_text = completion_response["choices"][0]["message"]["content"]
-                response_cost = litellm.completion_cost(
-                    model = kwargs["model"],
-                    messages = input_text,
-                    completion=output_text
-                )
-                print("streaming response_cost", response_cost)
-                logging.info(f"Model {kwargs['model']} Cost: ${response_cost:.8f}")
-
-            # for non streaming responses
-            else:
-                # we pass the completion_response obj
-                if kwargs["stream"] != True:
-                    response_cost = litellm.completion_cost(completion_response=completion_response)
-                    print("regular response_cost", response_cost)
-                    logging.info(f"Model {completion_response.model} Cost: ${response_cost:.8f}")
+            response_cost = litellm.completion_cost(completion_response=completion_response)
+            print("regular response_cost", response_cost)
+            logging.info(f"Model {completion_response.model} Cost: ${response_cost:.8f}")
         except:
             pass
 
-
-    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): 
-        print(f"On Async Failure")
-
 proxy_handler_instance = MyCustomHandler()
 
 # Set litellm.callbacks = [proxy_handler_instance] on the proxy