mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-07-14 09:06:10 +00:00
feat: rebase and implement file API methods
Signed-off-by: Varsha Prasad Narsing <varshaprasad96@gmail.com>
This commit is contained in:
Varsha Prasad Narsing
parent
61bddfe70e
commit
3d27b7054c
13 changed files with 195 additions and 141 deletions
|
|
@ -4,7 +4,7 @@
|
|||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
||||
|
||||
from typing import Any, Literal
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -23,7 +23,6 @@ class QdrantVectorIOConfig(BaseModel):
|
|||
prefix: str | None = None
|
||||
timeout: int | None = None
|
||||
host: str | None = None
|
||||
distance_metric: Literal["COSINE", "DOT", "EUCLID", "MANHATTAN"] = "COSINE"
|
||||
|
||||
@classmethod
|
||||
def sample_run_config(cls, **kwargs: Any) -> dict[str, Any]:
|
||||
|
|
|
|||
|
|
@ -18,17 +18,7 @@ from llama_stack.apis.vector_dbs import VectorDB
|
|||
from llama_stack.apis.vector_io import (
|
||||
Chunk,
|
||||
QueryChunksResponse,
|
||||
SearchRankingOptions,
|
||||
VectorIO,
|
||||
VectorStoreChunkingStrategy,
|
||||
VectorStoreDeleteResponse,
|
||||
VectorStoreFileContentsResponse,
|
||||
VectorStoreFileObject,
|
||||
VectorStoreFileStatus,
|
||||
VectorStoreListFilesResponse,
|
||||
VectorStoreListResponse,
|
||||
VectorStoreObject,
|
||||
VectorStoreSearchResponsePage,
|
||||
)
|
||||
from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
|
||||
from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig
|
||||
|
|
@ -57,24 +47,41 @@ def convert_id(_id: str) -> str:
|
|||
|
||||
|
||||
class QdrantIndex(EmbeddingIndex):
|
||||
def __init__(self, client: AsyncQdrantClient, collection_name: str, distance_metric: str = "COSINE"):
|
||||
def __init__(self, client: AsyncQdrantClient, collection_name: str):
|
||||
self.client = client
|
||||
self.collection_name = collection_name
|
||||
self.distance_metric = distance_metric
|
||||
self._distance_metric = None # Will be set when collection is created
|
||||
|
||||
async def add_chunks(self, chunks: list[Chunk], embeddings: NDArray):
|
||||
async def add_chunks(self, chunks: list[Chunk], embeddings: NDArray, metadata: dict[str, Any] | None = None):
|
||||
assert len(chunks) == len(embeddings), (
|
||||
f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
|
||||
)
|
||||
|
||||
# Extract distance_metric from metadata if provided, default to COSINE
|
||||
distance_metric = "COSINE" # Default
|
||||
if metadata is not None and "distance_metric" in metadata:
|
||||
distance_metric = metadata["distance_metric"]
|
||||
|
||||
if not await self.client.collection_exists(self.collection_name):
|
||||
# Get distance metric, defaulting to COSINE
|
||||
distance = getattr(models.Distance, self.distance_metric, models.Distance.COSINE)
|
||||
# Create collection with the specified distance metric
|
||||
distance = getattr(models.Distance, distance_metric, models.Distance.COSINE)
|
||||
self._distance_metric = distance_metric
|
||||
|
||||
await self.client.create_collection(
|
||||
self.collection_name,
|
||||
vectors_config=models.VectorParams(size=len(embeddings[0]), distance=distance),
|
||||
)
|
||||
else:
|
||||
# Collection already exists, warn if different distance metric was requested
|
||||
if self._distance_metric is None:
|
||||
# For now, assume COSINE as default since we can't easily extract it from collection info
|
||||
self._distance_metric = "COSINE"
|
||||
|
||||
if self._distance_metric != distance_metric:
|
||||
log.warning(
|
||||
f"Collection {self.collection_name} was created with distance metric '{self._distance_metric}', "
|
||||
f"but '{distance_metric}' was requested. Using existing distance metric."
|
||||
)
|
||||
|
||||
points = []
|
||||
for _i, (chunk, embedding) in enumerate(zip(chunks, embeddings, strict=False)):
|
||||
|
|
@ -90,6 +97,7 @@ class QdrantIndex(EmbeddingIndex):
|
|||
await self.client.upsert(collection_name=self.collection_name, points=points)
|
||||
|
||||
async def query_vector(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
|
||||
# Distance metric is set at collection creation and cannot be changed
|
||||
results = (
|
||||
await self.client.query_points(
|
||||
collection_name=self.collection_name,
|
||||
|
|
@ -170,9 +178,8 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
|
|||
|
||||
# Create metadata collection if it doesn't exist
|
||||
if not await self.client.collection_exists(metadata_collection):
|
||||
# Get distance metric from config, defaulting to COSINE for backward compatibility
|
||||
distance_metric = getattr(self.config, "distance_metric", "COSINE")
|
||||
distance = getattr(models.Distance, distance_metric, models.Distance.COSINE)
|
||||
# Use default distance metric for metadata collection
|
||||
distance = models.Distance.COSINE
|
||||
|
||||
await self.client.create_collection(
|
||||
collection_name=metadata_collection,
|
||||
|
|
@ -226,13 +233,101 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
|
|||
collection_name=metadata_collection, points_selector=models.PointIdsList(points=[convert_id(store_id)])
|
||||
)
|
||||
|
||||
async def _save_openai_vector_store_file(
|
||||
self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
|
||||
) -> None:
|
||||
"""Save vector store file metadata to Qdrant collection metadata."""
|
||||
# Store file metadata in a special collection for vector store file metadata
|
||||
file_metadata_collection = f"{OPENAI_VECTOR_STORES_METADATA_COLLECTION}_files"
|
||||
|
||||
# Create file metadata collection if it doesn't exist
|
||||
if not await self.client.collection_exists(file_metadata_collection):
|
||||
distance = models.Distance.COSINE
|
||||
await self.client.create_collection(
|
||||
collection_name=file_metadata_collection,
|
||||
vectors_config=models.VectorParams(size=1, distance=distance),
|
||||
)
|
||||
|
||||
# Store file metadata as a point with dummy vector
|
||||
file_key = f"{store_id}:{file_id}"
|
||||
await self.client.upsert(
|
||||
collection_name=file_metadata_collection,
|
||||
points=[
|
||||
models.PointStruct(
|
||||
id=convert_id(file_key),
|
||||
vector=[0.0], # Dummy vector
|
||||
payload={"file_info": file_info, "file_contents": file_contents},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
|
||||
"""Load vector store file metadata from Qdrant."""
|
||||
file_metadata_collection = f"{OPENAI_VECTOR_STORES_METADATA_COLLECTION}_files"
|
||||
|
||||
if not await self.client.collection_exists(file_metadata_collection):
|
||||
return {}
|
||||
|
||||
file_key = f"{store_id}:{file_id}"
|
||||
points = await self.client.retrieve(
|
||||
collection_name=file_metadata_collection,
|
||||
ids=[convert_id(file_key)],
|
||||
with_payload=True,
|
||||
)
|
||||
|
||||
if points and points[0].payload and "file_info" in points[0].payload:
|
||||
return points[0].payload["file_info"]
|
||||
return {}
|
||||
|
||||
async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
|
||||
"""Load vector store file contents from Qdrant."""
|
||||
file_metadata_collection = f"{OPENAI_VECTOR_STORES_METADATA_COLLECTION}_files"
|
||||
|
||||
if not await self.client.collection_exists(file_metadata_collection):
|
||||
return []
|
||||
|
||||
file_key = f"{store_id}:{file_id}"
|
||||
points = await self.client.retrieve(
|
||||
collection_name=file_metadata_collection,
|
||||
ids=[convert_id(file_key)],
|
||||
with_payload=True,
|
||||
)
|
||||
|
||||
if points and points[0].payload and "file_contents" in points[0].payload:
|
||||
return points[0].payload["file_contents"]
|
||||
return []
|
||||
|
||||
async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
|
||||
"""Update vector store file metadata in Qdrant."""
|
||||
file_metadata_collection = f"{OPENAI_VECTOR_STORES_METADATA_COLLECTION}_files"
|
||||
|
||||
if not await self.client.collection_exists(file_metadata_collection):
|
||||
return
|
||||
|
||||
# Get existing file contents
|
||||
existing_contents = await self._load_openai_vector_store_file_contents(store_id, file_id)
|
||||
|
||||
# Update with new file info but keep existing contents
|
||||
await self._save_openai_vector_store_file(store_id, file_id, file_info, existing_contents)
|
||||
|
||||
async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
|
||||
"""Delete vector store file metadata from Qdrant."""
|
||||
file_metadata_collection = f"{OPENAI_VECTOR_STORES_METADATA_COLLECTION}_files"
|
||||
|
||||
if await self.client.collection_exists(file_metadata_collection):
|
||||
file_key = f"{store_id}:{file_id}"
|
||||
await self.client.delete(
|
||||
collection_name=file_metadata_collection,
|
||||
points_selector=models.PointIdsList(points=[convert_id(file_key)]),
|
||||
)
|
||||
|
||||
async def register_vector_db(
|
||||
self,
|
||||
vector_db: VectorDB,
|
||||
) -> None:
|
||||
index = VectorDBWithIndex(
|
||||
vector_db=vector_db,
|
||||
index=QdrantIndex(self.client, vector_db.identifier, self.config.distance_metric),
|
||||
index=QdrantIndex(self.client, vector_db.identifier),
|
||||
inference_api=self.inference_api,
|
||||
)
|
||||
|
||||
|
|
@ -253,9 +348,7 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
|
|||
|
||||
index = VectorDBWithIndex(
|
||||
vector_db=vector_db,
|
||||
index=QdrantIndex(
|
||||
client=self.client, collection_name=vector_db.identifier, distance_metric=self.config.distance_metric
|
||||
),
|
||||
index=QdrantIndex(client=self.client, collection_name=vector_db.identifier),
|
||||
inference_api=self.inference_api,
|
||||
)
|
||||
self.cache[vector_db_id] = index
|
||||
|
|
@ -266,12 +359,23 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
|
|||
vector_db_id: str,
|
||||
chunks: list[Chunk],
|
||||
ttl_seconds: int | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
index = await self._get_and_cache_vector_db_index(vector_db_id)
|
||||
if not index:
|
||||
raise ValueError(f"Vector DB {vector_db_id} not found")
|
||||
|
||||
await index.insert_chunks(chunks)
|
||||
# Extract distance_metric from params if provided
|
||||
distance_metric = None
|
||||
if params is not None:
|
||||
distance_metric = params.get("distance_metric")
|
||||
|
||||
# Create metadata dict with distance_metric if provided
|
||||
metadata = None
|
||||
if distance_metric is not None:
|
||||
metadata = {"distance_metric": distance_metric}
|
||||
|
||||
await index.insert_chunks(chunks, metadata=metadata)
|
||||
|
||||
async def query_chunks(
|
||||
self,
|
||||
|
|
@ -284,108 +388,3 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
|
|||
raise ValueError(f"Vector DB {vector_db_id} not found")
|
||||
|
||||
return await index.query_chunks(query, params)
|
||||
|
||||
async def openai_create_vector_store(
|
||||
self,
|
||||
name: str,
|
||||
file_ids: list[str] | None = None,
|
||||
expires_after: dict[str, Any] | None = None,
|
||||
chunking_strategy: dict[str, Any] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
embedding_model: str | None = None,
|
||||
embedding_dimension: int | None = 384,
|
||||
provider_id: str | None = None,
|
||||
provider_vector_db_id: str | None = None,
|
||||
) -> VectorStoreObject:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_list_vector_stores(
|
||||
self,
|
||||
limit: int | None = 20,
|
||||
order: str | None = "desc",
|
||||
after: str | None = None,
|
||||
before: str | None = None,
|
||||
) -> VectorStoreListResponse:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_retrieve_vector_store(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
) -> VectorStoreObject:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_update_vector_store(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
name: str | None = None,
|
||||
expires_after: dict[str, Any] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> VectorStoreObject:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_delete_vector_store(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
) -> VectorStoreDeleteResponse:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_search_vector_store(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
query: str | list[str],
|
||||
filters: dict[str, Any] | None = None,
|
||||
max_num_results: int | None = 10,
|
||||
ranking_options: SearchRankingOptions | None = None,
|
||||
rewrite_query: bool | None = False,
|
||||
search_mode: str | None = "vector",
|
||||
) -> VectorStoreSearchResponsePage:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_attach_file_to_vector_store(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
file_id: str,
|
||||
attributes: dict[str, Any] | None = None,
|
||||
chunking_strategy: VectorStoreChunkingStrategy | None = None,
|
||||
) -> VectorStoreFileObject:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_list_files_in_vector_store(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
limit: int | None = 20,
|
||||
order: str | None = "desc",
|
||||
after: str | None = None,
|
||||
before: str | None = None,
|
||||
filter: VectorStoreFileStatus | None = None,
|
||||
) -> VectorStoreListFilesResponse:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_retrieve_vector_store_file(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
file_id: str,
|
||||
) -> VectorStoreFileObject:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_retrieve_vector_store_file_contents(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
file_id: str,
|
||||
) -> VectorStoreFileContentsResponse:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_update_vector_store_file(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
file_id: str,
|
||||
attributes: dict[str, Any] | None = None,
|
||||
) -> VectorStoreFileObject:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
||||
async def openai_delete_vector_store_file(
|
||||
self,
|
||||
vector_store_id: str,
|
||||
file_id: str,
|
||||
) -> VectorStoreFileObject:
|
||||
raise NotImplementedError("OpenAI Vector Stores API is not supported in Qdrant")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue