tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 18:24:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feat] - [Backend] Search APIs - Allow storing configured Search APIs in DB (#15862)

Browse Source 
* add LiteLLM_SearchToolsTable

* init SearchToolRegistry

* fix add SearchToolRegistry

* fix add SearchToolRegistry

* fix handling search tool management

* fix search imports

* fix registry

* init search tools in memory

* fix init tools in mem

* fix TypedDict def

* add new SCHEMA

* bump proxy extras

* add LiteLLM_SearchToolsTable_search_tool_name_key

* bump extras with migration

* fix working CRUD Ops

* fix: _init_search_tools_in_db
This commit is contained in:
Ishaan Jaff
2025-10-23 17:57:49 -07:00
committed by GitHub
parent bf47c25de0
commit d8ea1665c7
16 changed files with 1533 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.28-py3-none-any.whl
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.29-py3-none-any.whl
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.29.tar.gz
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/litellm_proxy_extras/migrations/20251023141814_add_search_tool_table/migration.sql
+15
View File
@@ -0,0 +1,15 @@
-- CreateTable
CREATE TABLE "LiteLLM_SearchToolsTable" (
"search_tool_id" TEXT NOT NULL,
"search_tool_name" TEXT NOT NULL,
"litellm_params" JSONB NOT NULL,
"search_tool_info" JSONB,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_SearchToolsTable_pkey" PRIMARY KEY ("search_tool_id")
);
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_SearchToolsTable_search_tool_name_key" ON "LiteLLM_SearchToolsTable"("search_tool_name");
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+10
View File
@@ -570,4 +570,14 @@ model LiteLLM_HealthCheckTable {
@@index([model_name])
@@index([checked_at])
@@index([status])
}
// Search Tools table for storing search tool configurations
model LiteLLM_SearchToolsTable {
search_tool_id String @id @default(uuid())
search_tool_name String @unique
litellm_params Json
search_tool_info Json?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}
litellm-proxy-extras/pyproject.toml
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-proxy-extras"
version = "0.2.27"
version = "0.2.29"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
authors = ["BerriAI"]
readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "0.2.27"
version = "0.2.29"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-proxy-extras==",
litellm/proxy/proxy_server.py
+28 -2
View File
@@ -336,6 +336,9 @@ from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router
from litellm.proxy.response_api_endpoints.endpoints import router as response_router
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.search_endpoints.endpoints import router as search_router
from litellm.proxy.search_endpoints.search_tool_management import (
router as search_tool_management_router,
)
from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router
from litellm.proxy.spend_tracking.spend_management_endpoints import (
router as spend_management_router,
@@ -397,7 +400,9 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
LiteLLM_UpperboundKeyGenerateParams,
)
from litellm.types.realtime import RealtimeQueryParams
from litellm.types.router import DeploymentTypedDict
from litellm.types.router import (
DeploymentTypedDict,
)
from litellm.types.router import ModelInfo as RouterModelInfo
from litellm.types.router import (
RouterGeneralSettings,
@@ -3320,6 +3325,9 @@ class ProxyConfig:
if self._should_load_db_object(object_type="prompts"):
await self._init_prompts_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="search_tools"):
await self._init_search_tools_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="model_cost_map"):
await self._check_and_reload_model_cost_map(prisma_client=prisma_client)
@@ -3513,7 +3521,24 @@ class ProxyConfig:
str(e)
)
)
async def _init_search_tools_in_db(self, prisma_client: PrismaClient):
from litellm.proxy.search_endpoints.search_tool_registry import (
IN_MEMORY_SEARCH_TOOL_HANDLER,
SearchToolRegistry,
)
from litellm.types.search import SearchTool
try:
search_tools = await SearchToolRegistry.get_all_search_tools_from_db(prisma_client=prisma_client)
for search_tool in search_tools:
IN_MEMORY_SEARCH_TOOL_HANDLER.add_search_tool(search_tool=cast(SearchTool, search_tool))
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_search_tools_in_db - {}".format(
str(e)
)
)
async def _init_pass_through_endpoints_in_db(self):
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
initialize_pass_through_endpoints_in_db,
@@ -9877,6 +9902,7 @@ app.include_router(cloudzero_router)
app.include_router(caching_router)
app.include_router(analytics_router)
app.include_router(guardrails_router)
app.include_router(search_tool_management_router)
app.include_router(prompts_router)
app.include_router(callback_management_endpoints_router)
app.include_router(debugging_endpoints_router)
litellm/proxy/schema.prisma
+10
View File
@@ -570,4 +570,14 @@ model LiteLLM_HealthCheckTable {
@@index([model_name])
@@index([checked_at])
@@index([status])
}
// Search Tools table for storing search tool configurations
model LiteLLM_SearchToolsTable {
search_tool_id String @id @default(uuid())
search_tool_name String @unique
litellm_params Json
search_tool_info Json?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}
litellm/proxy/search_endpoints/__init__.py
+12
View File
@@ -1,2 +1,14 @@
# litellm/proxy/search_endpoints/__init__.py
from .search_tool_registry import (
IN_MEMORY_SEARCH_TOOL_HANDLER,
InMemorySearchToolHandler,
SearchToolRegistry,
)
__all__ = [
"SearchToolRegistry",
"InMemorySearchToolHandler",
"IN_MEMORY_SEARCH_TOOL_HANDLER",
]
litellm/proxy/search_endpoints/search_tool_management.py
+521
View File
@@ -0,0 +1,521 @@
"""
CRUD ENDPOINTS FOR SEARCH TOOLS
"""
from datetime import datetime
from typing import List, Union, cast
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.search_endpoints.search_tool_registry import (
IN_MEMORY_SEARCH_TOOL_HANDLER,
SearchToolRegistry,
)
from litellm.types.search import (
AvailableSearchProvider,
ListSearchToolsResponse,
SearchTool,
SearchToolInfoResponse,
)
from litellm.types.utils import SearchProviders
#### SEARCH TOOLS ENDPOINTS ####
router = APIRouter()
SEARCH_TOOL_REGISTRY = SearchToolRegistry()
def _convert_datetime_to_str(value: Union[datetime, str, None]) -> Union[str, None]:
"""
Convert datetime object to ISO format string.
Args:
value: datetime object, string, or None
Returns:
ISO format string or original value if already string or None
"""
if value is None:
return None
if isinstance(value, datetime):
return value.isoformat()
return value
@router.get(
"/search_tools/list",
tags=["Search Tools"],
dependencies=[Depends(user_api_key_auth)],
response_model=ListSearchToolsResponse,
)
async def list_search_tools():
"""
List all search tools that are available in the database.
Example Request:
```bash
curl -X GET "http://localhost:4000/search_tools/list" -H "Authorization: Bearer <your_api_key>"
```
Example Response:
```json
{
"search_tools": [
{
"search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
"search_tool_name": "litellm-search",
"litellm_params": {
"search_provider": "perplexity",
"api_key": "sk-***",
"api_base": "https://api.perplexity.ai"
},
"search_tool_info": {
"description": "Perplexity search tool"
},
"created_at": "2023-11-09T12:34:56.789Z",
"updated_at": "2023-11-09T12:34:56.789Z"
}
]
}
```
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
try:
search_tools = await SEARCH_TOOL_REGISTRY.get_all_search_tools_from_db(
prisma_client=prisma_client
)
search_tool_configs: List[SearchToolInfoResponse] = []
for search_tool in search_tools:
search_tool_configs.append(
SearchToolInfoResponse(
search_tool_id=search_tool.get("search_tool_id"),
search_tool_name=search_tool.get("search_tool_name", ""),
litellm_params=dict(search_tool.get("litellm_params", {})),
search_tool_info=search_tool.get("search_tool_info"),
created_at=_convert_datetime_to_str(search_tool.get("created_at")),
updated_at=_convert_datetime_to_str(search_tool.get("updated_at")),
)
)
return ListSearchToolsResponse(search_tools=search_tool_configs)
except Exception as e:
verbose_proxy_logger.exception(f"Error getting search tools from db: {e}")
raise HTTPException(status_code=500, detail=str(e))
class CreateSearchToolRequest(BaseModel):
search_tool: SearchTool
@router.post(
"/search_tools",
tags=["Search Tools"],
dependencies=[Depends(user_api_key_auth)],
)
async def create_search_tool(request: CreateSearchToolRequest):
"""
Create a new search tool.
Example Request:
```bash
curl -X POST "http://localhost:4000/search_tools" \\
-H "Authorization: Bearer <your_api_key>" \\
-H "Content-Type: application/json" \\
-d '{
"search_tool": {
"search_tool_name": "litellm-search",
"litellm_params": {
"search_provider": "perplexity",
"api_key": "sk-..."
},
"search_tool_info": {
"description": "Perplexity search tool"
}
}
}'
```
Example Response:
```json
{
"search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
"search_tool_name": "litellm-search",
"litellm_params": {
"search_provider": "perplexity",
"api_key": "sk-..."
},
"search_tool_info": {
"description": "Perplexity search tool"
},
"created_at": "2023-11-09T12:34:56.789Z",
"updated_at": "2023-11-09T12:34:56.789Z"
}
```
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
try:
result = await SEARCH_TOOL_REGISTRY.add_search_tool_to_db(
search_tool=request.search_tool, prisma_client=prisma_client
)
# Add to in-memory cache
try:
IN_MEMORY_SEARCH_TOOL_HANDLER.add_search_tool(search_tool=cast(SearchTool, result))
verbose_proxy_logger.info(
f"Successfully added search tool '{result.get('search_tool_name')}' to in-memory cache"
)
except Exception as cache_error:
verbose_proxy_logger.warning(
f"Failed to add search tool to in-memory cache: {cache_error}"
)
return result
except Exception as e:
verbose_proxy_logger.exception(f"Error adding search tool to db: {e}")
raise HTTPException(status_code=500, detail=str(e))
class UpdateSearchToolRequest(BaseModel):
search_tool: SearchTool
@router.put(
"/search_tools/{search_tool_id}",
tags=["Search Tools"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_search_tool(search_tool_id: str, request: UpdateSearchToolRequest):
"""
Update an existing search tool.
Example Request:
```bash
curl -X PUT "http://localhost:4000/search_tools/123e4567-e89b-12d3-a456-426614174000" \\
-H "Authorization: Bearer <your_api_key>" \\
-H "Content-Type: application/json" \\
-d '{
"search_tool": {
"search_tool_name": "updated-search",
"litellm_params": {
"search_provider": "perplexity",
"api_key": "sk-new-key"
},
"search_tool_info": {
"description": "Updated search tool"
}
}
}'
```
Example Response:
```json
{
"search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
"search_tool_name": "updated-search",
"litellm_params": {
"search_provider": "perplexity",
"api_key": "sk-new-key"
},
"search_tool_info": {
"description": "Updated search tool"
},
"created_at": "2023-11-09T12:34:56.789Z",
"updated_at": "2023-11-09T13:45:12.345Z"
}
```
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
try:
# Check if search tool exists
existing_tool = await SEARCH_TOOL_REGISTRY.get_search_tool_by_id_from_db(
search_tool_id=search_tool_id, prisma_client=prisma_client
)
if existing_tool is None:
raise HTTPException(
status_code=404,
detail=f"Search tool with ID {search_tool_id} not found",
)
result = await SEARCH_TOOL_REGISTRY.update_search_tool_in_db(
search_tool_id=search_tool_id,
search_tool=request.search_tool,
prisma_client=prisma_client,
)
# Update in-memory cache
try:
IN_MEMORY_SEARCH_TOOL_HANDLER.update_search_tool(
search_tool_id=search_tool_id, search_tool=cast(SearchTool, result)
)
verbose_proxy_logger.info(
f"Successfully updated search tool '{result.get('search_tool_name')}' in in-memory cache"
)
except Exception as cache_error:
verbose_proxy_logger.warning(
f"Failed to update search tool in in-memory cache: {cache_error}"
)
return result
except HTTPException as e:
raise e
except Exception as e:
verbose_proxy_logger.exception(f"Error updating search tool: {e}")
raise HTTPException(status_code=500, detail=str(e))
@router.delete(
"/search_tools/{search_tool_id}",
tags=["Search Tools"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_search_tool(search_tool_id: str):
"""
Delete a search tool.
Example Request:
```bash
curl -X DELETE "http://localhost:4000/search_tools/123e4567-e89b-12d3-a456-426614174000" \\
-H "Authorization: Bearer <your_api_key>"
```
Example Response:
```json
{
"message": "Search tool 123e4567-e89b-12d3-a456-426614174000 deleted successfully",
"search_tool_name": "litellm-search"
}
```
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
try:
# Check if search tool exists
existing_tool = await SEARCH_TOOL_REGISTRY.get_search_tool_by_id_from_db(
search_tool_id=search_tool_id, prisma_client=prisma_client
)
if existing_tool is None:
raise HTTPException(
status_code=404,
detail=f"Search tool with ID {search_tool_id} not found",
)
result = await SEARCH_TOOL_REGISTRY.delete_search_tool_from_db(
search_tool_id=search_tool_id, prisma_client=prisma_client
)
# Delete from in-memory cache
try:
IN_MEMORY_SEARCH_TOOL_HANDLER.delete_search_tool(
search_tool_id=search_tool_id
)
verbose_proxy_logger.info(
f"Successfully removed search tool from in-memory cache"
)
except Exception as cache_error:
verbose_proxy_logger.warning(
f"Failed to remove search tool from in-memory cache: {cache_error}"
)
return result
except HTTPException as e:
raise e
except Exception as e:
verbose_proxy_logger.exception(f"Error deleting search tool: {e}")
raise HTTPException(status_code=500, detail=str(e))
@router.get(
"/search_tools/{search_tool_id}",
tags=["Search Tools"],
dependencies=[Depends(user_api_key_auth)],
)
async def get_search_tool_info(search_tool_id: str):
"""
Get detailed information about a specific search tool by ID.
Example Request:
```bash
curl -X GET "http://localhost:4000/search_tools/123e4567-e89b-12d3-a456-426614174000" \\
-H "Authorization: Bearer <your_api_key>"
```
Example Response:
```json
{
"search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
"search_tool_name": "litellm-search",
"litellm_params": {
"search_provider": "perplexity",
"api_key": "sk-***"
},
"search_tool_info": {
"description": "Perplexity search tool"
},
"created_at": "2023-11-09T12:34:56.789Z",
"updated_at": "2023-11-09T12:34:56.789Z"
}
```
"""
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
try:
result = await SEARCH_TOOL_REGISTRY.get_search_tool_by_id_from_db(
search_tool_id=search_tool_id, prisma_client=prisma_client
)
if result is None:
# Try in-memory cache
result = IN_MEMORY_SEARCH_TOOL_HANDLER.get_search_tool_by_id(
search_tool_id=search_tool_id
)
if result is None:
raise HTTPException(
status_code=404,
detail=f"Search tool with ID {search_tool_id} not found",
)
# Mask sensitive data
litellm_params_dict = dict(result.get("litellm_params", {}))
masked_litellm_params_dict = _get_masked_values(
litellm_params_dict,
unmasked_length=4,
number_of_asterisks=4,
)
return SearchToolInfoResponse(
search_tool_id=result.get("search_tool_id"),
search_tool_name=result.get("search_tool_name", ""),
litellm_params=masked_litellm_params_dict,
search_tool_info=result.get("search_tool_info"),
created_at=_convert_datetime_to_str(result.get("created_at")),
updated_at=_convert_datetime_to_str(result.get("updated_at")),
)
except HTTPException as e:
raise e
except Exception as e:
verbose_proxy_logger.exception(f"Error getting search tool info: {e}")
raise HTTPException(status_code=500, detail=str(e))
@router.get(
"/search_tools/ui/available_providers",
tags=["Search Tools"],
dependencies=[Depends(user_api_key_auth)],
)
async def get_available_search_providers():
"""
Get the list of available search providers with their configuration fields.
This auto-discovers search providers from the SearchProviders enum.
Example Request:
```bash
curl -X GET "http://localhost:4000/search_tools/ui/available_providers" \\
-H "Authorization: Bearer <your_api_key>"
```
Example Response:
```json
[
{
"provider": "perplexity",
"display_name": "Perplexity",
"fields": [
{
"name": "api_key",
"type": "string",
"required": false,
"description": "API key for Perplexity"
},
{
"name": "api_base",
"type": "string",
"required": false,
"description": "API base URL"
}
]
}
]
```
"""
try:
available_providers: List[AvailableSearchProvider] = []
# Common fields for all search providers
common_fields = [
{
"name": "api_key",
"type": "string",
"required": False,
"description": "API key for the search provider",
},
{
"name": "api_base",
"type": "string",
"required": False,
"description": "Custom API base URL (optional)",
},
{
"name": "timeout",
"type": "number",
"required": False,
"description": "Request timeout in seconds",
},
{
"name": "max_retries",
"type": "number",
"required": False,
"description": "Maximum number of retry attempts",
},
]
# Provider display name mapping
provider_display_names = {
SearchProviders.PERPLEXITY: "Perplexity",
SearchProviders.TAVILY: "Tavily",
SearchProviders.PARALLEL_AI: "Parallel AI",
SearchProviders.EXA_AI: "Exa AI",
SearchProviders.GOOGLE_PSE: "Google PSE",
SearchProviders.DATAFORSEO: "DataForSEO",
}
# Auto-discover providers from SearchProviders enum
for provider in SearchProviders:
available_providers.append(
AvailableSearchProvider(
provider=provider.value,
display_name=provider_display_names.get(provider, provider.value.title()),
fields=common_fields,
)
)
return available_providers
except Exception as e:
verbose_proxy_logger.exception(f"Error getting available search providers: {e}")
raise HTTPException(status_code=500, detail=str(e))
litellm/proxy/search_endpoints/search_tool_registry.py
+320
View File
@@ -0,0 +1,320 @@
"""
Search Tool Registry for managing search tool configurations.
"""
from datetime import datetime, timezone
from typing import Dict, List, Optional
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy.utils import PrismaClient
from litellm.types.search import SearchTool
class SearchToolRegistry:
"""
Handles adding, removing, and getting search tools in DB + in memory.
"""
def __init__(self):
pass
@staticmethod
def _convert_prisma_to_dict(prisma_obj) -> dict:
"""
Convert Prisma result to dict with datetime objects as ISO format strings.
Args:
prisma_obj: Prisma model instance
Returns:
Dict with datetime fields converted to ISO strings
"""
result = dict(prisma_obj)
# Convert datetime objects to ISO format strings
if "created_at" in result and result["created_at"]:
result["created_at"] = result["created_at"].isoformat()
if "updated_at" in result and result["updated_at"]:
result["updated_at"] = result["updated_at"].isoformat()
return result
###########################################################
########### DB management helpers for search tools ########
###########################################################
async def add_search_tool_to_db(
self, search_tool: SearchTool, prisma_client: PrismaClient
):
"""
Add a search tool to the database.
Args:
search_tool: Search tool configuration
prisma_client: Prisma client instance
Returns:
Dict with created search tool data
"""
try:
search_tool_name = search_tool.get("search_tool_name")
litellm_params: str = safe_dumps(dict(search_tool.get("litellm_params", {})))
search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {}))
# Create search tool in DB
created_search_tool = await prisma_client.db.litellm_searchtoolstable.create(
data={
"search_tool_name": search_tool_name,
"litellm_params": litellm_params,
"search_tool_info": search_tool_info,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
)
# Add search_tool_id to the returned search tool object
search_tool_dict = dict(search_tool)
search_tool_dict["search_tool_id"] = created_search_tool.search_tool_id
search_tool_dict["created_at"] = created_search_tool.created_at.isoformat()
search_tool_dict["updated_at"] = created_search_tool.updated_at.isoformat()
return search_tool_dict
except Exception as e:
verbose_proxy_logger.exception(f"Error adding search tool to DB: {str(e)}")
raise Exception(f"Error adding search tool to DB: {str(e)}")
async def delete_search_tool_from_db(
self, search_tool_id: str, prisma_client: PrismaClient
):
"""
Delete a search tool from the database.
Args:
search_tool_id: ID of search tool to delete
prisma_client: Prisma client instance
Returns:
Dict with success message
"""
try:
# Get search tool before deletion for response
existing_tool = await prisma_client.db.litellm_searchtoolstable.find_unique(
where={"search_tool_id": search_tool_id}
)
if not existing_tool:
raise Exception(f"Search tool with ID {search_tool_id} not found")
# Delete from DB
await prisma_client.db.litellm_searchtoolstable.delete(
where={"search_tool_id": search_tool_id}
)
return {
"message": f"Search tool {search_tool_id} deleted successfully",
"search_tool_name": existing_tool.search_tool_name,
}
except Exception as e:
verbose_proxy_logger.exception(f"Error deleting search tool from DB: {str(e)}")
raise Exception(f"Error deleting search tool from DB: {str(e)}")
async def update_search_tool_in_db(
self, search_tool_id: str, search_tool: SearchTool, prisma_client: PrismaClient
):
"""
Update a search tool in the database.
Args:
search_tool_id: ID of search tool to update
search_tool: Updated search tool configuration
prisma_client: Prisma client instance
Returns:
Dict with updated search tool data
"""
try:
search_tool_name = search_tool.get("search_tool_name")
litellm_params: str = safe_dumps(dict(search_tool.get("litellm_params", {})))
search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {}))
# Update in DB
updated_search_tool = await prisma_client.db.litellm_searchtoolstable.update(
where={"search_tool_id": search_tool_id},
data={
"search_tool_name": search_tool_name,
"litellm_params": litellm_params,
"search_tool_info": search_tool_info,
"updated_at": datetime.now(timezone.utc),
},
)
# Convert to dict with ISO formatted datetimes
return self._convert_prisma_to_dict(updated_search_tool)
except Exception as e:
verbose_proxy_logger.exception(f"Error updating search tool in DB: {str(e)}")
raise Exception(f"Error updating search tool in DB: {str(e)}")
@staticmethod
async def get_all_search_tools_from_db(
prisma_client: PrismaClient,
) -> List[SearchTool]:
"""
Get all search tools from the database.
Args:
prisma_client: Prisma client instance
Returns:
List of search tool configurations
"""
try:
search_tools_from_db = (
await prisma_client.db.litellm_searchtoolstable.find_many(
order={"created_at": "desc"},
)
)
search_tools: List[SearchTool] = []
for search_tool in search_tools_from_db:
# Convert Prisma result to dict with ISO formatted datetimes
search_tool_dict = SearchToolRegistry._convert_prisma_to_dict(search_tool)
search_tools.append(SearchTool(**search_tool_dict)) # type: ignore
return search_tools
except Exception as e:
verbose_proxy_logger.exception(f"Error getting search tools from DB: {str(e)}")
raise Exception(f"Error getting search tools from DB: {str(e)}")
async def get_search_tool_by_id_from_db(
self, search_tool_id: str, prisma_client: PrismaClient
) -> Optional[SearchTool]:
"""
Get a search tool by its ID from the database.
Args:
search_tool_id: ID of search tool to retrieve
prisma_client: Prisma client instance
Returns:
Search tool configuration or None if not found
"""
try:
search_tool = await prisma_client.db.litellm_searchtoolstable.find_unique(
where={"search_tool_id": search_tool_id}
)
if not search_tool:
return None
# Convert Prisma result to dict with ISO formatted datetimes
search_tool_dict = self._convert_prisma_to_dict(search_tool)
return SearchTool(**search_tool_dict) # type: ignore
except Exception as e:
verbose_proxy_logger.exception(f"Error getting search tool from DB: {str(e)}")
raise Exception(f"Error getting search tool from DB: {str(e)}")
async def get_search_tool_by_name_from_db(
self, search_tool_name: str, prisma_client: PrismaClient
) -> Optional[SearchTool]:
"""
Get a search tool by its name from the database.
Args:
search_tool_name: Name of search tool to retrieve
prisma_client: Prisma client instance
Returns:
Search tool configuration or None if not found
"""
try:
search_tool = await prisma_client.db.litellm_searchtoolstable.find_unique(
where={"search_tool_name": search_tool_name}
)
if not search_tool:
return None
# Convert Prisma result to dict with ISO formatted datetimes
search_tool_dict = self._convert_prisma_to_dict(search_tool)
return SearchTool(**search_tool_dict) # type: ignore
except Exception as e:
verbose_proxy_logger.exception(f"Error getting search tool from DB: {str(e)}")
raise Exception(f"Error getting search tool from DB: {str(e)}")
class InMemorySearchToolHandler:
"""
Class that handles caching search tools in memory.
"""
def __init__(self):
self.IN_MEMORY_SEARCH_TOOLS: Dict[str, SearchTool] = {}
"""
Search tool id to SearchTool object mapping
"""
def add_search_tool(self, search_tool: SearchTool) -> None:
"""
Add a search tool to in-memory cache.
Args:
search_tool: Search tool configuration
"""
search_tool_id = search_tool.get("search_tool_id")
if search_tool_id:
self.IN_MEMORY_SEARCH_TOOLS[search_tool_id] = search_tool
verbose_proxy_logger.debug(
f"Added search tool '{search_tool.get('search_tool_name')}' to in-memory cache"
)
def update_search_tool(self, search_tool_id: str, search_tool: SearchTool) -> None:
"""
Update a search tool in in-memory cache.
Args:
search_tool_id: ID of search tool to update
search_tool: Updated search tool configuration
"""
self.IN_MEMORY_SEARCH_TOOLS[search_tool_id] = search_tool
verbose_proxy_logger.debug(
f"Updated search tool '{search_tool.get('search_tool_name')}' in in-memory cache"
)
def delete_search_tool(self, search_tool_id: str) -> None:
"""
Delete a search tool from in-memory cache.
Args:
search_tool_id: ID of search tool to delete
"""
self.IN_MEMORY_SEARCH_TOOLS.pop(search_tool_id, None)
verbose_proxy_logger.debug(
f"Deleted search tool with ID '{search_tool_id}' from in-memory cache"
)
def list_search_tools(self) -> List[SearchTool]:
"""
List all search tools in in-memory cache.
Returns:
List of search tool configurations
"""
return list(self.IN_MEMORY_SEARCH_TOOLS.values())
def get_search_tool_by_id(self, search_tool_id: str) -> Optional[SearchTool]:
"""
Get a search tool by its ID from in-memory cache.
Args:
search_tool_id: ID of search tool to retrieve
Returns:
Search tool configuration or None if not found
"""
return self.IN_MEMORY_SEARCH_TOOLS.get(search_tool_id)
########################################################
# In Memory Search Tool Handler for LiteLLM Proxy
########################################################
IN_MEMORY_SEARCH_TOOL_HANDLER = InMemorySearchToolHandler()
########################################################
litellm/types/search.py
+63
View File
@@ -3,6 +3,9 @@ LiteLLM Search API Types
This module defines types for the unified search API across different providers.
"""
from typing import List, Optional, Required
from typing_extensions import TypedDict
from litellm.types.utils import SearchProviders
@@ -11,3 +14,63 @@ SearchProvider = SearchProviders
__all__ = ["SearchProvider", "SearchProviders"]
class SearchToolLiteLLMParams(TypedDict, total=False):
"""
LiteLLM params for search tools configuration.
"""
search_provider: Required[str]
api_key: Optional[str]
api_base: Optional[str]
timeout: Optional[float]
max_retries: Optional[int]
class SearchTool(TypedDict, total=False):
"""
Search tool configuration.
Example:
{
"search_tool_id": "123e4567-e89b-12d3-a456-426614174000",
"search_tool_name": "litellm-search",
"litellm_params": {
"search_provider": "perplexity",
"api_key": "sk-..."
},
"search_tool_info": {
"description": "Perplexity search tool"
}
}
"""
search_tool_id: Optional[str]
search_tool_name: Required[str]
litellm_params: Required[SearchToolLiteLLMParams]
search_tool_info: Optional[dict]
created_at: Optional[str]
updated_at: Optional[str]
class SearchToolInfoResponse(TypedDict, total=False):
"""Response model for search tool information."""
search_tool_id: Optional[str]
search_tool_name: str
litellm_params: dict
search_tool_info: Optional[dict]
created_at: Optional[str]
updated_at: Optional[str]
class ListSearchToolsResponse(TypedDict):
"""Response model for listing search tools."""
search_tools: List[SearchToolInfoResponse]
class AvailableSearchProvider(TypedDict):
"""Information about an available search provider."""
provider: str
display_name: str
fields: List[dict]
poetry.lock
Generated
+540 -93
View File
File diff suppressed because it is too large Load Diff
pyproject.toml
+1 -1
View File
@@ -59,7 +59,7 @@ websockets = {version = "^13.1.0", optional = true}
boto3 = {version = "1.36.0", optional = true}
redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
mcp = {version = "^1.10.0", optional = true, python = ">=3.10"}
litellm-proxy-extras = {version = "0.2.27", optional = true}
litellm-proxy-extras = {version = "0.2.29", optional = true}
rich = {version = "13.7.1", optional = true}
litellm-enterprise = {version = "0.1.20", optional = true}
diskcache = {version = "^5.6.1", optional = true}
requirements.txt
+1 -1
View File
@@ -43,7 +43,7 @@ sentry_sdk==2.21.0 # for sentry error handling
detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
cryptography==44.0.1
tzdata==2025.1 # IANA time zone database
litellm-proxy-extras==0.2.27 # for proxy extras - e.g. prisma migrations
litellm-proxy-extras==0.2.29 # for proxy extras - e.g. prisma migrations
### LITELLM PACKAGE DEPENDENCIES
python-dotenv==1.0.1 # for env
tiktoken==0.8.0 # for calculating usage
schema.prisma
+10
View File
@@ -570,4 +570,14 @@ model LiteLLM_HealthCheckTable {
@@index([model_name])
@@index([checked_at])
@@index([status])
}
// Search Tools table for storing search tool configurations
model LiteLLM_SearchToolsTable {
search_tool_id String @id @default(uuid())
search_tool_name String @unique
litellm_params Json
search_tool_info Json?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}