fix: Migrate from deprecated vector_dbs to vector_stores API #678

luis5tb · 2025-10-15T13:06:11Z

Description

This PR updates the codebase to use the new vector_stores API instead of the deprecated vector_dbs API, which was removed in llama-stack PR #3774.

Type of change

Related Tickets & Documents

Related Issue [BUG] VectorDB API removed on llamastack #677

Summary by CodeRabbit

Refactor
- Backend now consistently uses "vector store" identifiers instead of "vector database" identifiers across query and streaming flows.
- RAG toolgroup selection and construction updated to accept and pass vector store IDs; streaming flow also includes server names when building toolgroups.
- Conditional handling preserved so toolgroups remain disabled when no IDs are available.
Tests
- Unit tests updated to use vector_stores.list wrappers (data attribute) and vector store id fields; fixtures and assertions reflect vector_store_ids.

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

coderabbitai · 2025-10-15T13:06:39Z

Walkthrough

Replaces vector DB listing with vector store listing in query endpoints; updates get_rag_toolgroups signature and internals to accept vector_store_ids; streaming and non-streaming flows derive IDs from client.vector_stores.list() and tests/mocks updated to the .data response shape and id attribute.

Changes

Cohort / File(s)	Summary
RAG toolgroup parameter migration `src/app/endpoints/query.py`	Signature and internals updated: `vector_db_ids` → `vector_store_ids`; docstring, conditional checks, and `ToolgroupAgentToolGroupWithArgs` usage now pass `vector_store_ids`.
Streaming endpoint alignment `src/app/endpoints/streaming_query.py`	Replaced `vector_dbs.list()` usage with `vector_stores.list()`; compute `vector_store_ids` from response `.data`, pass to `get_rag_toolgroups`, preserve empty-toolgroup → `None` behavior.
Unit test updates `tests/unit/app/endpoints/test_query.py`, `tests/unit/app/endpoints/test_streaming_query.py`	Tests updated to mock `client.vector_stores.list()` returning a wrapper with `.data` list of stores exposing `id`; local vars renamed to `vector_store_ids`; assertions and mocks adjusted to use `.data` and `id`.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant Client as HTTP Client
  participant Endpoint as Query / Streaming Endpoint
  participant SDK as API Client (client.vector_stores)
  participant RAG as get_rag_toolgroups
  participant Agent as Toolgroup Agent

  Client->>Endpoint: Request (query / streaming)
  Endpoint->>SDK: client.vector_stores.list()
  note right of SDK `#D6EAF8`: returns response object with `.data` list
  SDK-->>Endpoint: list response (data -> vector_store objects)
  Endpoint->>Endpoint: extract vector_store_ids from .data (ids)
  Endpoint->>RAG: get_rag_toolgroups(vector_store_ids)
  RAG-->>Endpoint: Toolgroups or None
  Endpoint->>Agent: Build/run agent with toolgroups
  Agent-->>Client: Response / stream

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Verify all renames from vector_db_ids → vector_store_ids across code and docstrings.
Confirm client.vector_stores.list() handling of the .data wrapper and id attribute in endpoints and tests.
Review Toolgroup construction and ToolgroupAgentToolGroupWithArgs arg name changes, and empty-toolgroup → None behavior.

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically describes the main change: migrating from the deprecated vector_dbs API to the vector_stores API, which is the core focus of all file modifications.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

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

📜 Recent 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 721963f and b6ccb45.

📒 Files selected for processing (4)

src/app/endpoints/query.py (2 hunks)
src/app/endpoints/streaming_query.py (1 hunks)
tests/unit/app/endpoints/test_query.py (16 hunks)
tests/unit/app/endpoints/test_streaming_query.py (12 hunks)

🚧 Files skipped from review as they are similar to previous changes (2)

tests/unit/app/endpoints/test_query.py
src/app/endpoints/query.py

🧰 Additional context used

📓 Path-based instructions (5)

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/app/endpoints/streaming_query.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/**/{client,app/endpoints/**}.py

📄 CodeRabbit inference engine (CLAUDE.md)

Handle APIConnectionError from Llama Stack in integration code

Files:

src/app/endpoints/streaming_query.py

tests/{unit,integration}/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

tests/{unit,integration}/**/*.py: Use pytest for all unit and integration tests; do not use unittest framework
Unit tests must achieve 60% code coverage; integration tests must achieve 10% coverage

Files:

tests/unit/app/endpoints/test_streaming_query.py

tests/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

Use pytest-mock with AsyncMock objects for mocking in tests

Files:

tests/unit/app/endpoints/test_streaming_query.py

🧬 Code graph analysis (1)

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

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

get_rag_toolgroups (832-859)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

GitHub Check: build-pr
GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
GitHub Check: e2e_tests (ci)
GitHub Check: e2e_tests (azure)

🔇 Additional comments (7)

tests/unit/app/endpoints/test_streaming_query.py (6)

379-383: Mock structure correctly reflects the new vector_stores API.

The mock setup properly mirrors the production code's access pattern: (await client.vector_stores.list()).data yielding objects with .id attributes.

426-428: Consistent mock pattern for empty vector stores.

The empty .data = [] correctly reflects scenarios where no vector stores are available, resulting in toolgroups=None when no MCP servers are configured.

659-661: LGTM.

Consistent mock update for attachment-focused tests.

1128-1130: LGTM.

Empty vector stores correctly results in toolgroups containing only MCP server names.

1533-1537: LGTM.

Test correctly verifies that no_tools=True bypasses vector store toolgroups even when vector stores are available.

1585-1589: LGTM.

Test correctly verifies that with no_tools=False, the toolgroups include both RAG (from vector stores) and MCP server names.

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

1064-1067: Migration to vector_stores API is correct.

The code properly accesses .data from the list response and .id from each vector store, aligning with the new API structure. The "vector_db_ids" key in the get_rag_toolgroups args dict is the correct parameter name expected by the Llama Stack RAG tool and does not need to be updated.

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

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

coderabbitai

Actionable comments posted: 1

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 52c0f6d and a5f1f16.

📒 Files selected for processing (2)

src/app/endpoints/query.py (2 hunks)
src/app/endpoints/streaming_query.py (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

src/app/endpoints/streaming_query.py

🧰 Additional context used

📓 Path-based instructions (5)

src/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)

Files:

src/app/endpoints/query.py

src/app/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

Use standard FastAPI imports (from fastapi import APIRouter, HTTPException, Request, status, Depends) in FastAPI app code

Files:

src/app/endpoints/query.py

**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

**/*.py: All modules start with descriptive module-level docstrings explaining purpose
Use logger = logging.getLogger(name) for module logging after import logging
Define type aliases at module level for clarity
All functions require docstrings with brief descriptions
Provide complete type annotations for all function parameters and return types
Use typing_extensions.Self in model validators where appropriate
Use modern union syntax (str | int) and Optional[T] or T | None consistently
Function names use snake_case with descriptive, action-oriented prefixes (get_, validate_, check_)
Avoid in-place parameter modification; return new data structures instead of mutating arguments
Use appropriate logging levels: debug, info, warning, error with clear messages
All classes require descriptive docstrings explaining purpose
Class names use PascalCase with conventional suffixes (Configuration, Error/Exception, Resolver, Interface)
Abstract base classes should use abc.ABC and @AbstractMethod for interfaces
Provide complete type annotations for all class attributes
Follow Google Python docstring style for modules, classes, and functions, including Args, Returns, Raises, Attributes sections as needed

Files:

src/app/endpoints/query.py

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

📄 CodeRabbit inference engine (CLAUDE.md)

Use async def for I/O-bound operations and external API calls

Files:

src/app/endpoints/query.py

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

📄 CodeRabbit inference engine (CLAUDE.md)

In API endpoints, raise FastAPI HTTPException with appropriate status codes for error handling

Files:

src/app/endpoints/query.py

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

GitHub Check: e2e_tests (ci)
GitHub Check: e2e_tests (azure)

🔇 Additional comments (2)

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

686-688: LGTM! API migration is correct.

The migration from client.vector_dbs.list() to client.vector_stores.list() correctly implements the API update. The list comprehension properly extracts IDs from the response data.

689-689: LGTM! Function call updated correctly.

The call to get_rag_toolgroups now correctly passes vector_store_ids, maintaining consistency with the API migration.

coderabbitai

Actionable comments posted: 1

📜 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 ad70fef and 721963f.

📒 Files selected for processing (4)

src/app/endpoints/query.py (2 hunks)
src/app/endpoints/streaming_query.py (1 hunks)
tests/unit/app/endpoints/test_query.py (16 hunks)
tests/unit/app/endpoints/test_streaming_query.py (12 hunks)

🧰 Additional context used

📓 Path-based instructions (5)

tests/{unit,integration}/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

tests/{unit,integration}/**/*.py: Use pytest for all unit and integration tests; do not use unittest framework
Unit tests must achieve 60% code coverage; integration tests must achieve 10% coverage

Files:

tests/unit/app/endpoints/test_streaming_query.py
tests/unit/app/endpoints/test_query.py

tests/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

Use pytest-mock with AsyncMock objects for mocking in tests

Files:

tests/unit/app/endpoints/test_streaming_query.py
tests/unit/app/endpoints/test_query.py

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/app/endpoints/query.py
src/app/endpoints/streaming_query.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/query.py
src/app/endpoints/streaming_query.py

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

📄 CodeRabbit inference engine (CLAUDE.md)

Handle APIConnectionError from Llama Stack in integration code

Files:

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

🧬 Code graph analysis (2)

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

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

get_rag_toolgroups (832-859)

tests/unit/app/endpoints/test_query.py (1)

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

get_rag_toolgroups (832-859)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
GitHub Check: e2e_tests (ci)
GitHub Check: e2e_tests (azure)
GitHub Check: build-pr

🔇 Additional comments (10)

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

1064-1072: Migration to vector_stores API looks correct.

The change properly updates from the deprecated vector_dbs API to vector_stores:

Accesses client.vector_stores.list().data to get the list

Uses vector_store.id instead of the old identifier attribute

Passes vector_store_ids to get_rag_toolgroups() matching the updated function signature

The toolgroup construction logic is preserved (empty list converts to None).

tests/unit/app/endpoints/test_streaming_query.py (3)

379-383: Test mock structure correctly updated for vector_stores API.

The test properly mocks the new vector_stores API structure:

mock_vector_store.id instead of identifier

mock_list_response.data wrapper around the list

mock_client.vector_stores.list returns the wrapper

This pattern is consistently applied throughout the test file.

426-428: Consistent mock updates across all test scenarios.

All test cases properly update their mocks to use:

mock_list_response.data = [] for empty vector stores

The vector_stores.list() API instead of vector_dbs.list()

The consistency across tests ensures comprehensive coverage of the API migration.

Also applies to: 484-486, 543-545, 604-606, 659-661, 715-717, 1128-1130, 1209-1211, 1275-1277

1533-1537: Vector store mocks correctly configured for no_tools tests.

Tests for no_tools parameter correctly set up vector store mocks with:

mock_vector_store.id = "VectorDB-1"

mock_list_response.data = [mock_vector_store]

These tests verify that even when vector stores are available, they're bypassed when no_tools=True, which is the expected behavior.

Also applies to: 1585-1589

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

736-739: Vector stores API integration looks correct.

The code properly:

Calls client.vector_stores.list().data to get the list

Accesses vector_store.id attribute on each item

Passes vector_store_ids to get_rag_toolgroups() with the updated signature

The toolgroup construction logic remains consistent (empty list converts to None).

tests/unit/app/endpoints/test_query.py (5)

539-543: Test mocks correctly updated for vector_stores API.

All test cases properly update their mocks to reflect the new API structure:

mock_vector_store.id instead of identifier

mock_vector_stores_response.data wrapper around the list

mock_client.vector_stores.list returns the wrapper

This ensures tests accurately reflect the production code changes.

Also applies to: 579-583, 620-624

669-671: Empty vector stores correctly mocked across test scenarios.

Tests with no vector stores available correctly mock:

mock_vector_stores_response.data = []

Return value from mock_client.vector_stores.list()

These tests verify behavior when RAG is not available, which should result in toolgroups=None being passed to agent.create_turn.

Also applies to: 727-729, 788-790, 851-853, 908-910, 966-968, 1150-1152, 1233-1235, 1294-1296

1407-1411: Shield violation test properly configured with vector stores.

The test for shield violations correctly sets up:

Vector store with id = "VectorDB-1"

Wrapper with mock_vector_stores_response.data = [mock_vector_store]

Assertion at line 1442 verifies get_rag_toolgroups(["VectorDB-1"]) is called

This ensures shield handling works correctly with the vector_stores API.

1448-1457: Direct test of get_rag_toolgroups updated correctly.

The unit test for get_rag_toolgroups properly validates:

Empty vector_store_ids list returns None (lines 1448-1450)

Non-empty list returns a toolgroup with the correct structure (lines 1452-1457)

Args dictionary contains "vector_store_ids" key (line 1457)

This confirms the function signature change and args key update are correct.

1678-1682: No-tools feature tests correctly configured.

Tests verifying the no_tools parameter behavior properly set up vector stores:

Line 1678-1682: no_tools=True test with vector stores available

Line 1735-1739: no_tools=False test with vector stores available

Line 1784 assertion: Verifies RAG toolgroups are included when no_tools=False

These tests confirm that vector stores are correctly bypassed when tools are disabled, regardless of their availability.

Also applies to: 1735-1739

src/app/endpoints/query.py

luis5tb force-pushed the vector_stores branch 2 times, most recently from 93b4b1b to a5f1f16 Compare October 27, 2025 10:49

coderabbitai bot reviewed Oct 27, 2025

View reviewed changes

luis5tb force-pushed the vector_stores branch from a5f1f16 to 7fd8a8f Compare October 27, 2025 11:54

luis5tb force-pushed the vector_stores branch from 7fd8a8f to 8fb26da Compare November 6, 2025 14:24

luis5tb force-pushed the vector_stores branch from 8fb26da to ad70fef Compare November 14, 2025 10:56

luis5tb force-pushed the vector_stores branch from ad70fef to 721963f Compare December 1, 2025 06:27

coderabbitai bot reviewed Dec 1, 2025

View reviewed changes

src/app/endpoints/query.py Show resolved Hide resolved

Move query and streaming_query to vector_stores

b6ccb45

luis5tb force-pushed the vector_stores branch from 721963f to b6ccb45 Compare December 1, 2025 11:42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix: Migrate from deprecated vector_dbs to vector_stores API #678

fix: Migrate from deprecated vector_dbs to vector_stores API #678

luis5tb commented Oct 15, 2025 •

edited by coderabbitai bot

Loading

Uh oh!

coderabbitai bot commented Oct 15, 2025 •

edited

Loading

Uh oh!

coderabbitai bot left a comment

Uh oh!

coderabbitai bot left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

fix: Migrate from deprecated vector_dbs to vector_stores API #678

Are you sure you want to change the base?

fix: Migrate from deprecated vector_dbs to vector_stores API #678

Conversation

luis5tb commented Oct 15, 2025 • edited by coderabbitai bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Related Tickets & Documents

Summary by CodeRabbit

Uh oh!

coderabbitai bot commented Oct 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Estimated code review effort

Pre-merge checks and finishing touches

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

luis5tb commented Oct 15, 2025 •

edited by coderabbitai bot

Loading

coderabbitai bot commented Oct 15, 2025 •

edited

Loading