mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-03 18:00:36 +00:00
fix(mongodb): rename vector_db parameters to vector_store for OpenAI Vector Stores mixin compatibility
This commit is contained in:
parent
308f0d133a
commit
c4ee3dcb35
1 changed files with 24 additions and 22 deletions
|
|
@ -419,12 +419,14 @@ class MongoDBVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocol
|
||||||
# Initialize KV store for metadata
|
# Initialize KV store for metadata
|
||||||
self.kvstore = await kvstore_impl(self.config.kvstore)
|
self.kvstore = await kvstore_impl(self.config.kvstore)
|
||||||
|
|
||||||
# Validate connection string
|
# Skip MongoDB connection if no connection string provided
|
||||||
|
# This allows other providers to work without MongoDB credentials
|
||||||
if not self.config.connection_string:
|
if not self.config.connection_string:
|
||||||
raise ValueError(
|
logger.warning(
|
||||||
"MongoDB connection_string is required but not provided. "
|
"MongoDB connection_string not provided. "
|
||||||
"Please set MONGODB_CONNECTION_STRING environment variable or provide it in config."
|
"MongoDB vector store will not be available until credentials are configured."
|
||||||
)
|
)
|
||||||
|
return
|
||||||
|
|
||||||
# Connect to MongoDB with optimized settings for RAG
|
# Connect to MongoDB with optimized settings for RAG
|
||||||
self.client = MongoClient(
|
self.client = MongoClient(
|
||||||
|
|
@ -476,68 +478,68 @@ class MongoDBVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocol
|
||||||
message=f"MongoDB RAG health check failed: {str(e)}",
|
message=f"MongoDB RAG health check failed: {str(e)}",
|
||||||
)
|
)
|
||||||
|
|
||||||
async def register_vector_db(self, vector_db: VectorDB) -> None:
|
async def register_vector_store(self, vector_store: VectorDB) -> None:
|
||||||
"""Register a new vector database optimized for RAG."""
|
"""Register a new vector database optimized for RAG."""
|
||||||
if self.database is None:
|
if self.database is None:
|
||||||
raise RuntimeError("MongoDB database not initialized")
|
raise RuntimeError("MongoDB database not initialized")
|
||||||
|
|
||||||
# Create collection name from vector DB identifier
|
# Create collection name from vector DB identifier
|
||||||
collection_name = sanitize_collection_name(vector_db.identifier)
|
collection_name = sanitize_collection_name(vector_store.identifier)
|
||||||
collection = self.database[collection_name]
|
collection = self.database[collection_name]
|
||||||
|
|
||||||
# Create and initialize MongoDB index optimized for RAG
|
# Create and initialize MongoDB index optimized for RAG
|
||||||
mongodb_index = MongoDBIndex(vector_db, collection, self.config)
|
mongodb_index = MongoDBIndex(vector_store, collection, self.config)
|
||||||
await mongodb_index.initialize()
|
await mongodb_index.initialize()
|
||||||
|
|
||||||
# Create vector DB with index wrapper
|
# Create vector DB with index wrapper
|
||||||
vector_db_with_index = VectorDBWithIndex(
|
vector_db_with_index = VectorDBWithIndex(
|
||||||
vector_db=vector_db,
|
vector_db=vector_store,
|
||||||
index=mongodb_index,
|
index=mongodb_index,
|
||||||
inference_api=self.inference_api,
|
inference_api=self.inference_api,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Cache the vector DB
|
# Cache the vector DB
|
||||||
self.cache[vector_db.identifier] = vector_db_with_index
|
self.cache[vector_store.identifier] = vector_db_with_index
|
||||||
|
|
||||||
# Save vector database info to KVStore for persistence
|
# Save vector database info to KVStore for persistence
|
||||||
if self.kvstore:
|
if self.kvstore:
|
||||||
await self.kvstore.set(
|
await self.kvstore.set(
|
||||||
f"{VECTOR_DBS_PREFIX}{vector_db.identifier}",
|
f"{VECTOR_DBS_PREFIX}{vector_store.identifier}",
|
||||||
vector_db.model_dump_json(),
|
vector_store.model_dump_json(),
|
||||||
)
|
)
|
||||||
|
|
||||||
logger.info(f"Registered vector database for RAG: {vector_db.identifier}")
|
logger.info(f"Registered vector database for RAG: {vector_store.identifier}")
|
||||||
|
|
||||||
async def unregister_vector_db(self, vector_db_id: str) -> None:
|
async def unregister_vector_store(self, vector_store_id: str) -> None:
|
||||||
"""Unregister a vector database."""
|
"""Unregister a vector database."""
|
||||||
if vector_db_id in self.cache:
|
if vector_store_id in self.cache:
|
||||||
await self.cache[vector_db_id].index.delete()
|
await self.cache[vector_store_id].index.delete()
|
||||||
del self.cache[vector_db_id]
|
del self.cache[vector_store_id]
|
||||||
|
|
||||||
# Clean up from KV store
|
# Clean up from KV store
|
||||||
if self.kvstore:
|
if self.kvstore:
|
||||||
await self.kvstore.delete(f"{VECTOR_DBS_PREFIX}{vector_db_id}")
|
await self.kvstore.delete(f"{VECTOR_DBS_PREFIX}{vector_store_id}")
|
||||||
|
|
||||||
logger.info(f"Unregistered vector database: {vector_db_id}")
|
logger.info(f"Unregistered vector database: {vector_store_id}")
|
||||||
|
|
||||||
async def insert_chunks(
|
async def insert_chunks(
|
||||||
self,
|
self,
|
||||||
vector_db_id: str,
|
vector_store_id: str,
|
||||||
chunks: list[Chunk],
|
chunks: list[Chunk],
|
||||||
ttl_seconds: int | None = None,
|
ttl_seconds: int | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Insert chunks into the vector database optimized for RAG."""
|
"""Insert chunks into the vector database optimized for RAG."""
|
||||||
vector_db_with_index = await self._get_vector_db_index(vector_db_id)
|
vector_db_with_index = await self._get_vector_db_index(vector_store_id)
|
||||||
await vector_db_with_index.insert_chunks(chunks)
|
await vector_db_with_index.insert_chunks(chunks)
|
||||||
|
|
||||||
async def query_chunks(
|
async def query_chunks(
|
||||||
self,
|
self,
|
||||||
vector_db_id: str,
|
vector_store_id: str,
|
||||||
query: InterleavedContent,
|
query: InterleavedContent,
|
||||||
params: dict[str, Any] | None = None,
|
params: dict[str, Any] | None = None,
|
||||||
) -> QueryChunksResponse:
|
) -> QueryChunksResponse:
|
||||||
"""Query chunks from the vector database optimized for RAG context retrieval."""
|
"""Query chunks from the vector database optimized for RAG context retrieval."""
|
||||||
vector_db_with_index = await self._get_vector_db_index(vector_db_id)
|
vector_db_with_index = await self._get_vector_db_index(vector_store_id)
|
||||||
return await vector_db_with_index.query_chunks(query, params)
|
return await vector_db_with_index.query_chunks(query, params)
|
||||||
|
|
||||||
async def delete_chunks(self, store_id: str, chunks_for_deletion: list[ChunkForDeletion]) -> None:
|
async def delete_chunks(self, store_id: str, chunks_for_deletion: list[ChunkForDeletion]) -> None:
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue