From ad37273ba7ea23c9641c0935f02f06d9d85e301f Mon Sep 17 00:00:00 2001 From: Thomas Menard Date: Sun, 3 May 2026 21:30:19 +0200 Subject: [PATCH] feat: add fetch_tools DB table, registry, and API endpoints - Add LiteLLM_FetchToolsTable to schema.prisma - Create fetch_tool_registry.py for DB CRUD operations - Create fetch_tool_management.py for REST API endpoints - Add FetchProviders enum to types/utils.py - Register fetch_tool routes in proxy_server.py --- litellm/proxy/fetch_endpoints/__init__.py | 3 + .../fetch_endpoints/fetch_tool_management.py | 352 ++++++++++++++++++ .../fetch_endpoints/fetch_tool_registry.py | 200 ++++++++++ litellm/proxy/proxy_server.py | 4 + litellm/types/fetch.py | 81 ++++ litellm/types/utils.py | 13 + schema.prisma | 10 + 7 files changed, 663 insertions(+) create mode 100644 litellm/proxy/fetch_endpoints/__init__.py create mode 100644 litellm/proxy/fetch_endpoints/fetch_tool_management.py create mode 100644 litellm/proxy/fetch_endpoints/fetch_tool_registry.py create mode 100644 litellm/types/fetch.py diff --git a/litellm/proxy/fetch_endpoints/__init__.py b/litellm/proxy/fetch_endpoints/__init__.py new file mode 100644 index 00000000000..797ada25c80 --- /dev/null +++ b/litellm/proxy/fetch_endpoints/__init__.py @@ -0,0 +1,3 @@ +""" +Fetch Endpoints for LiteLLM Proxy +""" diff --git a/litellm/proxy/fetch_endpoints/fetch_tool_management.py b/litellm/proxy/fetch_endpoints/fetch_tool_management.py new file mode 100644 index 00000000000..6c99d92f6fd --- /dev/null +++ b/litellm/proxy/fetch_endpoints/fetch_tool_management.py @@ -0,0 +1,352 @@ +""" +CRUD ENDPOINTS FOR FETCH TOOLS +""" + +from datetime import datetime +from typing import Any, Dict, Union + +from fastapi import APIRouter, Depends, HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.fetch_endpoints.fetch_tool_registry import FetchToolRegistry +from litellm.types.fetch import ( + FetchProviders, + FetchTool, + FetchToolInfoResponse, + ListFetchToolsResponse, +) + +#### FETCH TOOLS ENDPOINTS #### + +router = APIRouter() +FETCH_TOOL_REGISTRY = FetchToolRegistry() + + +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( + "/fetch_tools/list", + tags=["Fetch Tools"], + dependencies=[Depends(user_api_key_auth)], + response_model=ListFetchToolsResponse, +) +async def list_fetch_tools(): + """ + List all fetch tools that are available in the database and config file. + + Example Request: + ```bash + curl -X GET "http://localhost:4000/fetch_tools/list" -H "Authorization: Bearer " + ``` + + Example Response: + ```json + { + "fetch_tools": [ + { + "fetch_tool_id": "123e4567-e89b-12d3-a456-426614174000", + "fetch_tool_name": "litellm-fetch", + "litellm_params": { + "fetch_provider": "firecrawl", + "api_key": "fc-***", + "api_base": "https://api.firecrawl.dev" + }, + "fetch_tool_info": { + "description": "Firecrawl fetch tool" + }, + "created_at": "2023-11-09T12:34:56.789Z", + "updated_at": "2023-11-09T12:34:56.789Z", + "is_from_config": false + }, + { + "fetch_tool_name": "config-fetch-tool", + "litellm_params": { + "fetch_provider": "firecrawl", + "api_key": "fc-***" + }, + "is_from_config": true + } + ] + } + ``` + """ + from litellm.litellm_core_utils.litellm_logging import _get_masked_values + from litellm.proxy.proxy_server import prisma_client, proxy_config + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Prisma client not initialized") + + try: + fetch_tools_from_db = await FETCH_TOOL_REGISTRY.get_all_fetch_tools_from_db( + prisma_client=prisma_client + ) + + # Mask API keys + for tool in fetch_tools_from_db: + litellm_params = tool.get("litellm_params", {}) + if isinstance(litellm_params, dict): + for key, value in litellm_params.items(): + if isinstance(value, str): + litellm_params[key] = _get_masked_values(key, value) + tool["litellm_params"] = litellm_params + + db_tool_names = {tool.get("fetch_tool_name") for tool in fetch_tools_from_db} + + # Get fetch tools from config + config_fetch_tools = [] + config = getattr(proxy_config, "config", None) + if config: + parsed_tools = proxy_config.parse_fetch_tools(config) + if parsed_tools: + config_fetch_tools = parsed_tools + + for fetch_tool in config_fetch_tools: + tool_name = fetch_tool.get("fetch_tool_name") + if tool_name not in db_tool_names: + # Add config tool to response + fetch_tool_info = fetch_tool.get("fetch_tool_info", {}) + litellm_params = fetch_tool.get("litellm_params", {}) + + # Mask API keys + if isinstance(litellm_params, dict): + for key, value in litellm_params.items(): + if isinstance(value, str): + litellm_params[key] = _get_masked_values(key, value) + + fetch_tools_from_db.append( + FetchToolInfoResponse( + fetch_tool_name=tool_name, + litellm_params=litellm_params, + fetch_tool_info=fetch_tool_info if fetch_tool_info else None, + is_from_config=True, + ) + ) + + return ListFetchToolsResponse(fetch_tools=fetch_tools_from_db) + except Exception as e: + verbose_proxy_logger.exception(f"Error listing fetch tools: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/fetch_tools", + tags=["Fetch Tools"], + dependencies=[Depends(user_api_key_auth)], +) +async def create_fetch_tool(request: Dict[str, Any]): + """ + Create a new fetch tool in the database. + + Example Request: + ```bash + curl -X POST "http://localhost:4000/fetch_tools" \\ + -H "Authorization: Bearer " \\ + -H "Content-Type: application/json" \\ + -d '{ + "fetch_tool": { + "fetch_tool_name": "litellm-fetch", + "litellm_params": { + "fetch_provider": "firecrawl", + "api_key": "fc-1234", + "api_base": "https://api.firecrawl.dev" + }, + "fetch_tool_info": { + "description": "Firecrawl fetch tool" + } + } + }' + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Prisma client not initialized") + + try: + fetch_tool_data = request.get("fetch_tool", {}) + fetch_tool = FetchTool(**fetch_tool_data) + + created_tool = await FETCH_TOOL_REGISTRY.add_fetch_tool_to_db( + fetch_tool=fetch_tool, + prisma_client=prisma_client, + ) + + return created_tool + except Exception as e: + verbose_proxy_logger.exception(f"Error creating fetch tool: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.put( + "/fetch_tools/{fetch_tool_id}", + tags=["Fetch Tools"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_fetch_tool(fetch_tool_id: str, request: Dict[str, Any]): + """ + Update an existing fetch tool in the database. + + Example Request: + ```bash + curl -X PUT "http://localhost:4000/fetch_tools/123e4567-e89b-12d3-a456-426614174000" \\ + -H "Authorization: Bearer " \\ + -H "Content-Type: application/json" \\ + -d '{ + "fetch_tool": { + "fetch_tool_name": "litellm-fetch", + "litellm_params": { + "fetch_provider": "firecrawl", + "api_key": "fc-5678", + "api_base": "https://api.firecrawl.dev" + }, + "fetch_tool_info": { + "description": "Updated Firecrawl fetch tool" + } + } + }' + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Prisma client not initialized") + + try: + fetch_tool_data = request.get("fetch_tool", {}) + fetch_tool = FetchTool(**fetch_tool_data) + + updated_tool = await FETCH_TOOL_REGISTRY.update_fetch_tool_in_db( + fetch_tool_id=fetch_tool_id, + fetch_tool=fetch_tool, + prisma_client=prisma_client, + ) + + return updated_tool + except Exception as e: + verbose_proxy_logger.exception(f"Error updating fetch tool: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.delete( + "/fetch_tools/{fetch_tool_id}", + tags=["Fetch Tools"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_fetch_tool(fetch_tool_id: str): + """ + Delete a fetch tool from the database. + + Example Request: + ```bash + curl -X DELETE "http://localhost:4000/fetch_tools/123e4567-e89b-12d3-a456-426614174000" \\ + -H "Authorization: Bearer " + ``` + """ + 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 FETCH_TOOL_REGISTRY.delete_fetch_tool_from_db( + fetch_tool_id=fetch_tool_id, + prisma_client=prisma_client, + ) + + return result + except Exception as e: + verbose_proxy_logger.exception(f"Error deleting fetch tool: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/fetch_tools/{fetch_tool_id}", + tags=["Fetch Tools"], + dependencies=[Depends(user_api_key_auth)], +) +async def get_fetch_tool(fetch_tool_id: str): + """ + Get a specific fetch tool by ID. + + Example Request: + ```bash + curl -X GET "http://localhost:4000/fetch_tools/123e4567-e89b-12d3-a456-426614174000" \\ + -H "Authorization: Bearer " + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Prisma client not initialized") + + try: + fetch_tool = await FETCH_TOOL_REGISTRY.get_fetch_tool_from_db( + fetch_tool_id=fetch_tool_id, + prisma_client=prisma_client, + ) + + if fetch_tool is None: + raise HTTPException(status_code=404, detail="Fetch tool not found") + + return fetch_tool + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception(f"Error getting fetch tool: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/fetch_tools/ui/available_providers", + tags=["Fetch Tools"], + dependencies=[Depends(user_api_key_auth)], +) +async def list_available_fetch_providers(): + """ + List all available fetch providers for the UI. + + Example Request: + ```bash + curl -X GET "http://localhost:4000/fetch_tools/ui/available_providers" \\ + -H "Authorization: Bearer " + ``` + + Example Response: + ```json + { + "providers": [ + {"provider_name": "firecrawl", "ui_friendly_name": "Firecrawl"} + ] + } + ``` + """ + try: + providers = [ + { + "provider_name": provider.value, + "ui_friendly_name": provider.value.capitalize(), + } + for provider in FetchProviders + ] + return {"providers": providers} + except Exception as e: + verbose_proxy_logger.exception( + f"Error listing available fetch providers: {str(e)}" + ) + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/fetch_endpoints/fetch_tool_registry.py b/litellm/proxy/fetch_endpoints/fetch_tool_registry.py new file mode 100644 index 00000000000..b40e886574f --- /dev/null +++ b/litellm/proxy/fetch_endpoints/fetch_tool_registry.py @@ -0,0 +1,200 @@ +""" +Fetch Tool Registry for managing fetch tool configurations. +""" + +from datetime import datetime, timezone +from typing import List + +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.fetch import FetchTool + + +class FetchToolRegistry: + """ + Handles adding, removing, and getting fetch 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 fetch tools ######## + ########################################################### + + async def add_fetch_tool_to_db( + self, fetch_tool: FetchTool, prisma_client: PrismaClient + ): + """ + Add a fetch tool to the database. + + Args: + fetch_tool: Fetch tool configuration + prisma_client: Prisma client instance + + Returns: + Dict with created fetch tool data + """ + try: + fetch_tool_name = fetch_tool.get("fetch_tool_name") + litellm_params: str = safe_dumps(dict(fetch_tool.get("litellm_params", {}))) + fetch_tool_info: str = safe_dumps(fetch_tool.get("fetch_tool_info", {})) + + # Create fetch tool in DB + created_fetch_tool = await prisma_client.db.litellm_fetchtoolstable.create( + data={ + "fetch_tool_name": fetch_tool_name, + "litellm_params": litellm_params, + "fetch_tool_info": fetch_tool_info, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + ) + + # Add fetch_tool_id to the returned fetch tool object + fetch_tool_dict = dict(fetch_tool) + fetch_tool_dict["fetch_tool_id"] = created_fetch_tool.fetch_tool_id + fetch_tool_dict["created_at"] = created_fetch_tool.created_at.isoformat() + fetch_tool_dict["updated_at"] = created_fetch_tool.updated_at.isoformat() + + return fetch_tool_dict + except Exception as e: + verbose_proxy_logger.exception(f"Error adding fetch tool to DB: {str(e)}") + raise Exception(f"Error adding fetch tool to DB: {str(e)}") + + async def delete_fetch_tool_from_db( + self, fetch_tool_id: str, prisma_client: PrismaClient + ): + """ + Delete a fetch tool from the database. + + Args: + fetch_tool_id: ID of fetch tool to delete + prisma_client: Prisma client instance + + Returns: + Dict with success message + """ + try: + await prisma_client.db.litellm_fetchtoolstable.delete( + where={"fetch_tool_id": fetch_tool_id} + ) + return {"message": f"Fetch tool {fetch_tool_id} deleted successfully"} + except Exception as e: + verbose_proxy_logger.exception( + f"Error deleting fetch tool from DB: {str(e)}" + ) + raise Exception(f"Error deleting fetch tool from DB: {str(e)}") + + async def get_fetch_tool_from_db( + self, fetch_tool_id: str, prisma_client: PrismaClient + ): + """ + Get a fetch tool from the database by ID. + + Args: + fetch_tool_id: ID of fetch tool to retrieve + prisma_client: Prisma client instance + + Returns: + Dict with fetch tool data or None if not found + """ + try: + fetch_tool = await prisma_client.db.litellm_fetchtoolstable.find_unique( + where={"fetch_tool_id": fetch_tool_id} + ) + if fetch_tool is None: + return None + return self._convert_prisma_to_dict(fetch_tool) + except Exception as e: + verbose_proxy_logger.exception( + f"Error getting fetch tool from DB: {str(e)}" + ) + raise Exception(f"Error getting fetch tool from DB: {str(e)}") + + async def get_all_fetch_tools_from_db(self, prisma_client: PrismaClient): + """ + Get all fetch tools from the database. + + Args: + prisma_client: Prisma client instance + + Returns: + List of dicts with fetch tool data + """ + try: + fetch_tools_from_db = ( + await prisma_client.db.litellm_fetchtoolstable.find_many() + ) + fetch_tools: List[FetchTool] = [] + for fetch_tool in fetch_tools_from_db: + fetch_tool_dict = self._convert_prisma_to_dict(fetch_tool) + fetch_tools.append(FetchTool(**fetch_tool_dict)) # type: ignore + return fetch_tools + except Exception as e: + verbose_proxy_logger.exception( + f"Error getting all fetch tools from DB: {str(e)}" + ) + raise Exception(f"Error getting all fetch tools from DB: {str(e)}") + + async def update_fetch_tool_in_db( + self, + fetch_tool_id: str, + fetch_tool: FetchTool, + prisma_client: PrismaClient, + ): + """ + Update a fetch tool in the database. + + Args: + fetch_tool_id: ID of fetch tool to update + fetch_tool: Updated fetch tool configuration + prisma_client: Prisma client instance + + Returns: + Dict with updated fetch tool data + """ + try: + fetch_tool_name = fetch_tool.get("fetch_tool_name") + litellm_params: str = safe_dumps(dict(fetch_tool.get("litellm_params", {}))) + fetch_tool_info: str = safe_dumps(fetch_tool.get("fetch_tool_info", {})) + + updated_fetch_tool = await prisma_client.db.litellm_fetchtoolstable.update( + where={"fetch_tool_id": fetch_tool_id}, + data={ + "fetch_tool_name": fetch_tool_name, + "litellm_params": litellm_params, + "fetch_tool_info": fetch_tool_info, + "updated_at": datetime.now(timezone.utc), + }, + ) + + fetch_tool_dict = dict(fetch_tool) + fetch_tool_dict["fetch_tool_id"] = updated_fetch_tool.fetch_tool_id + fetch_tool_dict["created_at"] = updated_fetch_tool.created_at.isoformat() + fetch_tool_dict["updated_at"] = updated_fetch_tool.updated_at.isoformat() + + return fetch_tool_dict + except Exception as e: + verbose_proxy_logger.exception(f"Error updating fetch tool in DB: {str(e)}") + raise Exception(f"Error updating fetch tool in DB: {str(e)}") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fc9855abfa5..daf6221b930 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -479,6 +479,9 @@ 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.fetch_endpoints.fetch_tool_management import ( + router as fetch_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, @@ -14342,6 +14345,7 @@ app.include_router(usage_ai_router) app.include_router(policy_crud_router) app.include_router(policy_resolve_router) app.include_router(search_tool_management_router) +app.include_router(fetch_tool_management_router) app.include_router(prompts_router) app.include_router(callback_management_endpoints_router) app.include_router(debugging_endpoints_router) diff --git a/litellm/types/fetch.py b/litellm/types/fetch.py new file mode 100644 index 00000000000..9c9151236fa --- /dev/null +++ b/litellm/types/fetch.py @@ -0,0 +1,81 @@ +""" +LiteLLM Fetch API Types + +This module defines types for the unified fetch API across different providers. +""" + +from typing import List, Optional + +from typing_extensions import Required, TypedDict + +from litellm.types.utils import FetchProviders + +# Re-export FetchProviders as FetchProvider for backwards compatibility +FetchProvider = FetchProviders + +__all__ = ["FetchProvider", "FetchProviders"] + + +class FetchToolLiteLLMParams(TypedDict, total=False): + """ + LiteLLM params for fetch tools configuration. + """ + + fetch_provider: Required[str] + api_key: Optional[str] + api_base: Optional[str] + timeout: Optional[float] + max_retries: Optional[int] + + +class FetchTool(TypedDict, total=False): + """ + Fetch tool configuration. + + Example: + { + "fetch_tool_id": "123e4567-e89b-12d3-a456-426614174000", + "fetch_tool_name": "litellm-fetch", + "litellm_params": { + "fetch_provider": "firecrawl", + "api_key": "fc-..." + }, + "fetch_tool_info": { + "description": "Firecrawl fetch tool" + } + } + """ + + fetch_tool_id: Optional[str] + fetch_tool_name: Required[str] + litellm_params: Required[FetchToolLiteLLMParams] + fetch_tool_info: Optional[dict] + created_at: Optional[str] + updated_at: Optional[str] + + +class FetchToolInfoResponse(TypedDict, total=False): + """Response model for fetch tool information.""" + + fetch_tool_id: Optional[str] + fetch_tool_name: str + litellm_params: dict + fetch_tool_info: Optional[dict] + created_at: Optional[str] + updated_at: Optional[str] + is_from_config: Optional[ + bool + ] # True if this tool is defined in config file, False if from DB + + +class ListFetchToolsResponse(TypedDict): + """Response model for listing fetch tools.""" + + fetch_tools: List[FetchToolInfoResponse] + + +class AvailableFetchProvider(TypedDict): + """Information about an available fetch provider.""" + + provider_name: str + ui_friendly_name: str diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ed29d49fc29..83bc8a06459 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3374,6 +3374,19 @@ class SearchProviders(str, Enum): SearchProvidersSet = {provider.value for provider in SearchProviders} +class FetchProviders(str, Enum): + """ + Enum for fetch provider types. + Separate from SearchProviders for semantic clarity. + """ + + FIRECRAWL = "firecrawl" + + +# Create a set of all fetch provider values for quick lookup +FetchProvidersSet = {provider.value for provider in FetchProviders} + + class LiteLLMLoggingBaseClass: """ Base class for logging pre and post call diff --git a/schema.prisma b/schema.prisma index a9d3911c07b..18ef4edfd3e 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1065,6 +1065,16 @@ model LiteLLM_SearchToolsTable { updated_at DateTime @updatedAt } +// Fetch Tools table for storing fetch tool configurations +model LiteLLM_FetchToolsTable { + fetch_tool_id String @id @default(uuid()) + fetch_tool_name String @unique + litellm_params Json + fetch_tool_info Json? + created_at DateTime @default(now()) + updated_at DateTime @updatedAt +} + // SSO configuration table model LiteLLM_SSOConfig { id String @id @default("sso_config")