Variable RAG DB setup for query, query v2, and streaming query

[JslYoon]

Description

Currently, the RAG implementation lacks granularity in retrieval. To improve response accuracy and support multi-tenant or multi-context use cases, we need the ability to dynamically select which knowledge base (Vector DB/Index) to query.
Added a new request param vector_store_ids for:
query, query_v2, and streaming_query endpoints

Type of change

• Refactor
• New feature
• Bug fix
• CVE fix
• Optimization
• Documentation Update
• Configuration Update
• Bump-up service version
• Bump-up dependent library
• Bump-up library or tool used for development (does not change the final image)
• CI configuration change
• Konflux configuration change
• Unit tests improvement
• Integration tests improvement
• End to end tests improvement

Tools used to create PR

Identify any AI code assistants used in this PR (for transparency and review context)

• Assisted-by: (e.g., Claude code)
• Generated by: (Claude-4.5-Sonnet)

Related Tickets & Documents

• Related Issue LCORE-1067
• Closes LCORE-1067

Checklist before requesting a review

• I have performed a self-review of my code.
• PR has passed all pre-merge test jobs.
• If it is a core feature, I have added thorough tests.

Testing

• Please provide detailed steps to perform tests related to this code change.
• How were the fix/results from this change verified? Please provide relevant screenshots or results.

Summary by CodeRabbit

• New Features
  • Query requests now support optional vector store selection for RAG operations
  • Specify particular vector stores to query; defaults to all stores if not provided
  • Fully backward compatible with existing requests

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Walkthrough

The changes extend the system to support selective vector store querying in RAG operations. A new optional vector_store_ids field is added to the QueryRequest model. Query endpoints (query.py, query_v2.py, streaming_query.py) are updated to use explicitly provided vector store IDs when available; otherwise, all available stores are fetched as before.

Changes

Cohort / File(s)	Summary
Query endpoints with vector store selection logic src/app/endpoints/query.py, src/app/endpoints/query_v2.py, src/app/endpoints/streaming_query.py	Added conditional logic to respect explicit vector_store_ids from query requests. If provided, uses the specified list; otherwise, fetches all available vector stores. Converts empty toolgroups list to None to maintain consistency.
Request model extension src/models/requests.py	Added optional vector_store_ids: Optional[list[str]] field to QueryRequest class with description and examples, enabling clients to specify particular vector stores for RAG queries. Defaults to None (queries all stores if not provided).

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

• Consistent, repetitive pattern of conditional logic applied across three endpoints reduces cognitive load
• New model field is straightforward addition with clear documentation
• Changes are self-contained with no cross-cutting concerns
• Areas for attention: Verify the conditional logic in all three endpoints handles None cases consistently; confirm backward compatibility when vector_store_ids is omitted from requests

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Variable RAG DB setup for query, query v2, and streaming query' directly describes the main change: enabling selective vector store configuration across three query endpoints instead of always fetching all stores.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[openshift-ci]

Hi @JslYoon. Thanks for your PR.

I'm waiting for a github.com member to verify that this patch is reasonable to test. If it is, they should reply with /ok-to-test on its own line. Until that is done, I will not automatically test new commits in this PR, but the usual testing commands by org members will still work. Regular contributors should join the org to skip this step.

Once the patch is verified, the new status will be reflected by the ok-to-test label.

I understand the commands that are listed here.

Details

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository.

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/app/endpoints/streaming_query.py (1)

1091-1104: Handle empty list edge case consistently across all endpoints.

This endpoint has the same empty list ambiguity as query.py and query_v2.py. For a unified user experience, all three query endpoints should handle vector_store_ids identically:

• None → fetch all available stores
• [] → ??? (currently treated as "use no stores", but should this fetch all instead?)
• ["id1", "id2"] → use specified stores

Recommend normalizing the behavior across all endpoints (query.py, query_v2.py, and streaming_query.py) by either:

1. Treating empty list as "fetch all" (normalize at model validation level), or
2. Treating empty list as "disable RAG" and documenting it clearly

-        # Use specified vector stores or fetch all available ones
-        if query_request.vector_store_ids:
-            vector_db_ids = query_request.vector_store_ids
-        else:
-            vector_db_ids = [
-                vector_store.id
-                for vector_store in (await client.vector_stores.list()).data
-            ]
+        # Use specified vector stores or fetch all available ones
+        if query_request.vector_store_ids is not None and len(query_request.vector_store_ids) > 0:
+            vector_db_ids = query_request.vector_store_ids
+        else:
+            # None or empty list means fetch all available stores
+            vector_db_ids = [
+                vector_store.id
+                for vector_store in (await client.vector_stores.list()).data
+            ]

🧹 Nitpick comments (1)

src/models/requests.py (1)

163-168: Clarify empty list vs None semantics and consider validation.

The field uses truthiness checks in endpoints (if query_request.vector_store_ids:), so an empty list [] would be treated as "use these specific stores" (resulting in no RAG tools), while None means "fetch all available stores". This distinction is subtle and not documented.

Consider:

1. Adding a @field_validator to normalize empty lists to None if they should mean "use all stores"
2. Validating that list items are non-empty strings
3. Documenting the empty list behavior in the field description

@field_validator("vector_store_ids")
@classmethod
def validate_vector_store_ids(
    cls, value: Optional[list[str]]
) -> Optional[list[str]]:
    """
    Validate and normalize vector store IDs.
    
    Parameters:
        value: List of vector store IDs or None
        
    Returns:
        Normalized list or None
        
    Raises:
        ValueError: If list contains empty strings
    """
    if value is None:
        return value
    
    # Normalize empty list to None (fetch all stores)
    if len(value) == 0:
        return None
    
    # Validate no empty strings
    if any(not id.strip() for id in value):
        raise ValueError("vector_store_ids cannot contain empty strings")
    
    return value

As per coding guidelines, Pydantic models should use @field_validator for custom validation.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 6bae4a2 and 9f077e8.

📒 Files selected for processing (4)

• src/app/endpoints/query.py (1 hunks)
• src/app/endpoints/query_v2.py (1 hunks)
• src/app/endpoints/streaming_query.py (1 hunks)
• src/models/requests.py (3 hunks)

🧰 Additional context used

📓 Path-based instructions (4)

src/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

src/**/*.py: Use absolute imports for internal modules in LCS project (e.g., from auth import get_auth_dependency)
All modules must start with descriptive docstrings explaining their purpose
Use logger = logging.getLogger(__name__) pattern for module logging
All functions must include complete type annotations for parameters and return types, using modern syntax (str | int) and Optional[Type] or Type | None
All functions must have docstrings with brief descriptions following Google Python docstring conventions
Function names must use snake_case with descriptive, action-oriented names (get_, validate_, check_)
Avoid in-place parameter modification anti-patterns; return new data structures instead of modifying input parameters
Use async def for I/O operations and external API calls
All classes must include descriptive docstrings explaining their purpose following Google Python docstring conventions
Class names must use PascalCase with descriptive names and standard suffixes: Configuration for config classes, Error/Exception for exceptions, Resolver for strategy patterns, Interface for abstract base classes
Abstract classes must use ABC with @abstractmethod decorators
Include complete type annotations for all class attributes in Python classes
Use import logging and module logger pattern with standard log levels: debug, info, warning, error

Files:

• src/models/requests.py
• src/app/endpoints/streaming_query.py
• src/app/endpoints/query_v2.py
• src/app/endpoints/query.py

src/models/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

src/models/**/*.py: Use @field_validator and @model_validator for custom validation in Pydantic models
Pydantic configuration classes must extend ConfigurationBase; data models must extend BaseModel

Files:

• src/models/requests.py

src/app/endpoints/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

Use FastAPI HTTPException with appropriate status codes for API endpoint error handling

Files:

• src/app/endpoints/streaming_query.py
• src/app/endpoints/query_v2.py
• src/app/endpoints/query.py

src/**/{client,app/endpoints/**}.py

📄 CodeRabbit inference engine (CLAUDE.md)

Handle APIConnectionError from Llama Stack in integration code

Files:

• src/app/endpoints/streaming_query.py
• src/app/endpoints/query_v2.py
• src/app/endpoints/query.py

🔇 Additional comments (1)

src/app/endpoints/query.py (1)

731-744: [Rewritten review comment]
[Classification tag]

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Handle empty list edge case consistently.

The same empty list ambiguity exists here as in query.py: if query_request.vector_store_ids: evaluates to False for both None and [], but only None is falsy. An empty list would skip the fetch-all-stores branch, potentially disabling RAG unintentionally.

For consistency across endpoints and clarity, consider explicit handling:

-    # Get vector stores for RAG tools - use specified ones or fetch all
-    if query_request.vector_store_ids:
-        vector_store_ids = query_request.vector_store_ids
-    else:
-        vector_store_ids = [
-            vector_store.id for vector_store in (await client.vector_stores.list()).data
-        ]
+    # Get vector stores for RAG tools - use specified ones or fetch all
+    if query_request.vector_store_ids is not None and len(query_request.vector_store_ids) > 0:
+        vector_store_ids = query_request.vector_store_ids
+    else:
+        # None or empty list means fetch all available stores
+        vector_store_ids = [
+            vector_store.id for vector_store in (await client.vector_stores.list()).data
+        ]

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	# Get vector stores for RAG tools - use specified ones or fetch all
	if query_request.vector_store_ids:
	vector_store_ids = query_request.vector_store_ids
	else:
	vector_store_ids = [
	vector_store.id for vector_store in (await client.vector_stores.list()).data
	]
	# Get vector stores for RAG tools - use specified ones or fetch all
	if query_request.vector_store_ids is not None and len(query_request.vector_store_ids) > 0:
	vector_store_ids = query_request.vector_store_ids
	else:
	# None or empty list means fetch all available stores
	vector_store_ids = [
	vector_store.id for vector_store in (await client.vector_stores.list()).data
	]

🤖 Prompt for AI Agents

In src/app/endpoints/query_v2.py around lines 746 to 752, the current condition
treats both None and empty list as falsy so an empty list unintentionally
triggers fetching all stores; change the check to explicitly test for None
(e.g., if query_request.vector_store_ids is not None) so that a provided empty
list is preserved (disables RAG) while only None causes the code to fetch all
vector stores; update assignment accordingly and ensure any downstream code
expects an empty list vs populated list.

[are-ces]

LGTM!

[tisnik]

LGTM

[tisnik]

/ok-to-test