Merge branch 'BerriAI:main' into main

2026-09-12 23:01:41 +00:00 · 2025-03-10 19:21:37 -07:00 · 2025-03-10 19:21:37 -07:00 · 45f0ca68c4
commit 45f0ca68c4
parent e5fefe5105 6dcf64918c
16 changed files with 437 additions and 154 deletions
--- a/.github/pull_request_template.md
+++ b/.github/pull_request_template.md
@ -6,6 +6,16 @@

 <!-- e.g. "Fixes #000" -->

+## 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

 <!-- Select the type of Pull Request -->
@ -20,10 +30,4 @@

 ## Changes

-<!-- List of changes -->
-
-## [REQUIRED] Testing - Attach a screenshot of any new tests passing locally
-If UI changes, send a screenshot/GIF of working UI fixes
-
-<!-- Test procedure -->

--- 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
--- a/docs/my-website/docs/extras/contributing_code.md
+++ 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
+```
--- 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
--- 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
--- 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",
--- 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,
--- 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",
        ]
--- 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)
    """
--- 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)
--- 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,
--- 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",
--- a/tests/litellm/integrations/test_athina.py
+++ 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()
--- 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)
--- 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<ProviderSpecificFieldsProps> = ({
  console.log(`type of selectedProviderEnum: ${typeof selectedProviderEnum}`);
  return (
    <>
-      {selectedProviderEnum === Providers.OpenAI && (
+      {selectedProviderEnum === Providers.OpenAI || selectedProviderEnum === Providers.OpenAI_Text && (
        <>
          <Form.Item
            label="API Base"
@ -99,7 +99,8 @@ const ProviderSpecificFields: React.FC<ProviderSpecificFieldsProps> = ({

      {(selectedProviderEnum === Providers.Azure ||
        selectedProviderEnum === Providers.Azure_AI_Studio ||
-        selectedProviderEnum === Providers.OpenAI_Compatible
+        selectedProviderEnum === Providers.OpenAI_Compatible ||
+        selectedProviderEnum === Providers.OpenAI_Text_Compatible
      ) && (
        <Form.Item
          rules={[{ required: true, message: "Required" }]}
--- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
+++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
@ -1,7 +1,11 @@
+import OpenAI from "openai";
 import React from "react";

 export enum Providers {
    OpenAI = "OpenAI",
+    OpenAI_Compatible = "OpenAI-Compatible Endpoints (Together AI, etc.)",
+    OpenAI_Text = "OpenAI Text Completion",
+    OpenAI_Text_Compatible = "OpenAI-Compatible Text Completion Models (Together AI, etc.)",
    Azure = "Azure",
    Azure_AI_Studio = "Azure AI Foundry (Studio)",
    Anthropic = "Anthropic",
@ -11,7 +15,6 @@ export enum Providers {
    Groq = "Groq",
    MistralAI = "Mistral AI",
    Deepseek = "Deepseek",
-    OpenAI_Compatible = "OpenAI-Compatible Endpoints (Together AI, etc.)",
    Cohere = "Cohere",
    Databricks = "Databricks",
    Ollama = "Ollama",
@ -28,6 +31,7 @@ export enum Providers {
  
 export const provider_map: Record<string, string> = {
    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<string, string> = {
    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<string, string> = {

 export const providerLogoMap: Record<string, string> = {
    [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<string, string> = {
    [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",