From 87cf6f3ffe606e17e0b5cf4ef7e04086bdec28b0 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 8 Dec 2025 20:29:50 +0530 Subject: [PATCH] Add fireworks rerank support --- .../my-website/docs/providers/fireworks_ai.md | 87 ++++- docs/my-website/docs/rerank.md | 5 +- litellm/__init__.py | 1 + litellm/llms/fireworks_ai/rerank/__init__.py | 2 + .../fireworks_ai/rerank/transformation.py | 262 +++++++++++++ litellm/rerank_api/main.py | 33 +- litellm/utils.py | 2 + ...test_fireworks_ai_rerank_transformation.py | 344 ++++++++++++++++++ 8 files changed, 731 insertions(+), 5 deletions(-) create mode 100644 litellm/llms/fireworks_ai/rerank/__init__.py create mode 100644 litellm/llms/fireworks_ai/rerank/transformation.py create mode 100644 tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py diff --git a/docs/my-website/docs/providers/fireworks_ai.md b/docs/my-website/docs/providers/fireworks_ai.md index b1b10cd71b..29168dce93 100644 --- a/docs/my-website/docs/providers/fireworks_ai.md +++ b/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' \ ``` - \ No newline at end of file + + +## Rerank + +### Quick Start + + + + +```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) + + + + +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 + }' +``` + + + + +### Supported Models + +| Model Name | Function Call | +|------------|---------------| +| fireworks/qwen3-reranker-8b | `rerank(model="fireworks_ai/fireworks/qwen3-reranker-8b", query=query, documents=documents)` | \ No newline at end of file diff --git a/docs/my-website/docs/rerank.md b/docs/my-website/docs/rerank.md index ec0592f31f..a0433cb7a2 100644 --- a/docs/my-website/docs/rerank.md +++ b/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) | \ No newline at end of file +| Vertex AI| [Usage](../docs/providers/vertex#rerank-api) | +| Fireworks AI| [Usage](../docs/providers/fireworks_ai#rerank-endpoint) | \ No newline at end of file diff --git a/litellm/__init__.py b/litellm/__init__.py index a4ade3bcca..5a441b99f6 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/llms/fireworks_ai/rerank/__init__.py b/litellm/llms/fireworks_ai/rerank/__init__.py new file mode 100644 index 0000000000..b8e99317a2 --- /dev/null +++ b/litellm/llms/fireworks_ai/rerank/__init__.py @@ -0,0 +1,2 @@ +# Fireworks AI Rerank + diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py new file mode 100644 index 0000000000..f5caff5686 --- /dev/null +++ b/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, + ) + diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index fc45266536..80360d994e 100644 --- a/litellm/rerank_api/main.py +++ b/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, diff --git a/litellm/utils.py b/litellm/utils.py index b283d5ae6e..9bd553b6f4 100644 --- a/litellm/utils.py +++ b/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 diff --git a/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py b/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py new file mode 100644 index 0000000000..e17123f8ae --- /dev/null +++ b/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" +