Add fireworks rerank support

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

@@ -13,7 +13,7 @@ import TabItem from '@theme/TabItem';
 | Description | The fastest and most efficient inference engine to build production-ready, compound AI systems. |
 | Provider Route on LiteLLM | `fireworks_ai/` |
 | Provider Doc | [Fireworks AI ↗](https://docs.fireworks.ai/getting-started/introduction) |
-| Supported OpenAI Endpoints | `/chat/completions`, `/embeddings`, `/completions`, `/audio/transcriptions` |
+| Supported OpenAI Endpoints | `/chat/completions`, `/embeddings`, `/completions`, `/audio/transcriptions`, `/rerank` |
 
 
 ## Overview
@@ -386,4 +386,87 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/audio/transcriptions' \
 ```
 
 </TabItem>
-</Tabs>
+</Tabs>
+
+## Rerank
+
+### Quick Start
+
+<Tabs>
+<TabItem value="sdk" label="SDK">
+
+```python
+from litellm import rerank
+import os
+
+os.environ["FIREWORKS_AI_API_KEY"] = "YOUR_API_KEY"
+
+query = "What is the capital of France?"
+documents = [
+    "Paris is the capital and largest city of France, home to the Eiffel Tower and the Louvre Museum.",
+    "France is a country in Western Europe known for its wine, cuisine, and rich history.",
+    "The weather in Europe varies significantly between northern and southern regions.",
+    "Python is a popular programming language used for web development and data science.",
+]
+
+response = rerank(
+    model="fireworks_ai/fireworks/qwen3-reranker-8b",
+    query=query,
+    documents=documents,
+    top_n=3,
+    return_documents=True,
+)
+print(response)
+```
+
+[Pass API Key/API Base in `.rerank`](../set_keys.md#passing-args-to-completion)
+
+</TabItem>
+<TabItem value="proxy" label="PROXY">
+
+1. Setup config.yaml
+
+```yaml
+model_list:
+  - model_name: qwen3-reranker-8b
+    litellm_params:
+      model: fireworks_ai/fireworks/qwen3-reranker-8b
+      api_key: os.environ/FIREWORKS_API_KEY
+    model_info:
+      mode: rerank
+```
+
+2. Start Proxy
+
+```
+litellm --config config.yaml
+```
+
+3. Test it
+
+```bash
+curl http://0.0.0.0:4000/rerank \
+  -H "Authorization: Bearer sk-1234" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "qwen3-reranker-8b",
+    "query": "What is the capital of France?",
+    "documents": [
+        "Paris is the capital and largest city of France, home to the Eiffel Tower and the Louvre Museum.",
+        "France is a country in Western Europe known for its wine, cuisine, and rich history.",
+        "The weather in Europe varies significantly between northern and southern regions.",
+        "Python is a popular programming language used for web development and data science."
+    ],
+    "top_n": 3,
+    "return_documents": true
+  }'
+```
+
+</TabItem>
+</Tabs>
+
+### Supported Models
+
+| Model Name | Function Call |
+|------------|---------------|
+| fireworks/qwen3-reranker-8b | `rerank(model="fireworks_ai/fireworks/qwen3-reranker-8b", query=query, documents=documents)` |

docs/my-website/docs/rerank.md

@@ -16,7 +16,7 @@ LiteLLM Follows the [cohere api request / response for the rerank api](https://c
 | Fallbacks | ✅ | Works between supported models |
 | Loadbalancing | ✅ | Works between supported models |
 | Guardrails | ✅ | Applies to input query only (not documents) |
-| Supported Providers | Cohere, Together AI, Azure AI, DeepInfra, Nvidia NIM, Infinity | |
+| Supported Providers | Cohere, Together AI, Azure AI, DeepInfra, Nvidia NIM, Infinity, Fireworks AI | |
 
 ## **LiteLLM Python SDK Usage**
 ### Quick Start 
@@ -134,4 +134,5 @@ curl http://0.0.0.0:4000/rerank \
 | Infinity|   [Usage](../docs/providers/infinity)                 |  
 | vLLM|   [Usage](../docs/providers/vllm#rerank-endpoint)                 |  
 | DeepInfra|   [Usage](../docs/providers/deepinfra#rerank-endpoint)                 |
-| Vertex AI|   [Usage](../docs/providers/vertex#rerank-api)                 |  
+| Vertex AI|   [Usage](../docs/providers/vertex#rerank-api)                 |  
+| Fireworks AI|   [Usage](../docs/providers/fireworks_ai#rerank-endpoint)                 |  

litellm/__init__.py

@@ -1111,6 +1111,7 @@ from .llms.deepinfra.rerank.transformation import DeepinfraRerankConfig
 from .llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig
 from .llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig
 from .llms.vertex_ai.rerank.transformation import VertexAIRerankConfig
+from .llms.fireworks_ai.rerank.transformation import FireworksAIRerankConfig
 from .llms.clarifai.chat.transformation import ClarifaiConfig
 from .llms.ai21.chat.transformation import AI21ChatConfig, AI21ChatConfig as AI21Config
 from .llms.meta_llama.chat.transformation import LlamaAPIConfig

litellm/llms/fireworks_ai/rerank/__init__.py

@@ -0,0 +1,2 @@
+# Fireworks AI Rerank
+

litellm/llms/fireworks_ai/rerank/transformation.py

@@ -0,0 +1,262 @@
+"""
+Fireworks AI Rerank API transformation
+
+Reference: https://docs.fireworks.ai/inference-api-reference/rerank
+"""
+
+from typing import Any, Dict, List, Optional, Union
+
+import httpx
+
+from litellm._uuid import uuid
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
+from litellm.llms.fireworks_ai.common_utils import FireworksAIMixin
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.rerank import (
+    RerankBilledUnits,
+    RerankResponse,
+    RerankResponseDocument,
+    RerankResponseMeta,
+    RerankResponseResult,
+    RerankTokens,
+)
+
+
+class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig):
+    """
+    Fireworks AI Rerank API configuration
+    """
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        model: str,
+        optional_params: Optional[dict] = None,
+    ) -> str:
+        if api_base:
+            # Remove trailing slashes and ensure clean base URL
+            api_base = api_base.rstrip("/")
+            if not api_base.endswith("/rerank"):
+                if api_base.endswith("/v1"):
+                    api_base = f"{api_base}/rerank"
+                elif api_base.endswith("/inference/v1"):
+                    api_base = f"{api_base}/rerank"
+                else:
+                    api_base = f"{api_base}/inference/v1/rerank"
+            return api_base
+        return "https://api.fireworks.ai/inference/v1/rerank"
+
+    def get_supported_cohere_rerank_params(self, model: str) -> list:
+        return [
+            "query",
+            "documents",
+            "top_n",
+            "return_documents",
+        ]
+
+    def map_cohere_rerank_params(
+        self,
+        non_default_params: Optional[dict],
+        model: str,
+        drop_params: bool,
+        query: str,
+        documents: List[Union[str, Dict[str, Any]]],
+        custom_llm_provider: Optional[str] = None,
+        top_n: Optional[int] = None,
+        rank_fields: Optional[List[str]] = None,
+        return_documents: Optional[bool] = True,
+        max_chunks_per_doc: Optional[int] = None,
+        max_tokens_per_doc: Optional[int] = None,
+    ) -> Dict[str, Any]:
+        """
+        Map Cohere rerank params to Fireworks AI rerank params
+        """
+        params: Dict[str, Any] = {
+            "query": query,
+            "documents": documents,
+        }
+        
+        if top_n is not None:
+            params["top_n"] = top_n
+        
+        if return_documents is not None:
+            params["return_documents"] = return_documents
+        
+        # Fireworks AI doesn't support these params
+        if rank_fields is not None:
+            # Silently ignore rank_fields as Fireworks AI doesn't support it
+            pass
+        
+        if max_chunks_per_doc is not None:
+            # Silently ignore max_chunks_per_doc as Fireworks AI doesn't support it
+            pass
+        
+        if max_tokens_per_doc is not None:
+            # Silently ignore max_tokens_per_doc as Fireworks AI doesn't support it
+            pass
+        
+        return params
+
+    def validate_environment(  # type: ignore[override]
+        self,
+        headers: dict,
+        model: str,
+        api_key: Optional[str] = None,
+        optional_params: Optional[dict] = None,
+    ) -> dict:
+        api_key = self._get_api_key(api_key)
+        if api_key is None:
+            raise ValueError(
+                "FIREWORKS_API_KEY is not set. Please set 'FIREWORKS_API_KEY' or 'FIREWORKS_AI_API_KEY' in your environment"
+            )
+
+        default_headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        }
+
+        # If 'Authorization' is provided in headers, it overrides the default.
+        if "Authorization" in headers:
+            default_headers["Authorization"] = headers["Authorization"]
+
+        # Merge other headers, overriding any default ones except Authorization
+        return {**default_headers, **headers}
+
+    def transform_rerank_request(
+        self,
+        model: str,
+        optional_rerank_params: Dict,
+        headers: dict,
+    ) -> dict:
+        """
+        Transform request to Fireworks AI rerank format
+        """
+        if "query" not in optional_rerank_params:
+            raise ValueError("query is required for Fireworks AI rerank")
+        if "documents" not in optional_rerank_params:
+            raise ValueError("documents is required for Fireworks AI rerank")
+
+        # Handle model name - Fireworks AI expects model name like "fireworks/qwen3-reranker-8b"
+        # Remove fireworks_ai/ prefix if present
+        if model.startswith("fireworks_ai/"):
+            model = model.replace("fireworks_ai/", "")
+        
+        # If model doesn't start with "fireworks/", add it
+        # But don't add if it already has the prefix
+        if not model.startswith("fireworks/"):
+            model = f"fireworks/{model}"
+
+        request_data = {
+            "model": model,
+            "query": optional_rerank_params["query"],
+            "documents": optional_rerank_params["documents"],
+        }
+
+        if "top_n" in optional_rerank_params and optional_rerank_params["top_n"] is not None:
+            request_data["top_n"] = optional_rerank_params["top_n"]
+
+        if "return_documents" in optional_rerank_params and optional_rerank_params["return_documents"] is not None:
+            request_data["return_documents"] = optional_rerank_params["return_documents"]
+
+        return request_data
+
+    def transform_rerank_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: RerankResponse,
+        logging_obj: LiteLLMLoggingObj,
+        api_key: Optional[str] = None,
+        request_data: dict = {},
+        optional_params: dict = {},
+        litellm_params: dict = {},
+    ) -> RerankResponse:
+        """
+        Transform Fireworks AI rerank response to LiteLLM RerankResponse format
+        """
+        try:
+            raw_response_json = raw_response.json()
+        except Exception as e:
+            raise self.get_error_class(
+                error_message=f"Failed to parse response: {str(e)}",
+                status_code=raw_response.status_code,
+                headers=raw_response.headers,
+            )
+
+        # Fireworks AI response format:
+        # {
+        #   "object": "list",
+        #   "model": "accounts/fireworks/models/qwen3-reranker-8b",
+        #   "data": [
+        #     {
+        #       "index": 0,
+        #       "relevance_score": 0.95,
+        #       "document": "..."  
+        #     }
+        #   ],
+        #   "usage": {
+        #     "total_tokens": 100,
+        #     "prompt_tokens": 50,
+        #     "completion_tokens": 50
+        #   }
+        # }
+
+        # Extract usage information
+        usage = raw_response_json.get("usage", {})
+        _billed_units = RerankBilledUnits(
+            search_units=usage.get("total_tokens", 0)
+        )
+        _tokens = RerankTokens(
+            input_tokens=usage.get("prompt_tokens", 0),
+            output_tokens=usage.get("completion_tokens", 0),
+        )
+        rerank_meta = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens)
+
+        # Extract results - Fireworks AI uses "data" instead of "results"
+        _results: Optional[List[dict]] = raw_response_json.get("data") or raw_response_json.get("results")
+
+        if _results is None:
+            raise ValueError(f"No results found in the response={raw_response_json}")
+
+        rerank_results: List[RerankResponseResult] = []
+
+        for result in _results:
+            # Validate required fields exist
+            if not all(key in result for key in ["index", "relevance_score"]):
+                raise ValueError(f"Missing required fields in the result={result}")
+
+            # Get document data - Fireworks AI returns document as a string directly
+            document_text = result.get("document")
+            document = None
+            if document_text:
+                # Handle both string and object formats
+                if isinstance(document_text, str):
+                    document = RerankResponseDocument(text=document_text)
+                elif isinstance(document_text, dict):
+                    # Handle object format if it exists
+                    text = document_text.get("text", "")
+                    if text:
+                        document = RerankResponseDocument(text=str(text))
+
+            # Create typed result
+            rerank_result = RerankResponseResult(
+                index=int(result["index"]),
+                relevance_score=float(result["relevance_score"]),
+            )
+
+            # Only add document if it exists
+            if document:
+                rerank_result["document"] = document
+
+            rerank_results.append(rerank_result)
+
+        # Use model name as id if no id is provided
+        response_id = raw_response_json.get("id") or raw_response_json.get("model") or str(uuid.uuid4())
+
+        return RerankResponse(
+            id=response_id,
+            results=rerank_results,
+            meta=rerank_meta,
+        )
+

litellm/rerank_api/main.py

@@ -29,7 +29,7 @@ async def arerank(
     model: str,
     query: str,
     documents: List[Union[str, Dict[str, Any]]],
-    custom_llm_provider: Optional[Literal["cohere", "together_ai", "deepinfra"]] = None,
+    custom_llm_provider: Optional[Literal["cohere", "together_ai", "deepinfra", "fireworks_ai"]] = None,
     top_n: Optional[int] = None,
     rank_fields: Optional[List[str]] = None,
     return_documents: Optional[bool] = None,
@@ -83,6 +83,7 @@ def rerank(  # noqa: PLR0915
             "litellm_proxy",
             "hosted_vllm",
             "deepinfra",
+            "fireworks_ai",
         ]
     ] = None,
     top_n: Optional[int] = None,
@@ -411,6 +412,36 @@ def rerank(  # noqa: PLR0915
                     "api_base must be provided for Deepinfra rerank. Set in call or via DEEPINFRA_API_BASE env var."
                 )
 
+            response = base_llm_http_handler.rerank(
+                model=model,
+                custom_llm_provider=_custom_llm_provider,
+                provider_config=rerank_provider_config,
+                optional_rerank_params=optional_rerank_params,
+                logging_obj=litellm_logging_obj,
+                timeout=optional_params.timeout,
+                api_key=api_key,
+                api_base=api_base,
+                _is_async=_is_async,
+                headers=headers or litellm.headers or {},
+                client=client,
+                model_response=model_response,
+            )
+        elif _custom_llm_provider == litellm.LlmProviders.FIREWORKS_AI:
+            api_key = (
+                dynamic_api_key
+                or optional_params.api_key
+                or get_secret_str("FIREWORKS_API_KEY")
+                or get_secret_str("FIREWORKS_AI_API_KEY")
+                or get_secret_str("FIREWORKSAI_API_KEY")
+                or get_secret_str("FIREWORKS_AI_TOKEN")
+            )
+
+            api_base = (
+                dynamic_api_base
+                or optional_params.api_base
+                or get_secret_str("FIREWORKS_AI_API_BASE")
+            )
+
             response = base_llm_http_handler.rerank(
                 model=model,
                 custom_llm_provider=_custom_llm_provider,

litellm/utils.py

@@ -7346,6 +7346,8 @@ class ProviderConfigManager:
             return litellm.NvidiaNimRerankConfig()
         elif litellm.LlmProviders.VERTEX_AI == provider:
             return litellm.VertexAIRerankConfig()
+        elif litellm.LlmProviders.FIREWORKS_AI == provider:
+            return litellm.FireworksAIRerankConfig()
         return litellm.CohereRerankConfig()
 
     @staticmethod

tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py

@@ -0,0 +1,344 @@
+"""
+Tests for Fireworks AI rerank transformation functionality.
+"""
+import json
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+
+from litellm.llms.fireworks_ai.rerank.transformation import FireworksAIRerankConfig
+from litellm.types.rerank import RerankResponse
+
+
+class TestFireworksAIRerankTransform:
+    def setup_method(self):
+        self.config = FireworksAIRerankConfig()
+        self.model = "fireworks_ai/fireworks/qwen3-reranker-8b"
+
+    def test_get_complete_url(self):
+        """Test URL generation for Fireworks AI rerank API."""
+        # Test basic URL generation
+        api_base = None
+        model = "fireworks/qwen3-reranker-8b"
+        url = self.config.get_complete_url(api_base, model)
+        assert url == "https://api.fireworks.ai/inference/v1/rerank"
+
+        # Test URL with custom api_base
+        api_base = "https://api.fireworks.ai/inference/v1"
+        url = self.config.get_complete_url(api_base, model)
+        assert url == "https://api.fireworks.ai/inference/v1/rerank"
+
+        # Test URL with trailing slash
+        api_base_with_slash = "https://api.fireworks.ai/inference/v1/"
+        url = self.config.get_complete_url(api_base_with_slash, model)
+        assert url == "https://api.fireworks.ai/inference/v1/rerank"
+
+    def test_map_cohere_rerank_params_basic(self):
+        """Test basic parameter mapping for Fireworks AI rerank."""
+        params = self.config.map_cohere_rerank_params(
+            non_default_params={},
+            model=self.model,
+            drop_params=False,
+            query="test query",
+            documents=["doc1", "doc2"],
+            top_n=3,
+            return_documents=True,
+        )
+        assert params["query"] == "test query"
+        assert params["documents"] == ["doc1", "doc2"]
+        assert params["top_n"] == 3
+        assert params["return_documents"] is True
+
+    def test_map_cohere_rerank_params_ignores_unsupported(self):
+        """Test that unsupported params are silently ignored."""
+        params = self.config.map_cohere_rerank_params(
+            non_default_params={},
+            model=self.model,
+            drop_params=False,
+            query="test query",
+            documents=["doc1", "doc2"],
+            rank_fields=["field1"],  # Not supported by Fireworks AI
+            max_chunks_per_doc=5,  # Not supported by Fireworks AI
+            max_tokens_per_doc=100,  # Not supported by Fireworks AI
+        )
+        assert params["query"] == "test query"
+        assert params["documents"] == ["doc1", "doc2"]
+        # Unsupported params should not be in the result
+        assert "rank_fields" not in params
+        assert "max_chunks_per_doc" not in params
+        assert "max_tokens_per_doc" not in params
+
+    def test_transform_rerank_request(self):
+        """Test request transformation for Fireworks AI format."""
+        optional_params = {
+            "query": "What is the capital of France?",
+            "documents": [
+                "Paris is the capital of France.",
+                "France is a country in Europe.",
+            ],
+            "top_n": 2,
+            "return_documents": True,
+        }
+
+        request_body = self.config.transform_rerank_request(
+            model=self.model, optional_rerank_params=optional_params, headers={}
+        )
+
+        # Model should be transformed to include "fireworks/" prefix
+        assert request_body["model"] == "fireworks/qwen3-reranker-8b"
+        assert request_body["query"] == "What is the capital of France?"
+        assert request_body["documents"] == optional_params["documents"]
+        assert request_body["top_n"] == 2
+        assert request_body["return_documents"] is True
+
+    def test_transform_rerank_request_model_prefix_handling(self):
+        """Test that model prefix is handled correctly."""
+        # Test with fireworks_ai/ prefix
+        optional_params = {
+            "query": "test",
+            "documents": ["doc1"],
+        }
+        request_body = self.config.transform_rerank_request(
+            model="fireworks_ai/fireworks/qwen3-reranker-8b",
+            optional_rerank_params=optional_params,
+            headers={},
+        )
+        assert request_body["model"] == "fireworks/qwen3-reranker-8b"
+
+        # Test with model already having fireworks/ prefix
+        request_body = self.config.transform_rerank_request(
+            model="fireworks/qwen3-reranker-8b",
+            optional_rerank_params=optional_params,
+            headers={},
+        )
+        assert request_body["model"] == "fireworks/qwen3-reranker-8b"
+
+    def test_transform_rerank_request_missing_query(self):
+        """Test that transform_rerank_request raises error for missing query."""
+        optional_params = {
+            "documents": ["doc1"],
+        }
+
+        with pytest.raises(ValueError, match="query is required"):
+            self.config.transform_rerank_request(
+                model=self.model, optional_rerank_params=optional_params, headers={}
+            )
+
+    def test_transform_rerank_request_missing_documents(self):
+        """Test that transform_rerank_request raises error for missing documents."""
+        optional_params = {
+            "query": "test query",
+        }
+
+        with pytest.raises(ValueError, match="documents is required"):
+            self.config.transform_rerank_request(
+                model=self.model, optional_rerank_params=optional_params, headers={}
+            )
+
+    def test_transform_rerank_response_success(self):
+        """Test successful response transformation."""
+        # Mock Fireworks AI response format (uses "data" not "results", and document is a string)
+        response_data = {
+            "object": "list",
+            "model": "accounts/fireworks/models/qwen3-reranker-8b",
+            "data": [
+                {
+                    "index": 0,
+                    "relevance_score": 0.95,
+                    "document": "Paris is the capital of France.",
+                },
+                {
+                    "index": 1,
+                    "relevance_score": 0.75,
+                    "document": "France is a country in Europe.",
+                },
+            ],
+            "usage": {
+                "total_tokens": 100,
+                "prompt_tokens": 50,
+                "completion_tokens": 50,
+            },
+        }
+
+        # Create mock httpx response
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.json.return_value = response_data
+        mock_response.status_code = 200
+        mock_response.headers = {}
+
+        # Create mock logging object
+        mock_logging = MagicMock()
+
+        model_response = RerankResponse()
+
+        result = self.config.transform_rerank_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=mock_logging,
+        )
+
+        # Verify response structure
+        # Fireworks AI doesn't return "id", so it uses "model" as the id
+        assert result.id == "accounts/fireworks/models/qwen3-reranker-8b"
+        assert len(result.results) == 2
+        assert result.results[0]["index"] == 0
+        assert result.results[0]["relevance_score"] == 0.95
+        assert result.results[0]["document"]["text"] == "Paris is the capital of France."
+        assert result.results[1]["index"] == 1
+        assert result.results[1]["relevance_score"] == 0.75
+        assert result.results[1]["document"]["text"] == "France is a country in Europe."
+
+        # Verify metadata
+        assert result.meta["tokens"]["input_tokens"] == 50
+        assert result.meta["tokens"]["output_tokens"] == 50
+        assert result.meta["billed_units"]["search_units"] == 100
+
+    def test_transform_rerank_response_without_documents(self):
+        """Test response transformation when return_documents is False."""
+        response_data = {
+            "object": "list",
+            "model": "accounts/fireworks/models/qwen3-reranker-8b",
+            "data": [
+                {"index": 0, "relevance_score": 0.95},
+                {"index": 1, "relevance_score": 0.75},
+            ],
+            "usage": {
+                "total_tokens": 50,
+                "prompt_tokens": 30,
+                "completion_tokens": 20,
+            },
+        }
+
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.json.return_value = response_data
+        mock_response.status_code = 200
+        mock_response.headers = {}
+
+        mock_logging = MagicMock()
+        model_response = RerankResponse()
+
+        result = self.config.transform_rerank_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=mock_logging,
+        )
+
+        # Fireworks AI doesn't return "id", so it uses "model" as the id
+        assert result.id == "accounts/fireworks/models/qwen3-reranker-8b"
+        assert len(result.results) == 2
+        assert result.results[0]["index"] == 0
+        assert result.results[0]["relevance_score"] == 0.95
+        # Document should not be present
+        assert "document" not in result.results[0]
+
+    def test_transform_rerank_response_missing_id(self):
+        """Test response transformation when id is missing (should use model name or generate UUID)."""
+        response_data = {
+            "object": "list",
+            "model": "accounts/fireworks/models/qwen3-reranker-8b",
+            "data": [
+                {"index": 0, "relevance_score": 0.95},
+            ],
+            "usage": {"total_tokens": 10},
+        }
+
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.json.return_value = response_data
+        mock_response.status_code = 200
+        mock_response.headers = {}
+
+        mock_logging = MagicMock()
+        model_response = RerankResponse()
+
+        result = self.config.transform_rerank_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=mock_logging,
+        )
+
+        # Should use model name when id is missing
+        assert result.id == "accounts/fireworks/models/qwen3-reranker-8b"
+
+    def test_transform_rerank_response_missing_results(self):
+        """Test that missing results raises ValueError."""
+        response_data = {
+            "object": "list",
+            "model": "accounts/fireworks/models/qwen3-reranker-8b",
+            "usage": {"total_tokens": 10},
+        }
+
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.json.return_value = response_data
+        mock_response.status_code = 200
+        mock_response.headers = {}
+
+        mock_logging = MagicMock()
+        model_response = RerankResponse()
+
+        with pytest.raises(ValueError, match="No results found"):
+            self.config.transform_rerank_response(
+                model=self.model,
+                raw_response=mock_response,
+                model_response=model_response,
+                logging_obj=mock_logging,
+            )
+
+    def test_transform_rerank_response_invalid_json(self):
+        """Test error handling for invalid JSON response."""
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "doc", 0)
+        mock_response.text = "Invalid JSON response"
+        mock_response.status_code = 500
+        mock_response.headers = {}
+
+        mock_logging = MagicMock()
+        model_response = RerankResponse()
+
+        with pytest.raises(Exception) as exc_info:
+            self.config.transform_rerank_response(
+                model=self.model,
+                raw_response=mock_response,
+                model_response=model_response,
+                logging_obj=mock_logging,
+            )
+
+        # Should raise an error with appropriate message
+        assert "Failed to parse response" in str(exc_info.value)
+
+    def test_get_supported_cohere_rerank_params(self):
+        """Test getting supported parameters for Fireworks AI rerank."""
+        supported_params = self.config.get_supported_cohere_rerank_params(self.model)
+        assert "query" in supported_params
+        assert "documents" in supported_params
+        assert "top_n" in supported_params
+        assert "return_documents" in supported_params
+        assert len(supported_params) == 4
+
+    def test_validate_environment_missing_api_key(self):
+        """Test that validate_environment raises error when API key is missing."""
+        from unittest.mock import patch
+
+        # Mock _get_api_key to return None
+        with patch.object(self.config, "_get_api_key", return_value=None):
+            with pytest.raises(ValueError, match="FIREWORKS_API_KEY is not set"):
+                self.config.validate_environment(
+                    headers={},
+                    model=self.model,
+                    api_key=None,
+                )
+
+    def test_validate_environment_with_api_key(self):
+        """Test that validate_environment works with API key."""
+        headers = self.config.validate_environment(
+            headers={},
+            model=self.model,
+            api_key="test-api-key",
+        )
+
+        assert headers["Authorization"] == "Bearer test-api-key"
+        assert headers["Content-Type"] == "application/json"
+