import hashlib import json from datetime import datetime, timezone from typing import Any, Dict, List, Optional import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.management_helpers.object_permission_utils import \ handle_update_object_permission_common from litellm.proxy.utils import PrismaClient from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest class AgentRegistry: def __init__(self): self.agent_list: List[AgentResponse] = [] def reset_agent_list(self): self.agent_list = [] def register_agent(self, agent_config: AgentResponse): self.agent_list.append(agent_config) def deregister_agent(self, agent_name: str): self.agent_list = [ agent for agent in self.agent_list if agent.agent_name != agent_name ] def get_agent_list(self, agent_names: Optional[List[str]] = None): if agent_names is not None: return [ agent for agent in self.agent_list if agent.agent_name in agent_names ] return self.agent_list def get_public_agent_list(self) -> List[AgentResponse]: public_agent_list: List[AgentResponse] = [] if litellm.public_agent_groups is None: return public_agent_list for agent in self.agent_list: if agent.agent_id in litellm.public_agent_groups: public_agent_list.append(agent) return public_agent_list def _create_agent_id(self, agent_config: AgentConfig) -> str: return hashlib.sha256( json.dumps(agent_config, sort_keys=True).encode() ).hexdigest() def load_agents_from_config(self, agent_config: Optional[List[AgentConfig]] = None): if agent_config is None: return None for agent_config_item in agent_config: if not isinstance(agent_config_item, dict): raise ValueError("agent_config must be a list of dictionaries") agent_name = agent_config_item.get("agent_name") agent_card_params = agent_config_item.get("agent_card_params") if not all([agent_name, agent_card_params]): continue # create a stable hash id for config item config_hash = self._create_agent_id(agent_config_item) self.register_agent(agent_config=AgentResponse(agent_id=config_hash, **agent_config_item)) # type: ignore def load_agents_from_db_and_config( self, agent_config: Optional[List[AgentConfig]] = None, db_agents: Optional[List[Dict[str, Any]]] = None, ): self.reset_agent_list() if agent_config: for agent_config_item in agent_config: if not isinstance(agent_config_item, dict): raise ValueError("agent_config must be a list of dictionaries") self.register_agent(agent_config=AgentResponse(agent_id=self._create_agent_id(agent_config_item), **agent_config_item)) # type: ignore if db_agents: for db_agent in db_agents: if not isinstance(db_agent, dict): raise ValueError("db_agents must be a list of dictionaries") self.register_agent(agent_config=AgentResponse(**db_agent)) # type: ignore return self.agent_list ########################################################### ########### DB management helpers for agents ########### ############################################################ async def add_agent_to_db( self, agent: AgentConfig, prisma_client: PrismaClient, created_by: str ) -> AgentResponse: """ Add an agent to the database """ try: agent_name = agent.get("agent_name") # Serialize litellm_params litellm_params_obj: Any = agent.get("litellm_params", {}) if hasattr(litellm_params_obj, "model_dump"): litellm_params_dict = litellm_params_obj.model_dump() else: litellm_params_dict = ( dict(litellm_params_obj) if litellm_params_obj else {} ) litellm_params: str = safe_dumps(litellm_params_dict) # Serialize agent_card_params agent_card_params_obj: Any = agent.get("agent_card_params", {}) if hasattr(agent_card_params_obj, "model_dump"): agent_card_params_dict = agent_card_params_obj.model_dump() else: agent_card_params_dict = ( dict(agent_card_params_obj) if agent_card_params_obj else {} ) agent_card_params: str = safe_dumps(agent_card_params_dict) # Handle object_permission (MCP tool access for agent) object_permission_id: Optional[str] = None if agent.get("object_permission") is not None: agent_copy = dict(agent) object_permission_id = await handle_update_object_permission_common( agent_copy, None, prisma_client ) # Serialize static_headers static_headers_obj = agent.get("static_headers") static_headers_val: Optional[str] = ( safe_dumps(dict(static_headers_obj)) if static_headers_obj else None ) extra_headers_val: Optional[List[str]] = agent.get("extra_headers") create_data: Dict[str, Any] = { "agent_name": agent_name, "litellm_params": litellm_params, "agent_card_params": agent_card_params, "created_by": created_by, "updated_by": created_by, "created_at": datetime.now(timezone.utc), "updated_at": datetime.now(timezone.utc), } if static_headers_val is not None: create_data["static_headers"] = static_headers_val if extra_headers_val is not None: create_data["extra_headers"] = extra_headers_val if object_permission_id is not None: create_data["object_permission_id"] = object_permission_id for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): _val = agent.get(rate_field) if _val is not None: create_data[rate_field] = _val # Create agent in DB created_agent = await prisma_client.db.litellm_agentstable.create( data=create_data, include={"object_permission": True}, ) created_agent_dict = created_agent.model_dump() if created_agent.object_permission is not None: try: created_agent_dict["object_permission"] = created_agent.object_permission.model_dump() except Exception: created_agent_dict["object_permission"] = created_agent.object_permission.dict() return AgentResponse(**created_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error adding agent to DB: {str(e)}") async def delete_agent_from_db( self, agent_id: str, prisma_client: PrismaClient ) -> Dict[str, Any]: """ Delete an agent from the database """ try: deleted_agent = await prisma_client.db.litellm_agentstable.delete( where={"agent_id": agent_id} ) return dict(deleted_agent) except Exception as e: raise Exception(f"Error deleting agent from DB: {str(e)}") async def patch_agent_in_db( self, agent_id: str, agent: PatchAgentRequest, prisma_client: PrismaClient, updated_by: str, ) -> AgentResponse: """ Patch an agent in the database. Get the existing agent from the database and patch it with the new values. Args: agent_id: The ID of the agent to patch agent: The new agent values to patch prisma_client: The Prisma client to use updated_by: The user ID of the user who is patching the agent Returns: The patched agent """ try: existing_agent = await prisma_client.db.litellm_agentstable.find_unique( where={"agent_id": agent_id} ) if existing_agent is not None: existing_agent = dict(existing_agent) if existing_agent is None: raise Exception(f"Agent with ID {agent_id} not found") augment_agent = {**existing_agent, **agent} update_data: Dict[str, Any] = {} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if augment_agent.get("litellm_params"): update_data["litellm_params"] = safe_dumps( augment_agent.get("litellm_params") ) if augment_agent.get("agent_card_params"): update_data["agent_card_params"] = safe_dumps( augment_agent.get("agent_card_params") ) for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): if rate_field in agent: update_data[rate_field] = agent.get(rate_field) if "static_headers" in agent: headers_value = agent.get("static_headers") update_data["static_headers"] = safe_dumps( dict(headers_value) if headers_value is not None else {} ) if "extra_headers" in agent: extra_headers_value = agent.get("extra_headers") update_data["extra_headers"] = ( extra_headers_value if extra_headers_value is not None else [] ) if agent.get("object_permission") is not None: agent_copy = dict(augment_agent) existing_object_permission_id = existing_agent.get( "object_permission_id" ) object_permission_id = ( await handle_update_object_permission_common( agent_copy, existing_object_permission_id, prisma_client, ) ) if object_permission_id is not None: update_data["object_permission_id"] = object_permission_id # Patch agent in DB patched_agent = await prisma_client.db.litellm_agentstable.update( where={"agent_id": agent_id}, data={ **update_data, "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), }, include={"object_permission": True}, ) patched_agent_dict = patched_agent.model_dump() if patched_agent.object_permission is not None: try: patched_agent_dict["object_permission"] = patched_agent.object_permission.model_dump() except Exception: patched_agent_dict["object_permission"] = patched_agent.object_permission.dict() return AgentResponse(**patched_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error patching agent in DB: {str(e)}") async def update_agent_in_db( self, agent_id: str, agent: AgentConfig, prisma_client: PrismaClient, updated_by: str, ) -> AgentResponse: """ Update an agent in the database """ try: agent_name = agent.get("agent_name") # Serialize litellm_params litellm_params_obj: Any = agent.get("litellm_params", {}) if hasattr(litellm_params_obj, "model_dump"): litellm_params_dict = litellm_params_obj.model_dump() else: litellm_params_dict = ( dict(litellm_params_obj) if litellm_params_obj else {} ) litellm_params: str = safe_dumps(litellm_params_dict) # Serialize agent_card_params agent_card_params_obj: Any = agent.get("agent_card_params", {}) if hasattr(agent_card_params_obj, "model_dump"): agent_card_params_dict = agent_card_params_obj.model_dump() else: agent_card_params_dict = ( dict(agent_card_params_obj) if agent_card_params_obj else {} ) agent_card_params: str = safe_dumps(agent_card_params_dict) # Serialize static_headers for update static_headers_obj_u = agent.get("static_headers") static_headers_val_u: str = ( safe_dumps(dict(static_headers_obj_u)) if static_headers_obj_u is not None else safe_dumps({}) ) extra_headers_val_u: List[str] = agent.get("extra_headers") or [] update_data: Dict[str, Any] = { "agent_name": agent_name, "litellm_params": litellm_params, "agent_card_params": agent_card_params, "static_headers": static_headers_val_u, "extra_headers": extra_headers_val_u, "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), } for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): _val = agent.get(rate_field) if _val is not None: update_data[rate_field] = _val if agent.get("object_permission") is not None: existing_agent = await prisma_client.db.litellm_agentstable.find_unique( where={"agent_id": agent_id} ) existing_object_permission_id = ( existing_agent.object_permission_id if existing_agent is not None else None ) agent_copy = dict(agent) object_permission_id = ( await handle_update_object_permission_common( agent_copy, existing_object_permission_id, prisma_client, ) ) if object_permission_id is not None: update_data["object_permission_id"] = object_permission_id # Update agent in DB updated_agent = await prisma_client.db.litellm_agentstable.update( where={"agent_id": agent_id}, data=update_data, include={"object_permission": True}, ) updated_agent_dict = updated_agent.model_dump() if updated_agent.object_permission is not None: try: updated_agent_dict["object_permission"] = updated_agent.object_permission.model_dump() except Exception: updated_agent_dict["object_permission"] = updated_agent.object_permission.dict() return AgentResponse(**updated_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error updating agent in DB: {str(e)}") @staticmethod async def get_all_agents_from_db( prisma_client: PrismaClient, ) -> List[Dict[str, Any]]: """ Get all agents from the database """ try: agents_from_db = await prisma_client.db.litellm_agentstable.find_many( order={"created_at": "desc"}, include={"object_permission": True}, ) agents: List[Dict[str, Any]] = [] for agent in agents_from_db: agent_dict = dict(agent) # object_permission is eagerly loaded via include above if agent.object_permission is not None: try: agent_dict["object_permission"] = agent.object_permission.model_dump() except Exception: agent_dict["object_permission"] = agent.object_permission.dict() agents.append(agent_dict) return agents except Exception as e: raise Exception(f"Error getting agents from DB: {str(e)}") def get_agent_by_id( self, agent_id: str, ) -> Optional[AgentResponse]: """ Get an agent by its ID from the database """ try: for agent in self.agent_list: if agent.agent_id == agent_id: return agent return None except Exception as e: raise Exception(f"Error getting agent from DB: {str(e)}") def get_agent_by_name(self, agent_name: str) -> Optional[AgentResponse]: """ Get an agent by its name from the database """ try: for agent in self.agent_list: if agent.agent_name == agent_name: return agent return None except Exception as e: raise Exception(f"Error getting agent from DB: {str(e)}") global_agent_registry = AgentRegistry()