From 375650c6b397b7a831edfe4899e6a8bd16ba0525 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Mon, 20 Oct 2025 19:16:01 -0700 Subject: [PATCH] 1: rename vector db to store, remove external APIs --- llama_stack/apis/resource.py | 4 +- llama_stack/apis/vector_dbs/vector_dbs.py | 67 ++++------------------- llama_stack/apis/vector_io/vector_io.py | 10 ++-- llama_stack/core/datatypes.py | 10 ++-- llama_stack/core/routers/vector_io.py | 19 ------- 5 files changed, 25 insertions(+), 85 deletions(-) diff --git a/llama_stack/apis/resource.py b/llama_stack/apis/resource.py index 7c4130f7d..dafdb28b0 100644 --- a/llama_stack/apis/resource.py +++ b/llama_stack/apis/resource.py @@ -13,7 +13,7 @@ from pydantic import BaseModel, Field class ResourceType(StrEnum): model = "model" shield = "shield" - vector_db = "vector_db" + vector_store = "vector_store" dataset = "dataset" scoring_function = "scoring_function" benchmark = "benchmark" @@ -34,4 +34,4 @@ class Resource(BaseModel): provider_id: str = Field(description="ID of the provider that owns this resource") - type: ResourceType = Field(description="Type of resource (e.g. 'model', 'shield', 'vector_db', etc.)") + type: ResourceType = Field(description="Type of resource (e.g. 'model', 'shield', 'vector_store', etc.)") diff --git a/llama_stack/apis/vector_dbs/vector_dbs.py b/llama_stack/apis/vector_dbs/vector_dbs.py index 0368095cb..5e618819d 100644 --- a/llama_stack/apis/vector_dbs/vector_dbs.py +++ b/llama_stack/apis/vector_dbs/vector_dbs.py @@ -4,90 +4,47 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -from typing import Literal, Protocol, runtime_checkable +from typing import Literal from pydantic import BaseModel from llama_stack.apis.resource import Resource, ResourceType -from llama_stack.schema_utils import json_schema_type -@json_schema_type -class VectorDB(Resource): +# Internal resource type for storing the vector store routing and other information +class VectorStore(Resource): """Vector database resource for storing and querying vector embeddings. - :param type: Type of resource, always 'vector_db' for vector databases + :param type: Type of resource, always 'vector_store' for vector stores :param embedding_model: Name of the embedding model to use for vector generation :param embedding_dimension: Dimension of the embedding vectors """ - type: Literal[ResourceType.vector_db] = ResourceType.vector_db + type: Literal[ResourceType.vector_store] = ResourceType.vector_store embedding_model: str embedding_dimension: int - vector_db_name: str | None = None @property - def vector_db_id(self) -> str: + def vector_store_id(self) -> str: return self.identifier @property - def provider_vector_db_id(self) -> str | None: + def provider_vector_store_id(self) -> str | None: return self.provider_resource_id -class VectorDBInput(BaseModel): +class VectorStoreInput(BaseModel): """Input parameters for creating or configuring a vector database. - :param vector_db_id: Unique identifier for the vector database + :param vector_store_id: Unique identifier for the vector store :param embedding_model: Name of the embedding model to use for vector generation :param embedding_dimension: Dimension of the embedding vectors - :param provider_vector_db_id: (Optional) Provider-specific identifier for the vector database + :param provider_vector_store_id: (Optional) Provider-specific identifier for the vector store """ - vector_db_id: str + vector_store_id: str embedding_model: str embedding_dimension: int provider_id: str | None = None - provider_vector_db_id: str | None = None - - -class ListVectorDBsResponse(BaseModel): - """Response from listing vector databases. - - :param data: List of vector databases - """ - - data: list[VectorDB] - - -@runtime_checkable -class VectorDBs(Protocol): - """Internal protocol for vector_dbs routing - no public API endpoints.""" - - async def list_vector_dbs(self) -> ListVectorDBsResponse: - """Internal method to list vector databases.""" - ... - - async def get_vector_db( - self, - vector_db_id: str, - ) -> VectorDB: - """Internal method to get a vector database by ID.""" - ... - - async def register_vector_db( - self, - vector_db_id: str, - embedding_model: str, - embedding_dimension: int | None = 384, - provider_id: str | None = None, - vector_db_name: str | None = None, - provider_vector_db_id: str | None = None, - ) -> VectorDB: - """Internal method to register a vector database.""" - ... - - async def unregister_vector_db(self, vector_db_id: str) -> None: - """Internal method to unregister a vector database.""" - ... + provider_vector_store_id: str | None = None diff --git a/llama_stack/apis/vector_io/vector_io.py b/llama_stack/apis/vector_io/vector_io.py index a309c47f9..2f3a11424 100644 --- a/llama_stack/apis/vector_io/vector_io.py +++ b/llama_stack/apis/vector_io/vector_io.py @@ -15,7 +15,7 @@ from fastapi import Body from pydantic import BaseModel, Field from llama_stack.apis.inference import InterleavedContent -from llama_stack.apis.vector_dbs import VectorDB +from llama_stack.apis.vector_dbs import VectorStore from llama_stack.apis.version import LLAMA_STACK_API_V1 from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol from llama_stack.providers.utils.vector_io.vector_utils import generate_chunk_id @@ -517,17 +517,18 @@ class OpenAICreateVectorStoreFileBatchRequestWithExtraBody(BaseModel, extra="all chunking_strategy: VectorStoreChunkingStrategy | None = None -class VectorDBStore(Protocol): - def get_vector_db(self, vector_db_id: str) -> VectorDB | None: ... +class VectorStoreTable(Protocol): + def get_vector_store(self, vector_store_id: str) -> VectorStore | None: ... @runtime_checkable @trace_protocol class VectorIO(Protocol): - vector_db_store: VectorDBStore | None = None + vector_store_table: VectorStoreTable | None = None # this will just block now until chunks are inserted, but it should # probably return a Job instance which can be polled for completion + # TODO: rename vector_db_id to vector_store_id once Stainless is working @webmethod(route="/vector-io/insert", method="POST", level=LLAMA_STACK_API_V1) async def insert_chunks( self, @@ -546,6 +547,7 @@ class VectorIO(Protocol): """ ... + # TODO: rename vector_db_id to vector_store_id once Stainless is working @webmethod(route="/vector-io/query", method="POST", level=LLAMA_STACK_API_V1) async def query_chunks( self, diff --git a/llama_stack/core/datatypes.py b/llama_stack/core/datatypes.py index e8cb36a02..abac10268 100644 --- a/llama_stack/core/datatypes.py +++ b/llama_stack/core/datatypes.py @@ -23,7 +23,7 @@ from llama_stack.apis.scoring import Scoring from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnInput from llama_stack.apis.shields import Shield, ShieldInput from llama_stack.apis.tools import ToolGroup, ToolGroupInput, ToolRuntime -from llama_stack.apis.vector_dbs import VectorDB, VectorDBInput +from llama_stack.apis.vector_dbs import VectorStore, VectorStoreInput from llama_stack.apis.vector_io import VectorIO from llama_stack.core.access_control.datatypes import AccessRule from llama_stack.core.storage.datatypes import ( @@ -71,7 +71,7 @@ class ShieldWithOwner(Shield, ResourceWithOwner): pass -class VectorDBWithOwner(VectorDB, ResourceWithOwner): +class VectorStoreWithOwner(VectorStore, ResourceWithOwner): pass @@ -91,12 +91,12 @@ class ToolGroupWithOwner(ToolGroup, ResourceWithOwner): pass -RoutableObject = Model | Shield | VectorDB | Dataset | ScoringFn | Benchmark | ToolGroup +RoutableObject = Model | Shield | VectorStore | Dataset | ScoringFn | Benchmark | ToolGroup RoutableObjectWithProvider = Annotated[ ModelWithOwner | ShieldWithOwner - | VectorDBWithOwner + | VectorStoreWithOwner | DatasetWithOwner | ScoringFnWithOwner | BenchmarkWithOwner @@ -427,7 +427,7 @@ class RegisteredResources(BaseModel): models: list[ModelInput] = Field(default_factory=list) shields: list[ShieldInput] = Field(default_factory=list) - vector_dbs: list[VectorDBInput] = Field(default_factory=list) + vector_stores: list[VectorStoreInput] = Field(default_factory=list) datasets: list[DatasetInput] = Field(default_factory=list) scoring_fns: list[ScoringFnInput] = Field(default_factory=list) benchmarks: list[BenchmarkInput] = Field(default_factory=list) diff --git a/llama_stack/core/routers/vector_io.py b/llama_stack/core/routers/vector_io.py index bfc5f7164..11f6b5dbf 100644 --- a/llama_stack/core/routers/vector_io.py +++ b/llama_stack/core/routers/vector_io.py @@ -71,25 +71,6 @@ class VectorIORouter(VectorIO): raise ValueError(f"Embedding model '{embedding_model_id}' not found or not an embedding model") - async def register_vector_db( - self, - vector_db_id: str, - embedding_model: str, - embedding_dimension: int | None = 384, - provider_id: str | None = None, - vector_db_name: str | None = None, - provider_vector_db_id: str | None = None, - ) -> None: - logger.debug(f"VectorIORouter.register_vector_db: {vector_db_id}, {embedding_model}") - await self.routing_table.register_vector_db( - vector_db_id, - embedding_model, - embedding_dimension, - provider_id, - vector_db_name, - provider_vector_db_id, - ) - async def insert_chunks( self, vector_db_id: str,