From e82b8ed7e2b0f664d72e1ba43509726123456ca9 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 14 Nov 2023 17:53:48 -0800
Subject: [PATCH] (feat) debug ollama POST request

---
 litellm/llms/ollama.py | 18 ++++++++++++++++--
 litellm/main.py        |  7 ++-----
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py
index e4639d0731..a24e47c074 100644
--- a/litellm/llms/ollama.py
+++ b/litellm/llms/ollama.py
@@ -113,7 +113,8 @@ def get_ollama_response_stream(
         api_base="http://localhost:11434",
         model="llama2",
         prompt="Why is the sky blue?", 
-        optional_params=None
+        optional_params=None,
+        logging_obj=None,
     ):
     if api_base.endswith("/api/generate"):
         url = api_base
@@ -131,6 +132,12 @@ def get_ollama_response_stream(
         "prompt": prompt,
         **optional_params
     }
+    ## LOGGING
+    logging_obj.pre_call(
+        input=None,
+        api_key=None,
+        additional_args={"api_base": url, "complete_input_dict": data},
+    )
     session = requests.Session()
 
     with session.post(url, json=data, stream=True) as resp:
@@ -169,7 +176,8 @@ if async_generator_imported:
             api_base="http://localhost:11434",
             model="llama2",
             prompt="Why is the sky blue?",
-            optional_params=None
+            optional_params=None,
+            logging_obj=None,
         ):
         url = f"{api_base}/api/generate"
         
@@ -184,6 +192,12 @@ if async_generator_imported:
             "prompt": prompt,
             **optional_params
         }
+        ## LOGGING
+        logging_obj.pre_call(
+            input=None,
+            api_key=None,
+            additional_args={"api_base": url, "complete_input_dict": data},
+        )
         session = requests.Session()
 
         with session.post(url, json=data, stream=True) as resp:
diff --git a/litellm/main.py b/litellm/main.py
index 2299a45701..d168422c80 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -1235,16 +1235,13 @@ def completion(
                 prompt = prompt_factory(model=model, messages=messages, custom_llm_provider=custom_llm_provider)
 
             ## LOGGING
-            logging.pre_call(
-                input=prompt, api_key=None, additional_args={"api_base": api_base, "custom_prompt_dict": custom_prompt_dict}
-            )
             if kwargs.get('acompletion', False) == True:    
                 if optional_params.get("stream", False) == True:
                 # assume all ollama responses are streamed
-                    async_generator = ollama.async_get_ollama_response_stream(api_base, model, prompt, optional_params)
+                    async_generator = ollama.async_get_ollama_response_stream(api_base, model, prompt, optional_params, logging_obj=logging)
                     return async_generator
 
-            generator = ollama.get_ollama_response_stream(api_base, model, prompt, optional_params)
+            generator = ollama.get_ollama_response_stream(api_base, model, prompt, optional_params, logging_obj=logging)
             if optional_params.get("stream", False) == True:
                 # assume all ollama responses are streamed
                 response = CustomStreamWrapper(