[FIX ERROR]: multiple args 'query'

Author: fahmiaziz98

src/api/routers/rerank.py

@@ -6,7 +6,7 @@
 """
 
 import time
-from typing import Union, List
+from typing import List
 from fastapi import APIRouter, Depends, HTTPException, status
 from loguru import logger
 
@@ -18,15 +18,17 @@
     RerankingDocumentError,
     ValidationError,
 )
-
 from src.api.dependencies import get_model_manager
 from src.utils.validators import extract_embedding_kwargs
 
-router = APIRouter(tags=["rerank"]) 
+router = APIRouter(tags=["rerank"])
 
 
 @router.post(
-    "/rerank", response_model=RerankResponse, summary="Rerank documents", description="Reranks the provided documents based on the given query."
+    "/rerank",
+    response_model=RerankResponse,
+    summary="Rerank documents",
+    description="Reranks the provided documents based on the given query.",
 )
 async def rerank_documents(
     request: RerankRequest,
@@ -35,54 +37,82 @@ async def rerank_documents(
     """
     Rerank documents based on a query.
 
-    This endpoint processes a list of documents and returns them ranked according to their relevance to the query.
-    
+    This endpoint processes a list of documents and returns them ranked
+    according to their relevance to the query.
+
     Args:
-        request (RerankRequest): The request object containing the query and documents to rank.
-        manager (ModelManager): The model manager dependency to access the model.
+        request: The request object containing the query and documents to rank
+        manager: The model manager dependency to access the model
 
     Returns:
-        RerankResponse: The response containing the ranked documents and processing time.
+        RerankResponse: The response containing the ranked documents and processing time
 
     Raises:
-        HTTPException: If there are validation errors, model loading issues, or unexpected errors.
+        HTTPException: If there are validation errors, model loading issues, or unexpected errors
     """
     # Filter out empty documents and keep original indices
     valid_docs = [
         (i, doc.strip()) for i, doc in enumerate(request.documents) if doc.strip()
     ]
 
     if not valid_docs:
-        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="No valid documents provided.")
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="No valid documents provided.",
+        )
 
     try:
+        # Extract kwargs but exclude rerank-specific fields
         kwargs = extract_embedding_kwargs(request)
+
+        # Remove fields that are already passed as positional arguments
+        # to avoid "got multiple values for argument" error
+        kwargs.pop("query", None)
+        kwargs.pop("documents", None)
+        kwargs.pop("top_k", None)
+
         model = manager.get_model(request.model_id)
         config = manager.model_configs[request.model_id]
 
-        start = time.time()
-        if config.type == "rerank":
-            scores = model.rank_document(
-                request.query, request.documents, request.top_k, **kwargs
+        if config.type != "rerank":
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail=f"Model '{request.model_id}' is not a rerank model. Type: {config.type}",
             )
-            processing_time = time.time() - start
 
-            original_indices, documents_list = zip(*valid_docs)
-            results: List[RerankResult] = []
+        start = time.time()
+
+        # Call rank_document with clean kwargs
+        scores = model.rank_document(
+            query=request.query,
+            documents=[doc for _, doc in valid_docs],  # Use filtered documents
+            top_k=request.top_k,
+            **kwargs,
+        )
+
+        processing_time = time.time() - start
 
-            for i, (orig_idx, doc) in enumerate(zip(original_indices, documents_list)):
-                results.append(RerankResult(text=doc, score=scores[i], index=orig_idx))
+        # Build results with original indices
+        original_indices, documents_list = zip(*valid_docs)
+        results: List[RerankResult] = []
 
-            # Sort results by score in descending order
-            results.sort(key=lambda x: x.score, reverse=True)
+        for i, (orig_idx, doc) in enumerate(zip(original_indices, documents_list)):
+            results.append(RerankResult(text=doc, score=scores[i], index=orig_idx))
 
-            logger.info(f"Reranked documents in {processing_time:.3f} seconds")
-            return RerankResponse(
-                model_id=request.model_id,
-                processing_time=processing_time,
-                query=request.query,
-                results=results,
-            )
+        # Sort results by score in descending order
+        results.sort(key=lambda x: x.score, reverse=True)
+
+        logger.info(
+            f"Reranked {len(results)} documents in {processing_time:.3f}s "
+            f"(model: {request.model_id})"
+        )
+
+        return RerankResponse(
+            model_id=request.model_id,
+            processing_time=processing_time,
+            query=request.query,
+            results=results,
+        )
 
     except (ValidationError, ModelNotFoundError) as e:
         raise HTTPException(status_code=e.status_code, detail=e.message)
@@ -94,5 +124,5 @@ async def rerank_documents(
         logger.exception("Unexpected error in rerank_documents")
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"Failed to create query embedding: {str(e)}",
+            detail=f"Failed to rerank documents: {str(e)}",
         )

src/models/schemas/requests.py

@@ -126,7 +126,6 @@ class Config:
                 "model_id": "jina-reranker-v3",
                 "query": "Python best programming languages for data science",
                 "top_k": 4,
-                "prompt": "Rerank document based user query",
                 "documents": [
                     "Python is a popular language for data science due to its extensive libraries.",
                     "R is widely used in statistical computing and data analysis.",

src/utils/validators.py

@@ -10,6 +10,26 @@
 from src.core.exceptions import TextTooLongError, BatchTooLargeError, ValidationError
 
 
+def validate_text(text: str, max_length: int = 8192, allow_empty: bool = False) -> None:
+    """
+    Validate a single text input.
+
+    Args:
+        text: Input text to validate
+        max_length: Maximum allowed text length
+        allow_empty: Whether to allow empty strings
+
+    Raises:
+        ValidationError: If text is empty and not allowed
+        TextTooLongError: If text exceeds max_length
+    """
+    if not allow_empty and not text.strip():
+        raise ValidationError("text", "Text cannot be empty")
+
+    if len(text) > max_length:
+        raise TextTooLongError(len(text), max_length)
+
+
 def validate_texts(
     texts: List[str],
     max_length: int = 8192,
@@ -71,30 +91,6 @@ def validate_model_id(model_id: str, available_models: List[str]) -> None:
         )
 
 
-def sanitize_text(text: str, max_length: int = 8192) -> str:
-    """
-    Sanitize text input by removing excessive whitespace and truncating.
-
-    Args:
-        text: Input text to sanitize
-        max_length: Maximum length to truncate to
-
-    Returns:
-        Sanitized text
-    """
-    # Remove leading/trailing whitespace
-    text = text.strip()
-
-    # Replace multiple whitespaces with single space
-    text = " ".join(text.split())
-
-    # Truncate if too long
-    if len(text) > max_length:
-        text = text[:max_length]
-
-    return text
-
-
 def extract_embedding_kwargs(request: BaseModel) -> Dict[str, Any]:
     """
     Extract embedding kwargs from a request object.
@@ -110,7 +106,7 @@ def extract_embedding_kwargs(request: BaseModel) -> Dict[str, Any]:
 
     Example:
         >>> request = EmbedRequest(
-        ...     text="hello",
+        ...     texts=["hello"],
         ...     model_id="qwen3-0.6b",
         ...     options=EmbeddingOptions(normalize_embeddings=True),
         ...     batch_size=32  # Extra field
@@ -125,7 +121,16 @@ def extract_embedding_kwargs(request: BaseModel) -> Dict[str, Any]:
         kwargs.update(request.options.to_kwargs())
 
     # Extract extra fields (excluding standard fields)
-    standard_fields = {"text", "texts", "model_id", "prompt", "options"}
+    standard_fields = {
+        "text",
+        "texts",
+        "model_id",
+        "prompt",
+        "options",
+        "query",
+        "documents",
+        "top_k",
+    }
     request_dict = request.model_dump()
 
     for key, value in request_dict.items():