chore: initial public snapshot for github upload
This commit is contained in:
Your Name
@@ -0,0 +1,4 @@
|
||||
from .main import acreate, asearch, create, search
|
||||
from .vector_store_registry import VectorStoreRegistry
|
||||
|
||||
__all__ = ["search", "asearch", "create", "acreate", "VectorStoreRegistry"]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,58 @@
|
||||
from typing import Any, Dict, cast, get_type_hints
|
||||
|
||||
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
|
||||
from litellm.types.vector_stores import (
|
||||
VectorStoreCreateOptionalRequestParams,
|
||||
VectorStoreSearchOptionalRequestParams,
|
||||
)
|
||||
|
||||
|
||||
class VectorStoreRequestUtils:
|
||||
"""Helper utils for constructing Vector Store search requests"""
|
||||
|
||||
@staticmethod
|
||||
def get_requested_vector_store_search_optional_param(
|
||||
params: Dict[str, Any],
|
||||
vector_store_provider_config: BaseVectorStoreConfig,
|
||||
) -> VectorStoreSearchOptionalRequestParams:
|
||||
"""
|
||||
Filter parameters to only include those defined in VectorStoreSearchOptionalRequestParams.
|
||||
|
||||
Args:
|
||||
params: Dictionary of parameters to filter
|
||||
|
||||
Returns:
|
||||
VectorStoreSearchOptionalRequestParams instance with only the valid parameters
|
||||
"""
|
||||
valid_keys = get_type_hints(VectorStoreSearchOptionalRequestParams).keys()
|
||||
filtered_params = {
|
||||
k: v for k, v in params.items() if k in valid_keys and v is not None
|
||||
}
|
||||
|
||||
optional_params = vector_store_provider_config.map_openai_params(
|
||||
non_default_params=params,
|
||||
optional_params=filtered_params,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
return cast(VectorStoreSearchOptionalRequestParams, optional_params)
|
||||
|
||||
@staticmethod
|
||||
def get_requested_vector_store_create_optional_param(
|
||||
params: Dict[str, Any],
|
||||
) -> VectorStoreCreateOptionalRequestParams:
|
||||
"""
|
||||
Filter parameters to only include those defined in VectorStoreCreateOptionalRequestParams.
|
||||
|
||||
Args:
|
||||
params: Dictionary of parameters to filter
|
||||
|
||||
Returns:
|
||||
VectorStoreCreateOptionalRequestParams instance with only the valid parameters
|
||||
"""
|
||||
valid_keys = get_type_hints(VectorStoreCreateOptionalRequestParams).keys()
|
||||
filtered_params = {
|
||||
k: v for k, v in params.items() if k in valid_keys and v is not None
|
||||
}
|
||||
|
||||
return cast(VectorStoreCreateOptionalRequestParams, filtered_params)
|
||||
@@ -0,0 +1,570 @@
|
||||
# litellm/proxy/vector_stores/vector_store_registry.py
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, get_args
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import remove_items_at_indices
|
||||
from litellm.types.vector_stores import (
|
||||
VECTOR_STORE_OPENAI_PARAMS,
|
||||
LiteLLM_ManagedVectorStore,
|
||||
LiteLLM_ManagedVectorStoreIndex,
|
||||
LiteLLM_ManagedVectorStoreListResponse,
|
||||
LiteLLM_VectorStoreConfig,
|
||||
VectorStoreToolParams,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
else:
|
||||
PrismaClient = Any
|
||||
|
||||
|
||||
class VectorStoreIndexRegistry:
|
||||
def __init__(
|
||||
self, vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = []
|
||||
):
|
||||
self.vector_store_indexes: List[
|
||||
LiteLLM_ManagedVectorStoreIndex
|
||||
] = vector_store_indexes
|
||||
|
||||
def get_vector_store_indexes(self) -> List[LiteLLM_ManagedVectorStoreIndex]:
|
||||
"""
|
||||
Returns the vector store indexes
|
||||
"""
|
||||
return self.vector_store_indexes
|
||||
|
||||
def get_vector_store_index_by_name(
|
||||
self, vector_store_index_name: str
|
||||
) -> Optional[LiteLLM_ManagedVectorStoreIndex]:
|
||||
"""
|
||||
Returns the vector store index by name
|
||||
"""
|
||||
for vector_store_index in self.vector_store_indexes:
|
||||
if vector_store_index.index_name == vector_store_index_name:
|
||||
return vector_store_index
|
||||
return None
|
||||
|
||||
def upsert_vector_store_index(
|
||||
self, vector_store_index: LiteLLM_ManagedVectorStoreIndex
|
||||
):
|
||||
"""
|
||||
Adds a vector store index to the registry.
|
||||
|
||||
If it already exists, it will be updated.
|
||||
"""
|
||||
for i, _vector_store_index in enumerate[LiteLLM_ManagedVectorStoreIndex](
|
||||
self.vector_store_indexes
|
||||
):
|
||||
if _vector_store_index.index_name == vector_store_index.index_name:
|
||||
self.vector_store_indexes[i] = vector_store_index
|
||||
return
|
||||
self.vector_store_indexes.append(vector_store_index)
|
||||
|
||||
def delete_vector_store_index(self, vector_store_index: str):
|
||||
"""
|
||||
Deletes a vector store index from the registry
|
||||
"""
|
||||
self.vector_store_indexes = [
|
||||
index for index in self.vector_store_indexes if index != vector_store_index
|
||||
]
|
||||
|
||||
def is_vector_store_index(self, vector_store_index_name: str) -> bool:
|
||||
"""
|
||||
Returns True if the vector store index is in the registry
|
||||
"""
|
||||
for vector_store_index in self.vector_store_indexes:
|
||||
if vector_store_index.index_name == vector_store_index_name:
|
||||
return True
|
||||
return False
|
||||
|
||||
#########################################################
|
||||
########### DB management helpers for vector stores ###########
|
||||
#########################################################
|
||||
|
||||
@staticmethod
|
||||
async def _get_vector_store_indexes_from_db(
|
||||
prisma_client: Optional[PrismaClient],
|
||||
) -> List[LiteLLM_ManagedVectorStoreIndex]:
|
||||
"""
|
||||
Get vector stores from the database
|
||||
"""
|
||||
vector_stores_from_db: List[LiteLLM_ManagedVectorStoreIndex] = []
|
||||
if prisma_client is not None:
|
||||
_vector_stores_from_db = (
|
||||
await prisma_client.db.litellm_managedvectorstoreindextable.find_many(
|
||||
order={"created_at": "desc"},
|
||||
)
|
||||
)
|
||||
for vector_store in _vector_stores_from_db:
|
||||
_dict_vector_store = dict(vector_store)
|
||||
_litellm_managed_vector_store = LiteLLM_ManagedVectorStoreIndex(
|
||||
**_dict_vector_store
|
||||
)
|
||||
vector_stores_from_db.append(_litellm_managed_vector_store)
|
||||
return vector_stores_from_db
|
||||
|
||||
|
||||
class VectorStoreRegistry:
|
||||
def __init__(self, vector_stores: List[LiteLLM_ManagedVectorStore] = []):
|
||||
self.vector_stores: List[LiteLLM_ManagedVectorStore] = vector_stores
|
||||
self.vector_store_ids_to_vector_store_map: Dict[
|
||||
str, LiteLLM_ManagedVectorStore
|
||||
] = {}
|
||||
|
||||
def _extract_tool_params(self, tool: Dict) -> VectorStoreToolParams:
|
||||
"""
|
||||
Extract supported parameters from a tool definition.
|
||||
|
||||
Dynamically extracts all parameters defined in VECTOR_STORE_OPENAI_PARAMS.
|
||||
"""
|
||||
# Get the list of supported param names from the Literal type
|
||||
supported_params = get_args(VECTOR_STORE_OPENAI_PARAMS)
|
||||
|
||||
# Extract only the params that exist in the tool
|
||||
kwargs = {param: tool.get(param) for param in supported_params if param in tool}
|
||||
|
||||
return VectorStoreToolParams(**kwargs)
|
||||
|
||||
def get_vector_store_ids_to_run(
|
||||
self, non_default_params: Dict, tools: Optional[List[Dict]] = None
|
||||
) -> List[str]:
|
||||
"""
|
||||
Returns the vector store ids to run
|
||||
|
||||
vector_store_ids can be provided in two ways:
|
||||
"""
|
||||
vector_store_ids: List[str] = []
|
||||
|
||||
# 1. check if vector_store_ids is provided in the non_default_params
|
||||
vector_store_ids_param = non_default_params.get("vector_store_ids")
|
||||
if isinstance(vector_store_ids_param, list):
|
||||
vector_store_ids.extend(vector_store_ids_param)
|
||||
|
||||
# 2. check if vector_store_ids is provided as a tool in the request
|
||||
vector_store_ids = self._get_vector_store_ids_from_tool_calls(
|
||||
tools=tools, vector_store_ids=vector_store_ids
|
||||
)
|
||||
|
||||
return list(dict.fromkeys(vector_store_ids))
|
||||
|
||||
def get_and_pop_recognised_vector_store_tools(
|
||||
self,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
vector_store_ids: Optional[List[str]] = None,
|
||||
) -> Dict[str, VectorStoreToolParams]:
|
||||
"""
|
||||
Returns and pops recognized vector store tools from the tools list.
|
||||
|
||||
Args:
|
||||
tools: The tools to extract and remove vector store IDs from
|
||||
vector_store_ids: Mutable list to append found vector_store_ids to
|
||||
|
||||
Returns:
|
||||
Dict mapping vector_store_id to its extracted tool parameters
|
||||
"""
|
||||
params_by_id: Dict[str, VectorStoreToolParams] = {}
|
||||
|
||||
if not tools:
|
||||
return params_by_id
|
||||
|
||||
if vector_store_ids is None:
|
||||
vector_store_ids = []
|
||||
|
||||
tools_to_remove: List[int] = []
|
||||
|
||||
for i, tool in enumerate(tools):
|
||||
tool_vector_store_ids = tool.get("vector_store_ids", [])
|
||||
if not tool_vector_store_ids:
|
||||
continue
|
||||
|
||||
# Check if all vector_store_ids are recognized in the registry
|
||||
recognised = all(
|
||||
any(vs.get("vector_store_id") == vs_id for vs in self.vector_stores)
|
||||
for vs_id in tool_vector_store_ids
|
||||
)
|
||||
|
||||
if recognised:
|
||||
tools_to_remove.append(i)
|
||||
vector_store_ids.extend(tool_vector_store_ids)
|
||||
|
||||
# Extract and store params for each vector store
|
||||
tool_params = self._extract_tool_params(tool)
|
||||
for vs_id in tool_vector_store_ids:
|
||||
params_by_id[vs_id] = tool_params
|
||||
|
||||
# Remove recognized tools from the original list
|
||||
remove_items_at_indices(items=tools, indices=tools_to_remove)
|
||||
|
||||
return params_by_id
|
||||
|
||||
def get_vector_store_to_run(
|
||||
self, non_default_params: Dict, tools: Optional[List[Dict]] = None
|
||||
) -> Optional[LiteLLM_ManagedVectorStore]:
|
||||
"""
|
||||
Returns the vector store to run
|
||||
|
||||
vectore_stores can be run in two ways:
|
||||
1. vector_store_ids is provided in the non_default_params
|
||||
2. vector_store_ids is provided as a tool in the request
|
||||
|
||||
|
||||
This will return the first vector store found in the registry.
|
||||
"""
|
||||
vector_store_ids = self.get_vector_store_ids_to_run(
|
||||
non_default_params=non_default_params, tools=tools
|
||||
)
|
||||
|
||||
# check if the vector store ids are in the registry
|
||||
if len(vector_store_ids) <= 0:
|
||||
return None
|
||||
|
||||
for vector_store_id in vector_store_ids:
|
||||
for vector_store in self.vector_stores:
|
||||
if vector_store.get("vector_store_id") == vector_store_id:
|
||||
return vector_store
|
||||
return None
|
||||
|
||||
def get_litellm_managed_vector_store_from_registry(
|
||||
self, vector_store_id: str
|
||||
) -> Optional[LiteLLM_ManagedVectorStore]:
|
||||
"""
|
||||
Returns the vector store from the registry
|
||||
"""
|
||||
for vector_store in self.vector_stores:
|
||||
if vector_store.get("vector_store_id") == vector_store_id:
|
||||
return vector_store
|
||||
return None
|
||||
|
||||
async def get_litellm_managed_vector_store_from_registry_or_db(
|
||||
self, vector_store_id: str, prisma_client: Optional[PrismaClient] = None
|
||||
) -> Optional[LiteLLM_ManagedVectorStore]:
|
||||
"""
|
||||
Returns the vector store from the registry, falling back to database if not found.
|
||||
This ensures synchronization across multiple instances.
|
||||
"""
|
||||
# First check in-memory registry
|
||||
vector_store = self.get_litellm_managed_vector_store_from_registry(
|
||||
vector_store_id
|
||||
)
|
||||
if vector_store is not None:
|
||||
return vector_store
|
||||
|
||||
# Fall back to database if not found in memory
|
||||
if prisma_client is not None:
|
||||
try:
|
||||
vector_stores_from_db = await self._get_vector_stores_from_db(
|
||||
prisma_client=prisma_client
|
||||
)
|
||||
for db_vector_store in vector_stores_from_db:
|
||||
if db_vector_store.get("vector_store_id") == vector_store_id:
|
||||
# Add to in-memory registry for future use
|
||||
self.add_vector_store_to_registry(vector_store=db_vector_store)
|
||||
return db_vector_store
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
f"Error fetching vector store from database: {str(e)}"
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
def get_litellm_managed_vector_store_from_registry_by_name(
|
||||
self, vector_store_name: str
|
||||
) -> Optional[LiteLLM_ManagedVectorStore]:
|
||||
"""
|
||||
Returns the vector store from the registry by name
|
||||
"""
|
||||
for vector_store in self.vector_stores:
|
||||
if vector_store.get("vector_store_name") == vector_store_name:
|
||||
return vector_store
|
||||
return None
|
||||
|
||||
def pop_vector_stores_to_run(
|
||||
self, non_default_params: Dict, tools: Optional[List[Dict]] = None
|
||||
) -> List[LiteLLM_ManagedVectorStore]:
|
||||
"""
|
||||
Pops the vector stores to run with their tool parameters merged.
|
||||
|
||||
Primary function to use for vector store pre call hook.
|
||||
|
||||
Args:
|
||||
non_default_params: Parameters dict to pop vector_store_ids from
|
||||
tools: Optional list of tools to extract vector store params from
|
||||
|
||||
Returns:
|
||||
List of vector stores with tool parameters merged into litellm_params
|
||||
"""
|
||||
# Pop vector_store_ids from params
|
||||
vector_store_ids: List[str] = (
|
||||
non_default_params.pop("vector_store_ids", None) or []
|
||||
)
|
||||
|
||||
# Extract params from tools and collect IDs
|
||||
params_by_id = self.get_and_pop_recognised_vector_store_tools(
|
||||
tools=tools, vector_store_ids=vector_store_ids
|
||||
)
|
||||
|
||||
vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = []
|
||||
|
||||
for vector_store_id in vector_store_ids:
|
||||
for vector_store in self.vector_stores:
|
||||
if vector_store.get("vector_store_id") == vector_store_id:
|
||||
# Create a copy to avoid modifying the registry
|
||||
vector_store_copy = vector_store.copy()
|
||||
|
||||
# Merge tool params if they exist
|
||||
if vector_store_id in params_by_id:
|
||||
existing_params = (
|
||||
vector_store_copy.get("litellm_params", {}) or {}
|
||||
)
|
||||
tool_params_dict = params_by_id[vector_store_id].to_dict()
|
||||
# Tool params take precedence over existing params
|
||||
tool_params_dict.update(existing_params)
|
||||
vector_store_copy["litellm_params"] = tool_params_dict
|
||||
|
||||
vector_stores_to_run.append(vector_store_copy)
|
||||
break
|
||||
|
||||
return vector_stores_to_run
|
||||
|
||||
async def pop_vector_stores_to_run_with_db_fallback(
|
||||
self,
|
||||
non_default_params: Dict,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
prisma_client: Optional[PrismaClient] = None,
|
||||
) -> List[LiteLLM_ManagedVectorStore]:
|
||||
"""
|
||||
Pops the vector stores to run with their tool parameters merged.
|
||||
Falls back to database if vector stores are not found in memory.
|
||||
This ensures synchronization across multiple instances.
|
||||
|
||||
Primary function to use for vector store pre call hook.
|
||||
|
||||
Args:
|
||||
non_default_params: Parameters dict to pop vector_store_ids from
|
||||
tools: Optional list of tools to extract vector store params from
|
||||
prisma_client: Optional database client for fallback lookup
|
||||
|
||||
Returns:
|
||||
List of vector stores with tool parameters merged into litellm_params
|
||||
"""
|
||||
# Pop vector_store_ids from params
|
||||
vector_store_ids: List[str] = (
|
||||
non_default_params.pop("vector_store_ids", None) or []
|
||||
)
|
||||
|
||||
# Extract params from tools and collect IDs
|
||||
params_by_id = self.get_and_pop_recognised_vector_store_tools(
|
||||
tools=tools, vector_store_ids=vector_store_ids
|
||||
)
|
||||
|
||||
vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = []
|
||||
|
||||
for vector_store_id in vector_store_ids:
|
||||
vector_store = None
|
||||
|
||||
# First check in-memory registry
|
||||
for vs in self.vector_stores:
|
||||
if vs.get("vector_store_id") == vector_store_id:
|
||||
vector_store = vs
|
||||
break
|
||||
|
||||
# Verify vector store still exists in database (if we have DB access)
|
||||
# This ensures deleted vector stores are removed from cache
|
||||
if vector_store is not None and prisma_client is not None:
|
||||
try:
|
||||
# Check if it still exists in database
|
||||
db_vector_store = await prisma_client.db.litellm_managedvectorstorestable.find_unique(
|
||||
where={"vector_store_id": vector_store_id}
|
||||
)
|
||||
if db_vector_store is None:
|
||||
# Vector store was deleted from database, remove from cache
|
||||
verbose_logger.debug(
|
||||
f"Vector store {vector_store_id} found in memory but deleted from database, removing from cache"
|
||||
)
|
||||
self.delete_vector_store_from_registry(
|
||||
vector_store_id=vector_store_id
|
||||
)
|
||||
vector_store = None
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
f"Error verifying vector store {vector_store_id} in database: {str(e)}"
|
||||
)
|
||||
|
||||
# Fall back to database if not found in memory (or was deleted)
|
||||
if vector_store is None and prisma_client is not None:
|
||||
try:
|
||||
vector_store = (
|
||||
await self.get_litellm_managed_vector_store_from_registry_or_db(
|
||||
vector_store_id=vector_store_id, prisma_client=prisma_client
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
f"Error fetching vector store {vector_store_id} from database: {str(e)}"
|
||||
)
|
||||
|
||||
if vector_store is not None:
|
||||
# Create a copy to avoid modifying the registry
|
||||
vector_store_copy = vector_store.copy()
|
||||
|
||||
# Merge tool params if they exist
|
||||
if vector_store_id in params_by_id:
|
||||
existing_params = vector_store_copy.get("litellm_params", {}) or {}
|
||||
tool_params_dict = params_by_id[vector_store_id].to_dict()
|
||||
# Tool params take precedence over existing params
|
||||
tool_params_dict.update(existing_params)
|
||||
vector_store_copy["litellm_params"] = tool_params_dict
|
||||
|
||||
vector_stores_to_run.append(vector_store_copy)
|
||||
|
||||
return vector_stores_to_run
|
||||
|
||||
def _get_vector_store_ids_from_tool_calls(
|
||||
self, tools: Optional[List[Dict]] = None, vector_store_ids: List[str] = []
|
||||
) -> List[str]:
|
||||
"""
|
||||
Returns the vector store ids from the tool calls
|
||||
"""
|
||||
if tools:
|
||||
for tool in tools:
|
||||
if "vector_store_ids" in tool:
|
||||
vector_store_ids.extend(tool["vector_store_ids"])
|
||||
return vector_store_ids
|
||||
|
||||
def load_vector_stores_from_config(self, vector_stores_config: List[Dict]):
|
||||
"""
|
||||
Loads vector stores from the litellm proxy config.yaml
|
||||
"""
|
||||
for vector_store_config in vector_stores_config:
|
||||
# cast to VectorStoreConfig
|
||||
litellm_vector_store_config = LiteLLM_VectorStoreConfig(
|
||||
**vector_store_config
|
||||
)
|
||||
vector_store_name = litellm_vector_store_config.get("vector_store_name")
|
||||
vector_store_litellm_params: Dict[str, Any] = (
|
||||
litellm_vector_store_config.get("litellm_params") or {}
|
||||
)
|
||||
|
||||
vector_store_id = vector_store_litellm_params.get("vector_store_id")
|
||||
if vector_store_id is None:
|
||||
raise ValueError(
|
||||
f"vector_store_id is required for initializing vector store, got vector_store_id={vector_store_id}"
|
||||
)
|
||||
custom_llm_provider = vector_store_litellm_params.get("custom_llm_provider")
|
||||
if custom_llm_provider is None:
|
||||
raise ValueError(
|
||||
f"custom_llm_provider is required for initializing vector store, got custom_llm_provider={custom_llm_provider}"
|
||||
)
|
||||
|
||||
litellm_managed_vector_store = LiteLLM_ManagedVectorStore(
|
||||
vector_store_id=vector_store_id,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=vector_store_litellm_params,
|
||||
vector_store_name=vector_store_name,
|
||||
vector_store_description=vector_store_litellm_params.get(
|
||||
"vector_store_description"
|
||||
),
|
||||
vector_store_metadata=vector_store_litellm_params.get(
|
||||
"vector_store_metadata"
|
||||
),
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
self.vector_stores.append(litellm_managed_vector_store)
|
||||
|
||||
verbose_logger.debug(
|
||||
"all loaded vector stores = %s",
|
||||
json.dumps(self.vector_stores, indent=4, default=str),
|
||||
)
|
||||
|
||||
def list_all_vector_stores(self) -> LiteLLM_ManagedVectorStoreListResponse:
|
||||
"""
|
||||
List all vector stores in the required format
|
||||
|
||||
Returns:
|
||||
LiteLLM_ManagedVectorStoreListResponse: A standardized response with vector store data
|
||||
"""
|
||||
# Prepare the response
|
||||
response = LiteLLM_ManagedVectorStoreListResponse(
|
||||
object="list",
|
||||
data=self.vector_stores,
|
||||
total_count=len(self.vector_stores),
|
||||
current_page=1,
|
||||
total_pages=1,
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
def add_vector_store_to_registry(self, vector_store: LiteLLM_ManagedVectorStore):
|
||||
"""
|
||||
Add a vector store to the registry
|
||||
|
||||
Only add the vector store if it is not already in the registry
|
||||
"""
|
||||
vector_store_id = vector_store.get("vector_store_id")
|
||||
for _vector_store in self.vector_stores:
|
||||
if _vector_store.get("vector_store_id") == vector_store_id:
|
||||
return
|
||||
self.vector_stores.append(vector_store)
|
||||
|
||||
def delete_vector_store_from_registry(self, vector_store_id: str):
|
||||
"""
|
||||
Delete a vector store from the registry
|
||||
"""
|
||||
self.vector_stores = [
|
||||
vector_store
|
||||
for vector_store in self.vector_stores
|
||||
if vector_store.get("vector_store_id") != vector_store_id
|
||||
]
|
||||
|
||||
def update_vector_store_in_registry(
|
||||
self, vector_store_id: str, updated_data: LiteLLM_ManagedVectorStore
|
||||
):
|
||||
"""Update or add a vector store in the registry"""
|
||||
for i, vector_store in enumerate(self.vector_stores):
|
||||
if vector_store.get("vector_store_id") == vector_store_id:
|
||||
self.vector_stores[i] = updated_data
|
||||
return
|
||||
self.vector_stores.append(updated_data)
|
||||
|
||||
#########################################################
|
||||
########### DB management helpers for vector stores ###########
|
||||
#########################################################
|
||||
|
||||
@staticmethod
|
||||
async def _get_vector_stores_from_db(
|
||||
prisma_client: Optional[PrismaClient],
|
||||
) -> List[LiteLLM_ManagedVectorStore]:
|
||||
"""
|
||||
Get vector stores from the database
|
||||
"""
|
||||
vector_stores_from_db: List[LiteLLM_ManagedVectorStore] = []
|
||||
if prisma_client is not None:
|
||||
_vector_stores_from_db = (
|
||||
await prisma_client.db.litellm_managedvectorstorestable.find_many(
|
||||
order={"created_at": "desc"},
|
||||
)
|
||||
)
|
||||
for vector_store in _vector_stores_from_db:
|
||||
_dict_vector_store = dict(vector_store)
|
||||
_litellm_managed_vector_store = LiteLLM_ManagedVectorStore(
|
||||
**_dict_vector_store
|
||||
)
|
||||
vector_stores_from_db.append(_litellm_managed_vector_store)
|
||||
return vector_stores_from_db
|
||||
|
||||
def get_credentials_for_vector_store(self, vector_store_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get the credentials for a vector store
|
||||
|
||||
Returns a dictionary of unpacked credentials for the vector store to use for the request
|
||||
"""
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
|
||||
for vector_store in self.vector_stores:
|
||||
if vector_store.get("vector_store_id") == vector_store_id:
|
||||
credentials = vector_store.get("litellm_credential_name")
|
||||
if credentials:
|
||||
return CredentialAccessor.get_credential_values(credentials)
|
||||
return {}
|
||||
Reference in New Issue
Block a user