From 0f489b68eb1da0d1db06dc25af302c1538cd3ddc Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 26 Jun 2024 17:09:30 -0700
Subject: [PATCH] test volcengine

---
 litellm/tests/test_completion.py | 62 +++++++++++++-------------------
 1 file changed, 24 insertions(+), 38 deletions(-)

diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index a3b0e6ea26..2ceb11a79b 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -1222,44 +1222,6 @@ def test_completion_fireworks_ai():
         pytest.fail(f"Error occurred: {e}")
 
 
-def test_fireworks_ai_tool_calling():
-    litellm.set_verbose = True
-    model_name = "fireworks_ai/accounts/fireworks/models/firefunction-v2"
-    tools = [
-        {
-            "type": "function",
-            "function": {
-                "name": "get_current_weather",
-                "description": "Get the current weather in a given location",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "location": {
-                            "type": "string",
-                            "description": "The city and state, e.g. San Francisco, CA",
-                        },
-                        "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
-                    },
-                    "required": ["location"],
-                },
-            },
-        }
-    ]
-    messages = [
-        {
-            "role": "user",
-            "content": "What's the weather like in Boston today in Fahrenheit?",
-        }
-    ]
-    response = completion(
-        model=model_name,
-        messages=messages,
-        tools=tools,
-        tool_choice="required",
-    )
-    print(response)
-
-
 @pytest.mark.skip(reason="this test is flaky")
 def test_completion_perplexity_api():
     try:
@@ -3508,6 +3470,30 @@ def test_completion_deep_infra_mistral():
 # test_completion_deep_infra_mistral()
 
 
+@pytest.mark.skip(reason="Local test - don't have a volcengine account as yet")
+def test_completion_volcengine():
+    litellm.set_verbose = True
+    model_name = "volcengine/<OUR_ENDPOINT_ID>"
+    try:
+        response = completion(
+            model=model_name,
+            messages=[
+                {
+                    "role": "user",
+                    "content": "What's the weather like in Boston today in Fahrenheit?",
+                }
+            ],
+            api_key="<OUR_API_KEY>",
+        )
+        # Add any assertions here to check the response
+        print(response)
+
+    except litellm.exceptions.Timeout as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
 def test_completion_nvidia_nim():
     model_name = "nvidia_nim/databricks/dbrx-instruct"
     try: