diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 3615d030bfa..4519228a343 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -6,6 +6,16 @@ +## Pre-Submission checklist + +**Please complete all items before asking a LiteLLM maintainer to review your PR** + +- [ ] I have Added testing in the `tests/litellm/` directory, **Adding at least 1 test is a hard requirement** - [see details](https://docs.litellm.ai/docs/extras/contributing_code) +- [ ] I have added a screenshot of my new test passing locally +- [ ] My PR passes all unit tests on (`make unit-test`)[https://docs.litellm.ai/docs/extras/contributing_code] +- [ ] My PR's scope is as isolated as possible, it only solves 1 specific problem + + ## Type @@ -20,10 +30,4 @@ ## Changes - - -## [REQUIRED] Testing - Attach a screenshot of any new tests passing locally -If UI changes, send a screenshot/GIF of working UI fixes - - diff --git a/README.md b/README.md index 014df0ccdd4..2d2f71e4d18 100644 --- a/README.md +++ b/README.md @@ -340,71 +340,7 @@ curl 'http://0.0.0.0:4000/key/generate' \ ## Contributing -To contribute: Clone the repo locally -> Make a change -> Submit a PR with the change. - -Here's how to modify the repo locally: - -Step 1: Clone the repo - -``` -git clone https://github.com/BerriAI/litellm.git -``` - -Step 2: Install dependencies: - -``` -pip install -r requirements.txt -``` - -Step 3: Test your change: - -a. Add a pytest test within `tests/litellm/` - -This folder follows the same directory structure as `litellm/`. - -If a corresponding test file does not exist, create one. - -b. Run the test - -``` -cd tests/litellm # pwd: Documents/litellm/litellm/tests/litellm -pytest /path/to/test_file.py -``` - -Step 4: Submit a PR with your changes! 🚀 - -- push your fork to your GitHub repo -- submit a PR from there - -### Building LiteLLM Docker Image - -Follow these instructions if you want to build / run the LiteLLM Docker Image yourself. - -Step 1: Clone the repo - -``` -git clone https://github.com/BerriAI/litellm.git -``` - -Step 2: Build the Docker Image - -Build using Dockerfile.non_root -``` -docker build -f docker/Dockerfile.non_root -t litellm_test_image . -``` - -Step 3: Run the Docker Image - -Make sure config.yaml is present in the root directory. This is your litellm proxy config file. -``` -docker run \ - -v $(pwd)/proxy_config.yaml:/app/config.yaml \ - -e DATABASE_URL="postgresql://xxxxxxxx" \ - -e LITELLM_MASTER_KEY="sk-1234" \ - -p 4000:4000 \ - litellm_test_image \ - --config /app/config.yaml --detailed_debug -``` +Interested in contributing? Contributions to LiteLLM Python SDK, Proxy Server, and contributing LLM integrations are both accepted and highly encouraged! [See our Contribution Guide for more details](https://docs.litellm.ai/docs/extras/contributing_code) # Enterprise For companies that need better security, user management and professional support diff --git a/docs/my-website/docs/extras/contributing_code.md b/docs/my-website/docs/extras/contributing_code.md new file mode 100644 index 00000000000..0fe7675ead6 --- /dev/null +++ b/docs/my-website/docs/extras/contributing_code.md @@ -0,0 +1,96 @@ +# Contributing Code + +## **Checklist before submitting a PR** + +Here are the core requirements for any PR submitted to LiteLLM + + +- [ ] Add testing, **Adding at least 1 test is a hard requirement** - [see details](#2-adding-testing-to-your-pr) +- [ ] Ensure your PR passes the following tests: + - [ ] [Unit Tests](#3-running-unit-tests) + - [ ] Formatting / Linting Tests +- [ ] Keep scope as isolated as possible. As a general rule, your changes should address 1 specific problem at a time + + + +## Quick start + +## 1. Setup your local dev environment + + +Here's how to modify the repo locally: + +Step 1: Clone the repo + +```shell +git clone https://github.com/BerriAI/litellm.git +``` + +Step 2: Install dev dependencies: + +```shell +poetry install --with dev --extras proxy +``` + +That's it, your local dev environment is ready! + +## 2. Adding Testing to your PR + +- Add your test to the [`tests/litellm/` directory](https://github.com/BerriAI/litellm/tree/main/tests/litellm) + +- This directory 1:1 maps the the `litellm/` directory, and can only contain mocked tests. +- Do not add real llm api calls to this directory. + +### 2.1 File Naming Convention for `tests/litellm/` + +The `tests/litellm/` directory follows the same directory structure as `litellm/`. + +- `litellm/proxy/test_caching_routes.py` maps to `litellm/proxy/caching_routes.py` +- `test_{filename}.py` maps to `litellm/{filename}.py` + +## 3. Running Unit Tests + +run the following command on the root of the litellm directory + +```shell +make test-unit +``` + +## 4. Submit a PR with your changes! + +- push your fork to your GitHub repo +- submit a PR from there + + +## Advanced +### Building LiteLLM Docker Image + +Some people might want to build the LiteLLM docker image themselves. Follow these instructions if you want to build / run the LiteLLM Docker Image yourself. + +Step 1: Clone the repo + +```shell +git clone https://github.com/BerriAI/litellm.git +``` + +Step 2: Build the Docker Image + +Build using Dockerfile.non_root + +```shell +docker build -f docker/Dockerfile.non_root -t litellm_test_image . +``` + +Step 3: Run the Docker Image + +Make sure config.yaml is present in the root directory. This is your litellm proxy config file. + +```shell +docker run \ + -v $(pwd)/proxy_config.yaml:/app/config.yaml \ + -e DATABASE_URL="postgresql://xxxxxxxx" \ + -e LITELLM_MASTER_KEY="sk-1234" \ + -p 4000:4000 \ + litellm_test_image \ + --config /app/config.yaml --detailed_debug +``` diff --git a/docs/my-website/docs/observability/athina_integration.md b/docs/my-website/docs/observability/athina_integration.md index 4994d553c62..ba93ea4c980 100644 --- a/docs/my-website/docs/observability/athina_integration.md +++ b/docs/my-website/docs/observability/athina_integration.md @@ -78,6 +78,9 @@ Following are the allowed fields in metadata, their types, and their description * `context: Optional[Union[dict, str]]` - This is the context used as information for the prompt. For RAG applications, this is the "retrieved" data. You may log context as a string or as an object (dictionary). * `expected_response: Optional[str]` - This is the reference response to compare against for evaluation purposes. This is useful for segmenting inference calls by expected response. * `user_query: Optional[str]` - This is the user's query. For conversational applications, this is the user's last message. +* `tags: Optional[list]` - This is a list of tags. This is useful for segmenting inference calls by tags. +* `user_feedback: Optional[str]` - The end user’s feedback. +* `model_options: Optional[dict]` - This is a dictionary of model options. This is useful for getting insights into how model behavior affects your end users. * `custom_attributes: Optional[dict]` - This is a dictionary of custom attributes. This is useful for additional information about the inference. ## Using a self hosted deployment of Athina diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index b90709ce5b7..10ac13ecaf0 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -404,14 +404,16 @@ curl http://localhost:4000/v1/chat/completions \ If this was your initial VertexAI Grounding code, ```python -import vertexai +import vertexai +from vertexai.generative_models import GenerativeModel, GenerationConfig, Tool, grounding + vertexai.init(project=project_id, location="us-central1") model = GenerativeModel("gemini-1.5-flash-001") # Use Google Search for grounding -tool = Tool.from_google_search_retrieval(grounding.GoogleSearchRetrieval(disable_attributon=False)) +tool = Tool.from_google_search_retrieval(grounding.GoogleSearchRetrieval()) prompt = "When is the next total solar eclipse in US?" response = model.generate_content( @@ -428,7 +430,7 @@ print(response) then, this is what it looks like now ```python -from litellm import completion +from litellm import completion # !gcloud auth application-default login - run this to add vertex credentials to your env diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 85b94182d7d..cf4f14b202a 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -437,6 +437,7 @@ const sidebars = { type: "category", label: "Contributing", items: [ + "extras/contributing_code", { type: "category", label: "Adding Providers", diff --git a/litellm/integrations/athina.py b/litellm/integrations/athina.py index 754e980c2ad..705dc11f1d3 100644 --- a/litellm/integrations/athina.py +++ b/litellm/integrations/athina.py @@ -23,6 +23,9 @@ class AthinaLogger: "context", "expected_response", "user_query", + "tags", + "user_feedback", + "model_options", "custom_attributes", ] @@ -81,7 +84,6 @@ class AthinaLogger: for key in self.additional_keys: if key in metadata: data[key] = metadata[key] - response = litellm.module_level_client.post( self.athina_logging_url, headers=self.headers, diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py index 8f71cc153f3..dab64283ec2 100644 --- a/litellm/llms/perplexity/chat/transformation.py +++ b/litellm/llms/perplexity/chat/transformation.py @@ -37,6 +37,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): "response_format", "stream", "temperature", - "top_p" "max_retries", + "top_p", + "max_retries", "extra_headers", ] diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py index 0cd69400637..0a65e216dfe 100644 --- a/litellm/llms/triton/completion/transformation.py +++ b/litellm/llms/triton/completion/transformation.py @@ -3,7 +3,7 @@ Translates from OpenAI's `/v1/chat/completions` endpoint to Triton's `/generate` """ import json -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, AsyncIterator, Dict, Iterator, List, Literal, Optional, Union from httpx import Headers, Response @@ -67,6 +67,18 @@ class TritonConfig(BaseConfig): optional_params[param] = value return optional_params + def get_complete_url( + self, + api_base: str, + model: str, + optional_params: dict, + stream: Optional[bool] = None, + ) -> str: + llm_type = self._get_triton_llm_type(api_base) + if llm_type == "generate" and stream: + return api_base + "_stream" + return api_base + def transform_response( self, model: str, @@ -149,6 +161,18 @@ class TritonConfig(BaseConfig): else: raise ValueError(f"Invalid Triton API base: {api_base}") + def get_model_response_iterator( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], + sync_stream: bool, + json_mode: Optional[bool] = False, + ) -> Any: + return TritonResponseIterator( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) + class TritonGenerateConfig(TritonConfig): """ @@ -204,7 +228,7 @@ class TritonGenerateConfig(TritonConfig): return model_response -class TritonInferConfig(TritonGenerateConfig): +class TritonInferConfig(TritonConfig): """ Transformations for triton /infer endpoint (his is an infer model with a custom model on triton) """ diff --git a/litellm/main.py b/litellm/main.py index 846a908a8e9..b90030a6bb6 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3900,42 +3900,19 @@ async def atext_completion( ctx = contextvars.copy_context() func_with_context = partial(ctx.run, func) - _, custom_llm_provider, _, _ = get_llm_provider( - model=model, api_base=kwargs.get("api_base", None) - ) - - if ( - custom_llm_provider == "openai" - or custom_llm_provider == "azure" - or custom_llm_provider == "azure_text" - or custom_llm_provider == "custom_openai" - or custom_llm_provider == "anyscale" - or custom_llm_provider == "mistral" - or custom_llm_provider == "openrouter" - or custom_llm_provider == "deepinfra" - or custom_llm_provider == "perplexity" - or custom_llm_provider == "groq" - or custom_llm_provider == "nvidia_nim" - or custom_llm_provider == "cerebras" - or custom_llm_provider == "sambanova" - or custom_llm_provider == "ai21_chat" - or custom_llm_provider == "ai21" - or custom_llm_provider == "volcengine" - or custom_llm_provider == "text-completion-codestral" - or custom_llm_provider == "deepseek" - or custom_llm_provider == "text-completion-openai" - or custom_llm_provider == "huggingface" - or custom_llm_provider == "ollama" - or custom_llm_provider == "vertex_ai" - or custom_llm_provider in litellm.openai_compatible_providers - ): # currently implemented aiohttp calls for just azure and openai, soon all. - # Await normally - response = await loop.run_in_executor(None, func_with_context) - if asyncio.iscoroutine(response): - response = await response + init_response = await loop.run_in_executor(None, func_with_context) + if isinstance(init_response, dict) or isinstance( + init_response, TextCompletionResponse + ): ## CACHING SCENARIO + if isinstance(init_response, dict): + response = TextCompletionResponse(**init_response) + else: + response = init_response + elif asyncio.iscoroutine(init_response): + response = await init_response else: - # Call the synchronous function using run_in_executor - response = await loop.run_in_executor(None, func_with_context) + response = init_response # type: ignore + if ( kwargs.get("stream", False) is True or isinstance(response, TextCompletionStreamWrapper) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b61ddd45620..2e740a3ca33 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1994,8 +1994,8 @@ "max_tokens": 8191, "max_input_tokens": 32000, "max_output_tokens": 8191, - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000003, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.0000003, "litellm_provider": "mistral", "supports_function_calling": true, "mode": "chat", @@ -2006,8 +2006,8 @@ "max_tokens": 8191, "max_input_tokens": 32000, "max_output_tokens": 8191, - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000003, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.0000003, "litellm_provider": "mistral", "supports_function_calling": true, "mode": "chat", @@ -6057,26 +6057,6 @@ "mode": "chat", "supports_tool_choice": true }, - "jamba-large-1.6": { - "max_tokens": 256000, - "max_input_tokens": 256000, - "max_output_tokens": 256000, - "input_cost_per_token": 0.000002, - "output_cost_per_token": 0.000008, - "litellm_provider": "ai21", - "mode": "chat", - "supports_tool_choice": true - }, - "jamba-mini-1.6": { - "max_tokens": 256000, - "max_input_tokens": 256000, - "max_output_tokens": 256000, - "input_cost_per_token": 0.0000002, - "output_cost_per_token": 0.0000004, - "litellm_provider": "ai21", - "mode": "chat", - "supports_tool_choice": true - }, "jamba-1.5-mini": { "max_tokens": 256000, "max_input_tokens": 256000, @@ -6097,6 +6077,26 @@ "mode": "chat", "supports_tool_choice": true }, + "jamba-large-1.6": { + "max_tokens": 256000, + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "input_cost_per_token": 0.000002, + "output_cost_per_token": 0.000008, + "litellm_provider": "ai21", + "mode": "chat", + "supports_tool_choice": true + }, + "jamba-mini-1.6": { + "max_tokens": 256000, + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "input_cost_per_token": 0.0000002, + "output_cost_per_token": 0.0000004, + "litellm_provider": "ai21", + "mode": "chat", + "supports_tool_choice": true + }, "j2-mid": { "max_tokens": 8192, "max_input_tokens": 8192, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e35afb0ab7c..0020ff6d8cd 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1994,8 +1994,8 @@ "max_tokens": 8191, "max_input_tokens": 32000, "max_output_tokens": 8191, - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000003, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.0000003, "litellm_provider": "mistral", "supports_function_calling": true, "mode": "chat", @@ -2006,8 +2006,8 @@ "max_tokens": 8191, "max_input_tokens": 32000, "max_output_tokens": 8191, - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000003, + "input_cost_per_token": 0.0000001, + "output_cost_per_token": 0.0000003, "litellm_provider": "mistral", "supports_function_calling": true, "mode": "chat", diff --git a/tests/litellm/integrations/test_athina.py b/tests/litellm/integrations/test_athina.py new file mode 100644 index 00000000000..fd660a036ed --- /dev/null +++ b/tests/litellm/integrations/test_athina.py @@ -0,0 +1,207 @@ +import unittest +from unittest.mock import patch, MagicMock, ANY +import json +import datetime +import sys +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system-path + +from litellm.integrations.athina import AthinaLogger + +class TestAthinaLogger(unittest.TestCase): + + def setUp(self): + # Set up environment variables for testing + self.env_patcher = patch.dict('os.environ', { + 'ATHINA_API_KEY': 'test-api-key', + 'ATHINA_BASE_URL': 'https://test.athina.ai' + }) + self.env_patcher.start() + self.logger = AthinaLogger() + + # Setup common test variables + self.start_time = datetime.datetime(2023, 1, 1, 12, 0, 0) + self.end_time = datetime.datetime(2023, 1, 1, 12, 0, 1) + self.print_verbose = MagicMock() + + def tearDown(self): + self.env_patcher.stop() + + def test_init(self): + """Test the initialization of AthinaLogger""" + self.assertEqual(self.logger.athina_api_key, 'test-api-key') + self.assertEqual(self.logger.athina_logging_url, 'https://test.athina.ai/api/v1/log/inference') + self.assertEqual(self.logger.headers, { + 'athina-api-key': 'test-api-key', + 'Content-Type': 'application/json' + }) + + @patch('litellm.module_level_client.post') + def test_log_event_success(self, mock_post): + """Test successful logging of an event""" + # Setup mock response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.text = "Success" + mock_post.return_value = mock_response + + # Create test data + kwargs = { + 'model': 'gpt-4', + 'messages': [{'role': 'user', 'content': 'Hello'}], + 'stream': False, + 'litellm_params': { + 'metadata': { + 'environment': 'test-environment', + 'prompt_slug': 'test-prompt', + 'customer_id': 'test-customer', + 'customer_user_id': 'test-user', + 'session_id': 'test-session', + 'external_reference_id': 'test-ext-ref', + 'context': 'test-context', + 'expected_response': 'test-expected', + 'user_query': 'test-query', + 'tags': ['test-tag'], + 'user_feedback': 'test-feedback', + 'model_options': {'test-opt': 'test-val'}, + 'custom_attributes': {'test-attr': 'test-val'} + } + } + } + + response_obj = MagicMock() + response_obj.model_dump.return_value = { + 'id': 'resp-123', + 'choices': [{'message': {'content': 'Hi there'}}], + 'usage': { + 'prompt_tokens': 10, + 'completion_tokens': 5, + 'total_tokens': 15 + } + } + + # Call the method + self.logger.log_event(kwargs, response_obj, self.start_time, self.end_time, self.print_verbose) + + # Verify the results + mock_post.assert_called_once() + call_args = mock_post.call_args + self.assertEqual(call_args[0][0], 'https://test.athina.ai/api/v1/log/inference') + self.assertEqual(call_args[1]['headers'], self.logger.headers) + + # Parse and verify the sent data + sent_data = json.loads(call_args[1]['data']) + self.assertEqual(sent_data['language_model_id'], 'gpt-4') + self.assertEqual(sent_data['prompt'], kwargs['messages']) + self.assertEqual(sent_data['prompt_tokens'], 10) + self.assertEqual(sent_data['completion_tokens'], 5) + self.assertEqual(sent_data['total_tokens'], 15) + self.assertEqual(sent_data['response_time'], 1000) # 1 second = 1000ms + self.assertEqual(sent_data['customer_id'], 'test-customer') + self.assertEqual(sent_data['session_id'], 'test-session') + self.assertEqual(sent_data['environment'], 'test-environment') + self.assertEqual(sent_data['prompt_slug'], 'test-prompt') + self.assertEqual(sent_data['external_reference_id'], 'test-ext-ref') + self.assertEqual(sent_data['context'], 'test-context') + self.assertEqual(sent_data['expected_response'], 'test-expected') + self.assertEqual(sent_data['user_query'], 'test-query') + self.assertEqual(sent_data['tags'], ['test-tag']) + self.assertEqual(sent_data['user_feedback'], 'test-feedback') + self.assertEqual(sent_data['model_options'], {'test-opt': 'test-val'}) + self.assertEqual(sent_data['custom_attributes'], {'test-attr': 'test-val'}) + # Verify the print_verbose was called + self.print_verbose.assert_called_once_with("Athina Logger Succeeded - Success") + + @patch('litellm.module_level_client.post') + def test_log_event_error_response(self, mock_post): + """Test handling of error response from the API""" + # Setup mock error response + mock_response = MagicMock() + mock_response.status_code = 400 + mock_response.text = "Bad Request" + mock_post.return_value = mock_response + + # Create test data + kwargs = { + 'model': 'gpt-4', + 'messages': [{'role': 'user', 'content': 'Hello'}], + 'stream': False + } + + response_obj = MagicMock() + response_obj.model_dump.return_value = { + 'id': 'resp-123', + 'choices': [{'message': {'content': 'Hi there'}}], + 'usage': { + 'prompt_tokens': 10, + 'completion_tokens': 5, + 'total_tokens': 15 + } + } + + # Call the method + self.logger.log_event(kwargs, response_obj, self.start_time, self.end_time, self.print_verbose) + + # Verify print_verbose was called with error message + self.print_verbose.assert_called_once_with("Athina Logger Error - Bad Request, 400") + + @patch('litellm.module_level_client.post') + def test_log_event_exception(self, mock_post): + """Test handling of exceptions during logging""" + # Setup mock to raise exception + mock_post.side_effect = Exception("Test exception") + + # Create test data + kwargs = { + 'model': 'gpt-4', + 'messages': [{'role': 'user', 'content': 'Hello'}], + 'stream': False + } + + response_obj = MagicMock() + response_obj.model_dump.return_value = {} + + # Call the method + self.logger.log_event(kwargs, response_obj, self.start_time, self.end_time, self.print_verbose) + + # Verify print_verbose was called with exception info + self.print_verbose.assert_called_once() + self.assertIn("Athina Logger Error - Test exception", self.print_verbose.call_args[0][0]) + + @patch('litellm.module_level_client.post') + def test_log_event_with_tools(self, mock_post): + """Test logging with tools/functions data""" + # Setup mock response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_post.return_value = mock_response + + # Create test data with tools + kwargs = { + 'model': 'gpt-4', + 'messages': [{'role': 'user', 'content': "What's the weather?"}], + 'stream': False, + 'optional_params': { + 'tools': [{'type': 'function', 'function': {'name': 'get_weather'}}] + } + } + + response_obj = MagicMock() + response_obj.model_dump.return_value = { + 'id': 'resp-123', + 'usage': {'prompt_tokens': 10, 'completion_tokens': 5, 'total_tokens': 15} + } + + # Call the method + self.logger.log_event(kwargs, response_obj, self.start_time, self.end_time, self.print_verbose) + + # Verify the results + sent_data = json.loads(mock_post.call_args[1]['data']) + self.assertEqual(sent_data['tools'], [{'type': 'function', 'function': {'name': 'get_weather'}}]) + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/tests/llm_translation/test_triton.py b/tests/llm_translation/test_triton.py index 0835d09fab6..7e4ba92f23b 100644 --- a/tests/llm_translation/test_triton.py +++ b/tests/llm_translation/test_triton.py @@ -49,16 +49,26 @@ def test_split_embedding_by_shape_fails_with_shape_value_error(): ) -def test_completion_triton_generate_api(): +@pytest.mark.parametrize("stream", [True, False]) +def test_completion_triton_generate_api(stream): try: mock_response = MagicMock() + if stream: + def mock_iter_lines(): + mock_output = ''.join([ + 'data: {"model_name":"ensemble","model_version":"1","sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"' + t + '"}\n\n' + for t in ["I", " am", " an", " AI", " assistant"] + ]) + for out in mock_output.split('\n'): + yield out + mock_response.iter_lines = mock_iter_lines + else: + def return_val(): + return { + "text_output": "I am an AI assistant", + } - def return_val(): - return { - "text_output": "I am an AI assistant", - } - - mock_response.json = return_val + mock_response.json = return_val mock_response.status_code = 200 with patch( @@ -71,6 +81,7 @@ def test_completion_triton_generate_api(): max_tokens=10, timeout=5, api_base="http://localhost:8000/generate", + stream=stream, ) # Verify the call was made @@ -81,7 +92,10 @@ def test_completion_triton_generate_api(): call_kwargs = mock_post.call_args.kwargs # Access kwargs directly # Verify URL - assert call_kwargs["url"] == "http://localhost:8000/generate" + if stream: + assert call_kwargs["url"] == "http://localhost:8000/generate_stream" + else: + assert call_kwargs["url"] == "http://localhost:8000/generate" # Parse the request data from the JSON string request_data = json.loads(call_kwargs["data"]) @@ -91,7 +105,15 @@ def test_completion_triton_generate_api(): assert request_data["parameters"]["max_tokens"] == 10 # Verify response - assert response.choices[0].message.content == "I am an AI assistant" + if stream: + tokens = ["I", " am", " an", " AI", " assistant", None] + idx = 0 + for chunk in response: + assert chunk.choices[0].delta.content == tokens[idx] + idx += 1 + assert idx == len(tokens) + else: + assert response.choices[0].message.content == "I am an AI assistant" except Exception as e: print("exception", e) diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index b3da80c7155..b7565b0494f 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -23,7 +23,7 @@ const ProviderSpecificFields: React.FC = ({ console.log(`type of selectedProviderEnum: ${typeof selectedProviderEnum}`); return ( <> - {selectedProviderEnum === Providers.OpenAI && ( + {selectedProviderEnum === Providers.OpenAI || selectedProviderEnum === Providers.OpenAI_Text && ( <> = ({ {(selectedProviderEnum === Providers.Azure || selectedProviderEnum === Providers.Azure_AI_Studio || - selectedProviderEnum === Providers.OpenAI_Compatible + selectedProviderEnum === Providers.OpenAI_Compatible || + selectedProviderEnum === Providers.OpenAI_Text_Compatible ) && ( = { OpenAI: "openai", + OpenAI_Text: "text-completion-openai", Azure: "azure", Azure_AI_Studio: "azure_ai", Anthropic: "anthropic", @@ -37,6 +41,7 @@ export const provider_map: Record = { MistralAI: "mistral", Cohere: "cohere_chat", OpenAI_Compatible: "openai", + OpenAI_Text_Compatible: "text-completion-openai", Vertex_AI: "vertex_ai", Databricks: "databricks", xAI: "xai", @@ -53,6 +58,9 @@ export const provider_map: Record = { export const providerLogoMap: Record = { [Providers.OpenAI]: "https://artificialanalysis.ai/img/logos/openai_small.svg", + [Providers.OpenAI_Text]: "https://artificialanalysis.ai/img/logos/openai_small.svg", + [Providers.OpenAI_Text_Compatible]: "https://artificialanalysis.ai/img/logos/openai_small.svg", + [Providers.OpenAI_Compatible]: "https://artificialanalysis.ai/img/logos/openai_small.svg", [Providers.Azure]: "https://upload.wikimedia.org/wikipedia/commons/a/a8/Microsoft_Azure_Logo.svg", [Providers.Azure_AI_Studio]: "https://upload.wikimedia.org/wikipedia/commons/a/a8/Microsoft_Azure_Logo.svg", [Providers.Anthropic]: "https://artificialanalysis.ai/img/logos/anthropic_small.svg", @@ -61,7 +69,6 @@ export const providerLogoMap: Record = { [Providers.Groq]: "https://artificialanalysis.ai/img/logos/groq_small.png", [Providers.MistralAI]: "https://artificialanalysis.ai/img/logos/mistral_small.png", [Providers.Cohere]: "https://artificialanalysis.ai/img/logos/cohere_small.png", - [Providers.OpenAI_Compatible]: "https://upload.wikimedia.org/wikipedia/commons/4/4e/OpenAI_Logo.svg", [Providers.Vertex_AI]: "https://artificialanalysis.ai/img/logos/google_small.svg", [Providers.Databricks]: "https://artificialanalysis.ai/img/logos/databricks_small.png", [Providers.Ollama]: "https://artificialanalysis.ai/img/logos/ollama_small.svg",