phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-10 04:08:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

updated batching and config

Browse source 
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
Francisco Javier Arceo 2025-02-18 00:09:22 -05:00
parent 663396ec5a
commit 9ee59443aa
2 changed files with 5 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

1
llama_stack/providers/inline/vector_io/sqlite_vec/config.py
View file

@ -18,7 +18,6 @@ from llama_stack.providers.utils.kvstore.config import (
class SQLiteVectorIOConfig(BaseModel):
db_path: str
kvstore: KVStoreConfig
batch_size: bool = 500
@classmethod
def sample_run_config(cls, __distro_dir__: str) -> Dict[str, Any]:

10
llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py
View file

@ -69,7 +69,7 @@ class SQLiteVecIndex(EmbeddingIndex):
cur.execute(f"DROP TABLE IF EXISTS {self.vector_table};")
self.connection.commit()
async def add_chunks(self, chunks: List[Chunk], embeddings: NDArray):
async def add_chunks(self, chunks: List[Chunk], embeddings: NDArray, batch_size: int = 500):
"""
Add new chunks along with their embeddings using batch inserts.
For each chunk, we insert its JSON into the metadata table and then insert its
@ -80,9 +80,9 @@ class SQLiteVecIndex(EmbeddingIndex):
try:
# Start transaction
cur.execute("BEGIN TRANSACTION")
for i in range(0, len(chunks), self.config.batch_size):
batch_chunks = chunks[i : i + self.config.batch_size]
batch_embeddings = embeddings[i : i + self.config.batch_size]
for i in range(0, len(chunks), batch_size):
batch_chunks = chunks[i : i + batch_size]
batch_embeddings = embeddings[i : i + batch_size]
# Prepare metadata inserts
metadata_data = [
(f"{chunk.metadata['document_id']}:chunk-{j}", chunk.model_dump_json())
@ -218,7 +218,7 @@ class SQLiteVecVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
if vector_db_id not in self.cache:
raise ValueError(f"Vector DB {vector_db_id} not found. Found: {list(self.cache.keys())}")
# The VectorDBWithIndex helper is expected to compute embeddings via the inference_api
# and then call our index’s add_chunks.
# and then call our index's add_chunks.
await self.cache[vector_db_id].insert_chunks(chunks)
async def query_chunks(