From 64d7db125d83ccaf9140d7822fe5b783ceff5a97 Mon Sep 17 00:00:00 2001 From: UltimateGamingShack Date: Thu, 10 Sep 2026 01:26:31 +0530 Subject: [PATCH] feat: add ADEPT deterministic template routing for adaptive SLM delegation in agentic workflows --- litellm/router.py | 163 +- .../router_strategy/adept_router/__init__.py | 0 .../adept_router/adept_router.py | 308 +++ .../router_strategy/adept_router/config.py | 3 + .../adept_router/store/__init__.py | 0 .../store/implementation/__init__.py | 0 .../store/implementation/prisma.py | 326 +++ .../adept_router/store/store_template.py | 71 + .../adept_router/template/__init__.py | 0 .../template/implementation/__init__.py | 0 .../implementation/adept_template_router.py | 305 +++ .../adept_router/template/router_template.py | 41 + litellm/types/router.py | 13 + litellm/types/utils.py | 13 + .../router_code_coverage.py | 1 + .../test_get_supported_openai_params.py | 48 +- ...test_batch_embed_content_transformation.py | 42 +- tests/test_litellm/test_adept_router.py | 1792 +++++++++++++++++ ui/litellm-dashboard/eslint-suppressions.json | 2 +- .../models-and-endpoints/page.test.tsx | 1 + .../(dashboard)/models-and-endpoints/page.tsx | 7 +- .../panels/AddAdeptRouterPanel.tsx | 40 + .../add_model/AddAdeptRouterTab.tsx | 291 +++ .../add_model/HandleAddAdeptRouterSubmit.tsx | 66 + .../AdeptRouterEditControl.tsx | 46 + .../EditAdeptRouterModal.tsx | 297 +++ .../src/components/model_info_view.tsx | 9 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 48 + 28 files changed, 3861 insertions(+), 72 deletions(-) create mode 100644 litellm/router_strategy/adept_router/__init__.py create mode 100644 litellm/router_strategy/adept_router/adept_router.py create mode 100644 litellm/router_strategy/adept_router/config.py create mode 100644 litellm/router_strategy/adept_router/store/__init__.py create mode 100644 litellm/router_strategy/adept_router/store/implementation/__init__.py create mode 100644 litellm/router_strategy/adept_router/store/implementation/prisma.py create mode 100644 litellm/router_strategy/adept_router/store/store_template.py create mode 100644 litellm/router_strategy/adept_router/template/__init__.py create mode 100644 litellm/router_strategy/adept_router/template/implementation/__init__.py create mode 100644 litellm/router_strategy/adept_router/template/implementation/adept_template_router.py create mode 100644 litellm/router_strategy/adept_router/template/router_template.py create mode 100644 tests/test_litellm/test_adept_router.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddAdeptRouterPanel.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/AddAdeptRouterTab.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/HandleAddAdeptRouterSubmit.tsx create mode 100644 ui/litellm-dashboard/src/components/edit_adept_router/AdeptRouterEditControl.tsx create mode 100644 ui/litellm-dashboard/src/components/edit_adept_router/EditAdeptRouterModal.tsx diff --git a/litellm/router.py b/litellm/router.py index 0582805943e..5184e9479e9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -300,6 +300,9 @@ if TYPE_CHECKING: from litellm.router_strategy.adaptive_router.adaptive_router import ( AdaptiveRouter, ) + from litellm.router_strategy.adept_router.adept_router import ( + AdeptRouter, + ) from litellm.router_strategy.auto_router.auto_router import ( AutoRouter, PreRoutingHookResponse, @@ -924,6 +927,7 @@ class Router: self.adaptive_routers: dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]] = {} self.quality_routers: dict[str, list[TaggedPreRoutingStrategy[QualityRouter]]] = {} self.routing_plugins: list[RoutingPlugin] = list(plugins) if plugins else [] + self.adept_routers: dict[str, AdeptRouter] = {} # mutable-ok: registry filled per deploy sync # Initialize model_group_alias early since it's used in set_model_list self.model_group_alias: dict[str, str | RouterModelGroupAliasItem] = ( @@ -9199,6 +9203,26 @@ class Router: if self._unregister_pre_routing_strategy(self.adaptive_routers, model_name, tags): self._sync_adaptive_router_hooks() + def _release_adept_router_for_deleted_deployment(self, item: object) -> None: + """Drop an ADEPT deployment's in-memory router, unregister its callbacks, and release + its shared PG client when no sibling deployment still references the same URL.""" + model_name: Final = ( + item.model_name + if isinstance(item, Deployment) + else (item.get("model_name") if isinstance(item, dict) else None) + ) + if model_name is None: + return + removed_adept: Final = self.adept_routers.pop(model_name, None) + if removed_adept is None: + return + litellm.logging_callback_manager.remove_callback_from_all_lists(removed_adept) + removed_adept.close() + if not any(r.pg_url == removed_adept.pg_url for r in self.adept_routers.values()): + from litellm.router_strategy.adept_router.store.implementation.prisma import schedule_disconnect + + schedule_disconnect(removed_adept.pg_url) + def _finalize_adaptive_router_if_configured(self) -> None: """Locate every adaptive-router deployment in the finalized model_list and build an AdaptiveRouter for each. Safe no-op when none are configured. @@ -9374,6 +9398,97 @@ class Router: strategy_label="Quality-router", ) + def _is_adept_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: + """Returns True when the model prefix is 'adept/'.""" + return litellm_params.model.startswith("adept/") + + def init_adept_router_deployment(self, deployment: Deployment) -> None: + """Initialize an ADEPT router deployment and register it in self.adept_routers. + + Idempotent: repeat calls with identical params are a no-op; a param change rebuilds + the in-memory router so the new value takes effect without a proxy restart. + """ + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + from litellm.router_strategy.adept_router.config import ( + DEFAULT_CONVERSATIONS_THRESHOLD, + ) + + lp: Final = deployment.litellm_params + default_model: Final = lp.adept_router_default_model + if default_model is None: + raise ValueError("adept_router_default_model is required for ADEPT router deployments.") + + if not lp.adept_router_pg_host: + raise ValueError( + "adept_router_pg_host is required for ADEPT router deployments. " + "Configure a PostgreSQL database so the trainer pipeline can access the data." + ) + + from urllib.parse import quote_plus + + password: Final = lp.adept_router_pg_password or "" + port: Final = lp.adept_router_pg_port or 5432 + user: Final = lp.adept_router_pg_user or "" + database: Final = lp.adept_router_pg_database or "" + ssl_mode: Final = lp.adept_router_pg_ssl_mode or "prefer" + pg_url: Final = f"postgresql://{quote_plus(user)}:{quote_plus(password)}@{lp.adept_router_pg_host}:{port}/{database}?sslmode={ssl_mode}" + + trainer_url: Final = lp.adept_router_trainer_url + if trainer_url: + from litellm.litellm_core_utils.url_utils import SSRFError, validate_url + + try: + validate_url(trainer_url) + except SSRFError as ssrf_err: + raise ValueError( + f"adept_router_trainer_url {trainer_url!r} rejected by SSRF guard: {ssrf_err}. " + "Add the host to `general_settings.user_url_allowed_hosts` if it is an internal " + "trainer the operator has explicitly cleared." + ) from ssrf_err + + threshold: Final = lp.adept_router_conversations_threshold or DEFAULT_CONVERSATIONS_THRESHOLD + tag_prefix: Final = lp.adept_router_tag_prefix or "" + + existing: Final = self.adept_routers.get(deployment.model_name) + if existing is not None: + params_changed: Final = ( + existing.default_model != default_model + or existing.pg_url != pg_url + or existing.template_router.trainer_url != trainer_url + or existing.template_router.conversations_threshold != threshold + or existing.template_router.tag_prefix != tag_prefix + ) + if not params_changed: + verbose_router_logger.debug( + "AdeptRouter: '%s' already registered with matching params — skipping re-init.", + deployment.model_name, + ) + return + verbose_router_logger.info( + "AdeptRouter: '%s' params changed — rebuilding in-memory router.", deployment.model_name + ) + litellm.logging_callback_manager.remove_callback_from_all_lists(existing) + existing.close() + if existing.pg_url != pg_url and not any( + name != deployment.model_name and r.pg_url == existing.pg_url for name, r in self.adept_routers.items() + ): + from litellm.router_strategy.adept_router.store.implementation.prisma import schedule_disconnect + + schedule_disconnect(existing.pg_url) + + adept_router: Final = AdeptRouter( + model_name=deployment.model_name, + default_model=default_model, + litellm_router_instance=self, + pg_url=pg_url, + tag_prefix=tag_prefix, + conversations_threshold=threshold, + trainer_url=trainer_url, + seed_config=lp.adept_router_seed_config, + ) + self.adept_routers[deployment.model_name] = adept_router + litellm.logging_callback_manager.add_litellm_callback(adept_router) + def deployment_is_active_for_environment(self, deployment: Deployment) -> bool: """ Function to check if a llm deployment is active for a given environment. Allows using the same config.yaml across multople environments @@ -9403,6 +9518,11 @@ class Router: self.complexity_routers = {} self.auto_routers = {} self._provider_unresolved_deployments = () + for stale_adept in tuple(self.adept_routers.values()): + litellm.logging_callback_manager.remove_callback_from_all_lists(stale_adept) + stale_adept.close() + pre_reload_adept_urls: Final = frozenset(r.pg_url for r in self.adept_routers.values()) + self.adept_routers = {} # mutable-ok: registry reset on model_list reload; populated incrementally per deployment sync self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works @@ -9480,6 +9600,15 @@ class Router: # deployments have been registered. self._finalize_adaptive_router_if_configured() + # Deferred until after re-init so we never disconnect a client a rebuilt deployment reused. + post_reload_adept_urls: Final = frozenset(r.pg_url for r in self.adept_routers.values()) + orphaned_adept_urls: Final = pre_reload_adept_urls - post_reload_adept_urls + if orphaned_adept_urls: + from litellm.router_strategy.adept_router.store.implementation.prisma import schedule_disconnect + + for orphaned_url in orphaned_adept_urls: + schedule_disconnect(orphaned_url) + def _add_deployment(self, deployment: Deployment) -> Deployment: import os @@ -9499,9 +9628,10 @@ class Router: if split_litellm_model in litellm._known_custom_logger_compatible_callbacks: is_prompt_management_model = True - if is_prompt_management_model: - # For prompt management models, skip LLM provider validation - # The actual model will be resolved at runtime from the prompt file + # ADEPT deployments skip standard get_llm_provider validation; init_adept_router_deployment handles them. + is_adept_router_model: Final = self._is_adept_router_deployment(litellm_params=deployment.litellm_params) + + if is_prompt_management_model or is_adept_router_model: _model = litellm_model custom_llm_provider = None dynamic_api_key = None @@ -9606,6 +9736,9 @@ class Router: if self._is_quality_router_deployment(litellm_params=deployment.litellm_params): self.init_quality_router_deployment(deployment=deployment) + if self._is_adept_router_deployment(litellm_params=deployment.litellm_params): + self.init_adept_router_deployment(deployment=deployment) + return deployment def _initialize_deployment_for_pass_through(self, deployment: Deployment, custom_llm_provider: str): @@ -10097,6 +10230,7 @@ class Router: "the deployment is out of the model_list and its indices are repaired", id, ) + self._release_adept_router_for_deleted_deployment(item) return item else: return None @@ -13511,6 +13645,14 @@ class Router: model=registered_model_name, request_kwargs=request_kwargs ) if selected_strategy is None: + if registered_model_name in self.adept_routers: + return await self.adept_routers[registered_model_name].async_pre_routing_hook( + model=registered_model_name, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) self._stamp_or_clear_metadata_key( request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None @@ -13617,6 +13759,15 @@ class Router: if newly_forwarded: request_kwargs.update(((_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, tuple(key for key, _ in newly_forwarded)),)) + if registered_model_name in self.adept_routers: + return await self.adept_routers[registered_model_name].async_pre_routing_hook( + model=registered_model_name, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) + return pre_routing_hook_response def _forwardable_alias_marker_params( @@ -13790,6 +13941,12 @@ class Router: "usage-based-routing-v2, cost-based-routing, latency-based-routing, least-busy), " "or remove `plugins` from the Router config." ) + if self.adept_routers: + raise ValueError( + "An ADEPT router is configured but this call resolved to the synchronous " + "deployment-selection path, which never runs the async pre-routing hook. " + "Use an async Router method." + ) # users need to explicitly call a specific deployment, by setting `specific_deployment = True` as completion()/embedding() kwarg # When this was no explicit we had several issues with fallbacks timing out diff --git a/litellm/router_strategy/adept_router/__init__.py b/litellm/router_strategy/adept_router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/router_strategy/adept_router/adept_router.py b/litellm/router_strategy/adept_router/adept_router.py new file mode 100644 index 00000000000..50cd0b56f5c --- /dev/null +++ b/litellm/router_strategy/adept_router/adept_router.py @@ -0,0 +1,308 @@ +""" +ADEPT (Adaptive Deployment via Prompt Templates) Router. + +Extracts a structural skeleton from each single-turn prompt, hashes it with the system prompt +for per-tool isolation, and routes to a task-specific SLM once one has been trained. Until then, +traffic falls back to the default model while conversations accumulate as training data. +""" + +import asyncio +import datetime +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm._logging import verbose_router_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_strategy.adept_router.config import DEFAULT_CONVERSATIONS_THRESHOLD +from litellm.router_strategy.adept_router.template.implementation.adept_template_router import ( + AdeptTemplateRouter, +) +from litellm.types.utils import ModelResponse, Usage + +if TYPE_CHECKING: + from litellm.router import Router + from litellm.types.router import PreRoutingHookResponse +else: + Router: Final = object + PreRoutingHookResponse: Final = object + + +class _MessageContentBlock(BaseModel): + model_config = ConfigDict(extra="ignore") + type: str | None = None + text: str | None = None + + +class _UsageEnvelope(BaseModel): + model_config = ConfigDict(extra="ignore") + usage: Usage | None = None + + +_MessageList: TypeAlias = list[dict[str, object]] +_MESSAGES_ADAPTER: Final = TypeAdapter(_MessageList) +_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[_MessageContentBlock]) +_METADATA_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +class AdeptRouter(CustomLogger): + def __init__( + self, + model_name: str, + default_model: str, + litellm_router_instance: "Router", + pg_url: str, + tag_prefix: str = "", + conversations_threshold: int = DEFAULT_CONVERSATIONS_THRESHOLD, + trainer_url: str | None = None, + seed_config: Sequence[Mapping[str, object]] | None = None, + ) -> None: + self.model_name = model_name + self.default_model = default_model + self.pg_url = pg_url + self.litellm_router_instance = litellm_router_instance + self.template_router = AdeptTemplateRouter( + model_name=model_name, + litellm_router_instance=litellm_router_instance, + pg_url=pg_url, + tag_prefix=tag_prefix, + conversations_threshold=conversations_threshold, + trainer_url=trainer_url, + ) + self._seed_config = seed_config + self._seeded = not seed_config + self._seed_lock = asyncio.Lock() + self._seed_task: asyncio.Task[None] | None = None # mutable-ok: task handle rotates on restart + + def _kick_off_seed(self) -> None: + if self._seeded: + return + existing: Final = self._seed_task + if existing is not None and not existing.done(): + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + self._seed_task = loop.create_task(self._run_seed()) + + def close(self) -> None: + """Cancel background tasks so a retired router cannot mutate state or write to the DB.""" + if self._seed_task is not None: + self._seed_task.cancel() + self._seed_task = None + self.template_router.stop_refresh() + + async def _run_seed(self) -> None: + async with self._seed_lock: + if self._seeded: + return + for entry in self._seed_config or (): + description = entry.get("description", "") + target_model = entry.get("target_model", self.default_model) + if not description: + verbose_router_logger.warning( + "AdeptRouter: seed_config entry missing 'description', skipping: %s", str(entry)[:100] + ) + continue + if await self.template_router.seed_template(str(description), str(target_model)): + verbose_router_logger.info("AdeptRouter: seeded template for target_model=%s", target_model) + self._seeded = True + + async def async_pre_routing_hook( + self, + model: str, + request_kwargs: Mapping[str, object], + messages: _MessageList | None = None, + input: str | Sequence[object] | None = None, + specific_deployment: bool | None = False, + ) -> "PreRoutingHookResponse | None": + from litellm.types.router import PreRoutingHookResponse + + if messages is None: + return None + + self._kick_off_seed() + + message_content: Final = self._extract_user_text(messages) + if not message_content: + await self._authorize_routed_model(request_kwargs, self.default_model) + return PreRoutingHookResponse(model=self.default_model, messages=messages) + + system_prompt: Final = self._extract_system_prompt(messages) + + template_match: Final = await self.template_router.route(message_content, system_prompt) + + target: Final = template_match.get("target_model") if template_match is not None else None + routed_model: Final = target or self.default_model + routed_to_slm: Final = bool(target) + await self._authorize_routed_model(request_kwargs, routed_model) + if template_match is not None: + verbose_router_logger.info( + "AdeptRouter: matched template %s, routing to %s", + template_match.get("template_id"), + routed_model, + ) + else: + verbose_router_logger.info("AdeptRouter: no template match, falling back to %s", self.default_model) + + for md_key in ("metadata", "litellm_metadata"): + candidate = request_kwargs.get(md_key) + if isinstance(candidate, dict): + candidate["adept_routed_to_slm"] = routed_to_slm + break + + return PreRoutingHookResponse(model=routed_model, messages=messages) + + async def _authorize_routed_model(self, request_kwargs: Mapping[str, object], routed_model: str) -> None: + """Re-run proxy model-access check against the swapped target so a caller cannot be silently + upgraded from the ADEPT alias to a SLM they lack access to.""" + if routed_model == self.model_name: + return + auth_from_litellm_params: Final = self._read_request_metadata( + request_kwargs.get("litellm_params"), "user_api_key_auth" + ) + auth_obj: Final = ( + auth_from_litellm_params + if auth_from_litellm_params is not None + else self._extract_user_api_key_auth(request_kwargs) + ) + if auth_obj is None: + return + try: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + except ImportError: + return + if not isinstance(auth_obj, UserAPIKeyAuth): + return + await can_key_call_resolved_model( + model=routed_model, + llm_model_list=self.litellm_router_instance.model_list, + valid_token=auth_obj, + llm_router=self.litellm_router_instance, + ) + + @staticmethod + def _extract_user_api_key_auth(request_kwargs: Mapping[str, object]) -> object: + for md_key in ("metadata", "litellm_metadata"): + candidate = request_kwargs.get(md_key) + if isinstance(candidate, Mapping): + value = candidate.get("user_api_key_auth") + if value is not None: + return value + return None + + @staticmethod + def _read_request_metadata(litellm_params: object, key: str) -> object: + try: + params: Final = _METADATA_ADAPTER.validate_python(litellm_params) + except ValidationError: + return None + for md_key in ("metadata", "litellm_metadata"): + try: + nested = _METADATA_ADAPTER.validate_python(params.get(md_key)) # rebind-ok: loop rebinds per iteration + except ValidationError: + continue + value = nested.get(key) # rebind-ok: loop rebinds per iteration + if value is not None: + return value + return None + + async def async_log_success_event( + self, + kwargs: Mapping[str, object], + response_obj: ModelResponse, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + # Callback fires for every request; gate on model_group to only log our own. + lp_raw: Final = kwargs.get("litellm_params") + request_model_group: Final = self._read_request_metadata(lp_raw, "model_group") + if request_model_group != self.model_name: + return + + try: + messages: Final = _MESSAGES_ADAPTER.validate_python(kwargs.get("messages")) + except ValidationError: + return + if not messages: + return + if messages[-1].get("role") == "tool": + return + + try: + prompt_text: Final = self._extract_user_text(messages) + if not prompt_text: + return + + usage: Final = _UsageEnvelope.model_validate(response_obj, from_attributes=True).usage + if usage is None: + return + + response_content: Final = self._response_text(response_obj) + if response_content is None: + return + + token_usage: Final[dict[str, object]] = { # mutable-ok: JSON payload persisted to a Postgres JSON column + "prompt_tokens": usage.prompt_tokens, + "completion_tokens": usage.completion_tokens, + "total_tokens": usage.total_tokens, + } + + cost_raw: Final = kwargs.get("response_cost") + cost_usd: Final = cost_raw if isinstance(cost_raw, (int, float)) else None + latency_ms: Final = (end_time - start_time).total_seconds() * 1000 + system_prompt: Final = self._extract_system_prompt(messages) + routed_to_slm_raw: Final = self._read_request_metadata(lp_raw, "adept_routed_to_slm") + routed_to_slm: Final = routed_to_slm_raw if isinstance(routed_to_slm_raw, bool) else None + actual_model: Final = str(kwargs.get("model", "unknown")) + + await self.template_router.store_conversation( + prompt_text, + response_content, + actual_model, + token_usage, + cost_usd, + latency_ms, + system_prompt, + routed_to_slm, + ) + verbose_router_logger.info("AdeptRouter: stored interaction.") + except (AttributeError, KeyError, TypeError, ValueError): + verbose_router_logger.exception("AdeptRouter: failed to log success event") + + @staticmethod + def _response_text(response_obj: ModelResponse) -> str | None: + if not response_obj.choices: + return None + choice: Final = response_obj.choices[0] + return choice.message.content + + @staticmethod + def _content_to_text(content: object) -> str: + if content is None: + return "" + if isinstance(content, str): + return content + try: + blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content) + except ValidationError: + return str(content) + return " ".join(block.text or "" for block in blocks if block.type == "text") + + @staticmethod + def _extract_system_prompt(messages: Sequence[Mapping[str, object]]) -> str | None: + for msg in messages: + if msg.get("role") == "system": + content = msg.get("content") + return str(content) if content else None + return None + + @staticmethod + def _extract_user_text(messages: Sequence[Mapping[str, object]]) -> str: + for msg in reversed(messages): + if msg.get("role") == "user": + return AdeptRouter._content_to_text(msg.get("content")) + return "" diff --git a/litellm/router_strategy/adept_router/config.py b/litellm/router_strategy/adept_router/config.py new file mode 100644 index 00000000000..727f5609b0a --- /dev/null +++ b/litellm/router_strategy/adept_router/config.py @@ -0,0 +1,3 @@ +from typing import Final + +DEFAULT_CONVERSATIONS_THRESHOLD: Final = 1000 diff --git a/litellm/router_strategy/adept_router/store/__init__.py b/litellm/router_strategy/adept_router/store/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/router_strategy/adept_router/store/implementation/__init__.py b/litellm/router_strategy/adept_router/store/implementation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/router_strategy/adept_router/store/implementation/prisma.py b/litellm/router_strategy/adept_router/store/implementation/prisma.py new file mode 100644 index 00000000000..306f8cd18df --- /dev/null +++ b/litellm/router_strategy/adept_router/store/implementation/prisma.py @@ -0,0 +1,326 @@ +import asyncio +import json +from collections.abc import AsyncGenerator, Mapping, Sequence +from contextlib import asynccontextmanager +from datetime import datetime +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +from prisma import Prisma +from prisma.errors import PrismaError +from prisma.types import DatasourceOverride +from pydantic import BaseModel, ConfigDict, TypeAdapter, field_validator + +from litellm._logging import verbose_router_logger +from litellm.router_strategy.adept_router.store.store_template import ( + AdeptTemplateStore, + StoredTemplate, +) + +_JSON_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_DRAIN_TIMEOUT_SECONDS: Final = 30.0 + + +class _IdRow(BaseModel): + id: str + + +class _CountRow(BaseModel): + c: int + + +class _TemplateRow(BaseModel): + model_config = ConfigDict(extra="ignore") + + id: str + template: str + template_hash: str | None = None + router_id: str + target_model: str | None = None + additional_information: Mapping[str, object] | None = None + created_at: datetime | None = None + + @field_validator("additional_information", mode="before") + @classmethod + def _coerce_json(cls, value: object) -> object: + if isinstance(value, str): + return _JSON_ADAPTER.validate_json(value) + return value + + +def _redact_url(url: str) -> str: + """Return the URL with any embedded password replaced by '***'.""" + try: + split: Final = urlsplit(url) + except ValueError: + return "" + if split.password is None: + return url + user: Final = split.username or "" + host: Final = split.hostname or "" + port_suffix: Final = f":{split.port}" if split.port is not None else "" + netloc: Final = f"{user}:***@{host}{port_suffix}" if user else f":***@{host}{port_suffix}" + return urlunsplit((split.scheme, netloc, split.path, split.query, split.fragment)) + + +async def _create_tables(client: Prisma) -> None: + await client.execute_raw( + "CREATE TABLE IF NOT EXISTS templates (" + "id TEXT PRIMARY KEY, template TEXT NOT NULL, template_hash VARCHAR(64) NOT NULL, " + "router_id TEXT NOT NULL, target_model TEXT, additional_information JSONB, " + "created_at TIMESTAMPTZ NOT NULL DEFAULT now())" + ) + await client.execute_raw( + "CREATE UNIQUE INDEX IF NOT EXISTS ix_templates_router_hash ON templates (router_id, template_hash)" + ) + await client.execute_raw( + "CREATE TABLE IF NOT EXISTS conversations (" + "id SERIAL PRIMARY KEY, template_id TEXT NOT NULL REFERENCES templates(id), " + "prompt TEXT NOT NULL, response TEXT NOT NULL, additional_information JSONB, " + "created_at TIMESTAMPTZ NOT NULL DEFAULT now())" + ) + await client.execute_raw("CREATE INDEX IF NOT EXISTS ix_conversations_template_id ON conversations (template_id)") + + +def _json_or_none(payload: Mapping[str, object] | None) -> str | None: + return json.dumps(payload) if payload is not None else None + + +class _ClientHandle: + """Wraps a shared Prisma client with an in-flight refcount so `disconnect_client` can drain + active operations before tearing the socket down.""" + + __slots__ = ("_closed", "_drained", "_inflight", "_lock", "client") + + def __init__(self, client: Prisma) -> None: + self.client: Final = client + self._lock: Final = asyncio.Lock() + self._inflight = 0 # mutable-ok: refcount for graceful drain + self._drained: Final = asyncio.Event() + self._drained.set() + self._closed = False # mutable-ok: one-way close flag rejecting new borrows + + @asynccontextmanager + async def borrow(self) -> AsyncGenerator[Prisma]: + async with self._lock: + if self._closed: + raise PrismaError("Prisma client has been disconnected") + self._inflight += 1 # rebind-ok: refcount increment guarded by _lock + self._drained.clear() + try: + yield self.client + finally: + async with self._lock: + self._inflight -= 1 # rebind-ok: refcount decrement guarded by _lock + if self._inflight == 0: + self._drained.set() + + async def close(self) -> None: + async with self._lock: + self._closed = True # rebind-ok: one-way transition guarded by _lock + already_drained: Final = self._inflight == 0 + if already_drained: + return + try: + await asyncio.wait_for(self._drained.wait(), timeout=_DRAIN_TIMEOUT_SECONDS) + except asyncio.TimeoutError: + verbose_router_logger.warning( + "AdeptPrismaRepo: drain timed out with %s in-flight operations; disconnecting anyway.", + self._inflight, + ) + + +_CLIENTS: Final[dict[str, _ClientHandle]] = {} # mutable-ok: connection registry keyed by database URL +_REGISTRY_LOCK: Final = asyncio.Lock() + + +async def _get_handle(db_url: str) -> _ClientHandle: + cached: Final = _CLIENTS.get(db_url) + if cached is not None: + return cached + async with _REGISTRY_LOCK: + existing: Final = _CLIENTS.get(db_url) + if existing is not None: + return existing + client: Final = Prisma(datasource=DatasourceOverride(url=db_url)) + await client.connect() + await _create_tables(client) + handle: Final = _ClientHandle(client) + _CLIENTS[db_url] = handle + return handle + + +async def disconnect_client(db_url: str) -> None: + async with _REGISTRY_LOCK: + existing: Final = _CLIENTS.pop(db_url, None) + if existing is None: + return + await existing.close() + try: + await existing.client.disconnect() + except Exception as e: # noqa: BLE001 # best-effort: a hung PG socket must not block the rebuild path + verbose_router_logger.warning("AdeptPrismaRepo: disconnect failed for %s: %s", _redact_url(db_url), e) + + +def schedule_disconnect(db_url: str) -> None: + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + loop.create_task(disconnect_client(db_url)) + + +class AdeptPrismaRepo(AdeptTemplateStore): + def __init__(self, db_url: str) -> None: + if not db_url: + raise ValueError( + "A PostgreSQL connection URL is required. Example: postgresql://user:password@host:5432/dbname" + ) + self._db_url = db_url + + async def match_by_hash(self, template_hash: str, router_id: str) -> str | None: + try: + handle: Final = await _get_handle(self._db_url) + async with handle.borrow() as client: + rows: Final = await client.query_raw( + "SELECT id FROM templates WHERE router_id = $1 AND template_hash = $2 LIMIT 1", + router_id, + template_hash, + model=_IdRow, + ) + except PrismaError as e: + verbose_router_logger.error("Error matching template by hash: %s", e) + return None + else: + return rows[0].id if rows else None + + async def load_all_for_router(self, router_id: str, limit: int) -> Sequence[StoredTemplate]: + try: + handle: Final = await _get_handle(self._db_url) + async with handle.borrow() as client: + rows: Final = await client.query_raw( + "SELECT id, template, template_hash, router_id, target_model, additional_information, created_at " + "FROM templates WHERE router_id = $1 ORDER BY created_at DESC LIMIT $2", + router_id, + limit, + model=_TemplateRow, + ) + except PrismaError as e: + verbose_router_logger.error("Error loading templates for router %s: %s", router_id, e) + return () + return tuple( + StoredTemplate( + id=row.id, + template=row.template, + template_hash=row.template_hash, + router_id=row.router_id, + target_model=row.target_model, + additional_information=row.additional_information, + created_at=row.created_at, + ) + for row in rows + ) + + async def store_conversation( + self, + prompt: str, + response: str, + template_id: str | None = None, + additional_information: Mapping[str, object] | None = None, + ) -> bool: + if not template_id: + verbose_router_logger.error("template_id is required to store a conversation.") + return False + try: + handle: Final = await _get_handle(self._db_url) + async with handle.borrow() as client: + await client.execute_raw( + "INSERT INTO conversations (template_id, prompt, response, additional_information) " + "VALUES ($1, $2, $3, $4::jsonb)", + template_id, + prompt, + response, + _json_or_none(additional_information), + ) + except PrismaError as e: + verbose_router_logger.error("Error storing conversation: %s", e) + return False + return True + + async def store_template( + self, + template_id: str, + template: str, + template_hash: str, + target_model: str, + router_id: str, + additional_information: Mapping[str, object] | None = None, + ) -> str | None: + """Insert a new template row, returning the surviving id (ours or a concurrent insert's). + + ON CONFLICT DO NOTHING on the (router_id, template_hash) unique index makes concurrent + inserts safe: the loser no-ops and we re-read the winner. + """ + try: + handle: Final = await _get_handle(self._db_url) + async with handle.borrow() as client: + await client.execute_raw( + "INSERT INTO templates (id, template, template_hash, target_model, router_id, additional_information) " + "VALUES ($1, $2, $3, $4, $5, $6::jsonb) ON CONFLICT (router_id, template_hash) DO NOTHING", + template_id, + template, + template_hash, + target_model, + router_id, + _json_or_none(additional_information), + ) + rows: Final = await client.query_raw( + "SELECT id FROM templates WHERE router_id = $1 AND template_hash = $2 LIMIT 1", + router_id, + template_hash, + model=_IdRow, + ) + except PrismaError as e: + verbose_router_logger.error("AdeptRouter: error storing template: %s", e) + return None + return rows[0].id if rows else template_id + + async def get_template(self, template_id: str) -> StoredTemplate | None: + try: + handle: Final = await _get_handle(self._db_url) + async with handle.borrow() as client: + rows: Final = await client.query_raw( + "SELECT id, template, template_hash, router_id, target_model, additional_information, created_at " + "FROM templates WHERE id = $1 LIMIT 1", + template_id, + model=_TemplateRow, + ) + except PrismaError as e: + verbose_router_logger.error("Error retrieving template: %s", e) + return None + if not rows: + return None + row: Final = rows[0] + return StoredTemplate( + id=row.id, + template=row.template, + template_hash=row.template_hash, + router_id=row.router_id, + target_model=row.target_model, + additional_information=row.additional_information, + created_at=row.created_at, + ) + + async def count_conversation_by_template_id(self, template_id: str) -> int | None: + try: + handle: Final = await _get_handle(self._db_url) + async with handle.borrow() as client: + rows: Final = await client.query_raw( + "SELECT count(*)::int AS c FROM conversations WHERE template_id = $1", + template_id, + model=_CountRow, + ) + except PrismaError as e: + verbose_router_logger.error("Error counting conversations: %s", e) + return None + return rows[0].c if rows else 0 diff --git a/litellm/router_strategy/adept_router/store/store_template.py b/litellm/router_strategy/adept_router/store/store_template.py new file mode 100644 index 00000000000..d2e23dcee06 --- /dev/null +++ b/litellm/router_strategy/adept_router/store/store_template.py @@ -0,0 +1,71 @@ +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime + + +@dataclass(frozen=True, slots=True) +class StoredTemplate: + """A template row read back from the store.""" + + id: str + template: str + template_hash: str | None + router_id: str + target_model: str | None + additional_information: Mapping[str, object] | None + created_at: datetime | None + + +class AdeptTemplateStore(ABC): + """Abstract interface for storing and retrieving ADEPT prompt templates and conversations. + + Implementations talk to the user's own database, so every method is async. + """ + + @abstractmethod + async def match_by_hash(self, template_hash: str, router_id: str) -> str | None: + """Look up a template ID by the SHA-256 hash of its masked template string.""" + ... + + @abstractmethod + async def load_all_for_router(self, router_id: str, limit: int) -> Sequence[StoredTemplate]: + """Bulk-load recent templates for a router so the request path can serve from memory.""" + ... + + @abstractmethod + async def store_conversation( + self, + prompt: str, + response: str, + template_id: str | None = None, + additional_information: Mapping[str, object] | None = None, + ) -> bool: + """Store a prompt-response pair linked to a template.""" + ... + + @abstractmethod + async def store_template( + self, + template_id: str, + template: str, + template_hash: str, + target_model: str, + router_id: str, + additional_information: Mapping[str, object] | None = None, + ) -> str | None: + """ + Store a new template row. Returns the surviving template_id (ours or a concurrent + insert's) so the caller can use it without a follow-up query. + """ + ... + + @abstractmethod + async def get_template(self, template_id: str) -> StoredTemplate | None: + """Retrieve metadata for a specific template by ID.""" + ... + + @abstractmethod + async def count_conversation_by_template_id(self, template_id: str) -> int | None: + """Count conversations associated with a template.""" + ... diff --git a/litellm/router_strategy/adept_router/template/__init__.py b/litellm/router_strategy/adept_router/template/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/router_strategy/adept_router/template/implementation/__init__.py b/litellm/router_strategy/adept_router/template/implementation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/router_strategy/adept_router/template/implementation/adept_template_router.py b/litellm/router_strategy/adept_router/template/implementation/adept_template_router.py new file mode 100644 index 00000000000..9f71c55d4f1 --- /dev/null +++ b/litellm/router_strategy/adept_router/template/implementation/adept_template_router.py @@ -0,0 +1,305 @@ +import asyncio +import hashlib +import re +import time +from collections import OrderedDict +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final, TypeAlias +from uuid import uuid4 + +import httpx + +from litellm._logging import verbose_router_logger +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.router_strategy.adept_router.config import DEFAULT_CONVERSATIONS_THRESHOLD +from litellm.router_strategy.adept_router.store.store_template import ( + AdeptTemplateStore, + StoredTemplate, +) +from litellm.router_strategy.adept_router.template.router_template import ( + AdeptTemplateMatch, + BaseTemplateRouter, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + +if TYPE_CHECKING: + from litellm.router import Router +else: + Router: Final = object + +_TEMPLATE_CACHE_MAX_SIZE: Final = 1024 +_TEMPLATE_CACHE_TTL_SECONDS: Final = 60.0 +_REFRESH_INTERVAL_SECONDS: Final = 60.0 +_TRAINER_HTTP_TIMEOUT_SECONDS: Final = 10.0 + +_TemplateCache: TypeAlias = OrderedDict[tuple[str, str], tuple[float, StoredTemplate]] + + +class AdeptTemplateRouter(BaseTemplateRouter): + ID_RE = re.compile(r"\b[A-Z]{2,}-\d{3,}\b") + EMAIL_RE = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}") + URL_RE = re.compile(r"https?://\S+|www\.\S+") + # UUID must be masked before NUM — UUID hex digits partially match NUM_RE. + UUID_RE = re.compile( + r"\b[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-5][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}\b" + ) + NUM_RE = re.compile(r"\b\d{1,4}([/-]\d{1,2}([/-]\d{1,4})?)?\b") + NORMALIZE_RE = re.compile(r"\s+") + + def __init__( + self, + model_name: str, + litellm_router_instance: "Router", + pg_url: str, + tag_prefix: str = "", + conversations_threshold: int = DEFAULT_CONVERSATIONS_THRESHOLD, + trainer_url: str | None = None, + ) -> None: + from litellm.router_strategy.adept_router.store.implementation.prisma import ( + AdeptPrismaRepo, + ) + + self.model_name = model_name + self.litellm_router_instance = litellm_router_instance + self.tag_prefix = tag_prefix + self.conversations_threshold = conversations_threshold + self.trainer_url = trainer_url + self._router_id_cache: str | None = None + self._template_cache: _TemplateCache = OrderedDict() # mutable-ok: LRU cache, bounded and TTL-gated + self._refresh_task: asyncio.Task[None] | None = None # mutable-ok: task handle rotates on restart + + escaped_prefix: Final = re.escape(self.tag_prefix) + self.TAG_CONTENT_RE = re.compile( + r"<" + escaped_prefix + r"([a-zA-Z0-9_ ]+)>([^<]*)" + ) + self.TAG_REPLACEMENT = r"<" + escaped_prefix + r"\1>" + + self.template_store: AdeptTemplateStore = AdeptPrismaRepo(pg_url) + + def get_router_id(self) -> str: + if self._router_id_cache is None: + self._router_id_cache = self.litellm_router_instance.get_model_ids(model_name=self.model_name)[0] + return self._router_id_cache + + def _normalize_text(self, text: str) -> str: + return self.NORMALIZE_RE.sub(" ", text.strip()) + + def _mask_text(self, text: str) -> str: + ids: Final = self.ID_RE.sub("{ID}", text) + emails: Final = self.EMAIL_RE.sub("{EMAIL}", ids) + urls: Final = self.URL_RE.sub("{URL}", emails) + uuids: Final = self.UUID_RE.sub("{UUID}", urls) + return self.NUM_RE.sub("{NUM}", uuids) + + def _extract_tag_content(self, text: str) -> Sequence[tuple[str, str]]: + return tuple((match.group(1), match.group(2)) for match in self.TAG_CONTENT_RE.finditer(text)) + + def _extract_template(self, prompt: str) -> tuple[str, Sequence[tuple[str, str]]]: + normalized: Final = self._normalize_text(prompt) + extractions: Final = self._extract_tag_content(normalized) + skeleton: Final = self.TAG_CONTENT_RE.sub(self.TAG_REPLACEMENT, normalized) + masked_template: Final = self._mask_text(skeleton) + verbose_router_logger.debug("Extracted template with %s tag(s)", len(extractions)) + return masked_template, extractions + + @staticmethod + def _hash_template(masked_template: str, system_prompt: str | None = None) -> str: + # Prepending the system prompt isolates two tools with identical user-message + # structure but different task definitions. + if system_prompt: + normalized_sys: Final = re.sub(r"\s+", " ", system_prompt.strip()) + payload: Final = normalized_sys + " | " + masked_template + return hashlib.sha256(payload.encode()).hexdigest() + return hashlib.sha256(masked_template.encode()).hexdigest() + + async def seed_template(self, description: str, target_model: str) -> bool: + masked: Final = self._mask_text(self._normalize_text(description)) + template_hash: Final = self._hash_template(masked) + router_id: Final = self.get_router_id() + if await self.template_store.match_by_hash(template_hash, router_id) is not None: + return False + await self.template_store.store_template( + template_id=str(uuid4()), + template=masked, + template_hash=template_hash, + target_model=target_model, + router_id=router_id, + ) + return True + + async def route(self, prompt: str, system_prompt: str | None = None) -> AdeptTemplateMatch | None: + try: + self._ensure_refresh_running() + masked_template, _ = self._extract_template(prompt) + template_hash: Final = self._hash_template(masked_template, system_prompt) + router_id: Final = self.get_router_id() + cached: Final = self._cache_get((router_id, template_hash)) + if cached is None: + verbose_router_logger.debug("No matching template in cache") + return None + verbose_router_logger.debug("Template cache hit for hash %s", template_hash[:8]) + return AdeptTemplateMatch( + template_id=cached.id, + template=cached.template, + target_model=cached.target_model, + metadata=cached.additional_information, + ) + except (AttributeError, IndexError, KeyError, TypeError, ValueError) as e: + verbose_router_logger.exception("Error matching template: %s", e) + return None + + def _ensure_refresh_running(self) -> None: + existing: Final = self._refresh_task + if existing is not None and not existing.done(): + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + self._refresh_task = loop.create_task(self._refresh_loop()) + + def stop_refresh(self) -> None: + if self._refresh_task is not None: + self._refresh_task.cancel() + self._refresh_task = None + + async def _refresh_loop(self) -> None: + while True: + try: + await self.refresh_cache() + except (AttributeError, IndexError, KeyError, TypeError, ValueError) as e: + verbose_router_logger.warning("AdeptRouter: cache refresh failed: %s", e) + await asyncio.sleep(_REFRESH_INTERVAL_SECONDS) + + async def refresh_cache(self) -> None: + router_id: Final = self.get_router_id() + templates: Final = await self.template_store.load_all_for_router(router_id, limit=_TEMPLATE_CACHE_MAX_SIZE) + now: Final = time.monotonic() + rebuilt: Final[_TemplateCache] = OrderedDict( # mutable-ok: cache rebuild populated in one pass then swapped in + ((router_id, t.template_hash), (now, t)) for t in templates if t.template_hash is not None + ) + self._template_cache = rebuilt # rebind-ok: full cache swap after background refresh + + def _cache_get(self, key: tuple[str, str]) -> StoredTemplate | None: + entry: Final = self._template_cache.get(key) + if entry is None: + return None + inserted_at, stored = entry + if time.monotonic() - inserted_at > _TEMPLATE_CACHE_TTL_SECONDS: + self._template_cache.pop(key, None) + return None + self._template_cache.move_to_end(key) + return stored + + def _cache_put(self, key: tuple[str, str], stored: StoredTemplate) -> None: + self._template_cache[key] = (time.monotonic(), stored) + self._template_cache.move_to_end(key) + while len(self._template_cache) > _TEMPLATE_CACHE_MAX_SIZE: + self._template_cache.popitem(last=False) + + def _cache_invalidate(self, key: tuple[str, str]) -> None: + self._template_cache.pop(key, None) + + async def _resolve_template_id( + self, masked_template: str, template_hash: str, router_id: str, system_prompt: str | None + ) -> str: + matched_id: Final = await self.template_store.match_by_hash(template_hash, router_id) + if matched_id is not None: + return matched_id + + verbose_router_logger.info("No existing template found, storing new template.") + sys_prompt_payload: Final = {"system_prompt": system_prompt} # mutable-ok: JSON column payload + template_additional_info: Final[Mapping[str, object] | None] = sys_prompt_payload if system_prompt else None + stored_id: Final = await self.template_store.store_template( + template_id=str(uuid4()), + template=masked_template, + template_hash=template_hash, + target_model="", + router_id=router_id, + additional_information=template_additional_info, + ) + self._cache_invalidate((router_id, template_hash)) + return stored_id or str(uuid4()) + + async def store_conversation( + self, + prompt: str, + response: str, + model: str | None = None, + token_usage: Mapping[str, object] | None = None, + cost_usd: float | None = None, + latency_ms: float | None = None, + system_prompt: str | None = None, + routed_to_slm: bool | None = None, + ) -> None: + try: + masked_template, extractions = self._extract_template(prompt) + template_hash: Final = self._hash_template(masked_template, system_prompt) + router_id: Final = self.get_router_id() + template_id: Final = await self._resolve_template_id( + masked_template, template_hash, router_id, system_prompt + ) + + additional_info: Final[dict[str, object]] = {"extractions": extractions} # mutable-ok: JSON column payload + if model is not None: + additional_info["model"] = model + if token_usage is not None: + additional_info["token_usage"] = token_usage + if cost_usd is not None: + additional_info["cost_usd"] = cost_usd + if latency_ms is not None: + additional_info["latency_ms"] = round(latency_ms, 2) + if routed_to_slm is not None: + additional_info["routed_to_slm"] = routed_to_slm + + await self.template_store.store_conversation( + prompt=prompt, + response=response, + template_id=template_id, + additional_information=additional_info, + ) + + conversation_count: Final = await self.template_store.count_conversation_by_template_id(template_id) + if ( + conversation_count is not None + and conversation_count >= self.conversations_threshold + and conversation_count % self.conversations_threshold == 0 + ): + # Trainer runs may flip target_model, so drop the cached row before firing. + self._cache_invalidate((router_id, template_hash)) + self._trigger_trainer(template_id) + + verbose_router_logger.info("Stored interaction for template %s", template_id) + except (AttributeError, IndexError, KeyError, TypeError, ValueError) as e: + verbose_router_logger.exception("Error storing interaction: %s", e) + + def _trigger_trainer(self, template_id: str) -> None: + if not self.trainer_url: + verbose_router_logger.info( + "AdeptRouter: threshold reached for template %s but no trainer_url configured — skipping notification.", + template_id, + ) + return + asyncio.create_task(self._trainer_post(f"{self.trainer_url}/run-workflow/{template_id}", template_id)) + + @staticmethod + async def _trainer_post(url: str, template_id: str) -> None: + from litellm.litellm_core_utils.url_utils import SSRFError, validate_url + + try: + pinned_url, host_header = validate_url(url) + except SSRFError as ssrf_err: + verbose_router_logger.warning( + "AdeptRouter: refusing trainer POST for template %s — %s", template_id, ssrf_err + ) + return + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + try: + await client.post( + url=pinned_url, + headers={"Host": host_header}, # mutable-ok: one-shot per-request headers dict handed to httpx + timeout=_TRAINER_HTTP_TIMEOUT_SECONDS, + ) + verbose_router_logger.info("Triggered trainer for template %s", template_id) + except httpx.HTTPError as e: + verbose_router_logger.warning("Failed to trigger trainer: %s", e) diff --git a/litellm/router_strategy/adept_router/template/router_template.py b/litellm/router_strategy/adept_router/template/router_template.py new file mode 100644 index 00000000000..2c1b11ddc84 --- /dev/null +++ b/litellm/router_strategy/adept_router/template/router_template.py @@ -0,0 +1,41 @@ +from abc import ABC, abstractmethod +from collections.abc import Mapping + +from typing_extensions import ReadOnly, TypedDict + + +class AdeptTemplateMatch(TypedDict): + """Result of matching a prompt to a stored template.""" + + template_id: ReadOnly[str] + template: ReadOnly[str] + target_model: ReadOnly[str | None] + metadata: ReadOnly[Mapping[str, object] | None] + + +class BaseTemplateRouter(ABC): + """Abstract base class for template-based prompt routing.""" + + @abstractmethod + async def route(self, prompt: str, system_prompt: str | None = None) -> AdeptTemplateMatch | None: + """ + Match a prompt to a stored template. + + Returns a dict with template details if matched, None otherwise. + """ + ... + + @abstractmethod + async def store_conversation( + self, + prompt: str, + response: str, + model: str | None = None, + token_usage: Mapping[str, object] | None = None, + cost_usd: float | None = None, + latency_ms: float | None = None, + system_prompt: str | None = None, + routed_to_slm: bool | None = None, + ) -> None: + """Persist a prompt-response pair with its template and per-call metrics.""" + ... diff --git a/litellm/types/router.py b/litellm/types/router.py index e8a3bb6c01f..cf66edf62d1 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -387,6 +387,19 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): quality_router_config: dict | None = None quality_router_default_model: str | None = None + # adept-router params + adept_router_default_model: str | None = None + adept_router_tag_prefix: str | None = None + adept_router_seed_config: list[dict] | None = None # mutable-ok: pydantic config field + adept_router_conversations_threshold: int | None = None + adept_router_trainer_url: str | None = None + adept_router_pg_host: str | None = None + adept_router_pg_port: int | None = None + adept_router_pg_database: str | None = None + adept_router_pg_user: str | None = None + adept_router_pg_password: str | None = None + adept_router_pg_ssl_mode: str | None = None + # Vector Store Params vector_store_id: str | None = None milvus_text_field: str | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 00c55b35182..20eb9b52342 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3852,6 +3852,19 @@ all_litellm_params = ( "adaptive_router_default_model", "quality_router_config", "quality_router_default_model", + "adept_router", + "adept_router_default_model", + "adept_router_pg_host", + "adept_router_pg_port", + "adept_router_pg_user", + "adept_router_pg_password", + "adept_router_pg_database", + "adept_router_tag_prefix", + "adept_router_conversations_threshold", + "adept_router_trainer_url", + "adept_router_seed_config", + "adept_router_pg_ssl_mode", + "adept_routed_to_slm", ] + list(StandardCallbackDynamicParams.__annotations__.keys()) + list(CustomPricingLiteLLMParams.model_fields.keys()) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index a11f015743b..881138a20ff 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -88,6 +88,7 @@ ignored_function_names = [ "_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py "_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py "_apply_updated_routing_strategy_args", # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name) + "_release_adept_router_for_deleted_deployment", # Tested indirectly via delete_deployment in test_adept_router.py (the two test_delete_deployment_* tests) ] diff --git a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py index 722818598af..7dadd0e0330 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py @@ -1,7 +1,5 @@ - import pytest - from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, ) @@ -33,9 +31,7 @@ def test_base_model_label_alone_lacks_bedrock_tools(): """The label by itself does not advertise tools; this is what made the union necessary. Guards against the discrepancy disappearing (and the regression test above silently passing for the wrong reason).""" - params = get_supported_openai_params( - model=BEDROCK_LABEL, custom_llm_provider="bedrock" - ) + params = get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock") assert params is not None assert "tools" not in params @@ -46,14 +42,8 @@ def test_base_model_is_additive_not_replacement(): Bedrock: real id supports ``tools`` but not the label's reasoning hint; the union must contain the real model's ``tools`` regardless of the label being a subset.""" - real_only = set( - get_supported_openai_params( - model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock" - ) - ) - label_only = set( - get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock") - ) + real_only = set(get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock")) + label_only = set(get_supported_openai_params(model=BEDROCK_LABEL, custom_llm_provider="bedrock")) combined = set( get_supported_openai_params( model=BEDROCK_REAL_MODEL, @@ -67,17 +57,13 @@ def test_base_model_is_additive_not_replacement(): assert real_only <= combined -def test_base_model_adds_capabilities_the_real_model_lacks(): +def test_base_model_adds_capabilities_the_real_model_lacks(local_model_cost_map): """Regression for #27717 (the behavior the union must preserve). ``gemini-3.1-pro`` isn't in the cost map so it advertises no reasoning support, but the registered ``gemini-3.1-pro-preview`` base_model does. The hint must add ``reasoning_effort``/``thinking`` without the call erroring.""" - real_only = set( - get_supported_openai_params( - model="gemini-3.1-pro", custom_llm_provider="gemini" - ) - ) + real_only = set(get_supported_openai_params(model="gemini-3.1-pro", custom_llm_provider="gemini")) assert "reasoning_effort" not in real_only combined = set( @@ -93,21 +79,15 @@ def test_base_model_adds_capabilities_the_real_model_lacks(): def test_no_base_model_is_unchanged(): """Omitting ``base_model`` must resolve purely from ``model``.""" - with_none = get_supported_openai_params( - model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock", base_model=None - ) - plain = get_supported_openai_params( - model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock" - ) + with_none = get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock", base_model=None) + plain = get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock") assert with_none == plain def test_base_model_equal_to_model_is_unchanged(): """A ``base_model`` identical to ``model`` must not double-resolve or reorder.""" - plain = get_supported_openai_params( - model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock" - ) + plain = get_supported_openai_params(model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock") same = get_supported_openai_params( model=BEDROCK_REAL_MODEL, custom_llm_provider="bedrock", @@ -152,14 +132,10 @@ def test_bedrock_converse_alias_resolves_like_bedrock(): params saw no Bedrock capabilities for a Converse model invoked via the alias.""" anthropic_model = "bedrock/converse/us.anthropic.claude-sonnet-4-6" - via_alias = get_supported_openai_params( - model=anthropic_model, custom_llm_provider="bedrock_converse" - ) + via_alias = get_supported_openai_params(model=anthropic_model, custom_llm_provider="bedrock_converse") assert via_alias is not None - assert via_alias == get_supported_openai_params( - model=anthropic_model, custom_llm_provider="bedrock" - ) + assert via_alias == get_supported_openai_params(model=anthropic_model, custom_llm_provider="bedrock") assert "web_search_options" not in via_alias assert "tools" in via_alias @@ -167,9 +143,7 @@ def test_bedrock_converse_alias_resolves_like_bedrock(): def test_bedrock_converse_alias_keeps_nova_web_search_options(): """Nova on the ``bedrock_converse`` alias still advertises web_search_options, proving the alias routes through the model-aware config rather than a blanket Bedrock default.""" - nova_params = get_supported_openai_params( - model="amazon.nova-pro-v1:0", custom_llm_provider="bedrock_converse" - ) + nova_params = get_supported_openai_params(model="amazon.nova-pro-v1:0", custom_llm_provider="bedrock_converse") assert nova_params is not None assert "web_search_options" in nova_params diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 86b3f0976ab..932414d5f09 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -22,7 +22,6 @@ from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject from litellm.types.utils import EmbeddingResponse - IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" GCS_URL = "gs://my-bucket/image.png" @@ -74,9 +73,7 @@ class TestBuildPartForInput: assert part["file_data"]["file_uri"] == GCS_URL def test_file_reference_resolved(self): - resolved = { - "files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"} - } + resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}} part = _build_part_for_input("files/abc", resolved_files=resolved) assert part["file_data"] is not None assert part["file_data"]["mime_type"] == "image/jpeg" @@ -115,10 +112,7 @@ class TestTransformOpenaiInputGeminiContent: ) assert len(result["requests"]) == 2 # First request is text - assert ( - result["requests"][0]["content"]["parts"][0]["text"] - == "The food was delicious" - ) + assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious" # Second request is image assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None @@ -217,9 +211,7 @@ class TestProcessResponse: """Test that process_response sets correct indices.""" def test_single_embedding_index(self): - predictions: VertexAIBatchEmbeddingsResponseObject = { - "embeddings": [{"values": [0.1, 0.2]}] - } + predictions: VertexAIBatchEmbeddingsResponseObject = {"embeddings": [{"values": [0.1, 0.2]}]} model_response = EmbeddingResponse() result = process_response( input="hello", @@ -270,9 +262,7 @@ class TestProcessResponse: def test_nested_input_token_counting(self): """Nested list: only plain-text sub-elements should be counted.""" - predictions: VertexAIBatchEmbeddingsResponseObject = { - "embeddings": [{"values": [0.1, 0.2]}] - } + predictions: VertexAIBatchEmbeddingsResponseObject = {"embeddings": [{"values": [0.1, 0.2]}]} result = process_response( input=[["a red shoe", IMAGE_DATA_URI]], model_response=EmbeddingResponse(), @@ -307,7 +297,7 @@ class TestProcessEmbedContentResponseUsage: MODEL = "gemini-embedding-2" - def test_multimodal_image_preserves_usage_metadata(self): + def test_multimodal_image_preserves_usage_metadata(self, local_model_cost_map): response_json = { "embedding": {"values": [0.1, 0.2, 0.3]}, "usageMetadata": { @@ -374,9 +364,7 @@ class TestProcessEmbedContentResponseUsage: response_json=response_json, ) assert result.usage.prompt_tokens == 516 - assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( - 2.0 - ) + assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx(2.0) assert result.usage.prompt_tokens_details.text_tokens == 1 def test_missing_usage_metadata_does_not_estimate_from_base64(self): @@ -400,7 +388,7 @@ class TestProcessEmbedContentResponseUsage: ) assert result.usage.prompt_tokens > 0 - def test_file_reference_image_billed_per_image_not_text(self): + def test_file_reference_image_billed_per_image_not_text(self, local_model_cost_map): """files/... image refs must bill per-image, not at the text token rate.""" response_json = { "embedding": {"values": [0.1, 0.2, 0.3]}, @@ -432,7 +420,7 @@ class TestProcessEmbedContentResponseUsage: ) assert prompt_cost == pytest.approx(0.00012) - def test_file_reference_non_image_not_counted_as_image(self): + def test_file_reference_non_image_not_counted_as_image(self, local_model_cost_map): """A files/... ref resolving to a non-image mime must not be image-counted.""" response_json = { "embedding": {"values": [0.1, 0.2]}, @@ -456,9 +444,7 @@ class TestProcessEmbedContentResponseUsage: ) assert result.usage.prompt_tokens_details.image_count == 0 assert result.usage.prompt_tokens_details.audio_tokens == 64 - assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( - 2.0 - ) + assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx(2.0) prompt_cost, _ = generic_cost_per_token( model=self.MODEL, @@ -467,7 +453,7 @@ class TestProcessEmbedContentResponseUsage: ) assert prompt_cost == pytest.approx(2.0 * 0.00016) - def test_video_plus_audio_does_not_double_bill_text(self): + def test_video_plus_audio_does_not_double_bill_text(self, local_model_cost_map): """Video+audio responses must not get video tokens reassigned to text.""" response_json = { "embedding": {"values": [0.1]}, @@ -487,12 +473,8 @@ class TestProcessEmbedContentResponseUsage: response_json=response_json, ) assert result.usage.prompt_tokens_details.text_tokens == 1 - assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( - 2.0 - ) - assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( - 2.0 - ) + assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx(2.0) + assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx(2.0) prompt_cost, _ = generic_cost_per_token( model=self.MODEL, diff --git a/tests/test_litellm/test_adept_router.py b/tests/test_litellm/test_adept_router.py new file mode 100644 index 00000000000..5f0f64445a3 --- /dev/null +++ b/tests/test_litellm/test_adept_router.py @@ -0,0 +1,1792 @@ +"""Unit tests for the ADEPT router.""" + +import asyncio +import hashlib +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +# --------------------------------------------------------------------------- +# AdeptTemplateRouter tests (mock the Prisma-backed store — no live DB needed) +# --------------------------------------------------------------------------- + + +def _make_template_router(mock_storage, conversations_threshold=10, trainer_url=None): + from litellm.router_strategy.adept_router.template.implementation.adept_template_router import ( + AdeptTemplateRouter, + ) + + mock_router_instance = MagicMock() + mock_router_instance.get_model_ids.return_value = ["router-id-1"] + + with patch( # test-quality-ok: helper constructs AdeptTemplateRouter without a live PG connection; no HTTP boundary to fake at construction time + "litellm.router_strategy.adept_router.store.implementation.prisma.AdeptPrismaRepo", + return_value=mock_storage, + ): + router = AdeptTemplateRouter( + model_name="adept_router_test", + litellm_router_instance=mock_router_instance, + pg_url="postgresql://user:pass@localhost:5432/db", + tag_prefix="var", + conversations_threshold=conversations_threshold, + trainer_url=trainer_url, + ) + router.template_store = mock_storage + return router + + +def test_adept_template_router_route_returns_none_when_cache_empty(): + """route() must never await a DB lookup on the request path; an empty cache falls back to + None. The background refresh may load templates asynchronously but the request path is + cache-only.""" + mock_storage = AsyncMock() + mock_storage.load_all_for_router.return_value = () + + router = _make_template_router(mock_storage) + + async def _run() -> object: + try: + return await router.route("What is 2 + 2?") + finally: + if router._refresh_task is not None: + router._refresh_task.cancel() + + result = asyncio.run(_run()) + assert result is None + mock_storage.match_by_hash.assert_not_called() + mock_storage.get_template.assert_not_called() + + +def test_adept_template_router_route_hit_from_preloaded_cache(): + """After refresh_cache() populates the cache from load_all_for_router, route() serves the + match without any request-path DB call.""" + from litellm.router_strategy.adept_router.store.store_template import StoredTemplate + + stored = StoredTemplate( + id="tmpl-abc", + template="Get order {ID} for {EMAIL}", + template_hash=hashlib.sha256(b"Get order {ID} for {EMAIL}").hexdigest(), + router_id="router-id-1", + target_model="gpt-4o", + additional_information=None, + created_at=None, + ) + mock_storage = AsyncMock() + mock_storage.load_all_for_router.return_value = (stored,) + + router = _make_template_router(mock_storage) + + async def _run(): + await router.refresh_cache() + try: + first = await router.route("Get order ORD-123 for user@example.com") + second = await router.route("Get order ORD-999 for other@example.com") + finally: + if router._refresh_task is not None: + router._refresh_task.cancel() + return first, second + + first, second = asyncio.run(_run()) + assert first is not None and second is not None + assert first["target_model"] == "gpt-4o" + assert second["target_model"] == "gpt-4o" + mock_storage.match_by_hash.assert_not_called() + mock_storage.get_template.assert_not_called() + + +def test_refresh_cache_replaces_cache_from_store(): + """refresh_cache rebuilds the cache from a bulk load and keeps templates keyed by hash.""" + from litellm.router_strategy.adept_router.store.store_template import StoredTemplate + + first = StoredTemplate( + id="t1", + template="", + template_hash="h1", + router_id="router-id-1", + target_model="slm-a", + additional_information=None, + created_at=None, + ) + second = StoredTemplate( + id="t2", + template="", + template_hash="h2", + router_id="router-id-1", + target_model="slm-b", + additional_information=None, + created_at=None, + ) + mock_storage = AsyncMock() + mock_storage.load_all_for_router.return_value = (first, second) + + router = _make_template_router(mock_storage) + asyncio.run(router.refresh_cache()) + + assert ("router-id-1", "h1") in router._template_cache + assert ("router-id-1", "h2") in router._template_cache + mock_storage.load_all_for_router.assert_awaited_once_with("router-id-1", limit=1024) + + +def test_threshold_modulo_triggers_at_multiples(): # test-quality-ok: trigger has no return value; count-based orchestration is only observable via trigger invocations + """Trainer should be called at 5, 10, 15... but not at 7.""" + mock_storage = AsyncMock() + mock_storage.match_by_hash.return_value = "tmpl-1" + mock_storage.store_conversation.return_value = True + mock_storage.store_template.return_value = "tmpl-1" + + router = _make_template_router(mock_storage, conversations_threshold=5, trainer_url="http://trainer.test") + + with patch.object(router, "_trigger_trainer") as mock_trigger: + # count=5 -> triggers + mock_storage.count_conversation_by_template_id.return_value = 5 + asyncio.run(router.store_conversation("prompt", "response")) + mock_trigger.assert_called_once_with("tmpl-1") + + mock_trigger.reset_mock() + + # count=7 -> does not trigger + mock_storage.count_conversation_by_template_id.return_value = 7 + asyncio.run(router.store_conversation("prompt", "response")) + mock_trigger.assert_not_called() + + # count=10 -> triggers again + mock_storage.count_conversation_by_template_id.return_value = 10 + asyncio.run(router.store_conversation("prompt", "response")) + mock_trigger.assert_called_once_with("tmpl-1") + + +def test_trainer_url_used_in_trigger(): # test-quality-ok: asserts scheduling was requested for the correct URL; create_task IS the observable boundary + """_trigger_trainer schedules a POST to trainer_url, and no-ops if not set.""" + mock_storage = AsyncMock() + router_with = _make_template_router(mock_storage, trainer_url="http://my-trainer.internal") + router_without = _make_template_router(mock_storage, trainer_url=None) + + with ( + patch( # test-quality-ok: create_task IS the scheduling boundary the test is verifying; no HTTP round-trip happens here + "litellm.router_strategy.adept_router.template.implementation.adept_template_router.asyncio.create_task" + ) as mock_create_task + ): + + async def _run() -> None: + router_with._trigger_trainer("tmpl-xyz") + router_without._trigger_trainer("tmpl-xyz") + + asyncio.run(_run()) + + assert mock_create_task.call_count == 1 + scheduled_coro = mock_create_task.call_args[0][0] + # Coroutine target holds the URL through its cr_frame locals — close to release it. + scheduled_coro.close() + + +# --------------------------------------------------------------------------- +# System prompt isolation tests +# --------------------------------------------------------------------------- + + +def test_different_system_prompts_produce_different_hashes(): + """Two tools with the same XML structure but different system prompts must not collide.""" + from litellm.router_strategy.adept_router.template.implementation.adept_template_router import ( + AdeptTemplateRouter, + ) + + hash_a = AdeptTemplateRouter._hash_template("", system_prompt="You are an invoice extractor.") + hash_b = AdeptTemplateRouter._hash_template("", system_prompt="You are a contract reviewer.") + assert hash_a != hash_b + + +def test_same_tool_always_produces_same_hash(): + """Identical system prompt + same tag structure must always hash to the same value.""" + from litellm.router_strategy.adept_router.template.implementation.adept_template_router import ( + AdeptTemplateRouter, + ) + + system = "You are a ticket classifier." + hash_1 = AdeptTemplateRouter._hash_template("", system_prompt=system) + hash_2 = AdeptTemplateRouter._hash_template("", system_prompt=system) + assert hash_1 == hash_2 + + +def test_no_system_prompt_falls_back_to_user_message_hash(): + """Without a system prompt the hash is identical to hashing the masked template alone.""" + from litellm.router_strategy.adept_router.template.implementation.adept_template_router import ( + AdeptTemplateRouter, + ) + + masked = "" + expected = hashlib.sha256(masked.encode()).hexdigest() + assert AdeptTemplateRouter._hash_template(masked, system_prompt=None) == expected + assert AdeptTemplateRouter._hash_template(masked) == expected + + +# --------------------------------------------------------------------------- +# Router.py integration: detection and registration +# --------------------------------------------------------------------------- + + +def _make_minimal_litellm_params(**kwargs): + from litellm.types.router import LiteLLM_Params + + return LiteLLM_Params(**kwargs) + + +def test_is_adept_router_deployment(): + from litellm.router import Router + + router = Router(model_list=[]) + lp = _make_minimal_litellm_params(model="adept/my_adept") + assert router._is_adept_router_deployment(lp) is True + + +def test_adept_router_excluded_from_auto_router(): + from litellm.router import Router + + router = Router(model_list=[]) + lp = _make_minimal_litellm_params(model="adept/my_adept") + assert router._is_auto_router_deployment(lp) is False + + +def test_adept_router_prefix_is_not_semantic_auto_router(): + from litellm.router import Router + + router = Router(model_list=[]) + lp = _make_minimal_litellm_params(model="auto_router/my_semantic_router") + assert router._is_adept_router_deployment(lp) is False + assert router._is_auto_router_deployment(lp) is True + + +def test_adept_routers_dict_exists_on_router(): + from litellm.router import Router + + router = Router(model_list=[]) + assert hasattr(router, "adept_routers") + assert isinstance(router.adept_routers, dict) + assert hasattr(router, "init_adept_router_deployment") + assert callable(router.init_adept_router_deployment) + + +def test_init_adept_router_deployment_requires_pg_host(): + """init_adept_router_deployment raises ValueError when pg_host is missing.""" + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + deployment = Deployment( + model_name="my_adept", + litellm_params=LiteLLM_Params( + model="adept/my_adept", + adept_router_default_model="gpt-4o", + # adept_router_pg_host intentionally omitted + ), + model_info=ModelInfo(), + ) + + with pytest.raises(ValueError, match="adept_router_pg_host"): + router.init_adept_router_deployment(deployment) + + +def test_init_adept_router_deployment_registers_router(): + """init_adept_router_deployment wires up an AdeptRouter with correct params.""" + from unittest.mock import patch as _patch + + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + deployment = Deployment( + model_name="my_adept", + litellm_params=LiteLLM_Params( + model="adept/my_adept", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.internal.com", + adept_router_pg_port=5432, + adept_router_pg_database="adept_db", + adept_router_pg_user="user", + adept_router_pg_password="pass", + adept_router_conversations_threshold=20, + adept_router_trainer_url="http://trainer.internal", + ), + model_info=ModelInfo(), + ) + + mock_adept = MagicMock() + with ( + _patch( + "litellm.router_strategy.adept_router.adept_router.AdeptRouter", + return_value=mock_adept, + ) as MockAdeptRouter, + _patch( + "litellm.litellm_core_utils.url_utils.validate_url", + return_value=("http://trainer.internal", "trainer.internal"), + ), + ): + router.init_adept_router_deployment(deployment) + + assert "my_adept" in router.adept_routers + call_kwargs = MockAdeptRouter.call_args[1] + assert "postgresql://user:pass@db.internal.com:5432/adept_db" in call_kwargs["pg_url"] + assert call_kwargs["conversations_threshold"] == 20 + assert call_kwargs["trainer_url"] == "http://trainer.internal" + + +# --------------------------------------------------------------------------- +# Callback registration, routing decision, URL encoding, caching +# --------------------------------------------------------------------------- + + +def test_callback_registered_after_init(): + """After init_adept_router_deployment, AdeptRouter must appear in the async success callbacks.""" + from unittest.mock import patch as _patch + + import litellm + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + deployment = Deployment( + model_name="cb_test_adept", + litellm_params=LiteLLM_Params( + model="adept/cb_test_adept", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.internal.com", + adept_router_pg_database="adept_db", + adept_router_pg_user="user", + adept_router_pg_password="pass", + ), + model_info=ModelInfo(), + ) + + mock_adept = MagicMock() + with _patch( + "litellm.router_strategy.adept_router.adept_router.AdeptRouter", + return_value=mock_adept, + ): + router.init_adept_router_deployment(deployment) + + assert mock_adept in litellm.callbacks + + +def test_model_list_reload_unregisters_stale_adept_callbacks(): + """set_model_list() must remove old AdeptRouter instances from every litellm callback list + before clearing self.adept_routers; otherwise a stale router keeps exporting conversations + (including new prompts and responses) to its old PostgreSQL destination after the operator + has replaced or removed the deployment.""" + import litellm + from litellm.router import Router + + router = Router(model_list=[]) + + stale_adept = MagicMock() + router.adept_routers["stale_adept"] = stale_adept + litellm.callbacks.append(stale_adept) + litellm._async_success_callback.append(stale_adept) + try: + router.set_model_list([]) + + assert stale_adept not in litellm.callbacks + assert stale_adept not in litellm._async_success_callback + assert "stale_adept" not in router.adept_routers + finally: + for cb_list in (litellm.callbacks, litellm._async_success_callback): + while stale_adept in cb_list: + cb_list.remove(stale_adept) + + +def test_model_list_reload_disconnects_orphaned_adept_urls(): + """set_model_list() must reclaim PG clients for ADEPT URLs that the new model_list no + longer references, otherwise repeated config churn leaks connection pools and password- + bearing URLs until process restart.""" + from unittest.mock import patch as _patch + + from litellm.router import Router + from litellm.router_strategy.adept_router.store.implementation import prisma as prisma_mod + + router = Router(model_list=[]) + orphan = MagicMock() + orphan.pg_url = "postgresql://u:p@host.internal:5432/adept_db" + router.adept_routers["will_be_removed"] = orphan + + with _patch.object(prisma_mod, "schedule_disconnect") as mock_disc: + router.set_model_list([]) + + mock_disc.assert_called_once_with(orphan.pg_url) + + +def test_model_list_reload_does_not_disconnect_url_still_in_use(): + """When the new model_list re-registers a deployment on the SAME pg_url, the reclaim step + must skip that URL: disconnecting it would immediately tear down the client the rebuilt + deployment just started reusing.""" + from unittest.mock import patch as _patch + + from litellm.router import Router + from litellm.router_strategy.adept_router.store.implementation import prisma as prisma_mod + + router = Router(model_list=[]) + surviving_url = "postgresql://u:p@host.internal:5432/adept_db" + stale = MagicMock() + stale.pg_url = surviving_url + router.adept_routers["survivor"] = stale + + def _fake_reregister(_router: Router) -> None: + replacement = MagicMock() + replacement.pg_url = surviving_url + _router.adept_routers["survivor"] = replacement + + with ( + _patch.object(prisma_mod, "schedule_disconnect") as mock_disc, + _patch.object(Router, "_finalize_adaptive_router_if_configured", autospec=True, side_effect=_fake_reregister), + ): + router.set_model_list([]) + + mock_disc.assert_not_called() + + +def _make_success_event_adept(model_name="adept/test", default_model="gpt-4o"): + """An AdeptRouter with mocked template_router and seeding disabled, for callback tests.""" + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + + adept = AdeptRouter.__new__(AdeptRouter) + adept.model_name = model_name + adept.default_model = default_model + adept.litellm_router_instance = MagicMock() + adept.template_router = AsyncMock() + adept._seeded = True + return adept + + +def _model_response(content="output", prompt_tokens=10, completion_tokens=20, total_tokens=30): + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + response = ModelResponse(choices=[Choices(message=Message(content=content))]) + response.usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens) + return response + + +def test_routing_decision_stored_in_conversation(): + """routed_to_slm=True is persisted in conversation additional_information.""" + import datetime + + adept = _make_success_event_adept() + start = datetime.datetime(2024, 1, 1, 0, 0, 0) + end = datetime.datetime(2024, 1, 1, 0, 0, 1) + + kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "model": "my-slm", + "response_cost": 0.001, + "litellm_params": {"metadata": {"model_group": "adept/test", "adept_routed_to_slm": True}}, + } + + asyncio.run(adept.async_log_success_event(kwargs, _model_response(), start, end)) + + call_args = adept.template_router.store_conversation.call_args + assert call_args is not None + # routed_to_slm is the last positional arg + assert call_args[0][-1] is True + + +def test_routing_decision_fallback_stored(): + """routed_to_slm=False is persisted when fallback was used.""" + import datetime + + adept = _make_success_event_adept() + start = datetime.datetime(2024, 1, 1, 0, 0, 0) + end = datetime.datetime(2024, 1, 1, 0, 0, 1) + + kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "model": "gpt-4o", + "response_cost": 0.005, + "litellm_params": {"metadata": {"model_group": "adept/test", "adept_routed_to_slm": False}}, + } + + asyncio.run(adept.async_log_success_event(kwargs, _model_response(), start, end)) + + call_args = adept.template_router.store_conversation.call_args + assert call_args is not None + assert call_args[0][-1] is False + + +def test_success_event_skips_foreign_and_untagged_requests(): + """The success callback is global, so it fires for every proxy request. It must store rows + only for requests routed through THIS adept model: a request whose model_group is absent or + belongs to another deployment is skipped, so non-ADEPT traffic and other ADEPT deployments + never pollute or duplicate this store's conversations.""" + import datetime + + adept = _make_success_event_adept() + start = datetime.datetime(2024, 1, 1, 0, 0, 0) + end = datetime.datetime(2024, 1, 1, 0, 0, 1) + base_kwargs = {"messages": [{"role": "user", "content": "hi"}], "model": "gpt-4o"} + + asyncio.run( + adept.async_log_success_event( + {**base_kwargs, "litellm_params": {"metadata": {"model_group": "other-model"}}}, + _model_response(), + start, + end, + ) + ) + asyncio.run( + adept.async_log_success_event( + {**base_kwargs, "litellm_params": {"metadata": {}}}, _model_response(), start, end + ) + ) + + adept.template_router.store_conversation.assert_not_called() + + +def test_pre_routing_hook_stashes_routed_to_slm_in_metadata(): + """async_pre_routing_hook records the SLM decision in the request metadata dict. + + A bare top-level request_kwargs key never reaches the logging callback, so the + decision must live in metadata (the channel model_group already travels through). + """ + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + + adept = AdeptRouter.__new__(AdeptRouter) + adept.model_name = "adept/test" + adept.default_model = "big-llm" + adept.template_router = AsyncMock() + adept._seeded = True + + messages = [{"role": "user", "content": "hello"}] + + adept.template_router.route.return_value = {"template_id": "t1", "target_model": "slm-x"} + matched_kwargs = {"metadata": {}} + matched_resp = asyncio.run( + adept.async_pre_routing_hook(model="adept/test", request_kwargs=matched_kwargs, messages=messages) + ) + assert matched_resp.model == "slm-x" + assert matched_kwargs["metadata"]["adept_routed_to_slm"] is True + + adept.template_router.route.return_value = None + miss_kwargs = {"metadata": {}} + miss_resp = asyncio.run( + adept.async_pre_routing_hook(model="adept/test", request_kwargs=miss_kwargs, messages=messages) + ) + assert miss_resp.model == "big-llm" + assert miss_kwargs["metadata"]["adept_routed_to_slm"] is False + + +def test_routed_to_slm_survives_pre_hook_to_success_event(): + """Regression: the SLM decision set in the pre-routing hook reaches + async_log_success_event through the shared request metadata dict and is persisted. + + Models how litellm threads request metadata into litellm_params.metadata. With the old + top-level kwargs key this handoff dropped the flag and routed_to_slm was never stored. + """ + import datetime + + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + + adept = AdeptRouter.__new__(AdeptRouter) + adept.model_name = "adept/test" + adept.default_model = "big-llm" + adept.template_router = AsyncMock() + adept._seeded = True + adept.template_router.route.return_value = {"template_id": "t1", "target_model": "slm-x"} + + messages = [{"role": "user", "content": "hello"}] + metadata = {"model_group": "adept/test"} + asyncio.run( + adept.async_pre_routing_hook(model="adept/test", request_kwargs={"metadata": metadata}, messages=messages) + ) + assert metadata["adept_routed_to_slm"] is True + + start = datetime.datetime(2024, 1, 1, 0, 0, 0) + end = datetime.datetime(2024, 1, 1, 0, 0, 1) + + success_kwargs = { + "messages": messages, + "model": "slm-x", + "litellm_params": {"metadata": metadata}, + } + asyncio.run(adept.async_log_success_event(success_kwargs, _model_response(content="out"), start, end)) + + call_args = adept.template_router.store_conversation.call_args + assert call_args is not None + assert call_args[0][-1] is True + + +def test_pg_url_special_chars_encoded(): + """Passwords with @, :, / must be percent-encoded in the PG URL.""" + from unittest.mock import patch as _patch + + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + deployment = Deployment( + model_name="url_enc_test", + litellm_params=LiteLLM_Params( + model="adept/url_enc_test", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.host", + adept_router_pg_database="mydb", + adept_router_pg_user="adept_user", + adept_router_pg_password="p@ss:w/rd", + ), + model_info=ModelInfo(), + ) + + captured_url = {} + + def capture_adept(model_name, default_model, litellm_router_instance, pg_url, **kwargs): + captured_url["pg_url"] = pg_url + return MagicMock() + + with _patch( + "litellm.router_strategy.adept_router.adept_router.AdeptRouter", + side_effect=capture_adept, + ): + router.init_adept_router_deployment(deployment) + + pg_url = captured_url["pg_url"] + assert "p%40ss%3Aw%2Frd" in pg_url, f"Expected encoded password in URL, got: {pg_url}" + assert "p@ss:w/rd" not in pg_url + + +def test_pre_routing_hook_does_not_await_seed_on_request_path(): + """The pre-routing hook must not await seed DB lookups: it kicks seeding off in the + background so the first request is not blocked on `seed_template` round-trips.""" + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + + adept = AdeptRouter.__new__(AdeptRouter) + adept.model_name = "adept/test" + adept.default_model = "gpt-4o" + adept.litellm_router_instance = MagicMock() + adept.template_router = AsyncMock() + adept.template_router.route.return_value = None + adept._seed_config = [{"description": "Extract order", "target_model": "slm-order"}] + adept._seeded = False + adept._seed_lock = asyncio.Lock() + adept._seed_task = None + + async def _slow_seed(_desc, _model): + await asyncio.sleep(60) # would block for a minute if awaited on the request path + return True + + adept.template_router.seed_template.side_effect = _slow_seed + + async def _run(): + try: + return await asyncio.wait_for( + adept.async_pre_routing_hook( + model="adept/test", + request_kwargs={"metadata": {}}, + messages=[{"role": "user", "content": "hi"}], + ), + timeout=1.0, + ) + finally: + if adept._seed_task is not None: + adept._seed_task.cancel() + + result = asyncio.run(_run()) + assert result is not None + assert result.model == "gpt-4o" + + +def test_close_cancels_seed_task_so_retired_router_stops_writing(): + """When an ADEPT router is retired (delete or rebuild), close() must cancel any in-flight + seed task so it cannot write templates against a replacement deployment.""" + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + + adept = AdeptRouter.__new__(AdeptRouter) + adept.model_name = "adept/test" + adept.default_model = "gpt-4o" + adept.litellm_router_instance = MagicMock() + adept.template_router = MagicMock() + adept.template_router.stop_refresh = MagicMock() + adept._seed_config = [{"description": "d", "target_model": "m"}] + adept._seeded = False + adept._seed_lock = asyncio.Lock() + adept._seed_task = None + + async def _run(): + adept._seed_task = asyncio.create_task(asyncio.sleep(60)) + seed_task = adept._seed_task + adept.close() + assert adept._seed_task is None + assert seed_task.cancelled() or seed_task.cancelling() + + asyncio.run(_run()) + adept.template_router.stop_refresh.assert_called_once() + + +def test_seed_config_missing_description_logs_warning(): + """_run_seed warns and skips entries without a description.""" + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + + adept = AdeptRouter.__new__(AdeptRouter) + adept.model_name = "adept/test" + adept.default_model = "gpt-4o" + adept.litellm_router_instance = MagicMock() + adept.template_router = AsyncMock() + adept._seed_config = [{"target_model": "my-slm"}] # missing description + adept._seeded = False + adept._seed_lock = asyncio.Lock() + + with patch( # test-quality-ok: warning log is the observable output of the misconfiguration guard + "litellm.router_strategy.adept_router.adept_router.verbose_router_logger" + ) as mock_log: + asyncio.run(adept._run_seed()) + warning_calls = [str(c) for c in mock_log.warning.call_args_list] + assert any("description" in w for w in warning_calls) + + adept.template_router.seed_template.assert_not_called() + + +def test_router_id_cached_after_first_call(): + """get_router_id() should call get_model_ids only once regardless of how many times it's called.""" + mock_storage = AsyncMock() + mock_storage.match_by_hash.return_value = None + + router = _make_template_router(mock_storage) + router._router_id_cache = None # ensure cache is clear + + router.get_router_id() + router.get_router_id() + router.get_router_id() + + assert router.litellm_router_instance.get_model_ids.call_count == 1 + + +# --------------------------------------------------------------------------- +# AdeptPrismaRepo store tests: row mapping/guards with a mocked Prisma client +# (no DB), plus a real-database integration test that runs only when a prisma +# engine and an ADEPT_TEST_DB_URL are configured. +# --------------------------------------------------------------------------- + + +def test_prisma_repo_row_mapping_and_guards(): + """The store maps raw rows to StoredTemplate, serializes JSON payloads, issues the right SQL, + and rejects a conversation with no template_id, all without a live database.""" + import litellm.router_strategy.adept_router.store.implementation.prisma as prisma_mod + from litellm.router_strategy.adept_router.store.implementation.prisma import ( + AdeptPrismaRepo, + _CountRow, + _IdRow, + _TemplateRow, + ) + + repo = AdeptPrismaRepo("postgresql://u:p@localhost:5432/mockdb") + client = MagicMock() + client.query_raw = AsyncMock() + client.execute_raw = AsyncMock() + prisma_mod._CLIENTS[repo._db_url] = prisma_mod._ClientHandle( + client + ) # inject a fake connected client into the per-URL registry + + client.query_raw.return_value = [_IdRow(id="tmpl-1")] + assert asyncio.run(repo.match_by_hash("h", "r")) == "tmpl-1" + client.query_raw.return_value = [] + assert asyncio.run(repo.match_by_hash("h", "r")) is None + + client.query_raw.return_value = [ + _TemplateRow(id="t", template="skel", router_id="r", target_model="m", additional_information={"a": 1}) + ] + stored = asyncio.run(repo.get_template("t")) + assert stored is not None + assert stored.id == "t" and stored.target_model == "m" + assert stored.additional_information == {"a": 1} + client.query_raw.return_value = [] + assert asyncio.run(repo.get_template("missing")) is None + + assert asyncio.run(repo.store_conversation("p", "resp", "t", {"routed_to_slm": True})) is True + assert "INSERT INTO conversations" in client.execute_raw.call_args[0][0] + # guard: no template_id -> False, and no SQL issued for it + client.execute_raw.reset_mock() + assert asyncio.run(repo.store_conversation("p", "resp", None)) is False + client.execute_raw.assert_not_called() + + client.query_raw.return_value = [_CountRow(c=3)] + assert asyncio.run(repo.count_conversation_by_template_id("t")) == 3 + + +def test_prisma_repo_rejects_empty_db_url(): + """A misconfigured (empty) connection URL fails fast with a clear error.""" + from litellm.router_strategy.adept_router.store.implementation.prisma import AdeptPrismaRepo + + with pytest.raises(ValueError, match="PostgreSQL connection URL"): + AdeptPrismaRepo("") + + +def test_prisma_store_real_db_roundtrip(): + """Real end-to-end against a live PostgreSQL via the actual Prisma client: covers table + creation, ON CONFLICT concurrency safety, JSON round-trip, and the counter. Skipped unless a + prisma engine and an ADEPT_TEST_DB_URL are configured (so it runs locally / in the E2E env, + not in the dependency-light unit CI where no database or engine is present).""" + from litellm._uuid import uuid + + if not os.environ.get("PRISMA_QUERY_ENGINE_BINARY"): + pytest.skip("prisma query engine not configured (set PRISMA_QUERY_ENGINE_BINARY)") + db_url = os.environ.get("ADEPT_TEST_DB_URL") + if not db_url: + pytest.skip("no ADEPT_TEST_DB_URL configured") + + import litellm.router_strategy.adept_router.store.implementation.prisma as prisma_mod + from litellm.router_strategy.adept_router.store.implementation.prisma import AdeptPrismaRepo + + repo = AdeptPrismaRepo(db_url) + router_id = "test-router-" + uuid.uuid4().hex[:8] + template_hash = uuid.uuid4().hex + + async def run() -> None: + surviving = await repo.store_template( + template_id=uuid.uuid4().hex, + template="skeleton", + template_hash=template_hash, + target_model="slm-a", + router_id=router_id, + additional_information={"system_prompt": "sys"}, + ) + assert surviving is not None + # A concurrent duplicate (same router_id + hash) no-ops and resolves to the same id. + again = await repo.store_template( + template_id=uuid.uuid4().hex, + template="skeleton", + template_hash=template_hash, + target_model="", + router_id=router_id, + ) + assert again == surviving + assert await repo.match_by_hash(template_hash, router_id) == surviving + + stored = await repo.get_template(surviving) + assert stored is not None + assert stored.target_model == "slm-a" + assert stored.additional_information == {"system_prompt": "sys"} + + assert await repo.count_conversation_by_template_id(surviving) == 0 + assert await repo.store_conversation("p", "resp", surviving, {"routed_to_slm": True, "model": "slm-a"}) is True + assert await repo.count_conversation_by_template_id(surviving) == 1 + + handle = await prisma_mod._get_handle(db_url) + await handle.client.execute_raw("DELETE FROM conversations WHERE template_id = $1", surviving) + await handle.client.execute_raw("DELETE FROM templates WHERE router_id = $1", router_id) + await handle.client.disconnect() + prisma_mod._CLIENTS.pop(db_url, None) + + asyncio.run(run()) + + +def test_prisma_repo_reuses_one_client_per_url(): + """A router rebuild drops the old repo and builds a new one for the same database URL; the + store must reuse the existing client instead of connecting a second one and orphaning the + first (the connection-leak guard).""" + import litellm.router_strategy.adept_router.store.implementation.prisma as prisma_mod + from litellm.router_strategy.adept_router.store.implementation.prisma import AdeptPrismaRepo + + url = "postgresql://u:p@localhost:5432/leaktest" + prisma_mod._CLIENTS.pop(url, None) + connected = [] + + class _FakeClient: + async def connect(self): + connected.append(self) + + async def execute_raw(self, *args, **kwargs): + return 0 + + with patch.object( # test-quality-ok: asserts connection reuse across repo rebuilds; the Prisma constructor IS the reuse boundary + prisma_mod, "Prisma", side_effect=lambda datasource: _FakeClient() + ): + + async def run(): + AdeptPrismaRepo(url) # first router + h1 = await prisma_mod._get_handle(url) + AdeptPrismaRepo(url) # simulate a rebuild: a fresh repo for the same URL + h2 = await prisma_mod._get_handle(url) + return h1, h2 + + h1, h2 = asyncio.run(run()) + + assert h1 is h2 # reused, not reconnected + assert len(connected) == 1 # connected exactly once across both repos -> no leak + prisma_mod._CLIENTS.pop(url, None) + + +def test_trigger_trainer_uses_shared_async_client(): + """_trainer_post reuses the litellm-managed shared async client (not a per-call httpx.AsyncClient).""" + mock_storage = AsyncMock() + router = _make_template_router(mock_storage, trainer_url="http://trainer.test") + + fake_client = MagicMock() + fake_client.post = AsyncMock() + + with ( + patch( # test-quality-ok: the whole point of this test is that the shared client is used instead of a per-call httpx.AsyncClient + "litellm.router_strategy.adept_router.template.implementation.adept_template_router.get_async_httpx_client", + return_value=fake_client, + ) as mock_get, + patch( # test-quality-ok: validate_url would fail DNS resolution for the fake hostname in unit tests; production behavior is covered by test_trainer_post_re_validates_url_to_defeat_dns_rebinding + "litellm.litellm_core_utils.url_utils.validate_url", + return_value=("http://trainer.test/run-workflow/tmpl-httpx-test", "trainer.test"), + ), + ): + + async def _run() -> None: + await router._trainer_post("http://trainer.test/run-workflow/tmpl-httpx-test", "tmpl-httpx-test") + + asyncio.run(_run()) + + mock_get.assert_called_once() + fake_client.post.assert_awaited_once() + call_kwargs = fake_client.post.await_args.kwargs + assert call_kwargs["url"] == "http://trainer.test/run-workflow/tmpl-httpx-test" + assert call_kwargs["timeout"] == 10.0 + + +def test_trigger_trainer_is_fire_and_forget(): # test-quality-ok: asserts scheduling is non-blocking; the only observable is that _trainer_post was scheduled but the caller returned before it awaited + """_trigger_trainer must not await the HTTP round-trip — it schedules a background task.""" + mock_storage = AsyncMock() + router = _make_template_router(mock_storage, trainer_url="http://slow-trainer.test") + + async def _run() -> None: + # A running loop is required for asyncio.create_task; store_conversation always + # runs inside one so we mirror that here. + with patch.object(router, "_trainer_post", new_callable=AsyncMock) as mock_post: + router._trigger_trainer("tmpl-fire-forget") + # Yield once so the scheduled task starts; then assert it was scheduled + # and store_conversation would have returned already. + await asyncio.sleep(0) + mock_post.assert_called_once() + # Drain scheduled tasks so pytest doesn't warn about an unawaited coroutine. + await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task())) + + asyncio.run(_run()) + + +# --------------------------------------------------------------------------- +# Rebuild-on-change tests: editing an ADEPT deployment in the DB should +# refresh the in-memory router without requiring a proxy restart. +# --------------------------------------------------------------------------- + + +def _make_adept_deployment( + model_name: str = "fin_agent", + trainer_url: str | None = None, + threshold: int | None = None, + tag_prefix: str | None = None, +): + """Helper: build a Deployment for the rebuild-on-change tests.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + return Deployment( + model_name=model_name, + litellm_params=LiteLLM_Params( + model=f"adept/{model_name}", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.internal.com", + adept_router_pg_database="adept_db", + adept_router_pg_user="user", + adept_router_pg_password="pass", + adept_router_trainer_url=trainer_url, + adept_router_conversations_threshold=threshold, + adept_router_tag_prefix=tag_prefix, + ), + model_info=ModelInfo(), + ) + + +def test_init_adept_router_idempotent_when_params_unchanged(): + """ + Calling init twice with identical params must not rebuild the AdeptRouter — + the second call is a no-op so the DB-sync loop doesn't churn callbacks. + """ + from unittest.mock import MagicMock + from unittest.mock import patch as _patch + + from litellm.router import Router + + router = Router(model_list=[]) + deployment = _make_adept_deployment(trainer_url="http://trainer.internal", threshold=10) + + existing_mock = MagicMock() + existing_mock.default_model = "gpt-4o" + existing_mock.pg_url = "postgresql://user:pass@db.internal.com:5432/adept_db?sslmode=prefer" + existing_mock.template_router = MagicMock( + trainer_url="http://trainer.internal", + conversations_threshold=10, + tag_prefix="", + ) + + with ( + _patch( + "litellm.router_strategy.adept_router.adept_router.AdeptRouter", + return_value=existing_mock, + ) as MockAdeptRouter, + _patch( + "litellm.litellm_core_utils.url_utils.validate_url", + return_value=("http://trainer.internal", "trainer.internal"), + ), + ): + router.init_adept_router_deployment(deployment) + first_instance = router.adept_routers["fin_agent"] + router.init_adept_router_deployment(deployment) + + assert MockAdeptRouter.call_count == 1 + assert router.adept_routers["fin_agent"] is first_instance + + +def test_init_adept_router_rebuilds_when_trainer_url_changes(): + """ + Editing trainer_url in the DB row must rebuild the in-memory AdeptRouter on + the next sync tick — otherwise edits silently never take effect (the bug + that hid 30 conversations' worth of trainer notifications). + """ + from unittest.mock import MagicMock + from unittest.mock import patch as _patch + + from litellm.router import Router + + router = Router(model_list=[]) + + initial_mock = MagicMock() + initial_mock.default_model = "gpt-4o" + initial_mock.pg_url = "postgresql://user:pass@db.internal.com:5432/adept_db?sslmode=prefer" + initial_mock.template_router = MagicMock(trainer_url=None, conversations_threshold=10, tag_prefix="") + with _patch( + "litellm.router_strategy.adept_router.adept_router.AdeptRouter", + return_value=initial_mock, + ): + router.init_adept_router_deployment(_make_adept_deployment(trainer_url=None)) + + assert router.adept_routers["fin_agent"] is initial_mock + + rebuilt_mock = MagicMock() + with ( + _patch( + "litellm.router_strategy.adept_router.adept_router.AdeptRouter", + return_value=rebuilt_mock, + ) as MockAdeptRouter, + _patch( + "litellm.litellm_core_utils.url_utils.validate_url", + return_value=("http://trainer.internal", "trainer.internal"), + ), + ): + router.init_adept_router_deployment(_make_adept_deployment(trainer_url="http://trainer.internal")) + + MockAdeptRouter.assert_called_once() + assert MockAdeptRouter.call_args[1]["trainer_url"] == "http://trainer.internal" + assert router.adept_routers["fin_agent"] is rebuilt_mock + assert router.adept_routers["fin_agent"] is not initial_mock + + +# --------------------------------------------------------------------------- +# Security: trainer_url SSRF validation, pg TLS, params_changed pg fields +# --------------------------------------------------------------------------- + + +def test_trainer_url_blocked_cloud_metadata_host(): + """Cloud-metadata addresses in trainer_url must be rejected at init time by the SSRF guard.""" + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + for blocked in ("http://169.254.169.254/latest/meta-data", "http://metadata.google.internal/"): + deployment = Deployment( + model_name="ssrf_test", + litellm_params=LiteLLM_Params( + model="adept/ssrf_test", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.internal.com", + adept_router_trainer_url=blocked, + ), + model_info=ModelInfo(), + ) + with pytest.raises(ValueError, match="rejected by SSRF guard"): + router.init_adept_router_deployment(deployment) + + +def test_trainer_url_invalid_scheme_rejected(): + """Non-http(s) schemes in trainer_url must be rejected by the SSRF guard.""" + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + deployment = Deployment( + model_name="scheme_test", + litellm_params=LiteLLM_Params( + model="adept/scheme_test", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.internal.com", + adept_router_trainer_url="file:///etc/passwd", + ), + model_info=ModelInfo(), + ) + with pytest.raises(ValueError, match="rejected by SSRF guard"): + router.init_adept_router_deployment(deployment) + + +def test_pg_ssl_mode_included_in_url(): + """adept_router_pg_ssl_mode is appended as ?sslmode=... in the pg_url.""" + from unittest.mock import patch as _patch + + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + deployment = Deployment( + model_name="ssl_test", + litellm_params=LiteLLM_Params( + model="adept/ssl_test", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.internal.com", + adept_router_pg_ssl_mode="verify-full", + ), + model_info=ModelInfo(), + ) + + captured: dict[str, str] = {} + + def capture(model_name, default_model, litellm_router_instance, pg_url, **kwargs): + captured["pg_url"] = pg_url + return MagicMock() + + with _patch("litellm.router_strategy.adept_router.adept_router.AdeptRouter", side_effect=capture): + router.init_adept_router_deployment(deployment) + + assert "sslmode=verify-full" in captured["pg_url"] + + +def test_pg_host_change_triggers_rebuild(): + """Changing pg_host must rebuild the in-memory AdeptRouter on the next sync tick.""" + from unittest.mock import patch as _patch + + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + + def _deployment(host: str) -> Deployment: + return Deployment( + model_name="pg_rebuild_test", + litellm_params=LiteLLM_Params( + model="adept/pg_rebuild_test", + adept_router_default_model="gpt-4o", + adept_router_pg_host=host, + adept_router_pg_database="adept_db", + adept_router_pg_user="user", + adept_router_pg_password="pass", + ), + model_info=ModelInfo(), + ) + + first_mock = MagicMock() + first_mock.default_model = "gpt-4o" + first_mock.pg_url = "postgresql://user:pass@db-old.internal.com:5432/adept_db?sslmode=prefer" + first_mock.template_router = MagicMock(trainer_url=None, conversations_threshold=1000, tag_prefix="") + with _patch("litellm.router_strategy.adept_router.adept_router.AdeptRouter", return_value=first_mock): + router.init_adept_router_deployment(_deployment("db-old.internal.com")) + + second_mock = MagicMock() + with _patch( + "litellm.router_strategy.adept_router.adept_router.AdeptRouter", return_value=second_mock + ) as MockAdeptRouter: + router.init_adept_router_deployment(_deployment("db-new.internal.com")) + + MockAdeptRouter.assert_called_once() + assert router.adept_routers["pg_rebuild_test"] is second_mock + assert router.adept_routers["pg_rebuild_test"] is not first_mock + + +# --------------------------------------------------------------------------- +# Model-access authorization on the routed target model +# --------------------------------------------------------------------------- + + +def _make_pre_hook_adept( + model_name: str = "adept/test", default_model: str = "gpt-4o", target_model: str | None = "slm-x" +): + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + + adept = AdeptRouter.__new__(AdeptRouter) + adept.model_name = model_name + adept.default_model = default_model + adept.litellm_router_instance = MagicMock() + adept.litellm_router_instance.model_list = [] + adept.template_router = AsyncMock() + adept.template_router.route.return_value = ( + {"template_id": "t1", "target_model": target_model} if target_model else None + ) + adept._seeded = True + return adept + + +def test_pre_routing_hook_rejects_when_caller_lacks_access_to_routed_model(): + """Regression: the swapped-in target model must go through the caller's model-access check. + + Without this, a key that can call the ADEPT alias but not the trained SLM (or the + default fallback) would be silently upgraded to a model it cannot legitimately call. + """ + from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth + + adept = _make_pre_hook_adept(target_model="slm-x") + messages = [{"role": "user", "content": "hello"}] + caller = UserAPIKeyAuth(api_key="sk-test", models=["adept/test"]) + request_kwargs = {"metadata": {"user_api_key_auth": caller}} + + denial = ProxyException( + message="Key not allowed", type=ProxyErrorTypes.key_model_access_denied.value, param=None, code="401" + ) + + with ( + patch( # test-quality-ok: verifies authz is delegated to the proxy's auth_checks with the resolved SLM model; that call IS the boundary + "litellm.proxy.auth.auth_checks.can_key_call_resolved_model", new_callable=AsyncMock, side_effect=denial + ) as check + ): + with pytest.raises(ProxyException): + asyncio.run( + adept.async_pre_routing_hook(model="adept/test", request_kwargs=request_kwargs, messages=messages) + ) + check.assert_awaited_once() + assert check.await_args.kwargs["model"] == "slm-x" + # The routing decision must NOT have been stamped on a rejected request. + assert "adept_routed_to_slm" not in request_kwargs["metadata"] + + +def test_pre_routing_hook_allows_when_caller_has_access_to_routed_model(): + """Happy path counterpart: an authorized caller reaches the SLM and metadata is stamped.""" + from litellm.proxy._types import UserAPIKeyAuth + + adept = _make_pre_hook_adept(target_model="slm-x") + messages = [{"role": "user", "content": "hello"}] + caller = UserAPIKeyAuth(api_key="sk-test", models=["adept/test", "slm-x"]) + request_kwargs = {"metadata": {"user_api_key_auth": caller}} + + with ( + patch( # test-quality-ok: verifies authz is delegated to the proxy's auth_checks with the resolved SLM model on the happy path + "litellm.proxy.auth.auth_checks.can_key_call_resolved_model", new_callable=AsyncMock, return_value=None + ) as check + ): + resp = asyncio.run( + adept.async_pre_routing_hook(model="adept/test", request_kwargs=request_kwargs, messages=messages) + ) + + check.assert_awaited_once() + assert check.await_args.kwargs["model"] == "slm-x" + assert resp is not None + assert resp.model == "slm-x" + assert request_kwargs["metadata"]["adept_routed_to_slm"] is True + + +def test_pre_routing_hook_authorizes_default_model_on_miss(): + """A miss falls back to the default model — that fallback must be authorized too.""" + from litellm.proxy._types import UserAPIKeyAuth + + adept = _make_pre_hook_adept(default_model="big-llm", target_model=None) + messages = [{"role": "user", "content": "hi"}] + caller = UserAPIKeyAuth(api_key="sk-test", models=["adept/test", "big-llm"]) + request_kwargs = {"metadata": {"user_api_key_auth": caller}} + + with ( + patch( # test-quality-ok: verifies authz is also enforced against the default fallback model, not just the SLM target + "litellm.proxy.auth.auth_checks.can_key_call_resolved_model", new_callable=AsyncMock, return_value=None + ) as check + ): + resp = asyncio.run( + adept.async_pre_routing_hook(model="adept/test", request_kwargs=request_kwargs, messages=messages) + ) + + check.assert_awaited_once() + assert check.await_args.kwargs["model"] == "big-llm" + assert resp is not None + assert resp.model == "big-llm" + + +def test_pre_routing_hook_skips_authz_when_no_user_api_key_auth(): + """Non-proxy paths (ADEPT used directly against a Router) have no auth object; + the hook must not crash, just skip the check.""" + adept = _make_pre_hook_adept(target_model="slm-x") + messages = [{"role": "user", "content": "hello"}] + request_kwargs: dict[str, object] = {"metadata": {}} + + with ( + patch( # test-quality-ok: verifies the authz check is SKIPPED when no proxy auth object is present; the check-not-called IS the observable + "litellm.proxy.auth.auth_checks.can_key_call_resolved_model", new_callable=AsyncMock + ) as check + ): + resp = asyncio.run( + adept.async_pre_routing_hook(model="adept/test", request_kwargs=request_kwargs, messages=messages) + ) + + check.assert_not_awaited() + assert resp is not None + assert resp.model == "slm-x" + + +# --------------------------------------------------------------------------- +# Trainer URL: operator-controlled host allowlist (general_settings.user_url_allowed_hosts) +# --------------------------------------------------------------------------- + + +def test_trainer_url_private_ip_rejected_by_ssrf_guard(): + """A private-range host in trainer_url must be rejected — a team admin cannot bypass this + by setting the URL directly on a team-scoped deployment because the SSRF guard is central.""" + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + deployment = Deployment( + model_name="private_ip", + litellm_params=LiteLLM_Params( + model="adept/private_ip", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.internal.com", + adept_router_trainer_url="http://10.0.0.5/hook", + ), + model_info=ModelInfo(), + ) + with pytest.raises(ValueError, match="rejected by SSRF guard"): + router.init_adept_router_deployment(deployment) + + +def test_trainer_url_private_host_allowed_when_operator_allowlists_it(): + """Operators can opt an internal trainer past the SSRF guard by adding it to + `litellm.user_url_allowed_hosts` (surfaced through general_settings) — team admins cannot.""" + from unittest.mock import patch as _patch + + import litellm + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + previous_allowed = list(getattr(litellm, "user_url_allowed_hosts", ()) or ()) + litellm.user_url_allowed_hosts = [ + *previous_allowed, + "trainer.internal.com", + ] # test-quality-ok: verifying operator-controlled allowlist behavior REQUIRES writing the process-wide setting; save/restore is done in the finally block + try: + router = Router(model_list=[]) + deployment = Deployment( + model_name="operator_allowlisted", + litellm_params=LiteLLM_Params( + model="adept/operator_allowlisted", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.internal.com", + adept_router_trainer_url="http://trainer.internal.com/hook", + ), + model_info=ModelInfo(), + ) + with ( + _patch( + "litellm.litellm_core_utils.url_utils.socket.getaddrinfo", + return_value=[(2, 1, 6, "", ("10.0.0.5", 80))], + ), + _patch("litellm.router_strategy.adept_router.adept_router.AdeptRouter", return_value=MagicMock()), + ): + router.init_adept_router_deployment(deployment) + assert "operator_allowlisted" in router.adept_routers + finally: + litellm.user_url_allowed_hosts = previous_allowed # test-quality-ok: restore the pre-test value so the write above does not leak into other tests + + +def test_trainer_post_re_validates_url_to_defeat_dns_rebinding(): # test-quality-ok: the observable behavior of the DNS-rebinding defense is exactly "no HTTP client is obtained when the URL fails re-validation"; nothing else to assert + """The trainer POST must re-run the SSRF guard so a hostname that resolved to a public IP + at deployment sync but rebinds to a private IP by call time is still refused.""" + from litellm.litellm_core_utils.url_utils import SSRFError + from litellm.router_strategy.adept_router.template.implementation.adept_template_router import ( + AdeptTemplateRouter, + ) + + async def _run() -> None: + with ( + patch( # test-quality-ok: validate_url IS the SSRF guard the test is verifying; simulating an SSRFError raise from it is the only way to force the re-validation branch + "litellm.litellm_core_utils.url_utils.validate_url", + side_effect=SSRFError("resolves to blocked network"), + ), + patch( # test-quality-ok: get_async_httpx_client is the only observable proxy for whether the POST would have been sent; asserting not-called IS the security invariant + "litellm.router_strategy.adept_router.template.implementation.adept_template_router.get_async_httpx_client" + ) as mock_get, + ): + await AdeptTemplateRouter._trainer_post("http://rebound.example.com/hook", "tmpl-rebind") + assert not mock_get.called + + asyncio.run(_run()) + + +# --------------------------------------------------------------------------- +# Coverage: cache TTL expiry, LRU eviction, seed_template, helpers +# --------------------------------------------------------------------------- + + +def test_cache_ttl_expiry_returns_none_and_evicts(): + """A cached entry older than the TTL must be dropped on read so the next call re-hits the store.""" + from litellm.router_strategy.adept_router.store.store_template import StoredTemplate + from litellm.router_strategy.adept_router.template.implementation import adept_template_router as mod + + router = _make_template_router(AsyncMock()) + key = ("router-id", "hash") + stored = StoredTemplate( + id="t", + template="", + template_hash="", + router_id="", + target_model="m", + additional_information=None, + created_at=None, + ) + router._template_cache[key] = (0.0, stored) + + with patch.object( # test-quality-ok: TTL expiry is only observable via a controlled clock; patching time.monotonic on the module is the standard way to fast-forward + mod.time, "monotonic", return_value=mod._TEMPLATE_CACHE_TTL_SECONDS + 1 + ): + assert router._cache_get(key) is None + assert key not in router._template_cache + + +def test_cache_put_evicts_lru_when_over_capacity(): + """Beyond _TEMPLATE_CACHE_MAX_SIZE, the oldest entry must be popped so memory stays bounded.""" + from litellm.router_strategy.adept_router.store.store_template import StoredTemplate + from litellm.router_strategy.adept_router.template.implementation import adept_template_router as mod + + router = _make_template_router(AsyncMock()) + stored = StoredTemplate( + id="t", + template="", + template_hash="", + router_id="", + target_model="m", + additional_information=None, + created_at=None, + ) + with patch.object( # test-quality-ok: shrinking the module-level bound is the only way to exercise LRU eviction without inserting 1024 real entries per test + mod, "_TEMPLATE_CACHE_MAX_SIZE", 2 + ): + router._cache_put(("r", "a"), stored) + router._cache_put(("r", "b"), stored) + router._cache_put(("r", "c"), stored) + assert ("r", "a") not in router._template_cache + assert ("r", "b") in router._template_cache + assert ("r", "c") in router._template_cache + + +def test_seed_template_stores_new_and_skips_existing(): + """seed_template writes on a miss and short-circuits on a hit (idempotent seed).""" + mock_storage = AsyncMock() + mock_storage.match_by_hash.return_value = None + mock_storage.store_template.return_value = "tmpl-seeded" + router = _make_template_router(mock_storage) + + assert asyncio.run(router.seed_template("Extract invoice fields", "slm-invoice")) is True + mock_storage.store_template.assert_awaited_once() + + mock_storage.reset_mock() + mock_storage.match_by_hash.return_value = "tmpl-existing" + assert asyncio.run(router.seed_template("Extract invoice fields", "slm-invoice")) is False + mock_storage.store_template.assert_not_awaited() + + +def test_response_text_returns_none_on_empty_choices(): + """`_response_text` guards against a response with no choices (the reject / no-content path).""" + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + from litellm.types.utils import ModelResponse + + response = ModelResponse() + response.choices = [] + assert AdeptRouter._response_text(response) is None + + +def test_content_to_text_flattens_list_content_and_handles_none(): + """Multimodal list content is flattened to text; None content returns empty string.""" + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + + assert AdeptRouter._content_to_text(None) == "" + assert AdeptRouter._content_to_text("plain string") == "plain string" + + blocks = [ + {"type": "text", "text": "hello"}, + {"type": "text", "text": "world"}, + {"type": "image_url", "image_url": {"url": "x"}}, + ] + assert AdeptRouter._content_to_text(blocks) == "hello world" + + +def test_extract_system_prompt_and_user_text_helpers(): + """`_extract_system_prompt` finds the system role; `_extract_user_text` reads the LAST user turn.""" + from litellm.router_strategy.adept_router.adept_router import AdeptRouter + + assert AdeptRouter._extract_system_prompt([{"role": "user", "content": "hi"}]) is None + assert AdeptRouter._extract_system_prompt([{"role": "system", "content": "you are..."}]) == "you are..." + + messages = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "reply"}, + {"role": "user", "content": "second"}, + ] + assert AdeptRouter._extract_user_text(messages) == "second" + assert AdeptRouter._extract_user_text([{"role": "assistant", "content": "only bot"}]) == "" + + +def test_resolve_template_id_returns_matched_id_on_hit(): + """When a template with the same hash already exists, `_resolve_template_id` reuses that id and does NOT re-store.""" + mock_storage = AsyncMock() + mock_storage.match_by_hash.return_value = "tmpl-existing" + router = _make_template_router(mock_storage) + + resolved = asyncio.run(router._resolve_template_id("masked", "hash-x", "router-1", None)) + assert resolved == "tmpl-existing" + mock_storage.store_template.assert_not_awaited() + + +def test_store_conversation_swallows_and_logs_expected_errors(): + """The success-event handler is fire-and-forget: expected raises (KeyError etc.) must be caught, not propagated.""" + mock_storage = AsyncMock() + mock_storage.match_by_hash.side_effect = KeyError("boom") + router = _make_template_router(mock_storage) + + asyncio.run(router.store_conversation("prompt", "response")) + + assert mock_storage.store_conversation.await_count == 0 + + +# --------------------------------------------------------------------------- +# Security: no prompt content in debug logs; stale PG clients get disconnected +# --------------------------------------------------------------------------- + + +def test_extract_template_debug_log_omits_prompt_content(caplog): + """The extract-template debug line must NOT include any prompt text — even 'masked' — because + `_mask_text` only rewrites tag content, so untagged prompt words would otherwise reach the log.""" + import logging + + router = _make_template_router(AsyncMock()) + secret_prompt = "USER_SECRET_TOKEN_abc123 please summarize this x" + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): + router._extract_template(secret_prompt) + + assert not any("USER_SECRET_TOKEN_abc123" in rec.getMessage() for rec in caplog.records) + + +def test_stale_pg_client_disconnected_on_rebuild_when_url_changes(): + """A rebuild whose pg_url differs from the existing router's must schedule a disconnect of the + old client so credentials do not remain in-memory for the process lifetime.""" + from unittest.mock import patch as _patch + + from litellm.router import Router + from litellm.router_strategy.adept_router.store.implementation import prisma as prisma_mod + + router = Router(model_list=[]) + stale = MagicMock() + stale.pg_url = "postgresql://old_user:old_pass@old.internal:5432/adept_db?sslmode=prefer" + stale.default_model = "gpt-4o" + stale.template_router = MagicMock(trainer_url=None, conversations_threshold=10, tag_prefix="") + router.adept_routers["cred_rotation"] = stale + + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + new_deployment = Deployment( + model_name="cred_rotation", + litellm_params=LiteLLM_Params( + model="adept/cred_rotation", + adept_router_default_model="gpt-4o", + adept_router_pg_host="new.internal.com", + adept_router_pg_user="new_user", + adept_router_pg_password="new_pass", + adept_router_pg_database="adept_db", + ), + model_info=ModelInfo(), + ) + + with ( + _patch.object(prisma_mod, "schedule_disconnect") as mock_disc, + _patch( + "litellm.router_strategy.adept_router.adept_router.AdeptRouter", + return_value=MagicMock(), + ), + ): + router.init_adept_router_deployment(new_deployment) + + mock_disc.assert_called_once_with(stale.pg_url) + + +def test_rebuild_skips_disconnect_when_another_deployment_still_uses_the_old_url(): + """If a second ADEPT deployment shares the OLD pg_url, the rebuild must NOT disconnect it — + the sibling deployment would immediately hit a torn-down client on its next query.""" + from unittest.mock import patch as _patch + + from litellm.router import Router + from litellm.router_strategy.adept_router.store.implementation import prisma as prisma_mod + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + shared_url = "postgresql://u:p@shared.internal:5432/adept_db?sslmode=prefer" + + outgoing = MagicMock() + outgoing.pg_url = shared_url + outgoing.default_model = "gpt-4o" + outgoing.template_router = MagicMock(trainer_url=None, conversations_threshold=10, tag_prefix="") + router.adept_routers["moving_deployment"] = outgoing + + sibling = MagicMock() + sibling.pg_url = shared_url + router.adept_routers["sibling_deployment"] = sibling + + new_deployment = Deployment( + model_name="moving_deployment", + litellm_params=LiteLLM_Params( + model="adept/moving_deployment", + adept_router_default_model="gpt-4o", + adept_router_pg_host="new.internal.com", + adept_router_pg_database="adept_db", + ), + model_info=ModelInfo(), + ) + + with ( + _patch.object(prisma_mod, "schedule_disconnect") as mock_disc, + _patch( + "litellm.router_strategy.adept_router.adept_router.AdeptRouter", + return_value=MagicMock(), + ), + ): + router.init_adept_router_deployment(new_deployment) + + mock_disc.assert_not_called() + + +def test_disconnect_client_evicts_from_registry_and_disconnects(): + """`disconnect_client` must pop the entry AND call `.disconnect()` so no zombie sockets or + passwords linger. If the URL is not present, it must not raise.""" + from litellm.router_strategy.adept_router.store.implementation import prisma as prisma_mod + + url = "postgresql://user:pw@host:5432/db" + fake_client = MagicMock() + fake_client.disconnect = AsyncMock() + prisma_mod._CLIENTS[url] = prisma_mod._ClientHandle(fake_client) + + asyncio.run(prisma_mod.disconnect_client(url)) + + assert url not in prisma_mod._CLIENTS + fake_client.disconnect.assert_awaited_once() + + asyncio.run(prisma_mod.disconnect_client(url)) + assert fake_client.disconnect.await_count == 1 + + +def test_disconnect_client_waits_for_in_flight_borrow_then_disconnects(): + """A borrowed client must not be disconnected mid-op: disconnect must wait for the borrow + to release before tearing the socket down.""" + from litellm.router_strategy.adept_router.store.implementation import prisma as prisma_mod + + url = "postgresql://user:pw@host:5432/db" + fake_client = MagicMock() + fake_client.disconnect = AsyncMock() + handle = prisma_mod._ClientHandle(fake_client) + prisma_mod._CLIENTS[url] = handle + + async def _run(): + borrow_started = asyncio.Event() + release_borrow = asyncio.Event() + + async def _hold_borrow(): + async with handle.borrow(): + borrow_started.set() + await release_borrow.wait() + + borrower = asyncio.create_task(_hold_borrow()) + await borrow_started.wait() + + disconnect_task = asyncio.create_task(prisma_mod.disconnect_client(url)) + await asyncio.sleep(0.05) + assert not disconnect_task.done(), "disconnect must wait for in-flight borrow" + assert fake_client.disconnect.await_count == 0 + + release_borrow.set() + await borrower + await disconnect_task + + asyncio.run(_run()) + assert url not in prisma_mod._CLIENTS + fake_client.disconnect.assert_awaited_once() + + +def test_redact_url_masks_password_in_logs(): + """_redact_url must replace embedded passwords with '***' so error logs cannot leak them.""" + from litellm.router_strategy.adept_router.store.implementation.prisma import _redact_url + + assert _redact_url("postgresql://user:secret@host:5432/db") == "postgresql://user:***@host:5432/db" + assert _redact_url("postgresql://user@host/db") == "postgresql://user@host/db" + + +def test_delete_deployment_releases_adept_router_and_schedules_disconnect(): + """delete_deployment must drop the ADEPT router from self.adept_routers, unregister its + callbacks, and schedule disconnect of its PG client when no sibling deployment shares the + URL. Otherwise repeated removals leak connection pools and password-bearing URLs.""" + from unittest.mock import patch as _patch + + import litellm + from litellm.router import Router + from litellm.router_strategy.adept_router.store.implementation import prisma as prisma_mod + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + to_remove = MagicMock() + to_remove.pg_url = "postgresql://u:p@removed.internal:5432/adept_db" + router.adept_routers["will_be_deleted"] = to_remove + litellm.callbacks.append(to_remove) + + deployment = Deployment( + model_name="will_be_deleted", + litellm_params=LiteLLM_Params(model="adept/will_be_deleted"), + model_info=ModelInfo(id="dep-id-1"), + ) + router.model_list = [deployment.model_dump(exclude_none=True)] + router.model_id_to_deployment_index_map = {"dep-id-1": 0} + + try: + with _patch.object(prisma_mod, "schedule_disconnect") as mock_disc: + router.delete_deployment("dep-id-1") + + assert "will_be_deleted" not in router.adept_routers + assert to_remove not in litellm.callbacks + mock_disc.assert_called_once_with(to_remove.pg_url) + finally: + while to_remove in litellm.callbacks: + litellm.callbacks.remove(to_remove) + + +def test_delete_deployment_skips_disconnect_when_sibling_shares_url(): + """A sibling ADEPT deployment on the same URL must keep its client: deleting the primary + should unregister only the primary and never disconnect the shared client.""" + from unittest.mock import patch as _patch + + import litellm + from litellm.router import Router + from litellm.router_strategy.adept_router.store.implementation import prisma as prisma_mod + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + shared_url = "postgresql://u:p@shared.internal:5432/adept_db" + to_remove = MagicMock() + to_remove.pg_url = shared_url + router.adept_routers["primary"] = to_remove + + sibling = MagicMock() + sibling.pg_url = shared_url + router.adept_routers["sibling"] = sibling + + litellm.callbacks.append(to_remove) + + deployment = Deployment( + model_name="primary", + litellm_params=LiteLLM_Params(model="adept/primary"), + model_info=ModelInfo(id="dep-id-2"), + ) + router.model_list = [deployment.model_dump(exclude_none=True)] + router.model_id_to_deployment_index_map = {"dep-id-2": 0} + + try: + with _patch.object(prisma_mod, "schedule_disconnect") as mock_disc: + router.delete_deployment("dep-id-2") + + assert "primary" not in router.adept_routers + assert "sibling" in router.adept_routers + mock_disc.assert_not_called() + finally: + while to_remove in litellm.callbacks: + litellm.callbacks.remove(to_remove) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index e2a08a40bcb..578d8cef1ca 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2455,4 +2455,4 @@ "count": 1 } } -} \ No newline at end of file +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 105f6ff3043..9a85b7f6170 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -128,6 +128,7 @@ describe("ModelsAndEndpointsPage", () => { "All Models", "Add Model", "Auto-Routers Beta", + "ADEPT Routers", "LLM Credentials", "Pass-Through Endpoints", "Health Status", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index 4d6a90fc56e..2243f26b441 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -17,6 +17,7 @@ import { useModelDashboardData } from "@/app/(dashboard)/models-and-endpoints/us import AllModelsPanel from "@/app/(dashboard)/models-and-endpoints/panels/AllModelsPanel"; import AutoRoutersTabPanel from "@/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel"; import AddModelPanel from "@/app/(dashboard)/models-and-endpoints/panels/AddModelPanel"; +import AddAdeptRouterPanel from "@/app/(dashboard)/models-and-endpoints/panels/AddAdeptRouterPanel"; import LlmCredentialsPanel from "@/app/(dashboard)/models-and-endpoints/panels/LlmCredentialsPanel"; import PassThroughPanel from "@/app/(dashboard)/models-and-endpoints/panels/PassThroughPanel"; import HealthStatusPanel from "@/app/(dashboard)/models-and-endpoints/panels/HealthStatusPanel"; @@ -30,6 +31,7 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; type ModelTabSlug = | "add" | "auto-routers" + | "adept-routers" | "llm-credentials" | "pass-through" | "health" @@ -43,6 +45,7 @@ const BASE_TAB_KEY = "all-models"; const TAB_LABELS: Record = { add: "Add Model", "auto-routers": "Auto-Routers", + "adept-routers": "ADEPT Routers", "llm-credentials": "LLM Credentials", "pass-through": "Pass-Through Endpoints", health: "Health Status", @@ -58,6 +61,8 @@ const renderPanel = (key: string) => { return ; case "auto-routers": return ; + case "adept-routers": + return ; case "add": return ; case "llm-credentials": @@ -105,7 +110,7 @@ export default function ModelsAndEndpointsPage() { () => [ "", ...(canCreate ? (["add"] as const) : []), - ...(isAdmin || canCreate ? (["auto-routers"] as const) : []), + ...(isAdmin || canCreate ? (["auto-routers", "adept-routers"] as const) : []), // effectiveSessionRole reports proxy_admin_viewer as "Admin", so isAdmin alone would show a // viewer these write-only panels; only the raw-role isViewOnly separates them. Health Status // stays: it is the bucket's one read view, and viewers keep read parity with admins. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddAdeptRouterPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddAdeptRouterPanel.tsx new file mode 100644 index 00000000000..58e74cdbcd4 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddAdeptRouterPanel.tsx @@ -0,0 +1,40 @@ +"use client"; + +import { useQueryClient } from "@tanstack/react-query"; +import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { internalUserRoles } from "@/utils/roles"; +import { modelCreationScope } from "@/utils/modelPermissions"; +import AddAdeptRouterTab from "@/components/add_model/AddAdeptRouterTab"; + +/** + * Owns the permission decision for the ADEPT Routers tab. Creating an ADEPT router is a + * POST /model/new like Add Model and Auto Router, so it takes the same audience rule: + * a proxy admin, or a team admin who scopes it to a team. + */ +export default function AddAdeptRouterPanel() { + const { accessToken, userRole, userId: userID, isViewOnly } = useAuthorized(); + const { data: teams } = useTeams(); + const { data: uiSettings } = useUISettings(); + const queryClient = useQueryClient(); + + const isInternalUser = userRole != null && internalUserRoles.includes(userRole); + const scope = modelCreationScope( + { userRole, userID, isViewOnly }, + { + teams: teams ?? null, + disabledForInternalUsers: isInternalUser && uiSettings?.values?.disable_model_add_for_internal_users === true, + }, + ); + + return ( + queryClient.invalidateQueries({ queryKey: ["models", "list"] })} + accessToken={accessToken ?? ""} + userRole={userRole ?? ""} + userId={userID ?? null} + createScope={scope} + /> + ); +} diff --git a/ui/litellm-dashboard/src/components/add_model/AddAdeptRouterTab.tsx b/ui/litellm-dashboard/src/components/add_model/AddAdeptRouterTab.tsx new file mode 100644 index 00000000000..ff47ae67931 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AddAdeptRouterTab.tsx @@ -0,0 +1,291 @@ +"use client"; + +import React, { useEffect, useState } from "react"; +import { useForm, useWatch, type UseFormReturn } from "react-hook-form"; +import { toast } from "@/lib/toast"; +import { FieldGroup } from "@/components/ui/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { Button } from "@/components/ui/button"; +import { Card } from "@/components/ui/card"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models"; +import { modelAvailableCall } from "../networking"; +import AccessGroupTagsCombobox from "./AccessGroupTagsCombobox"; +import ModelChoiceCombobox from "./ModelChoiceCombobox"; +import TeamDropdown from "../common_components/team_dropdown"; +import { type ModelWriteScope } from "@/utils/modelPermissions"; +import { handleAddAdeptRouterSubmit, type AddAdeptRouterValues } from "./HandleAddAdeptRouterSubmit"; + +// Fixed sslmode options mirror libpq's set. Kept in sync with the backend `LiteLLM_Params.adept_router_pg_ssl_mode`. +const PG_SSL_MODES = ["disable", "allow", "prefer", "require", "verify-ca", "verify-full"] as const; + +const EMPTY_FORM_VALUES: AddAdeptRouterValues = { + adept_router_name: "", + adept_router_default_model: "", + adept_router_tag_prefix: "", + adept_router_conversations_threshold: null, + adept_router_trainer_url: "", + adept_router_pg_host: "", + adept_router_pg_port: 5432, + adept_router_pg_database: "", + adept_router_pg_user: "", + adept_router_pg_password: "", + adept_router_pg_ssl_mode: "prefer", + team_id: "", + model_access_group: [], +}; + +interface AddAdeptRouterTabProps { + handleOk?: () => void; + accessToken: string; + userRole: string; + userId?: string | null; + createScope?: ModelWriteScope; +} + +const AddAdeptRouterTab: React.FC = ({ + handleOk, + accessToken, + userRole, + createScope = "unscoped-ok", +}) => { + const requiresTeamScope = createScope === "team-required"; + const form: UseFormReturn = useForm({ + defaultValues: EMPTY_FORM_VALUES, + }); + const [modelAccessGroups, setModelAccessGroups] = useState([]); + const [modelInfo, setModelInfo] = useState([]); + + const watchedName = useWatch({ control: form.control, name: "adept_router_name" }); + const watchedDefaultModel = useWatch({ control: form.control, name: "adept_router_default_model" }); + const watchedTeamId = useWatch({ control: form.control, name: "team_id" }); + + useEffect(() => { + const loadAccessGroups = async () => { + try { + const response = await modelAvailableCall(accessToken, "", "", false, null, true, true); + setModelAccessGroups(response["data"].map((model: { id: string }) => model["id"])); + } catch { + // access groups unavailable; the combobox falls back to a plain text entry + } + }; + loadAccessGroups(); + }, [accessToken]); + + useEffect(() => { + const loadModels = async () => { + try { + setModelInfo(await fetchAvailableModels(accessToken)); + } catch { + // model list unavailable; the combobox still renders a free-text entry + } + }; + loadModels(); + }, [accessToken]); + + const modelChoices = Array.from(new Set(modelInfo.map((m) => m.model_group))).map((g) => ({ value: g, label: g })); + + const computeSubmitBlockedReason = (): string | null => { + if (!watchedName?.trim()) return "Enter a router name"; + if (!watchedDefaultModel?.trim()) return "Select a default fallback model"; + if (requiresTeamScope && !watchedTeamId?.trim()) return "Select a team to create this router under"; + return null; + }; + const submitBlockedReason: string | null = computeSubmitBlockedReason(); + + const onSubmit = async (values: AddAdeptRouterValues) => { + if (submitBlockedReason !== null) { + toast.fromError(submitBlockedReason); + return; + } + await handleAddAdeptRouterSubmit(values, accessToken, () => form.reset(EMPTY_FORM_VALUES), handleOk); + }; + + return ( + +
+

Add ADEPT Router

+

+ Route XML-tagged agent prompts to task-specific SLMs. Requests fall back to the default model until a + template's target_model is trained. +

+
+
+ + + {({ ref, ...field }) => ( + + )} + + + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + + + {({ ref, ...field }) => } + + + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + + + + {({ ref, ...field }) => ( + + )} + + + + {({ ref, ...field }) => } + + + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + + + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => ( + + )} + + + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + + {requiresTeamScope && ( + + {({ id, value, onChange }) => ( + onChange(next ?? "")} /> + )} + + )} + + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + + +
+ +
+
+ + ); +}; + +export default AddAdeptRouterTab; diff --git a/ui/litellm-dashboard/src/components/add_model/HandleAddAdeptRouterSubmit.tsx b/ui/litellm-dashboard/src/components/add_model/HandleAddAdeptRouterSubmit.tsx new file mode 100644 index 00000000000..5795d994c38 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/HandleAddAdeptRouterSubmit.tsx @@ -0,0 +1,66 @@ +import { modelCreateCall, Model } from "../networking"; +import { toast } from "@/lib/toast"; + +export interface AddAdeptRouterValues { + adept_router_name: string; + adept_router_default_model: string; + adept_router_tag_prefix?: string; + adept_router_conversations_threshold?: number | null; + adept_router_trainer_url?: string; + adept_router_pg_host?: string; + adept_router_pg_port?: number | null; + adept_router_pg_database?: string; + adept_router_pg_user?: string; + adept_router_pg_password?: string; + adept_router_pg_ssl_mode?: string; + team_id?: string; + model_access_group?: string[]; +} + +// Drops undefined / empty-string entries so `litellm_params` on the backend only carries the +// keys the operator actually set; keeps the create payload aligned with the Pydantic defaults. +const dropEmpty = >(values: T): Partial => + Object.fromEntries( + Object.entries(values).filter(([, value]) => value !== undefined && value !== "" && value !== null), + ) as Partial; + +export const handleAddAdeptRouterSubmit = async ( + values: AddAdeptRouterValues, + accessToken: string, + resetForm: () => void, + callback?: () => void, +) => { + try { + const rawLitellmParams = { + model: `adept/${values.adept_router_name}`, + adept_router_default_model: values.adept_router_default_model, + adept_router_tag_prefix: values.adept_router_tag_prefix, + adept_router_conversations_threshold: values.adept_router_conversations_threshold ?? undefined, + adept_router_trainer_url: values.adept_router_trainer_url, + adept_router_pg_host: values.adept_router_pg_host, + adept_router_pg_port: values.adept_router_pg_port ?? undefined, + adept_router_pg_database: values.adept_router_pg_database, + adept_router_pg_user: values.adept_router_pg_user, + adept_router_pg_password: values.adept_router_pg_password, + adept_router_pg_ssl_mode: values.adept_router_pg_ssl_mode, + }; + const litellmParams = dropEmpty(rawLitellmParams); + + const adeptConfig = { + model_name: values.adept_router_name, + litellm_params: litellmParams, + model_info: { + ...(values.team_id ? { team_id: values.team_id } : {}), + ...(values.model_access_group?.length ? { access_groups: values.model_access_group } : {}), + }, + }; + + await modelCreateCall(accessToken, adeptConfig as unknown as Model); + toast.success(`Successfully created ADEPT Router: ${values.adept_router_name}`); + resetForm(); + callback?.(); + } catch (error) { + console.error("Failed to add ADEPT router:", error); + toast.fromError("Failed to add ADEPT router: " + error); + } +}; diff --git a/ui/litellm-dashboard/src/components/edit_adept_router/AdeptRouterEditControl.tsx b/ui/litellm-dashboard/src/components/edit_adept_router/AdeptRouterEditControl.tsx new file mode 100644 index 00000000000..b09c7fd137d --- /dev/null +++ b/ui/litellm-dashboard/src/components/edit_adept_router/AdeptRouterEditControl.tsx @@ -0,0 +1,46 @@ +"use client"; + +import { useState } from "react"; +import { Button } from "@/components/ui/button"; +import EditAdeptRouterModal, { type AdeptRouterModelData } from "./EditAdeptRouterModal"; + +interface AdeptRouterEditControlProps { + canEdit: boolean; + isEditing: boolean; + modelData: AdeptRouterModelData; + accessToken: string; + onUpdated: (updated: AdeptRouterModelData) => void; +} + +/** Composes the "Edit ADEPT Router" trigger + modal into one unit so model_info_view.tsx does + * not have to know about the ADEPT-specific state, mirroring how EditAutoRouterModal is wired + * next to it but without adding 15+ lines of view-model glue to the parent. */ +export const AdeptRouterEditControl: React.FC = ({ + canEdit, + isEditing, + modelData, + accessToken, + onUpdated, +}) => { + const [isOpen, setIsOpen] = useState(false); + const isAdept = (modelData?.litellm_params as { model?: string } | undefined)?.model?.startsWith("adept/") ?? false; + if (!isAdept) return null; + return ( + <> + {canEdit && !isEditing && ( + + )} + setIsOpen(false)} + onSuccess={onUpdated} + modelData={modelData} + accessToken={accessToken} + /> + + ); +}; + +export default AdeptRouterEditControl; diff --git a/ui/litellm-dashboard/src/components/edit_adept_router/EditAdeptRouterModal.tsx b/ui/litellm-dashboard/src/components/edit_adept_router/EditAdeptRouterModal.tsx new file mode 100644 index 00000000000..63ffa93f563 --- /dev/null +++ b/ui/litellm-dashboard/src/components/edit_adept_router/EditAdeptRouterModal.tsx @@ -0,0 +1,297 @@ +"use client"; + +import React, { useEffect, useState } from "react"; +import { useForm } from "react-hook-form"; +import { toast } from "@/lib/toast"; +import { FieldGroup } from "@/components/ui/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models"; +import { modelPatchUpdateCall } from "../networking"; +import ModelChoiceCombobox from "../add_model/ModelChoiceCombobox"; + +const PG_SSL_MODES = ["disable", "allow", "prefer", "require", "verify-ca", "verify-full"] as const; + +interface AdeptRouterLitellmParams { + adept_router_default_model?: string | null; + adept_router_tag_prefix?: string | null; + adept_router_conversations_threshold?: number | null; + adept_router_trainer_url?: string | null; + adept_router_pg_host?: string | null; + adept_router_pg_port?: number | null; + adept_router_pg_database?: string | null; + adept_router_pg_user?: string | null; + adept_router_pg_password?: string | null; + adept_router_pg_ssl_mode?: string | null; +} + +export interface AdeptRouterModelData { + model_name: string; + litellm_params?: AdeptRouterLitellmParams; + model_info: { id: string }; +} + +interface EditAdeptRouterModalProps { + isVisible: boolean; + onCancel: () => void; + onSuccess: (updatedModel: AdeptRouterModelData) => void; + modelData: AdeptRouterModelData; + accessToken: string; +} + +interface EditAdeptRouterFormValues { + adept_router_name: string; + adept_router_default_model: string; + adept_router_tag_prefix: string; + adept_router_conversations_threshold: number | null; + adept_router_trainer_url: string; + adept_router_pg_host: string; + adept_router_pg_port: number | null; + adept_router_pg_database: string; + adept_router_pg_user: string; + adept_router_pg_password: string; + adept_router_pg_ssl_mode: string; +} + +const toFormValues = (modelData: AdeptRouterModelData): EditAdeptRouterFormValues => { + const lp = modelData.litellm_params ?? {}; + return { + adept_router_name: modelData.model_name, + adept_router_default_model: lp.adept_router_default_model ?? "", + adept_router_tag_prefix: lp.adept_router_tag_prefix ?? "", + adept_router_conversations_threshold: lp.adept_router_conversations_threshold ?? null, + adept_router_trainer_url: lp.adept_router_trainer_url ?? "", + adept_router_pg_host: lp.adept_router_pg_host ?? "", + adept_router_pg_port: lp.adept_router_pg_port ?? 5432, + adept_router_pg_database: lp.adept_router_pg_database ?? "", + adept_router_pg_user: lp.adept_router_pg_user ?? "", + adept_router_pg_password: "", + adept_router_pg_ssl_mode: lp.adept_router_pg_ssl_mode ?? "prefer", + }; +}; + +const EditAdeptRouterModal: React.FC = ({ + isVisible, + onCancel, + onSuccess, + modelData, + accessToken, +}) => { + const form = useForm({ defaultValues: toFormValues(modelData) }); + const [modelInfo, setModelInfo] = useState([]); + const [saving, setSaving] = useState(false); + const pgPasswordAlreadySet = !!modelData.litellm_params?.adept_router_pg_password; + + useEffect(() => { + if (isVisible) { + form.reset(toFormValues(modelData)); + } + }, [isVisible, modelData, form]); + + useEffect(() => { + if (!isVisible) return; + const loadModels = async () => { + try { + setModelInfo(await fetchAvailableModels(accessToken)); + } catch { + // model list unavailable; combobox still accepts free-text + } + }; + loadModels(); + }, [isVisible, accessToken]); + + const modelChoices = Array.from(new Set(modelInfo.map((m) => m.model_group))).map((g) => ({ value: g, label: g })); + + const handleSave = async (values: EditAdeptRouterFormValues) => { + setSaving(true); + try { + const updatedLitellmParams: AdeptRouterLitellmParams & { model: string } = { + model: `adept/${values.adept_router_name}`, + adept_router_default_model: values.adept_router_default_model || null, + adept_router_tag_prefix: values.adept_router_tag_prefix || null, + adept_router_conversations_threshold: values.adept_router_conversations_threshold ?? null, + adept_router_trainer_url: values.adept_router_trainer_url || null, + adept_router_pg_host: values.adept_router_pg_host || null, + adept_router_pg_port: values.adept_router_pg_port ?? null, + adept_router_pg_database: values.adept_router_pg_database || null, + adept_router_pg_user: values.adept_router_pg_user || null, + adept_router_pg_ssl_mode: values.adept_router_pg_ssl_mode || null, + }; + // Only overwrite the password when the operator typed a new one; otherwise the + // stored value is preserved by omission. + if (values.adept_router_pg_password) { + updatedLitellmParams.adept_router_pg_password = values.adept_router_pg_password; + } + + const patchPayload = { + model_name: values.adept_router_name, + litellm_params: updatedLitellmParams, + }; + await modelPatchUpdateCall(accessToken, patchPayload, modelData.model_info.id); + + toast.success(`Updated ADEPT Router: ${values.adept_router_name}`); + const updatedModel: AdeptRouterModelData = { + ...modelData, + model_name: values.adept_router_name, + litellm_params: { + ...(modelData.litellm_params ?? {}), + ...updatedLitellmParams, + }, + }; + onSuccess(updatedModel); + onCancel(); + } catch (error) { + console.error("Failed to update ADEPT router:", error); + toast.fromError("Failed to update ADEPT router: " + error); + } finally { + setSaving(false); + } + }; + + return ( + !open && onCancel()}> + + + Edit ADEPT Router + + Update the default model, trainer, or Postgres connection. The password is only written on save when you + enter a new value. + + +
+ + + {({ ref, ...field }) => } + + + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + + + {({ ref, ...field }) => } + + + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + + + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => } + + + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + + + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => ( + + )} + + + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + + + + + + +
+ + + ); +}; + +export default EditAdeptRouterModal; diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index 35afcdb2985..50249ee33fe 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -26,6 +26,7 @@ import { canModifyModel } from "@/utils/modelPermissions"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import DeleteResourceModal from "./common_components/DeleteResourceModal"; import EditAutoRouterModal from "./edit_auto_router/edit_auto_router_modal"; +import AdeptRouterEditControl from "./edit_adept_router/AdeptRouterEditControl"; import ReuseCredentialsModal from "./model_add/reuse_credentials"; import { toast } from "@/lib/toast"; import { @@ -578,6 +579,7 @@ export default function ModelInfoView({ onModelUpdate(updatedModel); } }; + const isWildcardModel = modelData.litellm_model_name.includes("*"); const wildcardProvider = modelData.litellm_model_name.split("/")[0]; const healthCheckModelOptions = @@ -748,6 +750,13 @@ export default function ModelInfoView({ Edit Auto Router )} + {canEditModel ? ( !isEditing && (