Implement feedback on refactoring and unused code

Author: max-svistunov

src/app/endpoints/rags.py

@@ -11,8 +11,8 @@
 from authentication.interface import AuthTuple
 from authorization.middleware import authorize
 from client import AsyncLlamaStackClientHolder
-from configuration import LogicError, configuration
-from models.config import Action
+from configuration import configuration
+from models.config import Action, ByokRag
 from models.responses import (
     ForbiddenResponse,
     InternalServerErrorResponse,
@@ -104,7 +104,7 @@ async def rags_endpoint_handler(
         raise HTTPException(**response.model_dump()) from e
 
 
-def _resolve_rag_id_to_vector_db_id(rag_id: str) -> str:
+def _resolve_rag_id_to_vector_db_id(rag_id: str, byok_rags: list[ByokRag]) -> str:
     """Resolve a user-facing rag_id to the llama-stack vector_db_id.
 
     Checks if the given ID matches a rag_id in the BYOK config and returns
@@ -113,15 +113,11 @@ def _resolve_rag_id_to_vector_db_id(rag_id: str) -> str:
 
     Parameters:
         rag_id: The user-provided RAG identifier.
+        byok_rags: List of BYOK RAG config entries.
 
     Returns:
         The llama-stack vector_db_id, or the original ID if no mapping found.
     """
-    try:
-        byok_rags = configuration.configuration.byok_rag
-    except (AttributeError, RuntimeError, LogicError):
-        return rag_id
-
     for brag in byok_rags:
         if brag.rag_id == rag_id:
             return brag.vector_db_id
@@ -164,7 +160,9 @@ async def get_rag_endpoint_handler(
     logger.info("Llama stack config: %s", llama_stack_configuration)
 
     # Resolve user-facing rag_id to llama-stack vector_db_id
-    vector_db_id = _resolve_rag_id_to_vector_db_id(rag_id)
+    vector_db_id = _resolve_rag_id_to_vector_db_id(
+        rag_id, configuration.configuration.byok_rag
+    )
 
     try:
         # try to get Llama Stack client

src/models/context.py

@@ -1,8 +1,6 @@
 """Context objects for internal operations."""
 
 from dataclasses import dataclass, field
-from typing import Any
-
 from llama_stack_client import AsyncLlamaStackClient
 
 from models.requests import QueryRequest
@@ -24,7 +22,6 @@ class ResponseGeneratorContext:  # pylint: disable=too-many-instance-attributes
         query_request: The query request object
         started_at: Timestamp when the request started (ISO 8601 format)
         client: The Llama Stack client for API interactions
-        metadata_map: Dictionary for storing metadata from tool responses
         vector_store_ids: Vector store IDs used in the query for source resolution.
         rag_id_mapping: Mapping from vector_db_id to user-facing rag_id.
     """
@@ -43,7 +40,6 @@ class ResponseGeneratorContext:  # pylint: disable=too-many-instance-attributes
 
     # Dependencies & State
     client: AsyncLlamaStackClient
-    metadata_map: dict[str, dict[str, Any]] = field(default_factory=dict)
 
     # RAG index identification
     vector_store_ids: list[str] = field(default_factory=list)

src/utils/endpoints.py

@@ -196,27 +196,6 @@ def check_configuration_loaded(config: AppConfig) -> None:
         raise HTTPException(**response.model_dump()) from e
 
 
-def create_rag_chunks_dict(summary: TurnSummary) -> list[dict[str, Any]]:
-    """
-    Create dictionary representation of RAG chunks for streaming response.
-
-    Args:
-        summary: TurnSummary containing RAG chunks
-
-    Returns:
-        List of dictionaries with content, source, score, and attributes.
-    """
-    return [
-        {
-            "content": chunk.content,
-            "source": chunk.source,
-            "score": chunk.score,
-            "attributes": chunk.attributes,
-        }
-        for chunk in summary.rag_chunks
-    ]
-
-
 def _process_http_source(
     src: str, doc_urls: set[str]
 ) -> Optional[tuple[Optional[AnyUrl], str]]:

src/utils/responses.py

@@ -407,7 +407,7 @@ def _get_token_value(original: str, header: str) -> str | None:
     return tools
 
 
-def parse_referenced_documents(
+def parse_referenced_documents(  # pylint: disable=too-many-locals
     response: Optional[OpenAIResponseObject],
     vector_store_ids: Optional[list[str]] = None,
     rag_id_mapping: Optional[dict[str, str]] = None,
@@ -430,16 +430,17 @@ def parse_referenced_documents(
     if response is None or not response.output:
         return documents
 
-    resolved_source = _resolve_single_store_source(
-        vector_store_ids or [], rag_id_mapping or {}
-    )
+    vs_ids = vector_store_ids or []
+    id_mapping = rag_id_mapping or {}
 
     for output_item in response.output:
         item_type = getattr(output_item, "type", None)
 
         if item_type == "file_search_call":
             results = getattr(output_item, "results", []) or []
             for result in results:
+                resolved_source = _resolve_source_for_result(result, vs_ids, id_mapping)
+
                 # Handle both object and dict access
                 if isinstance(result, dict):
                     attributes = result.get("attributes", {})
@@ -472,26 +473,6 @@ def parse_referenced_documents(
     return documents
 
 
-def _resolve_single_store_source(
-    vector_store_ids: list[str],
-    rag_id_mapping: dict[str, str],
-) -> Optional[str]:
-    """Resolve source name when there is exactly one vector store.
-
-    Parameters:
-        vector_store_ids: The vector store IDs used in the query.
-        rag_id_mapping: Mapping from vector_db_id to user-facing rag_id.
-
-    Returns:
-        The resolved rag_id (or raw store_id as fallback) if exactly one
-        store is used, None otherwise.
-    """
-    if len(vector_store_ids) == 1:
-        store_id = vector_store_ids[0]
-        return rag_id_mapping.get(store_id, store_id)
-    return None
-
-
 def extract_token_usage(
     response: Optional[OpenAIResponseObject], model_id: str
 ) -> TokenCounter:

tests/unit/app/endpoints/test_rags.py

@@ -374,20 +374,16 @@ def __init__(self) -> None:
     assert response.id == "ocp-4.18-docs"
 
 
-def test_resolve_rag_id_to_vector_db_id_with_mapping(
-    mocker: MockerFixture, tmp_path: Path
-) -> None:
+def test_resolve_rag_id_to_vector_db_id_with_mapping(tmp_path: Path) -> None:
     """Test that _resolve_rag_id_to_vector_db_id maps rag_id to vector_db_id."""
     byok_config = _make_byok_config(str(tmp_path))
-    mocker.patch("app.endpoints.rags.configuration", byok_config)
-    assert _resolve_rag_id_to_vector_db_id("ocp-4.18-docs") == "vs_abc123"
-    assert _resolve_rag_id_to_vector_db_id("company-kb") == "vs_def456"
+    byok_rags = byok_config.configuration.byok_rag
+    assert _resolve_rag_id_to_vector_db_id("ocp-4.18-docs", byok_rags) == "vs_abc123"
+    assert _resolve_rag_id_to_vector_db_id("company-kb", byok_rags) == "vs_def456"
 
 
-def test_resolve_rag_id_to_vector_db_id_passthrough(
-    mocker: MockerFixture, tmp_path: Path
-) -> None:
+def test_resolve_rag_id_to_vector_db_id_passthrough(tmp_path: Path) -> None:
     """Test that unmapped IDs are passed through unchanged."""
     byok_config = _make_byok_config(str(tmp_path))
-    mocker.patch("app.endpoints.rags.configuration", byok_config)
-    assert _resolve_rag_id_to_vector_db_id("vs_unknown") == "vs_unknown"
+    byok_rags = byok_config.configuration.byok_rag
+    assert _resolve_rag_id_to_vector_db_id("vs_unknown", byok_rags) == "vs_unknown"