From a0cb56efd5e092f0b6219b77b8ee2e680f54fc0f Mon Sep 17 00:00:00 2001
From: adriensas <adrien.sas@gmail.com>
Date: Wed, 23 Aug 2023 10:59:36 +0200
Subject: [PATCH] fix

---
 litellm/llms/anthropic.py | 36 +++++++++++-------------------------
 1 file changed, 11 insertions(+), 25 deletions(-)

diff --git a/litellm/llms/anthropic.py b/litellm/llms/anthropic.py
index 983d4ca0e5..28623fdd6d 100644
--- a/litellm/llms/anthropic.py
+++ b/litellm/llms/anthropic.py
@@ -15,15 +15,11 @@ class AnthropicError(Exception):
     def __init__(self, status_code, message):
         self.status_code = status_code
         self.message = message
-        super().__init__(
-            self.message
-        )  # Call the base class constructor with the parameters it needs
+        super().__init__(self.message)  # Call the base class constructor with the parameters it needs
 
 
 class AnthropicLLM:
-    def __init__(
-        self, encoding, default_max_tokens_to_sample, logging_obj, api_key=None
-    ):
+    def __init__(self, encoding, default_max_tokens_to_sample, logging_obj, api_key=None):
         self.encoding = encoding
         self.default_max_tokens_to_sample = default_max_tokens_to_sample
         self.completion_url = "https://api.anthropic.com/v1/complete"
@@ -31,9 +27,7 @@ class AnthropicLLM:
         self.logging_obj = logging_obj
         self.validate_environment(api_key=api_key)
 
-    def validate_environment(
-        self, api_key
-    ):  # set up the environment required to run the model
+    def validate_environment(self, api_key):  # set up the environment required to run the model
         # set the api key
         if self.api_key == None:
             raise ValueError(
@@ -62,19 +56,13 @@ class AnthropicLLM:
         for message in messages:
             if "role" in message:
                 if message["role"] == "user":
-                    prompt += (
-                        f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}"
-                    )
+                    prompt += f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}"
                 else:
-                    prompt += (
-                        f"{AnthropicConstants.AI_PROMPT.value}{message['content']}"
-                    )
+                    prompt += f"{AnthropicConstants.AI_PROMPT.value}{message['content']}"
             else:
                 prompt += f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}"
         prompt += f"{AnthropicConstants.AI_PROMPT.value}"
-        if "max_tokens" in optional_params and optional_params["max_tokens"] != float(
-            "inf"
-        ):
+        if "max_tokens" in optional_params and optional_params["max_tokens"] != float("inf"):
             max_tokens = optional_params["max_tokens"]
         else:
             max_tokens = self.default_max_tokens_to_sample
@@ -93,9 +81,11 @@ class AnthropicLLM:
         )
         ## COMPLETION CALL
         response = requests.post(
-            self.completion_url, headers=self.headers, data=json.dumps(data)
+            self.completion_url, headers=self.headers, data=json.dumps(data), stream=optional_params["stream"]
         )
+        print(optional_params)
         if "stream" in optional_params and optional_params["stream"] == True:
+            print("IS STREAMING")
             return response.iter_lines()
         else:
             ## LOGGING
@@ -114,14 +104,10 @@ class AnthropicLLM:
                     status_code=response.status_code,
                 )
             else:
-                model_response["choices"][0]["message"][
-                    "content"
-                ] = completion_response["completion"]
+                model_response["choices"][0]["message"]["content"] = completion_response["completion"]
 
             ## CALCULATING USAGE
-            prompt_tokens = len(
-                self.encoding.encode(prompt)
-            )  ##[TODO] use the anthropic tokenizer here
+            prompt_tokens = len(self.encoding.encode(prompt))  ##[TODO] use the anthropic tokenizer here
             completion_tokens = len(
                 self.encoding.encode(model_response["choices"][0]["message"]["content"])
             )  ##[TODO] use the anthropic tokenizer here