fix: Handle suggestions from copilot.

Author: kalaspuffar

README.md

@@ -92,6 +92,8 @@ The following environment variables are required to run the application:
 - `AWS_SESSION_TOKEN`: (Optional) may be needed for bedrock embeddings
 - `GOOGLE_APPLICATION_CREDENTIALS`: (Optional) needed for Google VertexAI embeddings. This should be a path to a service account credential file in JSON format, as accepted by [langchain](https://python.langchain.com/api_reference/google_vertexai/index.html)
 - `RAG_CHECK_EMBEDDING_CTX_LENGTH` (Optional) Default is true, disabling this will send raw input to the embedder, use this for custom embedding models.
+- `SIMPLE_RERANKER_MODEL_NAME` (Optional) defaults to `mixedbread-ai/mxbai-rerank-large-v1`, more options at (https://github.com/AnswerDotAI/rerankers)
+- `SIMPLE_RERANKER_MODEL_TYPE` (Optional) defaults to `cross-encoder`, more options at (https://github.com/AnswerDotAI/rerankers)
 
 Make sure to set these environment variables before running the application. You can set them in a `.env` file or as system environment variables.
 

app/routes/document_routes.py

@@ -42,13 +42,11 @@
 from app.utils.health import is_health_ok
 
 router = APIRouter()
-reRankerInstance = Reranker(
-    model_name=os.getenv("SIMPLE_RERANKER_MODEL_NAME"),
-    model_type=os.getenv("SIMPLE_RERANKER_MODEL_TYPE"),
-    lang=os.getenv("SIMPLE_RERANKER_LANG"),
+reranker_instance = Reranker(
+    model_name=os.getenv("SIMPLE_RERANKER_MODEL_NAME", "mixedbread-ai/mxbai-rerank-large-v1"),
+    model_type=os.getenv("SIMPLE_RERANKER_MODEL_TYPE", "cross-encoder"),
 )
 
-
 def get_user_id(request: Request, entity_id: str = None) -> str:
     """Extract user ID from request or entity_id."""
     if not hasattr(request.state, "user"):
@@ -711,23 +709,27 @@ async def query_embeddings_by_file_ids(request: Request, body: QueryMultipleBody
 
 @router.post("/rerank")
 async def rerank_documents_by_query(request: Request, body: QueryMultipleDocs):
+    """
+    Rerank documents based on relevance to a query using a reranking model.
+
+    Args:
+        request: The FastAPI request object
+        body: Contains query string, list of documents, and optional k value
+
+    Returns:
+        List of ranked documents with their scores
+    """
+
     try:
+        if not body.docs:
+            raise HTTPException(status_code=400, detail="docs list cannot be empty")
         docs = []
         for i, d in enumerate(body.docs):
-            if isinstance(d, str):
-                docs.append(ReRankDocument(text=d, doc_id=i))
-            else:
-                docs.append(
-                    ReRankDocument(
-                        text=d.get("text", ""),
-                        doc_id=d.get("doc_id", i),
-                        metadata=d.get("metadata", {}) or {},
-                    )
-                )
+            docs.append(ReRankDocument(text=d, doc_id=i))
 
         top_k = body.k
 
-        results = reRankerInstance.rank(query=body.query, docs=docs)
+        results = reranker_instance.rank(query=body.query, docs=docs)
         items = results.top_k(top_k) if top_k else results
 
         return [

docker-compose.yaml

@@ -17,7 +17,7 @@ services:
       - DB_PORT=5432
     ports:
       - "8000:8000"
-    runtime: nvidia
+    runtime: ${DOCKER_RUNTIME:-runc}
     volumes:
       - ./uploads:/app/uploads
       - ~/.cache/huggingface:/root/.cache/huggingface:rw

tests/test_main.py

@@ -263,3 +263,41 @@ def test_extract_text_from_file(tmp_path, auth_headers):
     assert json_data["file_id"] == "test_text_123"
     assert json_data["filename"] == "test_text_extraction.txt"
     assert json_data["known_type"] is True  # text files are known types
+
+def test_query_rerank(auth_headers):
+    # Successful reranking with string documents
+    data = {
+        "query": "I love you",
+        "docs": ["I hate you", "I really like you"],
+        "k": 1
+    }
+    response = client.post("/rerank", json=data, headers=auth_headers)
+    assert response.status_code == 200
+    json_data = response.json()
+    assert isinstance(json_data, list)
+    assert len(json_data) == 1
+    doc = json_data[0]
+    assert doc["text"] == "I really like you"
+
+    # Handling of the k parameter (top_k filtering)
+    data = {
+        "query": "I love you",
+        "docs": ["I hate you", "I really like you", "I love you too"],
+        "k": 2
+    }
+    response = client.post("/rerank", json=data, headers=auth_headers)
+    assert response.status_code == 200
+    json_data = response.json()
+    assert isinstance(json_data, list)
+    assert len(json_data) == 2
+    assert json_data[0]["text"] == "I really like you"
+    assert json_data[1]["text"] == "I love you too"
+
+    # Error handling for invalid inputs
+    data = {
+        "query": "I love you",
+        "docs": [123, 456],
+        "k": 1
+    }
+    response = client.post("/rerank", json=data, headers=auth_headers)
+    assert response.status_code == 422