Merge pull request #2877 from BerriAI/litellm_fix_text_completion

[Feat] Text-Completion-OpenAI - Re-use OpenAI Client

docs/my-website/docs/providers/openai.md

@@ -44,7 +44,11 @@ export OPENAI_API_KEY=""
 model_list:
   - model_name: gpt-3.5-turbo
     litellm_params:
-      model: gpt-3.5-turbo
+      model: openai/gpt-3.5-turbo                          # The `openai/` prefix will call openai.chat.completions.create
+      api_key: os.environ/OPENAI_API_KEY
+  - model_name: gpt-3.5-turbo-instruct
+    litellm_params:
+      model: text-completion-openai/gpt-3.5-turbo-instruct # The `text-completion-openai/` prefix will call openai.completions.create
       api_key: os.environ/OPENAI_API_KEY
 ```
 </TabItem>

litellm/main.py

@@ -1060,6 +1060,7 @@ def completion(
                 api_key=api_key,
                 api_base=api_base,
                 acompletion=acompletion,
+                client=client,  # pass AsyncOpenAI, OpenAI client
                 logging_obj=logging,
                 optional_params=optional_params,
                 litellm_params=litellm_params,

litellm/router.py

@@ -1769,6 +1769,7 @@ class Router:
             or custom_llm_provider == "azure"
             or custom_llm_provider == "custom_openai"
             or custom_llm_provider == "openai"
+            or custom_llm_provider == "text-completion-openai"
             or "ft:gpt-3.5-turbo" in model_name
             or model_name in litellm.open_ai_embedding_models
         ):

litellm/tests/test_router.py

@@ -1228,3 +1228,35 @@ def test_router_add_deployment():
     assert len(new_model_id_list) > len(init_model_id_list)
 
     assert new_model_id_list[1] != new_model_id_list[0]
+
+
+@pytest.mark.asyncio
+async def test_router_text_completion_client():
+    # This tests if we re-use the Async OpenAI client
+    # This test fails when we create a new Async OpenAI client per request
+    try:
+        model_list = [
+            {
+                "model_name": "fake-openai-endpoint",
+                "litellm_params": {
+                    "model": "text-completion-openai/gpt-3.5-turbo-instruct",
+                    "api_key": os.getenv("OPENAI_API_KEY", None),
+                    "api_base": "https://exampleopenaiendpoint-production.up.railway.app/",
+                },
+            }
+        ]
+        router = Router(model_list=model_list, debug_level="DEBUG", set_verbose=True)
+        tasks = []
+        for _ in range(300):
+            tasks.append(
+                router.atext_completion(
+                    model="fake-openai-endpoint",
+                    prompt="hello from litellm test",
+                )
+            )
+
+        # Execute all coroutines concurrently
+        responses = await asyncio.gather(*tasks)
+        print(responses)
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")

litellm/tests/test_router_init.py

@@ -490,3 +490,63 @@ def test_init_clients_azure_command_r_plus():
     except Exception as e:
         traceback.print_exc()
         pytest.fail(f"Error occurred: {e}")
+
+
+@pytest.mark.asyncio
+async def test_text_completion_with_organization():
+    try:
+        print("Testing Text OpenAI with organization")
+        model_list = [
+            {
+                "model_name": "openai-bad-org",
+                "litellm_params": {
+                    "model": "text-completion-openai/gpt-3.5-turbo-instruct",
+                    "api_key": os.getenv("OPENAI_API_KEY", None),
+                    "organization": "org-ikDc4ex8NB",
+                },
+            },
+            {
+                "model_name": "openai-good-org",
+                "litellm_params": {
+                    "model": "text-completion-openai/gpt-3.5-turbo-instruct",
+                    "api_key": os.getenv("OPENAI_API_KEY", None),
+                    "organization": os.getenv("OPENAI_ORGANIZATION", None),
+                },
+            },
+        ]
+
+        router = Router(model_list=model_list)
+
+        print(router.model_list)
+        print(router.model_list[0])
+
+        openai_client = router._get_client(
+            deployment=router.model_list[0],
+            kwargs={"input": ["hello"], "model": "openai-bad-org"},
+        )
+        print(vars(openai_client))
+
+        assert openai_client.organization == "org-ikDc4ex8NB"
+
+        # bad org raises error
+
+        try:
+            response = await router.atext_completion(
+                model="openai-bad-org",
+                prompt="this is a test",
+            )
+            pytest.fail("Request should have failed - This organization does not exist")
+        except Exception as e:
+            print("Got exception: " + str(e))
+            assert "No such organization: org-ikDc4ex8NB" in str(e)
+
+        # good org works
+        response = await router.atext_completion(
+            model="openai-good-org",
+            prompt="this is a test",
+            max_tokens=5,
+        )
+        print("working response: ", response)
+
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")