From c7eed3ef803bbfe63b514f03d850b41c3ca8c778 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Mon, 20 Oct 2025 19:27:38 -0700 Subject: [PATCH] format --- llama_stack/core/datatypes.py | 2 +- llama_stack/core/resolver.py | 2 +- llama_stack/providers/inline/vector_io/faiss/faiss.py | 2 +- .../inline/vector_io/sqlite_vec/sqlite_vec.py | 6 ++++-- .../providers/remote/vector_io/chroma/chroma.py | 2 +- .../providers/remote/vector_io/milvus/milvus.py | 2 +- .../providers/remote/vector_io/pgvector/pgvector.py | 2 +- .../providers/remote/vector_io/qdrant/qdrant.py | 10 +++++++--- .../providers/remote/vector_io/weaviate/weaviate.py | 2 +- .../utils/memory/openai_vector_store_mixin.py | 2 +- llama_stack/providers/utils/memory/vector_store.py | 2 +- tests/integration/vector_io/test_vector_io.py | 8 ++++++-- tests/unit/providers/vector_io/conftest.py | 2 +- tests/unit/providers/vector_io/test_faiss.py | 2 +- .../vector_io/test_vector_io_openai_vector_stores.py | 2 +- tests/unit/registry/test_registry.py | 8 ++++++-- 16 files changed, 35 insertions(+), 21 deletions(-) diff --git a/llama_stack/core/datatypes.py b/llama_stack/core/datatypes.py index 38c6cc263..5f4775d87 100644 --- a/llama_stack/core/datatypes.py +++ b/llama_stack/core/datatypes.py @@ -23,8 +23,8 @@ 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_stores import VectorStore, VectorStoreInput from llama_stack.apis.vector_io import VectorIO +from llama_stack.apis.vector_stores import VectorStore, VectorStoreInput from llama_stack.core.access_control.datatypes import AccessRule from llama_stack.core.storage.datatypes import ( KVStoreReference, diff --git a/llama_stack/core/resolver.py b/llama_stack/core/resolver.py index 14a628282..681b7d7b6 100644 --- a/llama_stack/core/resolver.py +++ b/llama_stack/core/resolver.py @@ -29,8 +29,8 @@ from llama_stack.apis.scoring_functions import ScoringFunctions from llama_stack.apis.shields import Shields from llama_stack.apis.telemetry import Telemetry from llama_stack.apis.tools import ToolGroups, ToolRuntime -from llama_stack.apis.vector_stores import VectorStores from llama_stack.apis.vector_io import VectorIO +from llama_stack.apis.vector_stores import VectorStores from llama_stack.apis.version import LLAMA_STACK_API_V1ALPHA from llama_stack.core.client import get_client_impl from llama_stack.core.datatypes import ( diff --git a/llama_stack/providers/inline/vector_io/faiss/faiss.py b/llama_stack/providers/inline/vector_io/faiss/faiss.py index 1ba9be89c..cc7614ce8 100644 --- a/llama_stack/providers/inline/vector_io/faiss/faiss.py +++ b/llama_stack/providers/inline/vector_io/faiss/faiss.py @@ -17,8 +17,8 @@ from numpy.typing import NDArray from llama_stack.apis.common.errors import VectorStoreNotFoundError from llama_stack.apis.files import Files from llama_stack.apis.inference import Inference, InterleavedContent -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO +from llama_stack.apis.vector_stores import VectorStore from llama_stack.log import get_logger from llama_stack.providers.datatypes import HealthResponse, HealthStatus, VectorStoresProtocolPrivate from llama_stack.providers.utils.kvstore import kvstore_impl diff --git a/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py b/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py index 839e387fc..393bbf9c5 100644 --- a/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py +++ b/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py @@ -17,8 +17,8 @@ from numpy.typing import NDArray from llama_stack.apis.common.errors import VectorStoreNotFoundError from llama_stack.apis.files import Files from llama_stack.apis.inference import Inference -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO +from llama_stack.apis.vector_stores import VectorStore from llama_stack.log import get_logger from llama_stack.providers.datatypes import VectorStoresProtocolPrivate from llama_stack.providers.utils.kvstore import kvstore_impl @@ -412,7 +412,9 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresPro return [v.vector_store for v in self.cache.values()] async def register_vector_store(self, vector_store: VectorStore) -> None: - index = await SQLiteVecIndex.create(vector_store.embedding_dimension, self.config.db_path, vector_store.identifier) + index = await SQLiteVecIndex.create( + vector_store.embedding_dimension, self.config.db_path, vector_store.identifier + ) self.cache[vector_store.identifier] = VectorStoreWithIndex(vector_store, index, self.inference_api) async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> VectorStoreWithIndex | None: diff --git a/llama_stack/providers/remote/vector_io/chroma/chroma.py b/llama_stack/providers/remote/vector_io/chroma/chroma.py index 799e2d721..a4fd15f77 100644 --- a/llama_stack/providers/remote/vector_io/chroma/chroma.py +++ b/llama_stack/providers/remote/vector_io/chroma/chroma.py @@ -13,8 +13,8 @@ from numpy.typing import NDArray from llama_stack.apis.files import Files from llama_stack.apis.inference import Inference, InterleavedContent -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO +from llama_stack.apis.vector_stores import VectorStore from llama_stack.log import get_logger from llama_stack.providers.datatypes import VectorStoresProtocolPrivate from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py index eaca88aee..ace9ab1c4 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/llama_stack/providers/remote/vector_io/milvus/milvus.py @@ -14,8 +14,8 @@ from pymilvus import AnnSearchRequest, DataType, Function, FunctionType, MilvusC from llama_stack.apis.common.errors import VectorStoreNotFoundError from llama_stack.apis.files import Files from llama_stack.apis.inference import Inference, InterleavedContent -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO +from llama_stack.apis.vector_stores import VectorStore from llama_stack.log import get_logger from llama_stack.providers.datatypes import VectorStoresProtocolPrivate from llama_stack.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py index 61556912f..29cfd673f 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py +++ b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py @@ -16,8 +16,8 @@ from pydantic import BaseModel, TypeAdapter from llama_stack.apis.common.errors import VectorStoreNotFoundError from llama_stack.apis.files import Files from llama_stack.apis.inference import Inference, InterleavedContent -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO +from llama_stack.apis.vector_stores import VectorStore from llama_stack.log import get_logger from llama_stack.providers.datatypes import VectorStoresProtocolPrivate from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str diff --git a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py index 27804f125..266e9bf58 100644 --- a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py +++ b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py @@ -16,7 +16,6 @@ from qdrant_client.models import PointStruct from llama_stack.apis.common.errors import VectorStoreNotFoundError from llama_stack.apis.files import Files from llama_stack.apis.inference import Inference, InterleavedContent -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import ( Chunk, QueryChunksResponse, @@ -24,6 +23,7 @@ from llama_stack.apis.vector_io import ( VectorStoreChunkingStrategy, VectorStoreFileObject, ) +from llama_stack.apis.vector_stores import VectorStore from llama_stack.log import get_logger from llama_stack.providers.datatypes import VectorStoresProtocolPrivate from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig @@ -171,7 +171,9 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtoc for vector_store_data in stored_vector_stores: vector_store = VectorStore.model_validate_json(vector_store_data) - index = VectorStoreWithIndex(vector_store, QdrantIndex(self.client, vector_store.identifier), self.inference_api) + index = VectorStoreWithIndex( + vector_store, QdrantIndex(self.client, vector_store.identifier), self.inference_api + ) self.cache[vector_store.identifier] = index self.openai_vector_stores = await self._load_openai_vector_stores() @@ -186,7 +188,9 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtoc await self.kvstore.set(key=key, value=vector_store.model_dump_json()) index = VectorStoreWithIndex( - vector_store=vector_store, index=QdrantIndex(self.client, vector_store.identifier), inference_api=self.inference_api + vector_store=vector_store, + index=QdrantIndex(self.client, vector_store.identifier), + inference_api=self.inference_api, ) self.cache[vector_store.identifier] = index diff --git a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py index 111ec5cb4..7813f6e5c 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py +++ b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py @@ -16,8 +16,8 @@ from llama_stack.apis.common.content_types import InterleavedContent from llama_stack.apis.common.errors import VectorStoreNotFoundError from llama_stack.apis.files import Files from llama_stack.apis.inference import Inference -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO +from llama_stack.apis.vector_stores import VectorStore from llama_stack.core.request_headers import NeedsRequestProviderData from llama_stack.log import get_logger from llama_stack.providers.datatypes import VectorStoresProtocolPrivate diff --git a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py index 5d53bf172..bccb95871 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -17,7 +17,6 @@ from pydantic import TypeAdapter from llama_stack.apis.common.errors import VectorStoreNotFoundError from llama_stack.apis.files import Files, OpenAIFileObject -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import ( Chunk, OpenAICreateVectorStoreFileBatchRequestWithExtraBody, @@ -43,6 +42,7 @@ from llama_stack.apis.vector_io import ( VectorStoreSearchResponse, VectorStoreSearchResponsePage, ) +from llama_stack.apis.vector_stores import VectorStore from llama_stack.core.id_generation import generate_object_id from llama_stack.log import get_logger from llama_stack.providers.utils.kvstore.api import KVStore diff --git a/llama_stack/providers/utils/memory/vector_store.py b/llama_stack/providers/utils/memory/vector_store.py index 1b331e8fe..6c8746e92 100644 --- a/llama_stack/providers/utils/memory/vector_store.py +++ b/llama_stack/providers/utils/memory/vector_store.py @@ -23,8 +23,8 @@ from llama_stack.apis.common.content_types import ( ) from llama_stack.apis.inference import OpenAIEmbeddingsRequestWithExtraBody from llama_stack.apis.tools import RAGDocument -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import Chunk, ChunkMetadata, QueryChunksResponse +from llama_stack.apis.vector_stores import VectorStore from llama_stack.log import get_logger from llama_stack.models.llama.llama3.tokenizer import Tokenizer from llama_stack.providers.datatypes import Api diff --git a/tests/integration/vector_io/test_vector_io.py b/tests/integration/vector_io/test_vector_io.py index e61532442..a312456b9 100644 --- a/tests/integration/vector_io/test_vector_io.py +++ b/tests/integration/vector_io/test_vector_io.py @@ -49,7 +49,9 @@ def client_with_empty_registry(client_with_models): @vector_provider_wrapper -def test_vector_store_retrieve(client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_id): +def test_vector_store_retrieve( + client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_id +): vector_store_name = "test_vector_store" create_response = client_with_empty_registry.vector_stores.create( name=vector_store_name, @@ -69,7 +71,9 @@ def test_vector_store_retrieve(client_with_empty_registry, embedding_model_id, e @vector_provider_wrapper -def test_vector_store_register(client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_id): +def test_vector_store_register( + client_with_empty_registry, embedding_model_id, embedding_dimension, vector_io_provider_id +): vector_store_name = "test_vector_store" response = client_with_empty_registry.vector_stores.create( name=vector_store_name, diff --git a/tests/unit/providers/vector_io/conftest.py b/tests/unit/providers/vector_io/conftest.py index 63461e752..2951ca2e5 100644 --- a/tests/unit/providers/vector_io/conftest.py +++ b/tests/unit/providers/vector_io/conftest.py @@ -10,8 +10,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import numpy as np import pytest -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import Chunk, ChunkMetadata, QueryChunksResponse +from llama_stack.apis.vector_stores import VectorStore from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig from llama_stack.providers.inline.vector_io.faiss.faiss import FaissIndex, FaissVectorIOAdapter diff --git a/tests/unit/providers/vector_io/test_faiss.py b/tests/unit/providers/vector_io/test_faiss.py index 9b94f4a1a..7b870d16e 100644 --- a/tests/unit/providers/vector_io/test_faiss.py +++ b/tests/unit/providers/vector_io/test_faiss.py @@ -11,8 +11,8 @@ import numpy as np import pytest from llama_stack.apis.files import Files -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import Chunk, QueryChunksResponse +from llama_stack.apis.vector_stores import VectorStore from llama_stack.providers.datatypes import HealthStatus from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig from llama_stack.providers.inline.vector_io.faiss.faiss import ( diff --git a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py index 76958bf17..65d7b7602 100644 --- a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py +++ b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py @@ -12,7 +12,6 @@ import numpy as np import pytest from llama_stack.apis.common.errors import VectorStoreNotFoundError -from llama_stack.apis.vector_stores import VectorStore from llama_stack.apis.vector_io import ( Chunk, OpenAICreateVectorStoreFileBatchRequestWithExtraBody, @@ -21,6 +20,7 @@ from llama_stack.apis.vector_io import ( VectorStoreChunkingStrategyAuto, VectorStoreFileObject, ) +from llama_stack.apis.vector_stores import VectorStore from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import VECTOR_DBS_PREFIX # This test is a unit test for the inline VectorIO providers. This should only contain diff --git a/tests/unit/registry/test_registry.py b/tests/unit/registry/test_registry.py index bccca9379..d4c9786d1 100644 --- a/tests/unit/registry/test_registry.py +++ b/tests/unit/registry/test_registry.py @@ -134,7 +134,9 @@ async def test_duplicate_provider_registration(cached_disk_dist_registry): provider_resource_id="test_vector_store_2", provider_id="baz", # Same provider_id ) - with pytest.raises(ValueError, match="Object of type 'vector_store' and identifier 'test_vector_store_2' already exists"): + with pytest.raises( + ValueError, match="Object of type 'vector_store' and identifier 'test_vector_store_2' already exists" + ): await cached_disk_dist_registry.register(duplicate_vector_store) result = await cached_disk_dist_registry.get("vector_store", "test_vector_store_2") @@ -289,7 +291,9 @@ async def test_double_registration_different_objects(disk_dist_registry): assert result1 is True # Second registration with different data should fail - with pytest.raises(ValueError, match="Object of type 'vector_store' and identifier 'test_vector_store' already exists"): + with pytest.raises( + ValueError, match="Object of type 'vector_store' and identifier 'test_vector_store' already exists" + ): await disk_dist_registry.register(vector_store2) # Verify original object is unchanged