Mock reranking

Author: mliu-cloudera

llm-service/app/services/models/reranking.py

@@ -67,22 +67,21 @@ def get_noop() -> BaseNodePostprocessor:
     def list_available() -> list[ModelResponse]:
         return get_provider_class().list_reranking_models()
 
+    _TEST_NODES = [
+        NodeWithScore(node=TextNode(text="test node"), score=0.5),
+        NodeWithScore(node=TextNode(text="another test node"), score=0.4),
+    ]
+
     @classmethod
     def test(cls, model_name: str) -> str:
         models = cls.list_available()
         for model in models:
             if model.model_id == model_name:
-                node = NodeWithScore(node=TextNode(text="test"), score=0.5)
-                another_test_node = NodeWithScore(
-                    node=TextNode(text="another test node"), score=0.4
-                )
                 reranking_model: BaseNodePostprocessor | None = cls.get(
                     model_name=model_name
                 )
                 if reranking_model:
-                    reranking_model.postprocess_nodes(
-                        [node, another_test_node], None, "test"
-                    )
+                    reranking_model.postprocess_nodes(cls._TEST_NODES, None, "test")
                     return "ok"
         raise HTTPException(status_code=404, detail="Model not found")
 

llm-service/app/tests/model_provider_mocks/bedrock.py

@@ -52,6 +52,7 @@
 from app.config import settings
 from app.services.caii.types import ModelResponse
 from app.services.models.providers import BedrockModelProvider
+from app.services.models import Reranking
 from .testing_chat_history_manager import (
     patch_get_chat_history_manager,
 )
@@ -182,10 +183,8 @@ def mock_make_api_call(
         elif operation_name == "Rerank":
             return {
                 "results": [
-                    # TODO: Is the document store checked prior to this? Do I need to mock that too?
-                    {"index": 0, "relevanceScore": random.random()},
-                    {"index": 1, "relevanceScore": random.random()},
-                    {"index": 2, "relevanceScore": random.random()},
+                    {"index": i, "relevanceScore": random.random()}
+                    for i in range(len(Reranking._TEST_NODES))
                 ]
             }
         else:
@@ -234,7 +233,7 @@ def test_bedrock_models(client: TestClient) -> None:
     ] == available_embedding_models
     for model_id in available_embedding_models:
         response = client.get(f"/llm-service/models/embedding/{model_id}/test")
-        assert response.status_code == 200  # TODO
+        assert response.status_code == 200
 
     available_text_models = [
         model_id
@@ -258,9 +257,9 @@ def test_bedrock_models(client: TestClient) -> None:
     assert [
         model["model_id"] for model in response.json()
     ] == available_reranking_models
-    # for model_id in available_reranking_models:
-    #     response = client.get(f"/llm-service/models/reranking/{model_id}/test")
-    #     assert response.status_code == 200  # TODO
+    for model_id in available_reranking_models:
+        response = client.get(f"/llm-service/models/reranking/{model_id}/test")
+        assert response.status_code == 200
 
 
 def test_bedrock_sessions(client: TestClient) -> None: