phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2026-01-07 09:39:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

pre-commit

Browse source
This commit is contained in:
Ashwin Bharambe 2025-03-06 20:55:42 -08:00
parent 2313b1bab5
commit 5cab79646d
2 changed files with 15 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

2
docs/source/providers/vector_io/mivus.md
View file

@ -3,7 +3,7 @@ orphan: true
--- ---
# Milvus # Milvus
[Milvus](https://milvus.io/) is an inline and remote vector database provider for Llama Stack. It [Milvus](https://milvus.io/) is an inline and remote vector database provider for Llama Stack. It
allows you to store and query vectors directly within a Milvus database. allows you to store and query vectors directly within a Milvus database.
That means you're not limited to storing vectors in memory or in a separate service. That means you're not limited to storing vectors in memory or in a separate service.

19
llama_stack/providers/remote/vector_io/milvus/milvus.py
View file

@ -4,12 +4,12 @@
# This source code is licensed under the terms described in the LICENSE file in # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # the root directory of this source tree.
import hashlib
import logging import logging
import os import os
import uuid
from typing import Any, Dict, List, Optional, Union from typing import Any, Dict, List, Optional, Union
import hashlib
import uuid
from numpy.typing import NDArray from numpy.typing import NDArray
from pymilvus import MilvusClient from pymilvus import MilvusClient
@ -43,10 +43,15 @@ class MilvusIndex(EmbeddingIndex):
f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}" f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
) )
if not self.client.has_collection(self.collection_name): if not self.client.has_collection(self.collection_name):
self.client.create_collection(self.collection_name, dimension=len(embeddings[0]), auto_id=True, consistency_level=self.consistency_level) self.client.create_collection(
self.collection_name,
dimension=len(embeddings[0]),
auto_id=True,
consistency_level=self.consistency_level,
)
data = [] data = []
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings, strict=False)): for chunk, embedding in zip(chunks, embeddings, strict=False):
chunk_id = generate_chunk_id(chunk.metadata["document_id"], chunk.content) chunk_id = generate_chunk_id(chunk.metadata["document_id"], chunk.content)
data.append( data.append(
@ -79,7 +84,9 @@ class MilvusIndex(EmbeddingIndex):
class MilvusVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate): class MilvusVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
def __init__(self, config: Union[RemoteMilvusVectorIOConfig, InlineMilvusVectorIOConfig], inference_api: Api.inference) -> None: def __init__(
self, config: Union[RemoteMilvusVectorIOConfig, InlineMilvusVectorIOConfig], inference_api: Api.inference
) -> None:
self.config = config self.config = config
self.cache = {} self.cache = {}
self.client = None self.client = None
@ -158,9 +165,11 @@ class MilvusVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
return await index.query_chunks(query, params) return await index.query_chunks(query, params)
def generate_chunk_id(document_id: str, chunk_text: str) -> str: def generate_chunk_id(document_id: str, chunk_text: str) -> str:
"""Generate a unique chunk ID using a hash of document ID and chunk text.""" """Generate a unique chunk ID using a hash of document ID and chunk text."""
hash_input = f"{document_id}:{chunk_text}".encode("utf-8") hash_input = f"{document_id}:{chunk_text}".encode("utf-8")
return str(uuid.UUID(hashlib.md5(hash_input).hexdigest())) return str(uuid.UUID(hashlib.md5(hash_input).hexdigest()))
# TODO: refactor this generate_chunk_id along with the `sqlite-vec` implementation into a separate utils file # TODO: refactor this generate_chunk_id along with the `sqlite-vec` implementation into a separate utils file