llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/remote/tool_runtime/model_context_protocol/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .config import MCPProviderConfig
+
+
+async def get_adapter_impl(config: MCPProviderConfig, _deps):
+    from .model_context_protocol import ModelContextProtocolToolRuntimeImpl
+
+    impl = ModelContextProtocolToolRuntimeImpl(config, _deps)
+    await impl.initialize()
+    return impl

llama_stack/providers/remote/tool_runtime/model_context_protocol/config.py

@@ -0,0 +1,20 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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
+
+from pydantic import BaseModel
+
+
+class MCPProviderDataValidator(BaseModel):
+    # mcp_endpoint => dict of headers to send
+    mcp_headers: dict[str, dict[str, str]] | None = None
+
+
+class MCPProviderConfig(BaseModel):
+    @classmethod
+    def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
+        return {}

llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py

@@ -0,0 +1,73 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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
+from urllib.parse import urlparse
+
+from llama_stack.apis.common.content_types import URL
+from llama_stack.apis.datatypes import Api
+from llama_stack.apis.tools import (
+    ListToolDefsResponse,
+    ToolGroup,
+    ToolInvocationResult,
+    ToolRuntime,
+)
+from llama_stack.core.request_headers import NeedsRequestProviderData
+from llama_stack.log import get_logger
+from llama_stack.providers.datatypes import ToolGroupsProtocolPrivate
+from llama_stack.providers.utils.tools.mcp import invoke_mcp_tool, list_mcp_tools
+
+from .config import MCPProviderConfig
+
+logger = get_logger(__name__, category="tools")
+
+
+class ModelContextProtocolToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, NeedsRequestProviderData):
+    def __init__(self, config: MCPProviderConfig, _deps: dict[Api, Any]):
+        self.config = config
+
+    async def initialize(self):
+        pass
+
+    async def register_toolgroup(self, toolgroup: ToolGroup) -> None:
+        pass
+
+    async def unregister_toolgroup(self, toolgroup_id: str) -> None:
+        return
+
+    async def list_runtime_tools(
+        self, tool_group_id: str | None = None, mcp_endpoint: URL | None = None
+    ) -> ListToolDefsResponse:
+        # this endpoint should be retrieved by getting the tool group right?
+        if mcp_endpoint is None:
+            raise ValueError("mcp_endpoint is required")
+        headers = await self.get_headers_from_request(mcp_endpoint.uri)
+        return await list_mcp_tools(mcp_endpoint.uri, headers)
+
+    async def invoke_tool(self, tool_name: str, kwargs: dict[str, Any]) -> ToolInvocationResult:
+        tool = await self.tool_store.get_tool(tool_name)
+        if tool.metadata is None or tool.metadata.get("endpoint") is None:
+            raise ValueError(f"Tool {tool_name} does not have metadata")
+        endpoint = tool.metadata.get("endpoint")
+        if urlparse(endpoint).scheme not in ("http", "https"):
+            raise ValueError(f"Endpoint {endpoint} is not a valid HTTP(S) URL")
+
+        headers = await self.get_headers_from_request(endpoint)
+        return await invoke_mcp_tool(endpoint, headers, tool_name, kwargs)
+
+    async def get_headers_from_request(self, mcp_endpoint_uri: str) -> dict[str, str]:
+        def canonicalize_uri(uri: str) -> str:
+            return f"{urlparse(uri).netloc or ''}/{urlparse(uri).path or ''}"
+
+        headers = {}
+
+        provider_data = self.get_request_provider_data()
+        if provider_data and provider_data.mcp_headers:
+            for uri, values in provider_data.mcp_headers.items():
+                if canonicalize_uri(uri) != canonicalize_uri(mcp_endpoint_uri):
+                    continue
+                headers.update(values)
+        return headers

llama_stack/providers/remote/tool_runtime/tavily_search/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from pydantic import BaseModel
+
+from .config import TavilySearchToolConfig
+from .tavily_search import TavilySearchToolRuntimeImpl
+
+
+class TavilySearchToolProviderDataValidator(BaseModel):
+    tavily_search_api_key: str
+
+
+async def get_adapter_impl(config: TavilySearchToolConfig, _deps):
+    impl = TavilySearchToolRuntimeImpl(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/remote/tool_runtime/tavily_search/config.py

@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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
+
+from pydantic import BaseModel, Field
+
+
+class TavilySearchToolConfig(BaseModel):
+    api_key: str | None = Field(
+        default=None,
+        description="The Tavily Search API Key",
+    )
+    max_results: int = Field(
+        default=3,
+        description="The maximum number of results to return",
+    )
+
+    @classmethod
+    def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]:
+        return {
+            "api_key": "${env.TAVILY_SEARCH_API_KEY:=}",
+            "max_results": 3,
+        }

llama_stack/providers/remote/tool_runtime/tavily_search/tavily_search.py

@@ -0,0 +1,84 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import json
+from typing import Any
+
+import httpx
+
+from llama_stack.apis.common.content_types import URL
+from llama_stack.apis.tools import (
+    ListToolDefsResponse,
+    ToolDef,
+    ToolGroup,
+    ToolInvocationResult,
+    ToolRuntime,
+)
+from llama_stack.core.request_headers import NeedsRequestProviderData
+from llama_stack.providers.datatypes import ToolGroupsProtocolPrivate
+
+from .config import TavilySearchToolConfig
+
+
+class TavilySearchToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, NeedsRequestProviderData):
+    def __init__(self, config: TavilySearchToolConfig):
+        self.config = config
+
+    async def initialize(self):
+        pass
+
+    async def register_toolgroup(self, toolgroup: ToolGroup) -> None:
+        pass
+
+    async def unregister_toolgroup(self, toolgroup_id: str) -> None:
+        return
+
+    def _get_api_key(self) -> str:
+        if self.config.api_key:
+            return self.config.api_key
+
+        provider_data = self.get_request_provider_data()
+        if provider_data is None or not provider_data.tavily_search_api_key:
+            raise ValueError(
+                'Pass Search provider\'s API Key in the header X-LlamaStack-Provider-Data as { "tavily_search_api_key": <your api key>}'
+            )
+        return provider_data.tavily_search_api_key
+
+    async def list_runtime_tools(
+        self, tool_group_id: str | None = None, mcp_endpoint: URL | None = None
+    ) -> ListToolDefsResponse:
+        return ListToolDefsResponse(
+            data=[
+                ToolDef(
+                    name="web_search",
+                    description="Search the web for information",
+                    input_schema={
+                        "type": "object",
+                        "properties": {
+                            "query": {
+                                "type": "string",
+                                "description": "The query to search for",
+                            }
+                        },
+                        "required": ["query"],
+                    },
+                )
+            ]
+        )
+
+    async def invoke_tool(self, tool_name: str, kwargs: dict[str, Any]) -> ToolInvocationResult:
+        api_key = self._get_api_key()
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                "https://api.tavily.com/search",
+                json={"api_key": api_key, "query": kwargs["query"]},
+            )
+            response.raise_for_status()
+
+        return ToolInvocationResult(content=json.dumps(self._clean_tavily_response(response.json())))
+
+    def _clean_tavily_response(self, search_response, top_k=3):
+        return {"query": search_response["query"], "top_k": search_response["results"]}

llama_stack/providers/remote/tool_runtime/wolfram_alpha/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from pydantic import BaseModel
+
+from .config import WolframAlphaToolConfig
+from .wolfram_alpha import WolframAlphaToolRuntimeImpl
+
+__all__ = ["WolframAlphaToolConfig", "WolframAlphaToolRuntimeImpl"]
+
+
+class WolframAlphaToolProviderDataValidator(BaseModel):
+    wolfram_alpha_api_key: str
+
+
+async def get_adapter_impl(config: WolframAlphaToolConfig, _deps):
+    impl = WolframAlphaToolRuntimeImpl(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/remote/tool_runtime/wolfram_alpha/config.py

@@ -0,0 +1,21 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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
+
+from pydantic import BaseModel
+
+
+class WolframAlphaToolConfig(BaseModel):
+    """Configuration for WolframAlpha Tool Runtime"""
+
+    api_key: str | None = None
+
+    @classmethod
+    def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
+        return {
+            "api_key": "${env.WOLFRAM_ALPHA_API_KEY:=}",
+        }

llama_stack/providers/remote/tool_runtime/wolfram_alpha/wolfram_alpha.py

@@ -0,0 +1,140 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import json
+from typing import Any
+
+import httpx
+
+from llama_stack.apis.common.content_types import URL
+from llama_stack.apis.tools import (
+    ListToolDefsResponse,
+    ToolDef,
+    ToolGroup,
+    ToolInvocationResult,
+    ToolRuntime,
+)
+from llama_stack.core.request_headers import NeedsRequestProviderData
+from llama_stack.providers.datatypes import ToolGroupsProtocolPrivate
+
+from .config import WolframAlphaToolConfig
+
+
+class WolframAlphaToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, NeedsRequestProviderData):
+    def __init__(self, config: WolframAlphaToolConfig):
+        self.config = config
+        self.url = "https://api.wolframalpha.com/v2/query"
+
+    async def initialize(self):
+        pass
+
+    async def register_toolgroup(self, toolgroup: ToolGroup) -> None:
+        pass
+
+    async def unregister_toolgroup(self, toolgroup_id: str) -> None:
+        return
+
+    def _get_api_key(self) -> str:
+        if self.config.api_key:
+            return self.config.api_key
+
+        provider_data = self.get_request_provider_data()
+        if provider_data is None or not provider_data.wolfram_alpha_api_key:
+            raise ValueError(
+                'Pass WolframAlpha API Key in the header X-LlamaStack-Provider-Data as { "wolfram_alpha_api_key": <your api key>}'
+            )
+        return provider_data.wolfram_alpha_api_key
+
+    async def list_runtime_tools(
+        self, tool_group_id: str | None = None, mcp_endpoint: URL | None = None
+    ) -> ListToolDefsResponse:
+        return ListToolDefsResponse(
+            data=[
+                ToolDef(
+                    name="wolfram_alpha",
+                    description="Query WolframAlpha for computational knowledge",
+                    input_schema={
+                        "type": "object",
+                        "properties": {
+                            "query": {
+                                "type": "string",
+                                "description": "The query to compute",
+                            }
+                        },
+                        "required": ["query"],
+                    },
+                )
+            ]
+        )
+
+    async def invoke_tool(self, tool_name: str, kwargs: dict[str, Any]) -> ToolInvocationResult:
+        api_key = self._get_api_key()
+        params = {
+            "input": kwargs["query"],
+            "appid": api_key,
+            "format": "plaintext",
+            "output": "json",
+        }
+        async with httpx.AsyncClient() as client:
+            response = await client.get(params=params, url=self.url)
+            response.raise_for_status()
+        return ToolInvocationResult(content=json.dumps(self._clean_wolfram_alpha_response(response.json())))
+
+    def _clean_wolfram_alpha_response(self, wa_response):
+        remove = {
+            "queryresult": [
+                "datatypes",
+                "error",
+                "timedout",
+                "timedoutpods",
+                "numpods",
+                "timing",
+                "parsetiming",
+                "parsetimedout",
+                "recalculate",
+                "id",
+                "host",
+                "server",
+                "related",
+                "version",
+                {
+                    "pods": [
+                        "scanner",
+                        "id",
+                        "error",
+                        "expressiontypes",
+                        "states",
+                        "infos",
+                        "position",
+                        "numsubpods",
+                    ]
+                },
+                "assumptions",
+            ],
+        }
+        for main_key in remove:
+            for key_to_remove in remove[main_key]:
+                try:
+                    if key_to_remove == "assumptions":
+                        if "assumptions" in wa_response[main_key]:
+                            del wa_response[main_key][key_to_remove]
+                    if isinstance(key_to_remove, dict):
+                        for sub_key in key_to_remove:
+                            if sub_key == "pods":
+                                for i in range(len(wa_response[main_key][sub_key])):
+                                    if wa_response[main_key][sub_key][i]["title"] == "Result":
+                                        del wa_response[main_key][sub_key][i + 1 :]
+                                        break
+                            sub_items = wa_response[main_key][sub_key]
+                            for i in range(len(sub_items)):
+                                for sub_key_to_remove in key_to_remove[sub_key]:
+                                    if sub_key_to_remove in sub_items[i]:
+                                        del sub_items[i][sub_key_to_remove]
+                    elif key_to_remove in wa_response[main_key]:
+                        del wa_response[main_key][key_to_remove]
+                except KeyError:
+                    pass
+        return wa_response

llama_stack/providers/remote/vector_io/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_stack/providers/remote/vector_io/chroma/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from llama_stack.providers.datatypes import Api, ProviderSpec
+
+from .config import ChromaVectorIOConfig
+
+
+async def get_adapter_impl(config: ChromaVectorIOConfig, deps: dict[Api, ProviderSpec]):
+    from .chroma import ChromaVectorIOAdapter
+
+    impl = ChromaVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
+    await impl.initialize()
+    return impl

llama_stack/providers/remote/vector_io/chroma/chroma.py

@@ -0,0 +1,215 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+import asyncio
+import json
+from typing import Any
+from urllib.parse import urlparse
+
+import chromadb
+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_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
+from llama_stack.providers.utils.kvstore import kvstore_impl
+from llama_stack.providers.utils.kvstore.api import KVStore
+from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
+from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
+
+from .config import ChromaVectorIOConfig as RemoteChromaVectorIOConfig
+
+log = get_logger(name=__name__, category="vector_io::chroma")
+
+ChromaClientType = chromadb.api.AsyncClientAPI | chromadb.api.ClientAPI
+
+VERSION = "v3"
+VECTOR_DBS_PREFIX = f"vector_stores:chroma:{VERSION}::"
+VECTOR_INDEX_PREFIX = f"vector_index:chroma:{VERSION}::"
+OPENAI_VECTOR_STORES_PREFIX = f"openai_vector_stores:chroma:{VERSION}::"
+OPENAI_VECTOR_STORES_FILES_PREFIX = f"openai_vector_stores_files:chroma:{VERSION}::"
+OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX = f"openai_vector_stores_files_contents:chroma:{VERSION}::"
+
+
+# this is a helper to allow us to use async and non-async chroma clients interchangeably
+async def maybe_await(result):
+    if asyncio.iscoroutine(result):
+        return await result
+    return result
+
+
+class ChromaIndex(EmbeddingIndex):
+    def __init__(self, client: ChromaClientType, collection, kvstore: KVStore | None = None):
+        self.client = client
+        self.collection = collection
+        self.kvstore = kvstore
+
+    async def initialize(self):
+        pass
+
+    async def add_chunks(self, chunks: list[Chunk], embeddings: NDArray):
+        assert len(chunks) == len(embeddings), (
+            f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
+        )
+
+        ids = [f"{c.metadata.get('document_id', '')}:{c.chunk_id}" for c in chunks]
+        await maybe_await(
+            self.collection.add(documents=[chunk.model_dump_json() for chunk in chunks], embeddings=embeddings, ids=ids)
+        )
+
+    async def query_vector(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
+        results = await maybe_await(
+            self.collection.query(
+                query_embeddings=[embedding.tolist()], n_results=k, include=["documents", "distances"]
+            )
+        )
+        distances = results["distances"][0]
+        documents = results["documents"][0]
+
+        chunks = []
+        scores = []
+        for dist, doc in zip(distances, documents, strict=False):
+            try:
+                doc = json.loads(doc)
+                chunk = Chunk(**doc)
+            except Exception:
+                log.exception(f"Failed to parse document: {doc}")
+                continue
+
+            score = 1.0 / float(dist) if dist != 0 else float("inf")
+            if score < score_threshold:
+                continue
+
+            chunks.append(chunk)
+            scores.append(score)
+
+        return QueryChunksResponse(chunks=chunks, scores=scores)
+
+    async def delete(self):
+        await maybe_await(self.client.delete_collection(self.collection.name))
+
+    async def query_keyword(self, query_string: str, k: int, score_threshold: float) -> QueryChunksResponse:
+        raise NotImplementedError("Keyword search is not supported in Chroma")
+
+    async def delete_chunks(self, chunks_for_deletion: list[ChunkForDeletion]) -> None:
+        """Delete a single chunk from the Chroma collection by its ID."""
+        ids = [f"{chunk.document_id}:{chunk.chunk_id}" for chunk in chunks_for_deletion]
+        await maybe_await(self.collection.delete(ids=ids))
+
+    async def query_hybrid(
+        self,
+        embedding: NDArray,
+        query_string: str,
+        k: int,
+        score_threshold: float,
+        reranker_type: str,
+        reranker_params: dict[str, Any] | None = None,
+    ) -> QueryChunksResponse:
+        raise NotImplementedError("Hybrid search is not supported in Chroma")
+
+
+class ChromaVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtocolPrivate):
+    def __init__(
+        self,
+        config: RemoteChromaVectorIOConfig | InlineChromaVectorIOConfig,
+        inference_api: Inference,
+        files_api: Files | None,
+    ) -> None:
+        super().__init__(files_api=files_api, kvstore=None)
+        log.info(f"Initializing ChromaVectorIOAdapter with url: {config}")
+        self.config = config
+        self.inference_api = inference_api
+        self.client = None
+        self.cache = {}
+        self.vector_store_table = None
+
+    async def initialize(self) -> None:
+        self.kvstore = await kvstore_impl(self.config.persistence)
+
+        if isinstance(self.config, RemoteChromaVectorIOConfig):
+            log.info(f"Connecting to Chroma server at: {self.config.url}")
+            url = self.config.url.rstrip("/")
+            parsed = urlparse(url)
+
+            if parsed.path and parsed.path != "/":
+                raise ValueError("URL should not contain a path")
+
+            self.client = await chromadb.AsyncHttpClient(host=parsed.hostname, port=parsed.port)
+        else:
+            log.info(f"Connecting to Chroma local db at: {self.config.db_path}")
+            self.client = chromadb.PersistentClient(path=self.config.db_path)
+        self.openai_vector_stores = await self._load_openai_vector_stores()
+
+    async def shutdown(self) -> None:
+        # Clean up mixin resources (file batch tasks)
+        await super().shutdown()
+
+    async def register_vector_store(self, vector_store: VectorStore) -> None:
+        collection = await maybe_await(
+            self.client.get_or_create_collection(
+                name=vector_store.identifier, metadata={"vector_store": vector_store.model_dump_json()}
+            )
+        )
+        self.cache[vector_store.identifier] = VectorStoreWithIndex(
+            vector_store, ChromaIndex(self.client, collection), self.inference_api
+        )
+
+    async def unregister_vector_store(self, vector_store_id: str) -> None:
+        if vector_store_id not in self.cache:
+            log.warning(f"Vector DB {vector_store_id} not found")
+            return
+
+        await self.cache[vector_store_id].index.delete()
+        del self.cache[vector_store_id]
+
+    async def insert_chunks(self, vector_db_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None:
+        index = await self._get_and_cache_vector_store_index(vector_db_id)
+        if index is None:
+            raise ValueError(f"Vector DB {vector_db_id} not found in Chroma")
+
+        await index.insert_chunks(chunks)
+
+    async def query_chunks(
+        self, vector_db_id: str, query: InterleavedContent, params: dict[str, Any] | None = None
+    ) -> QueryChunksResponse:
+        index = await self._get_and_cache_vector_store_index(vector_db_id)
+
+        if index is None:
+            raise ValueError(f"Vector DB {vector_db_id} not found in Chroma")
+
+        return await index.query_chunks(query, params)
+
+    async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> VectorStoreWithIndex:
+        if vector_store_id in self.cache:
+            return self.cache[vector_store_id]
+
+        # Try to load from kvstore
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before using vector stores.")
+
+        key = f"{VECTOR_DBS_PREFIX}{vector_store_id}"
+        vector_store_data = await self.kvstore.get(key)
+        if not vector_store_data:
+            raise ValueError(f"Vector DB {vector_store_id} not found in Llama Stack")
+
+        vector_store = VectorStore.model_validate_json(vector_store_data)
+        collection = await maybe_await(self.client.get_collection(vector_store_id))
+        if not collection:
+            raise ValueError(f"Vector DB {vector_store_id} not found in Chroma")
+        index = VectorStoreWithIndex(vector_store, ChromaIndex(self.client, collection), self.inference_api)
+        self.cache[vector_store_id] = index
+        return index
+
+    async def delete_chunks(self, store_id: str, chunks_for_deletion: list[ChunkForDeletion]) -> None:
+        """Delete chunks from a Chroma vector store."""
+        index = await self._get_and_cache_vector_store_index(store_id)
+        if not index:
+            raise ValueError(f"Vector DB {store_id} not found")
+
+        await index.index.delete_chunks(chunks_for_deletion)