diff --git a/cookbook/adept_router_e2e_testing.md b/cookbook/adept_router_e2e_testing.md new file mode 100644 index 00000000000..0c70f99bb98 --- /dev/null +++ b/cookbook/adept_router_e2e_testing.md @@ -0,0 +1,508 @@ +# ADEPT Router — End-to-End Testing Guide + +This guide walks through every layer of the ADEPT system in a logical sequence. +Each stage builds on the previous one. Run them in order the first time. + +--- + +## Prerequisites + +- LiteLLM proxy running (default `http://localhost:4000`) +- PostgreSQL running and reachable +- A master key set (e.g. `LITELLM_MASTER_KEY=sk-1234`) +- `curl` and `jq` available in your terminal + +Set a convenience alias for all curl commands: + +```bash +export LITELLM_URL=http://localhost:4000 +export LITELLM_KEY=sk-1234 +``` + +--- + +## Stage 1 — Basic Health and Inference + +Confirm the proxy is alive and can call a model before touching ADEPT at all. + +### 1.1 Health check + +```bash +curl -s $LITELLM_URL/health | jq . +``` + +Expected: `{"status": "healthy", ...}` + +### 1.2 List available models + +```bash +curl -s $LITELLM_URL/v1/models \ + -H "Authorization: Bearer $LITELLM_KEY" | jq '.data[].id' +``` + +You should see the model names from your `config.yaml`. + +### 1.3 Basic inference (non-ADEPT model) + +```bash +curl -s $LITELLM_URL/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "Say hello."}] + }' | jq '.choices[0].message.content' +``` + +Expected: any sensible reply. This confirms the proxy and provider credentials work. + +--- + +## Stage 2 — ADEPT Router Config and Visibility + +### 2.0 Dependencies + +ADEPT's Postgres-backed template store reaches your database through litellm's existing +Prisma client (one client instance per ADEPT deployment, pointed at your database), so no +extra dependencies are needed beyond the proxy: + +```bash +pip install "litellm[proxy]" +``` + +The store connects to the database you configure with the `adept_router_pg_*` params below, +never to litellm's own database. + +### 2.1 Sample config.yaml for ADEPT + +Add a deployment with the `adept_router_*` params: + +```yaml +model_list: + - model_name: gpt-4o-mini + litellm_params: + model: openai/gpt-4o-mini + api_key: os.environ/OPENAI_API_KEY + + - model_name: invoice-slm # trained SLM (leave as default model initially) + litellm_params: + model: openai/gpt-4o-mini + api_key: os.environ/OPENAI_API_KEY + adept_router: true + adept_router_default_model: gpt-4o-mini + adept_router_pg_host: localhost + adept_router_pg_port: 5432 + adept_router_pg_user: litellm + adept_router_pg_password: yourpassword + adept_router_pg_database: adept + adept_router_tag_prefix: "" # or "tool:" for namespaced tags + adept_router_conversations_threshold: 5 + adept_router_trainer_url: "" # set to trainer URL when ready +``` + +Restart LiteLLM after saving the config. + +### 2.2 Confirm ADEPT router is registered + +On startup you should see in the LiteLLM logs: + +``` +AdeptRouter: initialized PostgreSQL template store. +``` + +If you see a Postgres connection error instead, check your `adept_router_pg_*` credentials. + +### 2.3 Confirm callback is wired + +```bash +# In a Python shell or script: +python3 -c " +import litellm +from litellm import Router +# load your config... +# Check: +adept_instances = [cb for cb in litellm.callbacks if type(cb).__name__ == 'AdeptRouter'] +print('AdeptRouter callbacks registered:', len(adept_instances)) +" +``` + +Expected: `AdeptRouter callbacks registered: 1` (one per ADEPT deployment). + +--- + +## Stage 3 — First ADEPT Inference (Cold — No Template Yet) + +### 3.1 Send a tagged prompt (simulating a tool call) + +This uses your ADEPT-enabled model name. The first request will: +- **Miss** the template store (no template exists yet) +- Fall back to the default model +- Store a new template + conversation in Postgres + +```bash +curl -s $LITELLM_URL/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "invoice-slm", + "messages": [ + { + "role": "system", + "content": "You are an invoice processor. Extract the total amount from the invoice." + }, + { + "role": "user", + "content": "Invoice #INV-001\nDate: 2024-01-15\nTotal: $1,234.56\nVendor: Acme Corp" + } + ] + }' | jq '{model: .model, content: .choices[0].message.content}' +``` + +Expected: a response from the default model (gpt-4o-mini). Check the logs for: + +``` +AdeptRouter: no template match, falling back to gpt-4o-mini +AdeptRouter: stored interaction. +``` + +### 3.2 Verify the template was stored in Postgres + +```bash +psql -U litellm -d adept -c "SELECT id, template_hash, target_model FROM adept_templates LIMIT 5;" +``` + +You should see one row. The `target_model` will be empty (`""`) — it gets populated after training. + +```bash +psql -U litellm -d adept -c "SELECT id, template_id FROM adept_conversations LIMIT 5;" +``` + +You should see one conversation row linked to the template. + +--- + +## Stage 4 — Template Matching (Same Tool, Different Data) + +### 4.1 Send the same tool's prompt with different runtime values + +The system prompt is the same. The user message has the same XML structure but different content. +ADEPT should strip the tag values and match the skeleton to the stored template. + +```bash +curl -s $LITELLM_URL/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "invoice-slm", + "messages": [ + { + "role": "system", + "content": "You are an invoice processor. Extract the total amount from the invoice." + }, + { + "role": "user", + "content": "Invoice #INV-002\nDate: 2024-02-20\nTotal: $9,876.00\nVendor: Beta Ltd" + } + ] + }' | jq '{model: .model, content: .choices[0].message.content}' +``` + +Logs should show: +``` +Matched template +AdeptRouter: matched template , routing to gpt-4o-mini +``` + +Even though the data is different, it matched the same template (same skeleton structure). + +### 4.2 Verify the conversation count increased + +```bash +psql -U litellm -d adept -c \ + "SELECT template_id, COUNT(*) as conversations FROM adept_conversations GROUP BY template_id;" +``` + +Expected: 2 conversations for the same template_id. + +### 4.3 Different system prompt = different template + +Send a different tool (different system prompt). This should create a **new** template: + +```bash +curl -s $LITELLM_URL/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "invoice-slm", + "messages": [ + { + "role": "system", + "content": "You are a contract reviewer. Summarize the key terms." + }, + { + "role": "user", + "content": "This agreement is between Party A and Party B effective 2024-03-01." + } + ] + }' | jq .model' +``` + +```bash +psql -U litellm -d adept -c "SELECT COUNT(*) FROM adept_templates;" +``` + +Expected: 2 templates now (one per unique tool/system prompt combo). + +--- + +## Stage 5 — Hit the Training Threshold + +Send enough requests to trigger the trainer notification. +With `adept_router_conversations_threshold: 5`, the 5th, 10th, 15th... conversation triggers it. + +### 5.1 Loop to threshold + +```bash +for i in $(seq 1 5); do + echo "=== Request $i ===" + curl -s $LITELLM_URL/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d "{ + \"model\": \"invoice-slm\", + \"messages\": [ + {\"role\": \"system\", \"content\": \"You are an invoice processor. Extract the total amount from the invoice.\"}, + {\"role\": \"user\", \"content\": \"Invoice #INV-00$i\nTotal: \$${i}00.00\nVendor: Vendor$i\"} + ] + }" | jq '.choices[0].message.content' -r + sleep 1 +done +``` + +### 5.2 Verify conversation count + +```bash +psql -U litellm -d adept -c \ + "SELECT template_id, COUNT(*) FROM adept_conversations GROUP BY template_id ORDER BY COUNT(*) DESC;" +``` + +### 5.3 Trainer trigger (if configured) + +If `adept_router_trainer_url` is set, on the 5th request you'll see in LiteLLM logs: + +``` +Triggered trainer for template +``` + +If no `trainer_url` is set, you'll see: + +``` +No trainer_url configured, skipping trainer notification +``` + +Both are correct — the skip is intentional when trainer is not deployed. + +--- + +## Stage 6 — Post-Training Routing (After SLM is Trained) + +Once the trainer has fine-tuned a model and loaded it into vLLM, update the template's +`target_model` in Postgres to point to the trained adapter: + +```bash +psql -U litellm -d adept -c \ + "UPDATE adept_templates SET target_model = 'invoice-slm-v1' WHERE target_model = '';" +``` + +Then also add `invoice-slm-v1` to your LiteLLM `config.yaml` pointing to your vLLM endpoint: + +```yaml + - model_name: invoice-slm-v1 + litellm_params: + model: openai/invoice-slm-lora + api_base: http://localhost:8000/v1 + api_key: fake-key +``` + +### 6.1 Verify routing switches to SLM + +```bash +curl -s $LITELLM_URL/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "invoice-slm", + "messages": [ + { + "role": "system", + "content": "You are an invoice processor. Extract the total amount from the invoice." + }, + { + "role": "user", + "content": "Invoice #INV-010\nTotal: $500.00\nVendor: TestCo" + } + ] + }' | jq '{model: .model, content: .choices[0].message.content}' +``` + +Logs should now show: +``` +AdeptRouter: matched template , routing to invoice-slm-v1 +``` + +### 6.2 Confirm `routed_to_slm` flag in Postgres + +```bash +psql -U litellm -d adept -c \ + "SELECT additional_information->>'routed_to_slm', additional_information->>'model' + FROM adept_conversations ORDER BY created_at DESC LIMIT 3;" +``` + +Expected: most recent row shows `routed_to_slm: true` and `model: invoice-slm-v1`. + +### 6.3 (Optional) Test vLLM LoRA adapter directly + +If you have vLLM running with a base model: + +```bash +# Check vLLM health +curl http://localhost:8000/health + +# List loaded models +curl http://localhost:8000/v1/models | jq '.data[].id' + +# Load a LoRA adapter (requires VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 env var on vLLM startup) +curl -X POST http://localhost:8000/v1/load_lora_adapter \ + -H "Content-Type: application/json" \ + -d '{ + "lora_name": "invoice-slm-lora", + "lora_path": "/path/to/lora/adapter" + }' + +# Verify the adapter appears in models list +curl http://localhost:8000/v1/models | jq '.data[].id' + +# Inference through the adapter +curl http://localhost:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "invoice-slm-lora", + "messages": [{"role": "user", "content": "Invoice #INV-001 Total: $500"}] + }' | jq '.choices[0].message.content' +``` + +> **Important:** vLLM must be started with `VLLM_ALLOW_RUNTIME_LORA_UPDATING=1` in its environment, +> otherwise the `/v1/load_lora_adapter` endpoint is not registered and returns 404. + +--- + +## Stage 7 — Edge Case Tests + +### 7.1 Non-tagged prompt (no XML structure) + +ADEPT should still work — it just won't strip any tag values. The skeleton = the full (normalized) prompt. + +```bash +curl -s $LITELLM_URL/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "invoice-slm", + "messages": [ + {"role": "user", "content": "What is 2 + 2?"} + ] + }' | jq '.choices[0].message.content' +``` + +This creates its own template (hash of "What is {NUM} + {NUM}?"). Works fine. + +### 7.2 Multi-turn conversation — only last user message used + +Send a conversation with multiple turns: + +```bash +curl -s $LITELLM_URL/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "invoice-slm", + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "Paris."}, + {"role": "user", "content": "Tell me more about Paris."} + ] + }' | jq '.choices[0].message.content' +``` + +ADEPT uses only the last `user` message for template matching. + +### 7.3 Tool-result turn is skipped + +If the last message is `role: tool`, ADEPT skips storing the conversation (the preceding assistant turn already captured the exchange): + +```bash +curl -s $LITELLM_URL/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "invoice-slm", + "messages": [ + {"role": "user", "content": "Run this tool"}, + {"role": "assistant", "content": null, "tool_calls": [{"id": "tc1", "type": "function", "function": {"name": "get_data", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "tc1", "content": "tool result here"} + ] + }' | jq '.choices[0].message.content' +``` + +Logs should show no "stored interaction" for this turn — that's correct behavior. + +### 7.4 Special characters in Postgres password + +The PG URL is URL-encoded, so passwords with `@`, `:`, `/` work. Verify by starting with: + +```yaml +adept_router_pg_password: "p@ss:w/rd" +``` + +LiteLLM should connect without error. The encoded URL will have `p%40ss%3Aw%2Frd` in it. + +--- + +## Quick Reference: What to Check at Each Step + +| Stage | What to verify | Where to check | +|-------|---------------|----------------| +| 1 | Proxy alive, model available, basic inference works | `curl /health`, `/v1/models`, one completion | +| 2 | ADEPT registered, Postgres connected | Startup logs, `adept_templates` table exists | +| 3 | Cold miss: fallback model used, template created | Logs, `adept_templates` row | +| 4 | Same skeleton → template hit; different sys prompt → new template | Logs, `adept_templates` count | +| 5 | Conversation counter grows; trainer triggered at threshold | `adept_conversations` count, logs | +| 6 | After `target_model` updated: routes to SLM; `routed_to_slm=true` in DB | Logs, `adept_conversations.additional_information` | +| 7 | Untagged prompt, multi-turn, tool-result skip, special-char password all work | Logs, no errors | + +--- + +## Common Problems + +**Template never matches (always "no template match")** +- Check that the system prompt is identical across requests (whitespace-normalized, same content). +- Check that the XML tag names are consistent (e.g. `` must match exactly). +- Query `adept_templates` — if empty, `async_log_success_event` is not firing. Check that `AdeptRouter` appears in `litellm.callbacks`. + +**Conversations not stored** +- `AdeptRouter` is not registered in the callback pipeline. Look for `"AdeptRouter: stored interaction"` in logs. +- Verify startup logs show `"AdeptRouter: initialized PostgreSQL template store."`. + +**Postgres connection error on startup** +- Check `adept_router_pg_*` values in config.yaml. +- Test connectivity: `psql -U -h -d `. +- If password has special characters, they are URL-encoded automatically — you do not need to encode them yourself. + +**vLLM `/v1/load_lora_adapter` returns 404** +- vLLM was not started with `VLLM_ALLOW_RUNTIME_LORA_UPDATING=1`. +- Restart vLLM: `VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 vllm serve --enable-lora ...` + +**Trainer never triggered** +- `adept_router_trainer_url` is empty or not set — that's fine in dev, logs will say "No trainer_url configured, skipping". +- If it IS set and still not triggered, check that `conversation_count % conversations_threshold == 0` is reached. Use `COUNT(*)` query above. diff --git a/litellm/router.py b/litellm/router.py index 843b916d90f..18dadd2e2db 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -280,6 +280,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, @@ -888,6 +891,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] = ( @@ -9256,6 +9260,104 @@ 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. + + Called from _add_deployment on every 30-second DB sync tick. If the deployment + is already registered with identical params, this is a no-op. If a param + changed (e.g. operator edited trainer_url in the UI), the in-memory router + is rebuilt 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, urlparse + + 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: + _parsed: Final = urlparse(trainer_url) + if _parsed.scheme not in ("http", "https"): + raise ValueError(f"adept_router_trainer_url must use http or https, got: {_parsed.scheme!r}") + _host: Final = (_parsed.hostname or "").lower() + _CLOUD_METADATA_HOSTS: Final = frozenset( + { + "169.254.169.254", + "100.100.100.200", + "metadata.google.internal", + "fd00:ec2::254", + } + ) + if _host in _CLOUD_METADATA_HOSTS: + raise ValueError(f"adept_router_trainer_url host {_host!r} is a blocked cloud-metadata endpoint") + _allowed_hosts: Final = lp.adept_router_trainer_url_allowed_hosts + if _allowed_hosts is not None: + _allowed_lower: Final = frozenset(h.lower() for h in _allowed_hosts) + if _host not in _allowed_lower: + raise ValueError( + f"adept_router_trainer_url host {_host!r} is not in adept_router_trainer_url_allowed_hosts" + ) + + 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) + + 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 @@ -9285,6 +9387,9 @@ 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) + 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 @@ -9375,9 +9480,15 @@ 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 route through their own strategy (initialized later in + # this method at `init_adept_router_deployment`) and don't have a standard + # LLM provider. Skip the get_llm_provider validation that would otherwise + # reject `adept/` as an unsupported provider. + 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: + # Skip LLM provider validation — actual routing happens via a + # strategy-specific init path further down. _model = litellm_model custom_llm_provider = None dynamic_api_key = None @@ -9482,6 +9593,12 @@ class Router: if self._is_quality_router_deployment(litellm_params=deployment.litellm_params): self.init_quality_router_deployment(deployment=deployment) + ######################################################### + # Check if this is an ADEPT router 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): @@ -13143,6 +13260,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 @@ -13248,6 +13373,18 @@ class Router: if newly_forwarded: request_kwargs.update(((_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, tuple(key for key, _ in newly_forwarded)),)) + ######################################################### + # Check if any ADEPT router should be used + ######################################################### + 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( @@ -13420,6 +13557,14 @@ 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 (e.g. Router.acompletion()) with a supported " + "routing_strategy (simple-shuffle, usage-based-routing-v2, cost-based-routing, " + "latency-based-routing, least-busy)." + ) # 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..22cf8a030f0 --- /dev/null +++ b/litellm/router_strategy/adept_router/adept_router.py @@ -0,0 +1,327 @@ +""" +ADEPT (Adaptive Deployment via Prompt Templates) Router. + +Designed for single-turn, task-specific routing. An agent or tool sends a fixed system +prompt (the task definition) plus XML-tagged variable user content (the runtime input). +ADEPT extracts a structural skeleton from each prompt, hashes it together with the system +prompt for per-tool isolation, and routes to a task-specific SLM once one has been trained. + +Until a trained SLM exists for a template, all traffic falls back to the default model +while conversations accumulate as training data in Postgres. +""" + +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): + """One block of an OpenAI multimodal message content list; non-text blocks parse with text=None.""" + + model_config = ConfigDict(extra="ignore") + type: str | None = None + text: str | None = None + + +class _UsageEnvelope(BaseModel): + """Reads the dynamically-set `usage` attribute off a ModelResponse in a typed way.""" + + 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): + """ + ADEPT routing strategy — matches incoming prompts to known templates via + SHA-256 hashing of the masked template string, with no external vector DB. + """ + + 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, + ) + # Seeding hits the store, which is async, so it runs lazily on first use (the store + # connects lazily too). __init__ stays sync for the router's construction path. + self._seed_config = seed_config + self._seeded = not seed_config + self._seed_lock = asyncio.Lock() + + async def _ensure_seeded(self) -> None: + """Pre-populate templates from seed_config once, on first use.""" + if self._seeded: + return + 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 + + await self._ensure_seeded() + + 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) + # Authorize the *actual* target before recording the routing decision — a caller with + # access to the ADEPT alias but not the trained SLM must not be silently upgraded. + 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) + + # Stash routing decision so async_log_success_event can record it without re-querying. + 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 the proxy model-access check against the swapped-in target model. + + The proxy's auth layer only sees the ADEPT alias (e.g. ``adept/my-tool``) and + authorizes that. Without this second check, a key that can call the alias but not + the underlying default/SLM would be silently escalated. When no ``user_api_key_auth`` + is on the request (ADEPT used outside the proxy), skip — the proxy is the only + surface that ever populates it. + """ + 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: + """Read ``user_api_key_auth`` off request_kwargs.metadata / litellm_metadata directly.""" + 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: + """Read a key from the request's metadata / litellm_metadata dict, if present.""" + try: + params = _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)) + except ValidationError: + continue + value = nested.get(key) + 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: + # Multiple ADEPT-router deployments each register an instance of this class + # as a global LiteLLM callback. Without a model_group filter, every instance + # would log every successful request — including requests routed through a + # *different* ADEPT deployment — duplicating conversation rows across + # router_ids. Gate on the requested model_group so only the router that + # actually handled this request logs it. + lp_raw: Final = kwargs.get("litellm_params") + request_model_group: Final = self._read_request_metadata(lp_raw, "model_group") + # Only log requests routed through *this* ADEPT model. A request whose model_group is + # absent or belongs to a different deployment is not ours, so skipping it avoids both + # storing non-ADEPT traffic and duplicating rows when several ADEPT deployments exist. + 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 + # Skip tool-result turns — the preceding assistant turn already captured this exchange. + 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: + """Flatten a message's content (string or OpenAI content-block list) to plain text.""" + 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..1aa998cdc01 --- /dev/null +++ b/litellm/router_strategy/adept_router/store/implementation/prisma.py @@ -0,0 +1,261 @@ +import asyncio +import json +from collections.abc import Mapping +from datetime import datetime +from typing import Final + +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]) + + +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: + # Prisma's raw-query model path returns JSONB as a JSON string; parse it back to a mapping. + if isinstance(value, str): + return _JSON_ADAPTER.validate_json(value) + return value + + +async def _create_tables(client: Prisma) -> None: + """Create ADEPT's tables in the user's database if they are absent. Column names match the + prior SQLAlchemy schema so existing ADEPT databases stay compatible.""" + 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 + + +# One long-lived Prisma client per database URL, connected once and reused. Keying on the URL +# (rather than per repo instance) means a router rebuild — which drops the old repo and builds a +# new one for the same database — reuses the existing client instead of connecting a second one +# and orphaning the first, so no connection or engine process leaks across rebuilds. +_CLIENTS: Final[dict[str, Prisma]] = {} # mutable-ok: connection registry keyed by database URL +_REGISTRY_LOCK: Final = asyncio.Lock() + + +async def _get_client(db_url: str) -> Prisma: + """Return the connected client for db_url (one per URL, shared), creating tables once.""" + 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) + _CLIENTS[db_url] = client + return client + + +class AdeptPrismaRepo(AdeptTemplateStore): + """ADEPT template/conversation store backed by the user's own PostgreSQL. + + Reaches the user's tables through litellm's Prisma client pointed at the user's database via a + datasource override, using parameterized raw SQL. One client is kept per database URL (see + `_CLIENTS`), connected once and reused for the app's lifetime, so a router rebuild does not + leak a connection. Nothing is added to litellm's own Prisma schema and litellm's database is + never opened; `auto_register` stays off so litellm's global client is never affected. Tables + are created on first use. + """ + + 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: + client: Final = await _get_client(self._db_url) + 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 get_template_by_hash(self, template_hash: str, router_id: str) -> StoredTemplate | None: + try: + client: Final = await _get_client(self._db_url) + 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 AND template_hash = $2 LIMIT 1", + router_id, + template_hash, + model=_TemplateRow, + ) + except PrismaError as e: + verbose_router_logger.error("Error retrieving template by hash: %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 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: + client: Final = await _get_client(self._db_url) + 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 + else: + verbose_router_logger.debug("Stored conversation for template %s", template_id) + 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 template_id (ours or a concurrent + insert's). ON CONFLICT DO NOTHING on the (router_id, template_hash) unique index makes + two concurrent requests with the same hash safe: the loser no-ops and we re-read the winner. + """ + try: + client: Final = await _get_client(self._db_url) + 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 + else: + surviving_id: Final = rows[0].id if rows else template_id + verbose_router_logger.debug("AdeptRouter: stored template %s", surviving_id) + return surviving_id + + async def get_template(self, template_id: str) -> StoredTemplate | None: + try: + client: Final = await _get_client(self._db_url) + 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: + client: Final = await _get_client(self._db_url) + 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 + else: + 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..2d443807d16 --- /dev/null +++ b/litellm/router_strategy/adept_router/store/store_template.py @@ -0,0 +1,84 @@ +from abc import ABC, abstractmethod +from collections.abc import Mapping +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. + + Args: + template_hash: SHA-256 hex digest of the masked template. + router_id: The router that owns this template. + + Returns: + The template ID if found, None otherwise. + """ + ... + + @abstractmethod + async def get_template_by_hash(self, template_hash: str, router_id: str) -> "StoredTemplate | None": + """Return the full template row for a hash in a single query. + + Kept alongside :meth:`match_by_hash` because the write path only needs the surviving id. + Read path uses this to avoid a second round-trip on every request. + """ + ... + + @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..51bb0a4c949 --- /dev/null +++ b/litellm/router_strategy/adept_router/template/implementation/adept_template_router.py @@ -0,0 +1,323 @@ +import asyncio +import hashlib +import re +import time +from collections import OrderedDict +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final +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 +# Fire-and-forget trainer notification: bounded so a stuck trainer can't leak sockets. +_TRAINER_HTTP_TIMEOUT_SECONDS: Final = 10.0 + + +class AdeptTemplateRouter(BaseTemplateRouter): + """ + Routes single-turn prompts to task-specific SLMs by matching their structural template. + + Intended use case: an agent/tool sends a fixed system prompt (the task definition) and + XML-tagged variable user content (the runtime input). ADEPT strips the tag values, leaving + a stable structural skeleton, and uses SHA-256(system_prompt | skeleton) as a routing key. + Each unique tool gets its own template family, its own training dataset, and — after enough + conversations — its own trained SLM. + + Flow: + 1. User message is normalized (whitespace) and XML tag values are stripped. + 2. Remaining variable spans (IDs, emails, URLs, numbers, UUIDs) are masked to placeholders. + 3. The masked skeleton is hashed together with the system prompt for per-tool isolation. + 4. Hash is looked up in Postgres — a hit routes to the template's target_model. + 5. On a miss the skeleton is stored; the default model handles the request. + 6. Every response is stored as a training conversation linked to the template. + 7. At every multiple of conversations_threshold, the external trainer is notified. + """ + + # Compiled at class level — shared across all instances, never recompiled per call. + 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 would otherwise 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 + # Bounded LRU: cache hot template lookups so `route()` skips the DB on repeat hits. + # (router_id, template_hash) -> (inserted_at_monotonic, StoredTemplate). OrderedDict + # gives O(1) LRU eviction; entries older than the TTL are dropped on read to bound + # staleness after a trainer update changes `target_model`. + self._template_cache: OrderedDict[ # mutable-ok: LRU cache, bounded and TTL-gated + tuple[str, str], tuple[float, StoredTemplate] + ] = OrderedDict() + + # Both the match pattern and its replacement string depend on tag_prefix, so they are + # built once here rather than on every _extract_template call. + 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) # before NUM — see UUID_RE comment above + return self.NUM_RE.sub("{NUM}", uuids) + + def _extract_tag_content(self, text: str) -> Sequence[tuple[str, str]]: + """Return (tag_name, value) pairs for all XML-tagged spans in text.""" + 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: %s... (%s tags)", masked_template[:100], len(extractions)) + return masked_template, extractions + + @staticmethod + def _hash_template(masked_template: str, system_prompt: str | None = None) -> str: + """ + Produce a routing key from the masked template skeleton. + + When a system prompt is provided it is prepended so that two tools with identical + user-message structure but different task definitions hash to different templates. + This is the per-tool isolation guarantee: same tool → same hash, different tool → different hash. + """ + 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: + """Pre-populate one template from a seed description. Returns True if a new template was stored.""" + masked: Final = self._mask_text(self._normalize_text(description)) + # Use the shared hash function so seeding stays consistent with live routing. + 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: + masked_template, _ = self._extract_template(prompt) + template_hash: Final = self._hash_template(masked_template, system_prompt) + router_id: Final = self.get_router_id() + cache_key: Final = (router_id, template_hash) + cached: Final = self._cache_get(cache_key) + if cached is not 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, + ) + stored: Final = await self.template_store.get_template_by_hash(template_hash, router_id) + if stored is None: + verbose_router_logger.debug("No matching template found") + return None + self._cache_put(cache_key, stored) + verbose_router_logger.info("Matched template %s", stored.id) + return AdeptTemplateMatch( + template_id=stored.id, + template=stored.template, + target_model=stored.target_model, + metadata=stored.additional_information, + ) + except (AttributeError, IndexError, KeyError, TypeError, ValueError) as e: + verbose_router_logger.exception("Error matching template: %s", e) + return None + + 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) # LRU refresh + 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: + """Return the id of the template for this hash, storing a new one on a miss.""" + 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 + # store_template returns the surviving id (handles concurrent inserts safely). + 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, + ) + # A new template row invalidates any negative-cache observation for the same key, + # so future route() calls read the freshly-stored row from the store. + 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) + # Resolve router_id once — used for both the hash lookup and, on a miss, template insert. + 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) + # Modulo check re-triggers at N, 2N, 3N... so training improves as traffic grows. + 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: + """Fire a background notification to the external trainer; never blocks the caller. + + A slow trainer previously parked a thread-pool worker for the full ``httpx`` timeout + because the sync client was awaited via :func:`asyncio.to_thread`. Under threshold + traffic this can starve every other :func:`asyncio.to_thread` call. Scheduling an + :class:`~httpx.AsyncClient` POST as a detached task keeps the notification bounded + by its own timeout without back-pressuring the request path. + """ + if not self.trainer_url: + # Visible by default — operators need to know they hit the threshold + # but no trainer is wired (common during initial setup / when + # `adept_router_trainer_url` was edited but the proxy not restarted). + 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: + # Reuses the litellm-managed shared async HTTP client rather than instantiating + # a per-call `httpx.AsyncClient`, which the code-quality gate blocks: creating a + # client per call re-does TLS/handshake work and adds +500ms per request. + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + try: + await client.post(url=url, 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 0db482d8a58..f1055e5e388 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -376,6 +376,20 @@ 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_trainer_url_allowed_hosts: list[str] | None = None # mutable-ok: pydantic config field + 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 9b5fb08a45f..d362fd17b25 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3803,6 +3803,20 @@ 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_trainer_url_allowed_hosts", + "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/test_litellm/test_adept_router.py b/tests/test_litellm/test_adept_router.py new file mode 100644 index 00000000000..288919ac96d --- /dev/null +++ b/tests/test_litellm/test_adept_router.py @@ -0,0 +1,1179 @@ +"""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_miss(): + mock_storage = AsyncMock() + mock_storage.get_template_by_hash.return_value = None + + router = _make_template_router(mock_storage) + result = asyncio.run(router.route("What is 2 + 2?")) + assert result is None + + +def test_adept_template_router_route_hit(): + from litellm.router_strategy.adept_router.store.store_template import StoredTemplate + + mock_storage = AsyncMock() + mock_storage.get_template_by_hash.return_value = StoredTemplate( + id="tmpl-abc", + template="Get order {ID} for {EMAIL}", + template_hash="hash-abc", + router_id="router-id-1", + target_model="gpt-4o", + additional_information=None, + created_at=None, + ) + + router = _make_template_router(mock_storage) + result = asyncio.run(router.route("Get order ORD-123 for user@example.com")) + assert result is not None + assert result["target_model"] == "gpt-4o" + # match_by_hash/get_template are the old two-query path; route() must not hit them. + mock_storage.match_by_hash.assert_not_called() + mock_storage.get_template.assert_not_called() + + +def test_route_uses_cache_on_second_call(): + """A repeated route() for the same prompt must serve from cache and skip the store.""" + from litellm.router_strategy.adept_router.store.store_template import StoredTemplate + + mock_storage = AsyncMock() + mock_storage.get_template_by_hash.return_value = StoredTemplate( + id="tmpl-abc", + template="Get order {ID} for {EMAIL}", + template_hash="hash-abc", + router_id="router-id-1", + target_model="gpt-4o", + additional_information=None, + created_at=None, + ) + + router = _make_template_router(mock_storage) + + async def _run(): + first = await router.route("Get order ORD-123 for user@example.com") + second = await router.route("Get order ORD-999 for other@example.com") + return first, second + + first, second = asyncio.run(_run()) + assert first is not None and second is not None + assert second["target_model"] == "gpt-4o" + # Second call must have been served from cache — store hit only once. + assert mock_storage.get_template_by_hash.await_count == 1 + + +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: + 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 _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_seed_config_missing_description_logs_warning(): + """_ensure_seeded 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("litellm.router_strategy.adept_router.adept_router.verbose_router_logger") as mock_log: # test-quality-ok: warning log is the observable output of the misconfiguration guard + asyncio.run(adept._ensure_seeded()) + 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] = 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 + + client = await prisma_mod._get_client(db_url) + await client.execute_raw("DELETE FROM conversations WHERE template_id = $1", surviving) + await client.execute_raw("DELETE FROM templates WHERE router_id = $1", router_id) + await 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(prisma_mod, "Prisma", side_effect=lambda datasource: _FakeClient()): # test-quality-ok: asserts connection reuse across repo rebuilds; the Prisma constructor IS the reuse boundary + + async def run(): + AdeptPrismaRepo(url) # first router + c1 = await prisma_mod._get_client(url) + AdeptPrismaRepo(url) # simulate a rebuild: a fresh repo for the same URL + c2 = await prisma_mod._get_client(url) + return c1, c2 + + c1, c2 = asyncio.run(run()) + + assert c1 is c2 # 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: + + 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: + 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: + 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.""" + 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="blocked cloud-metadata"): + router.init_adept_router_deployment(deployment) + + +def test_trainer_url_invalid_scheme_rejected(): + """Non-http(s) schemes in trainer_url must raise ValueError.""" + 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="http or https"): + 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("litellm.proxy.auth.auth_checks.can_key_call_resolved_model", new_callable=AsyncMock) as check: # test-quality-ok: verifies the authz check is SKIPPED when no proxy auth object is present; the check-not-called IS the observable + 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 +# --------------------------------------------------------------------------- + + +def test_trainer_url_rejected_when_not_in_operator_allowlist(): + """When adept_router_trainer_url_allowed_hosts is set, only listed hosts are accepted.""" + from litellm.router import Router + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = Router(model_list=[]) + deployment = Deployment( + model_name="allowlist_deny", + litellm_params=LiteLLM_Params( + model="adept/allowlist_deny", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.internal.com", + adept_router_trainer_url="http://rogue.example.com/hook", + adept_router_trainer_url_allowed_hosts=["trainer.internal.com"], + ), + model_info=ModelInfo(), + ) + with pytest.raises(ValueError, match="not in adept_router_trainer_url_allowed_hosts"): + router.init_adept_router_deployment(deployment) + + +def test_trainer_url_allowed_when_in_operator_allowlist(): + """A trainer host present in the allowlist must pass init.""" + 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="allowlist_allow", + litellm_params=LiteLLM_Params( + model="adept/allowlist_allow", + adept_router_default_model="gpt-4o", + adept_router_pg_host="db.internal.com", + adept_router_trainer_url="http://Trainer.Internal.com/hook", + adept_router_trainer_url_allowed_hosts=["trainer.internal.com"], + ), + model_info=ModelInfo(), + ) + with _patch("litellm.router_strategy.adept_router.adept_router.AdeptRouter", return_value=MagicMock()): + router.init_adept_router_deployment(deployment) + assert "allowlist_allow" in router.adept_routers diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 76ac60a6453..de4130a5b8f 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2460,4 +2460,4 @@ "count": 1 } } -} \ No newline at end of file +} 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 34c9d87004e..fc760b3d909 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) : []), ...(isAdmin ? ([ "llm-credentials", 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..14b6d4458a5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AddAdeptRouterTab.tsx @@ -0,0 +1,292 @@ +"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_trainer_url_allowed_hosts: [], + 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..16e0535f460 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/HandleAddAdeptRouterSubmit.tsx @@ -0,0 +1,70 @@ +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_trainer_url_allowed_hosts?: 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_trainer_url_allowed_hosts: values.adept_router_trainer_url_allowed_hosts?.length + ? values.adept_router_trainer_url_allowed_hosts + : undefined, + 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 && (