From 75bcb37cb2d138ac39fd4240162dac2e1e62aade Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 13 Dec 2023 12:27:31 -0800
Subject: [PATCH] fix(factory.py): fix tgai rendering template

---
 litellm/llms/prompt_templates/factory.py | 14 +++++++++++---
 litellm/tests/test_completion.py         | 23 +++++++++++++++++++++--
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py
index 7f68ca77bf7..a88ba02a6b0 100644
--- a/litellm/llms/prompt_templates/factory.py
+++ b/litellm/llms/prompt_templates/factory.py
@@ -161,6 +161,8 @@ def phind_codellama_pt(messages):
 
 def hf_chat_template(model: str, messages: list, chat_template: Optional[Any]=None):
     ## get the tokenizer config from huggingface
+    bos_token = ""
+    eos_token = ""
     if chat_template is None: 
         def _get_tokenizer_config(hf_model_name):
             url = f"https://huggingface.co/{hf_model_name}/raw/main/tokenizer_config.json"
@@ -187,7 +189,10 @@ def hf_chat_template(model: str, messages: list, chat_template: Optional[Any]=No
     # Create a template object from the template text
     env = Environment()
     env.globals['raise_exception'] = raise_exception
-    template = env.from_string(chat_template)
+    try:
+        template = env.from_string(chat_template)
+    except Exception as e:
+        raise e
 
     def _is_system_in_template():
         try:
@@ -227,8 +232,8 @@ def hf_chat_template(model: str, messages: list, chat_template: Optional[Any]=No
                     new_messages.append(reformatted_messages[-1])
                     rendered_text = template.render(bos_token=bos_token, eos_token=eos_token, messages=new_messages)
         return rendered_text
-    except: 
-        raise Exception("Error rendering template")
+    except Exception as e: 
+        raise Exception(f"Error rendering template - {str(e)}")
 
 # Anthropic template 
 def claude_2_1_pt(messages: list): # format - https://docs.anthropic.com/claude/docs/how-to-use-system-prompts
@@ -283,6 +288,9 @@ def get_model_info(token, model):
         return None, None
 
 def format_prompt_togetherai(messages, prompt_format, chat_template):
+    if prompt_format is None:
+        return default_pt(messages)
+    
     human_prompt, assistant_prompt = prompt_format.split('{prompt}')
 
     if chat_template is not None:
diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 7f2b20c67dd..3046dfe31bf 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -294,7 +294,7 @@ def hf_test_completion_tgi():
         print(response)
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
-hf_test_completion_tgi()
+# hf_test_completion_tgi()
 
 # ################### Hugging Face Conversational models ########################
 # def hf_test_completion_conv():
@@ -708,7 +708,7 @@ def test_completion_azure():
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
-test_completion_azure()
+# test_completion_azure()
 
 def test_azure_openai_ad_token():
     # this tests if the azure ad token is set in the request header
@@ -1026,6 +1026,25 @@ def test_completion_together_ai():
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
+def test_completion_together_ai_mixtral():
+    model_name = "together_ai/DiscoResearch/DiscoLM-mixtral-8x7b-v2"
+    try:
+        messages =[
+            {"role": "user", "content": "Who are you"},
+            {"role": "assistant", "content": "I am your helpful assistant."},
+            {"role": "user", "content": "Tell me a joke"},
+        ]
+        response = completion(model=model_name, messages=messages, max_tokens=256, n=1, logger_fn=logger_fn)
+        # Add any assertions here to check the response
+        print(response)
+        cost = completion_cost(completion_response=response)
+        assert cost > 0.0   
+        print("Cost for completion call together-computer/llama-2-70b: ", f"${float(cost):.10f}")
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+test_completion_together_ai_mixtral()
+
 def test_completion_together_ai_yi_chat():
     model_name = "together_ai/zero-one-ai/Yi-34B-Chat"
     try: