feat: simple/complex model routing and fallback chains (P-e) #27

Merged
KofTwentyTwo merged 2 commits from feature/p5-routing-fallback into develop 2026-10-14 00:13:27 +00:00
31 changed files with 2351 additions and 47 deletions

No files matched your search

+9
View File
@@ -1,5 +1,14 @@
# Changelog
## Unreleased - P-e
- Add per-profile heuristic/classifier model routing, with budgeted ROUTE usage.
- Add bounded cross-provider fallback chains, PHI-safe selection, neutral history
rebuilding, and a 60-second connection circuit breaker with cached local probes.
- Show provider health, fallback totals and turn routing; include additive schema
migration and admin record-screen controls.
## Unreleased - P-d
- Admin budgets, expiring overrides, deduplicated warning/downgrade/stop events.
+1
View File
@@ -69,6 +69,7 @@ Always use the wrapper: Maven 3.10 cannot read the parent POM. The committed def
- **Links:** only `/app` or `/app/...` hrefs are emitted (`TurnEngine.isAppHref`).
- **PHI routing:** PHI only to `hipaaEligible` providers; checked before every model call. Conversation classification never decreases.
- **HIPAA storage:** Conversation storage encrypted in HIPAA mode, only in `AssistantStore`; encryption keys are environment references and never logged or stored in tables.
- **Fallback:** fallback never crosses the PHI rule; never local-to-cloud by default. No retry after partial output from that call; provider switches rebuild neutral history without providerState.
- **Budgets:** budgets are checked before every model call; downgrade never crosses the PHI rule.
- **Kill switch:** read fresh before every model call and every tool step.
+27
View File
@@ -460,3 +460,30 @@ The Assistant home shows budgets, overrides, daily cost and tokens, and breakdow
by user, model, provider and tool. Existing `AssistantBudgets`/policy limits and
provider monthly caps continue to apply independently of these rows and overrides.
See [UPGRADE.md](UPGRADE.md) for the required additive schema.
### Profile routing and provider fallback
On an Assistant Profile, **Model routing** defaults to OFF. HEURISTIC selects the
complex slot for long questions (>600 characters), two mentioned app tables,
analysis/comparison/explanation/trend/planning requests, pending writes, more than
eight previous turns, or a previous tool limit; otherwise it selects the simple
slot. CLASSIFIER uses an optional registry model with just the question and a
fixed instruction, four output tokens and a three-second timeout. An invalid reply,
error, unavailable classifier or PHI-ineligible classifier uses HEURISTIC. ROUTE
usage counts toward budgets. Empty model slots keep the provider/policy default;
budgets may downgrade either choice. Without a profile, enabled code-level
`AssistantModelRouting.escalation*` supplies the complex default.
A model's **Fallback model** points to its next candidate, including another
provider. Chains are cycle-free and bounded to four hops. Only retryable failures
before any text or tool start from that call permit fallback. Disabled, down,
PHI-ineligible or budget-blocked candidates are skipped. Without explicit links,
only another LOCAL model in the same adapter family is eligible; local requests
never automatically move to cloud. Provider state/signatures are removed when
switching, preserving neutral tool history. Each attempted call is audited as usage.
NETWORK/TIMEOUT marks a provider DOWN for 60 seconds; success marks it UP. Local
health is probed when unknown/stale with a two-second timeout. Provider cards show
health and down-until, failures show fallback totals, and turn records show the
route and final destination. See [UPGRADE.md](UPGRADE.md) for the additive schema.
+26
View File
@@ -1,5 +1,31 @@
# Upgrading qbit-assistant
## P-e: profile routing and provider fallback
Apply the additive **P-e** ALTER TABLE statements at the end of
[docs/schema-postgres.sql](docs/schema-postgres.sql) before deploying. Existing
profiles default to OFF; optional model references and health columns start null.
No routing or fallback configuration is seeded. Existing turns retain unknown
routing/final-provider values; fallback_count defaults to zero.
Profiles now have a **Model routing** section. Select HEURISTIC or CLASSIFIER,
then optional simple/complex registry models and a classifier model. Empty route
slots retain provider/policy selection. CLASSIFIER with no usable classifier
uses HEURISTIC. OFF preserves prior provider selection. With no effective profile,
code-configured escalation supplies the complex model/effort when enabled.
Set a model's **Fallback model** to opt into a cross-provider chain. Saves reject
cycles, more than four edges, or disabled/missing targets. Provider health is
read-only; NETWORK/TIMEOUT opens a 60-second circuit, and successful calls close
it. Local endpoints with stale health get a bounded two-second probe. Automatic
fallback without an explicit link is confined to local models of the same adapter
family. PHI/HIPAA eligibility and budgets apply to every destination.
Classifier usage is recorded with purpose ROUTE and included in budgets. Reports
that filter usage purpose should include ROUTE. Turn screens show route decision,
source, fallback count and final provider; the home shows health and fallback totals.
All configuration remains assistant.admin-only, with AI exposure NONE.
## P-d: data-driven budgets
Create `assistant_budget`, `assistant_budget_override`, `assistant_budget_event`
+25
View File
@@ -395,6 +395,31 @@ AI exposure NONE is not enough, because the assistant acts as the user and an ad
- The home owns the budget cards and 30-day cost/token bars, with user, model,
provider and tool breakdowns, inherited colors and empty states.
### P-e as built
- Profiles select OFF (the migration default), HEURISTIC or CLASSIFIER and optional
simple/complex/classifier registry models. Selection uses ProfileService's existing
specificity. Without a profile, enabled escalation config supplies the complex
model and effort. Empty slots retain provider row/default/policy selection.
- The documented heuristic is a pure rule table: >600 characters, two distinct app
table names, complexity keywords, pending write context, >8 previous turns or a
previous tool limit. This branch has no conversation write-mode field.
- Classifier transport is non-streaming, bounded to four output tokens/three seconds,
and contains only the fixed instruction and user question. PHI guards apply before
dispatch. Errors/unexpected replies use the heuristic. ROUTE usage enters budget
accounting; normal dispatch budgets may override the selected route.
- FallbackChain follows at most four edges on retryable failures, only before any
text/tool start of that call. Candidates are checked for enablement, health,
PHI/HIPAA eligibility and budgets. Without explicit links only the same LOCAL
adapter family is considered. Provider switches rebuild neutral messages and
discard providerState; failed attempts remain visible in usage.
- ProviderHealth persists UP/DOWN, checked-at and a 60-second connection-failure
circuit. Local stale/unknown health uses a bounded two-second probe. Successful
completions close the breaker. Turn rows persist route/source, fallback count and
final provider/model; notices name only the unavailable and answering providers.
- Admin record screens, additive Postgres DDL and the flow-layout home expose these
settings and counters without granting AI access to assistant configuration.
### Phases (each a green PR)
| | Scope | Depends on |
|---|---|---|
+16
View File
@@ -461,3 +461,19 @@ INSERT INTO assistant_budget_usage
SELECT id, user_id, provider_id, model, create_date, cost_micros, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens
FROM assistant_usage
ON CONFLICT (usage_id) DO NOTHING;
-- P-e: profile routing, bounded fallback chains, provider health and turn outcome.
ALTER TABLE assistant_profile ADD COLUMN routing_mode VARCHAR(16) DEFAULT 'OFF'
CHECK (routing_mode IN ('OFF', 'HEURISTIC', 'CLASSIFIER'));
ALTER TABLE assistant_profile ADD COLUMN simple_model_id INTEGER REFERENCES assistant_model(id);
ALTER TABLE assistant_profile ADD COLUMN complex_model_id INTEGER REFERENCES assistant_model(id);
ALTER TABLE assistant_profile ADD COLUMN classifier_model_id INTEGER REFERENCES assistant_model(id);
ALTER TABLE assistant_model ADD COLUMN fallback_model_id INTEGER REFERENCES assistant_model(id);
ALTER TABLE assistant_provider ADD COLUMN health_state VARCHAR(8) CHECK (health_state IN ('UP', 'DOWN'));
ALTER TABLE assistant_provider ADD COLUMN health_checked_at TIMESTAMP WITH TIME ZONE;
ALTER TABLE assistant_provider ADD COLUMN down_until TIMESTAMP WITH TIME ZONE;
ALTER TABLE assistant_turn ADD COLUMN route_decision VARCHAR(16);
ALTER TABLE assistant_turn ADD COLUMN route_source VARCHAR(16);
ALTER TABLE assistant_turn ADD COLUMN fallback_count INTEGER DEFAULT 0;
ALTER TABLE assistant_turn ADD COLUMN provider_id INTEGER REFERENCES assistant_provider(id);
@@ -278,6 +278,26 @@ public class AnthropicModelProvider implements AssistantModelProvider
return (handle);
}
if(Boolean.FALSE.equals(request.getStreaming()))
{
java.util.concurrent.Future<?> completion = executor.submit(() ->
{
try
{
handle.complete(AnthropicMapping.toResponse(client.messages().create(params)));
}
catch(Throwable error)
{
if(!handle.isCancelled())
{
handle.fail(AnthropicMapping.toProviderException(error));
}
}
});
handle.setOnCancel(() -> completion.cancel(true));
return (handle);
}
AtomicReference<ModelUsage> seen = new AtomicReference<>();
handle.setPartialUsage(seen::get);
@@ -51,6 +51,39 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
*******************************************************************************/
class AnthropicModelProviderTest
{
/*******************************************************************************
** Routing requests use a single JSON reply without text deltas.
*******************************************************************************/
@Test
void nonStreamingClassifier() throws Exception
{
server.thenError(200, "classifier-response.json");
ModelRequest request = new ModelRequest().withModel("claude-sonnet-5-5").withMaxTokens(4).withStreaming(false)
.withMessages(List.of(ModelMessage.user("hello")));
List<String> deltas = new ArrayList<>();
ModelResponse response = provider.stream(request, deltas::add).await(Duration.ofSeconds(3));
assertThat(response.getText()).isEqualTo("SIMPLE");
assertThat(response.getUsage().getOutputTokens()).isEqualTo(1L);
assertThat(deltas).isEmpty();
assertThat(sentBody(0).get("stream")).isIn(null, false);
assertThat(sentBody(0).get("max_tokens")).isEqualTo(4);
}
/*******************************************************************************
** Non-streaming errors retain the normal exception mapping.
*******************************************************************************/
@Test
void nonStreamingError() throws Exception
{
server.thenError(401, "unauthorized.json");
assertThatThrownBy(() -> provider.stream(request().withStreaming(false), new StringBuilder()::append).await(Duration.ofSeconds(3)))
.isInstanceOfSatisfying(ModelProviderException.class, error -> assertThat(error.getKind()).isEqualTo(ModelProviderException.Kind.AUTHENTICATION));
}
private static final String TEST_KEY = "sk-ant-test-0000000000000000000000";
private RecordedAnthropicServer server;
@@ -0,0 +1 @@
{"id":"msg_classify","type":"message","role":"assistant","model":"claude-sonnet-5-5","content":[{"type":"text","text":"SIMPLE"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":18,"output_tokens":1}}
@@ -157,7 +157,7 @@ public class OpenAiCompatibleModelProvider implements AssistantModelProvider
{
throw (http.error(response.statusCode(), new String(response.body().readNBytes(64 * 1024), StandardCharsets.UTF_8), result.requestId, strict));
}
if(Boolean.TRUE.equals(caps.getStreaming()))
if(!Boolean.FALSE.equals(request.getStreaming()) && Boolean.TRUE.equals(caps.getStreaming()))
{
readEvents(response.body(), lastRead, result, strict, handle);
}
@@ -52,7 +52,7 @@ final class OpenAiMapping
static ObjectNode request(ModelRequest request, String kind, ModelCapabilities capabilities) throws JsonProcessingException
{
ObjectNode body = OpenAiHttp.JSON.createObjectNode();
boolean streaming = Boolean.TRUE.equals(capabilities.getStreaming());
boolean streaming = !Boolean.FALSE.equals(request.getStreaming()) && Boolean.TRUE.equals(capabilities.getStreaming());
boolean nativeOpenAi = List.of("OPENAI", "AZURE_OPENAI").contains(kind);
String model = request.getModel();
body.put("model", "VERTEX".equals(kind) && model != null && !model.contains("/") ? "google/" + model : model);
@@ -45,6 +45,27 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
*******************************************************************************/
class OpenAiContractTest
{
/*******************************************************************************
** Classifiers opt out of streaming even when the model supports it.
*******************************************************************************/
@Test
void nonStreamingClassifier() throws Exception
{
try(MockOpenAiServer server = new MockOpenAiServer())
{
server.enqueue(200, "{\"model\":\"test\",\"choices\":[{\"message\":{\"content\":\"SIMPLE\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":18,\"completion_tokens\":1}}");
var provider = new OpenAiCompatibleModelProvider(settings(server, "OPENAI"));
ModelResponse response = provider.stream(request().withStreaming(false).withMaxTokens(4), new StringBuilder()::append).await(Duration.ofSeconds(3));
assertThat(response.getText()).isEqualTo("SIMPLE");
JsonNode body = OpenAiHttp.JSON.readTree(server.requests.getFirst().body());
assertThat(body.path("stream").asBoolean()).isFalse();
assertThat(body.path("max_completion_tokens").asInt()).isEqualTo(4);
assertThat(body.has("stream_options")).isFalse();
}
}
static final String TEXT = "data: {\"model\":\"test\",\"choices\":[{\"delta\":{\"content\":\"OK\"},\"finish_reason\":\"stop\"}]}\r\n\r\ndata: [DONE]\r\n\r\n";
@@ -28,9 +28,9 @@ import com.kingsrook.qqq.backend.core.utils.StringUtils;
** from the provider (no date suffix). Effort is sent explicitly on every
** request, because the models' defaults differ (Sonnet 5.5 defaults to high).
**
** Phase 1 uses only the main slot. Escalation (Opus) and the summary model
** (Haiku: titles, compaction) arrive in phase 2; they are configurable now
** so the admin policy shape is stable.
** With no effective profile, enabled escalation supplies the complex route's
** model and effort. A profile's routing mode takes precedence. The summary
** slot remains available for titles and compaction.
*******************************************************************************/
public class AssistantModelRouting implements Serializable
{
@@ -0,0 +1,362 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.engine;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import java.util.function.Supplier;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.classification.DataClassification;
import com.kingsrook.qbits.assistant.model.AssistantModel;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.model.AssistantTurn;
import com.kingsrook.qbits.assistant.provider.ModelEventListener;
import com.kingsrook.qbits.assistant.provider.ModelMessage;
import com.kingsrook.qbits.assistant.provider.ModelProviderException;
import com.kingsrook.qbits.assistant.provider.ModelRequest;
import com.kingsrook.qbits.assistant.provider.ModelResponse;
import com.kingsrook.qqq.backend.core.exceptions.QException;
/*******************************************************************************
** At most four fallback edges per call, never after that call emitted output.
** Explicit chains may cross providers. Implicit alternatives stay local and in
** the same adapter family; PHI, enablement, health and budgets always narrow them.
*******************************************************************************/
public class FallbackChain
{
public static final Integer MAX_HOPS = 4;
private final AssistantRuntime runtime;
private final TurnClassification classification;
private final AssistantTurn turn;
private final Consumer<String> notice;
private final AtomicBoolean output = new AtomicBoolean();
private final ProviderHealth health;
/*******************************************************************************
** The engine owns dispatch, usage, cancellation and strict-schema retries.
*******************************************************************************/
@FunctionalInterface
public interface Dispatch
{
/*******************************************************************************
** Dispatch one selected model through the existing guards.
*******************************************************************************/
ModelResponse call(ProviderService.Active selected) throws QException, ModelProviderException, InterruptedException;
}
/*******************************************************************************
** One chain tracker per turn.
*******************************************************************************/
public FallbackChain(AssistantRuntime runtime, TurnClassification classification, AssistantTurn turn, Consumer<String> notice)
{
this.runtime = runtime;
this.classification = classification;
this.turn = turn;
this.notice = notice;
this.health = new ProviderHealth(runtime);
turn.setFallbackCount(0);
}
/*******************************************************************************
** Start a distinct attempt, including a strict-schema retry.
*******************************************************************************/
public void beforeCall(ProviderService.Active active) throws QException, ModelProviderException, InterruptedException
{
output.set(false);
health.beforeCall(active);
}
/*******************************************************************************
** Count text and tool starts conservatively, even when text is still buffered.
*******************************************************************************/
public ModelEventListener listener(ModelEventListener delegate)
{
return (new ModelEventListener()
{
/*******************************************************************************
** Visible text prevents another provider from replacing this attempt.
*******************************************************************************/
@Override
public void onTextDelta(String text)
{
if(text != null && !text.isEmpty())
{
output.set(true);
}
delegate.onTextDelta(text);
}
/*******************************************************************************
** A partial tool call also prevents retry.
*******************************************************************************/
@Override
public void onToolUseStart(String id, String name)
{
output.set(true);
delegate.onToolUseStart(id, name);
}
});
}
/*******************************************************************************
** Used by strict retries as well as provider fallback.
*******************************************************************************/
public Boolean hasOutput()
{
return (output.get());
}
/*******************************************************************************
** Retain the original route's budget scopes across downgrades and retries.
** A fallback may change the destination, never evade an original STOP.
*******************************************************************************/
public BudgetService.Decision budget(ProviderService.Active origin, ProviderService.Active candidate) throws QException
{
Boolean phi = classification.getClassification() == DataClassification.PHI;
BudgetService.Decision original = runtime.getBudgetService().evaluate(turn.getUserId(), origin, phi, false);
if(original.stopped() || (Objects.equals(origin.providerId(), candidate.providerId()) && Objects.equals(origin.model(), candidate.model())))
{
return (original);
}
BudgetService.Decision destination = runtime.getBudgetService().evaluate(turn.getUserId(), candidate, phi, false);
List<String> notices = new ArrayList<>(original.notices());
notices.addAll(destination.notices());
return (new BudgetService.Decision(destination.stopped(), destination.reason(), destination.active(), notices));
}
/*******************************************************************************
** Resolve candidates lazily after a failure; refresh rows before every attempt.
*******************************************************************************/
public ModelResponse call(ModelRequest request, Supplier<ProviderService.Active> current, Dispatch dispatch) throws QException, ModelProviderException, InterruptedException
{
ProviderService.Active selected = current.get();
String unavailable = null;
Set<Integer> visited = new HashSet<>();
Integer hops = 0;
Boolean explicitChain = null;
while(true)
{
try
{
ModelResponse response = dispatch.call(selected);
if(!ModelResponse.STOP_CANCELLED.equals(response.getStopReason()))
{
health.succeeded(current.get());
if(unavailable != null)
{
notice.accept("Provider " + unavailable + " is unavailable; answered with " + name(current.get()) + ".");
}
}
return (response);
}
catch(ModelProviderException error)
{
ProviderService.Active failed = current.get();
health.failed(failed, error);
if(!error.isRetryable() || output.get() || hops >= MAX_HOPS)
{
throw (error);
}
List<AssistantModel> models = runtime.getStore().listModels();
AssistantModel origin = models.stream().filter(m -> Objects.equals(m.getProviderId(), failed.providerId()) && Objects.equals(m.getModelId(), failed.model())).findFirst().orElse(null);
if(origin != null)
{
visited.add(origin.getId());
}
if(explicitChain == null)
{
explicitChain = origin != null && origin.getFallbackModelId() != null;
}
selected = null;
List<AssistantModel> candidates = explicitChain && (origin == null || origin.getFallbackModelId() == null) ? List.of() : candidates(origin, failed, models);
for(AssistantModel candidate : candidates)
{
if(hops >= MAX_HOPS)
{
break;
}
if(!visited.add(candidate.getId()))
{
continue;
}
hops++;
ProviderService.Active resolved = runtime.getProviderService().resolveModel(candidate.getId(), classification.getClassification() == DataClassification.PHI);
if(resolved == null || health.isDown(resolved.row()) || PhiRoutingGuard.check(classification, resolved.row(), runtime.getConfig().getHipaaMode()).isPresent())
{
continue;
}
BudgetService.Decision budget = runtime.getBudgetService().evaluate(turn.getUserId(), resolved, classification.getClassification() == DataClassification.PHI, false);
if(budget.stopped())
{
continue;
}
if(!explicitChain && !"LOCAL".equals(runtime.getProviderService().location(budget.active().row())))
{
continue;
}
Long cap = ProviderService.capMicros(budget.active().row());
if(health.isDown(budget.active().row()) || (cap != null && runtime.getProviderService().monthToDateCostMicros(budget.active().providerId()) >= cap))
{
continue;
}
selected = budget.active();
budget.notices().forEach(notice);
break;
}
if(selected == null)
{
throw (error);
}
if(unavailable == null)
{
unavailable = name(failed);
}
turn.setFallbackCount(turn.getFallbackCount() + 1);
neutralHistory(request);
}
}
}
/*******************************************************************************
** Explicit edges win. The automatic pool is LOCAL only, same adapter family.
*******************************************************************************/
private List<AssistantModel> candidates(AssistantModel origin, ProviderService.Active failed, List<AssistantModel> models) throws QException
{
List<AssistantModel> result = new ArrayList<>();
if(origin != null && origin.getFallbackModelId() != null)
{
Integer next = origin.getFallbackModelId();
Set<Integer> seen = new HashSet<>();
while(next != null && seen.add(next) && result.size() < MAX_HOPS)
{
Integer id = next;
AssistantModel row = models.stream().filter(m -> id.equals(m.getId())).findFirst().orElse(null);
if(row == null)
{
break;
}
result.add(row);
next = row.getFallbackModelId();
}
}
else if(failed.row() != null && "LOCAL".equals(runtime.getProviderService().location(failed.row())))
{
for(AssistantProvider provider : runtime.getStore().listProviders())
{
if("LOCAL".equals(runtime.getProviderService().location(provider)) && Objects.equals(provider.getProviderType(), failed.providerType()))
{
models.stream().filter(m -> Objects.equals(m.getProviderId(), provider.getId())).forEach(result::add);
}
}
}
return (result);
}
/*******************************************************************************
** Rebuild messages from neutral parts; signatures never cross provider boundaries.
*******************************************************************************/
public static void neutralHistory(ModelRequest request)
{
List<ModelMessage> messages = new ArrayList<>();
for(ModelMessage message : request.getMessages())
{
messages.add(new ModelMessage().withRole(message.getRole()).withParts(new ArrayList<>(message.getParts())).withCacheHint(message.getCacheHint()));
}
request.setMessages(messages);
request.getTools().forEach(tool -> tool.setStrict(false));
}
/*******************************************************************************
** User-visible names only.
*******************************************************************************/
private static String name(ProviderService.Active active)
{
return (active.row() == null ? active.providerType() : active.row().getName());
}
/*******************************************************************************
** Validate the proposed graph (including every row of a batch update).
*******************************************************************************/
public static String validate(AssistantModel origin, Map<Integer, AssistantModel> models)
{
return (validate(origin, models, true));
}
/*******************************************************************************
** Recheck ancestor lengths without preventing an admin from disabling a
** destination. Runtime selection skips disabled models and follows the chain.
*******************************************************************************/
public static String validate(AssistantModel origin, Map<Integer, AssistantModel> models, Boolean requireEnabled)
{
Set<Integer> seen = new HashSet<>();
seen.add(origin.getId());
Integer next = origin.getFallbackModelId();
Integer hops = 0;
while(next != null)
{
if(!seen.add(next))
{
return ("Fallback chain must not contain a cycle.");
}
if(++hops > MAX_HOPS)
{
return ("Fallback chain must have at most four hops.");
}
AssistantModel candidate = models.get(next);
if(candidate == null || (requireEnabled && !Boolean.TRUE.equals(candidate.getEnabled())))
{
return ("Fallback model must exist and be enabled.");
}
next = candidate.getFallbackModelId();
}
return (null);
}
}
@@ -0,0 +1,281 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.engine;
import java.util.stream.Collectors;
import java.util.Objects;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import java.time.Duration;
import java.time.Instant;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Semaphore;
import java.util.function.Consumer;
import java.util.regex.Pattern;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.classification.DataClassification;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
import com.kingsrook.qbits.assistant.model.AssistantPendingAction;
import com.kingsrook.qbits.assistant.model.AssistantProfile;
import com.kingsrook.qbits.assistant.model.AssistantTurn;
import com.kingsrook.qbits.assistant.provider.CacheHint;
import com.kingsrook.qbits.assistant.provider.ModelMessage;
import com.kingsrook.qbits.assistant.provider.ModelProviderException;
import com.kingsrook.qbits.assistant.provider.ModelRequest;
import com.kingsrook.qbits.assistant.provider.ModelResponse;
import com.kingsrook.qbits.assistant.provider.ModelUsage;
import com.kingsrook.qbits.assistant.provider.StreamHandle;
import com.kingsrook.qbits.assistant.provider.SystemBlock;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QCriteriaOperator;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QFilterOrderBy;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter;
import com.kingsrook.qqq.backend.core.model.session.QSession;
/*******************************************************************************
** Selects the initial model from the effective profile. Classifier content is
** strictly the fixed instruction plus the question; no context is attached.
*******************************************************************************/
public class ModelRouter
{
public static final String CLASSIFIER_PROMPT = "Classify the user's request as SIMPLE or COMPLEX. Reply with one word.";
private static final Pattern COMPLEX_WORDS = Pattern.compile("\\b(analyze|compare|why|trend|explain|summarize\\s+across|plan)\\b", Pattern.CASE_INSENSITIVE);
private final AssistantRuntime runtime;
private final ModelUsage usage = new ModelUsage();
private Long costMicros = 0L;
private Integer calls = 0;
/*******************************************************************************
** Routing returns accounting as well as the selected model.
*******************************************************************************/
public record Route(ProviderService.Active active, String effort, ModelUsage usage, Long costMicros, Integer calls)
{
}
/*******************************************************************************
** One router per turn; accounting is never shared between users.
*******************************************************************************/
public ModelRouter(AssistantRuntime runtime)
{
this.runtime = runtime;
}
/*******************************************************************************
** Preserve classifier accounting even when later route resolution fails.
*******************************************************************************/
public Route accounting()
{
return (new Route(null, null, usage, costMicros, calls));
}
/*******************************************************************************
** Each rule is independently sufficient. Boundaries are >600 characters,
** at least two distinct table names, and >8 previous turns. Match words,
** not substrings (for example, planet does not imply plan).
*******************************************************************************/
public static String heuristic(String question, Set<String> tableNames, Boolean writeContext, Integer priorTurns, Boolean previousToolLimit)
{
String text = question == null ? "" : question;
Long mentioned = tableNames.stream().filter(name -> Pattern.compile("(?<![\\p{L}\\p{N}_])" + Pattern.quote(name) + "(?![\\p{L}\\p{N}_])", Pattern.CASE_INSENSITIVE).matcher(text).find()).count();
return (Boolean.TRUE.equals(writeContext) || text.length() > 600 || mentioned >= 2 || COMPLEX_WORDS.matcher(text).find()
|| priorTurns > 8 || Boolean.TRUE.equals(previousToolLimit) ? "COMPLEX" : "SIMPLE");
}
/*******************************************************************************
** OFF preserves row/default/policy selection. With no profile, configured
** escalation supplies the complex slot and its effort when enabled.
*******************************************************************************/
public Route route(QSession session, AssistantConversation conversation, AssistantTurn turn, String question, Map<String, Object> pageContext,
PolicyService.EffectivePolicy policy, TurnClassification classification, RunningTurn running, Instant deadline, Consumer<String> notice) throws QException, InterruptedException
{
AssistantProfile profile = runtime.getProfileService().pick(session, runtime.appNameFor(pageContext));
String mode = profile == null || profile.getRoutingMode() == null ? "OFF" : profile.getRoutingMode();
Boolean escalation = profile == null && Boolean.TRUE.equals(runtime.getConfig().getModels().getEscalationEnabled());
String decision = "OFF";
String source = null;
Integer selectedId = null;
String effort = policy.mainEffort();
if(!"OFF".equals(mode) || escalation)
{
List<AssistantTurn> previous = runtime.getStore().query(AssistantTurn.TABLE_NAME, new QQueryFilter()
.withCriteria("userId", QCriteriaOperator.EQUALS, turn.getUserId()).withCriteria("conversationId", QCriteriaOperator.EQUALS, conversation.getId())
.withCriteria("id", QCriteriaOperator.LESS_THAN, turn.getId()).withOrderBy(new QFilterOrderBy("id", false)).withLimit(9), AssistantTurn.class);
Boolean pending = !runtime.getStore().query(AssistantPendingAction.TABLE_NAME, new QQueryFilter()
.withCriteria("userId", QCriteriaOperator.EQUALS, turn.getUserId()).withCriteria("conversationId", QCriteriaOperator.EQUALS, conversation.getId())
.withCriteria("status", QCriteriaOperator.EQUALS, AssistantPendingAction.STATUS_PENDING).withLimit(1), AssistantPendingAction.class).isEmpty();
String digest = runtime.getContextBuilder().appDigest();
Set<String> tables = runtime.getQInstance().getTables().keySet().stream().filter(digest::contains).collect(Collectors.toSet());
decision = heuristic(question, tables, pending, previous.size(), !previous.isEmpty() && "tool_limit".equals(previous.get(0).getStopReason()));
source = "HEURISTIC";
if("CLASSIFIER".equals(mode) && profile.getClassifierModelId() != null)
{
String classified = classify(profile.getClassifierModelId(), question, turn, policy, classification, running, deadline, notice);
if(classified != null)
{
decision = classified;
source = "CLASSIFIER";
}
}
selectedId = profile == null ? null : ("COMPLEX".equals(decision) ? profile.getComplexModelId() : profile.getSimpleModelId());
}
turn.withRouteDecision(decision).withRouteSource(source);
ProviderService.Active active = selectedId == null ? null : runtime.getProviderService().resolveModel(selectedId, classification.getClassification() == DataClassification.PHI);
if(selectedId != null && active == null)
{
throw (new ProviderService.ProviderSetupException("The routed model is disabled, unavailable or not eligible for this conversation."));
}
if(active == null)
{
TurnDataRouting.beforeResolve(runtime, classification, turn, policy.mainModel());
active = runtime.getProviderService().resolve(policy.mainModel());
if(active != null && escalation && "COMPLEX".equals(decision))
{
String model = runtime.getConfig().getModels().getEscalationModel();
if(active.row() == null)
{
active = new ProviderService.Active(null, active.providerType(), model, active.provider());
}
else
{
active = runtime.getProviderService().forRow(new AssistantProvider(active.row().toQRecord()).withModel(model), model);
}
effort = runtime.getConfig().getModels().getEscalationEffort();
}
}
return (new Route(active, effort, usage, costMicros, calls));
}
/*******************************************************************************
** Buffer a classifier reply, with a three-second deadline and four output
** tokens. Errors and non-exact answers leave the heuristic decision intact.
*******************************************************************************/
private String classify(Integer modelId, String question, AssistantTurn turn, PolicyService.EffectivePolicy policy, TurnClassification classification,
RunningTurn running, Instant deadline, Consumer<String> notice) throws QException, InterruptedException
{
classification.inherit(runtime.getStore().getConversation(turn.getUserId(), turn.getConversationId()));
ProviderService.Active active = runtime.getProviderService().resolveModel(modelId, classification.getClassification() == DataClassification.PHI, false);
if(active == null || PhiRoutingGuard.check(classification, active.row(), runtime.getConfig().getHipaaMode()).isPresent())
{
return (null);
}
BudgetService.Decision budget = runtime.getBudgetService().evaluate(turn.getUserId(), active, classification.getClassification() == DataClassification.PHI, false);
budget.notices().forEach(notice);
PolicyService.BudgetStatus configured = runtime.getPolicyService().checkConfiguredBudget(policy, turn.getUserId(), 0, 0);
if(budget.stopped() || configured.hardExceeded() || !Objects.equals(budget.active().model(), active.model()) || !Objects.equals(budget.active().providerId(), active.providerId()))
{
return (null);
}
Long cap = ProviderService.capMicros(active.row());
if((cap != null && runtime.getProviderService().monthToDateCostMicros(active.providerId()) >= cap) || running.isCancelled() || runtime.getPolicyService().isKillSwitchOn() || !runtime.getClock().instant().isBefore(deadline))
{
return (null);
}
try
{
TurnDataRouting.beforeModelCall(runtime, classification, turn, active);
}
catch(TurnDataRouting.Blocked e)
{
return (null);
}
ModelRequest request = new ModelRequest().withModel(active.model()).withMaxTokens(4).withStreaming(false)
.withSystem(List.of(new SystemBlock(CLASSIFIER_PROMPT, CacheHint.NONE))).withMessages(List.of(ModelMessage.user(question)));
ModelResponse response = null;
ModelProviderException error = null;
StreamHandle handle = null;
Long started = System.nanoTime();
ProviderHealth health = new ProviderHealth(runtime);
Semaphore gate = runtime.getProviderService().gate(active.row());
Boolean acquired = false;
try
{
health.beforeCall(active);
if(gate != null)
{
acquired = gate.tryAcquire();
if(!acquired)
{
throw (new ModelProviderException(ModelProviderException.Kind.RATE_LIMITED, null, "Classifier concurrency limit reached."));
}
}
calls++;
handle = active.provider().stream(request, text ->
{
});
running.setCurrentHandle(handle);
Duration remaining = Duration.between(runtime.getClock().instant(), deadline);
response = handle.await(remaining.compareTo(Duration.ofSeconds(3)) < 0 ? remaining : Duration.ofSeconds(3));
if(!ModelResponse.STOP_CANCELLED.equals(response.getStopReason()))
{
health.succeeded(active);
String answer = response.getText().strip().toUpperCase(Locale.ROOT);
if(List.of("SIMPLE", "COMPLEX").contains(answer) && response.getToolUses().isEmpty())
{
return (answer);
}
}
return (null);
}
catch(ModelProviderException e)
{
error = e;
health.failed(active, e);
return (null);
}
catch(RuntimeException e)
{
error = new ModelProviderException(ModelProviderException.Kind.OTHER, null, "Classifier did not return a usable answer.");
return (null);
}
finally
{
if(acquired)
{
gate.release();
}
if(handle != null)
{
handle.cancel();
}
running.setCurrentHandle(null);
ModelUsage measured = response == null ? (error == null ? null : error.getPartialUsage()) : response.getUsage();
if(measured != null)
{
usage.add(measured);
costMicros += runtime.getProviderService().pricing(active, active.model()).costMicros(measured);
}
UsageRecorder.save(runtime.getStore(), UsageRecorder.row(new UsageRecorder.Call(turn.getUserId(), turn.getConversationId(), turn.getId(), active, UsageRecorder.PURPOSE_ROUTE),
response, error, (System.nanoTime() - started) / 1_000_000, runtime.getProviderService().pricing(active, active.model())));
}
}
}
@@ -215,7 +215,7 @@ public class ProfileService
/*******************************************************************************
**
*******************************************************************************/
private AssistantProfile pick(QSession session, String appName)
public AssistantProfile pick(QSession session, String appName)
{
List<AssistantProfile> enabled = load().profiles().stream().filter(p -> !Boolean.FALSE.equals(p.getEnabled())).toList();
return (enabled.stream()
@@ -0,0 +1,144 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.engine;
import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.FutureTask;
import java.util.concurrent.TimeUnit;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.provider.HealthResult;
import com.kingsrook.qbits.assistant.provider.ModelProviderException;
import com.kingsrook.qbits.assistant.provider.ProviderSettings;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
/*******************************************************************************
** Persisted sixty-second circuit breaker; only local endpoints are probed.
*******************************************************************************/
public class ProviderHealth
{
private final AssistantRuntime runtime;
/*******************************************************************************
** Use the runtime clock for all cache and breaker boundaries.
*******************************************************************************/
public ProviderHealth(AssistantRuntime runtime)
{
this.runtime = runtime;
}
/*******************************************************************************
** Refuse an open circuit; probe unknown or stale local endpoints within two seconds.
*******************************************************************************/
public void beforeCall(ProviderService.Active active) throws QException, ModelProviderException, InterruptedException
{
if(active.providerId() == null)
{
return;
}
AssistantProvider row = runtime.getStore().getProvider(active.providerId());
if(row == null || !Boolean.TRUE.equals(row.getEnabled()))
{
throw (new ModelProviderException(ModelProviderException.Kind.AUTHENTICATION, null, "Provider is disabled."));
}
if(isDown(row))
{
throw (new ModelProviderException(ModelProviderException.Kind.OVERLOADED, null, "Provider circuit is open."));
}
Instant now = runtime.getClock().instant();
if(!"LOCAL".equals(runtime.getProviderService().location(row)) || (row.getHealthCheckedAt() != null && row.getHealthCheckedAt().plusSeconds(60).isAfter(now)))
{
return;
}
ProviderSettings settings = runtime.getProviderService().settings(row).withConnectTimeoutSeconds(2).withRequestTimeoutSeconds(2);
FutureTask<HealthResult> probe = new FutureTask<>(() -> runtime.getConfig().getProviderTypes().get(row.getProviderType()).checkHealth(settings));
Thread.ofVirtual().name("assistant-health").start(probe);
try
{
HealthResult result = probe.get(Duration.ofSeconds(2).toMillis(), TimeUnit.MILLISECONDS);
if(Boolean.FALSE.equals(result.healthy()))
{
throw (new ModelProviderException(ModelProviderException.Kind.NETWORK, null, "Provider health check failed."));
}
save(row.getId(), Boolean.TRUE.equals(result.healthy()) ? "UP" : null, null);
}
catch(java.util.concurrent.ExecutionException | java.util.concurrent.TimeoutException e)
{
throw (new ModelProviderException(ModelProviderException.Kind.TIMEOUT, null, "Provider health check did not complete."));
}
finally
{
probe.cancel(true);
}
}
/*******************************************************************************
** A future downUntil alone keeps the circuit open; expiry permits recovery.
*******************************************************************************/
public Boolean isDown(AssistantProvider row)
{
return (row != null && row.getDownUntil() != null && row.getDownUntil().isAfter(runtime.getClock().instant()));
}
/*******************************************************************************
** Connection errors open the breaker; model errors do not poison the endpoint.
*******************************************************************************/
public void failed(ProviderService.Active active, ModelProviderException error) throws QException
{
if(active.providerId() != null && (error.getKind() == ModelProviderException.Kind.NETWORK || error.getKind() == ModelProviderException.Kind.TIMEOUT))
{
save(active.providerId(), "DOWN", runtime.getClock().instant().plusSeconds(60));
}
}
/*******************************************************************************
** A successful completion closes the breaker.
*******************************************************************************/
public void succeeded(ProviderService.Active active) throws QException
{
if(active.providerId() != null)
{
save(active.providerId(), "UP", null);
}
}
/*******************************************************************************
** Only health columns are changed, under the requesting session.
*******************************************************************************/
private void save(Integer providerId, String state, Instant downUntil) throws QException
{
runtime.getStore().updateFields(AssistantProvider.TABLE_NAME, new QRecord().withValue("id", providerId)
.withValue("healthState", state).withValue("healthCheckedAt", runtime.getClock().instant()).withValue("downUntil", downUntil));
runtime.getProviderService().invalidate();
}
}
@@ -740,6 +740,16 @@ public class ProviderService
** the default. Copy the row so the shared provider cache is never mutated.
*******************************************************************************/
public Active resolveModel(Integer modelId, boolean requiresHipaa) throws QException
{
return (resolveModel(modelId, requiresHipaa, true));
}
/*******************************************************************************
** Classifiers may use models without tool support.
*******************************************************************************/
public Active resolveModel(Integer modelId, Boolean requiresHipaa, Boolean requireTools) throws QException
{
AssistantModel model = store.listModels().stream().filter(m -> modelId.equals(m.getId()) && Boolean.TRUE.equals(m.getEnabled())).findFirst().orElse(null);
if(model == null)
@@ -755,7 +765,7 @@ public class ProviderService
try
{
Active active = forRow(selected, model.getModelId());
return (active.provider().getCapabilities(active.model()) != null && Boolean.FALSE.equals(active.provider().getCapabilities(active.model()).getTools()) ? null : active);
return (requireTools && active.provider().getCapabilities(active.model()) != null && Boolean.FALSE.equals(active.provider().getCapabilities(active.model()).getTools()) ? null : active);
}
catch(ProviderSetupException e)
{
@@ -130,11 +130,10 @@ public class TurnEngine implements Runnable
private long costMicros = 0;
private TurnClassification classification;
private FallbackChain fallback;
private ProviderService.Active active;
private String model;
private ProviderService.Active budgetOrigin;
private Long capMicros;
private long monthCostAtStart = 0;
private volatile boolean finished = false;
@@ -193,19 +192,26 @@ public class TurnEngine implements Runnable
persistPendingEvents();
model = policy.mainModel();
TurnDataRouting.beforeResolve(runtime, classification, turn, policy.mainModel());
active = runtime.getProviderService().resolve(policy.mainModel());
ModelRouter router = new ModelRouter(runtime);
ModelRouter.Route route;
try
{
route = router.route(session, conversation, turn, userText, pageContext, policy, classification, running, deadline, text -> notice("warning", text));
}
finally
{
usage.add(router.accounting().usage());
costMicros += router.accounting().costMicros();
modelCallCount += router.accounting().calls();
}
active = route.active();
fallback = new FallbackChain(runtime, classification, turn, text -> notice("warning", text));
if(active == null)
{
throw (new ModelProviderException(ModelProviderException.Kind.AUTHENTICATION, null, "No usable assistant provider is configured."));
}
model = active.model();
budgetOrigin = active;
capMicros = ProviderService.capMicros(active.row());
if(capMicros != null)
{
monthCostAtStart = runtime.getProviderService().monthToDateCostMicros(active.providerId());
}
AssistantModelProvider provider = active.provider();
ModelCapabilities caps = provider.getCapabilities(model);
@@ -221,7 +227,7 @@ public class TurnEngine implements Runnable
String appName = runtime.appNameFor(pageContext);
ProfileService.Persona persona = runtime.getProfileService().persona(session, appName);
String profileText = ProfileService.instructions(persona, runtime.getProfileService().relevantSkills(persona.profileId(), userText, appName));
ModelRequest request = runtime.getContextBuilder().build(session, model, policy.mainEffort(), history, userText, pageContext, policy.tools(), profileText);
ModelRequest request = runtime.getContextBuilder().build(session, model, route.effort(), history, userText, pageContext, policy.tools(), profileText);
Integer maxPerRequest = active.row() == null ? null : active.row().getMaxOutputTokensPerRequest();
if(maxPerRequest != null && (request.getMaxTokens() == null || request.getMaxTokens() > maxPerRequest))
{
@@ -246,15 +252,7 @@ public class TurnEngine implements Runnable
break;
}
if(capMicros != null && monthCostAtStart + costMicros >= capMicros)
{
status = "ERROR";
errorCode = "BUDGET";
errorText = "The assistant has reached this provider's monthly cost cap. Ask an administrator.";
break;
}
ModelResponse response = callModelWithStrictFallback(provider, request);
ModelResponse response = fallback.call(request, () -> active, selected -> callModelWithStrictFallback(selected, request));
stopReason = response.getStopReason();
if(ModelResponse.STOP_CANCELLED.equals(stopReason))
@@ -351,11 +349,6 @@ public class TurnEngine implements Runnable
}
catch(ModelProviderException e)
{
if(e.getPartialUsage() != null)
{
usage.add(e.getPartialUsage());
costMicros += pricing().costMicros(e.getPartialUsage());
}
status = "ERROR";
errorCode = e.getKind() == ModelProviderException.Kind.TIMEOUT ? "TIMEOUT" : "PROVIDER";
errorText = e.getKind() == ModelProviderException.Kind.TIMEOUT
@@ -369,7 +362,7 @@ public class TurnEngine implements Runnable
errorText += " The provider said: " + UsageRecorder.errorText(e);
}
LOG.warn("Assistant model call failed", logPair("turnId", turn.getId()), logPair("kind", e.getKind()), logPair("httpStatus", e.getHttpStatus()),
logPair("errorType", hipaaSafe(e).getProviderErrorType()), logPair("errorMessage", hipaaSafe(e).getProviderMessage()), logPair("requestId", hipaaSafe(e).getRequestId()));
logPair("errorType", hipaaSafe(e.getProviderErrorType())), logPair("errorMessage", hipaaSafe(e.getProviderMessage())), logPair("requestId", hipaaSafe(e.getRequestId())));
}
catch(InterruptedException e)
{
@@ -514,23 +507,23 @@ public class TurnEngine implements Runnable
** tool schemas, retry once (and for the rest of the turn) with no strict
** tools. The kernel validates every tool input either way.
*******************************************************************************/
private ModelResponse callModelWithStrictFallback(AssistantModelProvider provider, ModelRequest request) throws QException, ModelProviderException, InterruptedException
private ModelResponse callModelWithStrictFallback(ProviderService.Active candidate, ModelRequest request) throws QException, ModelProviderException, InterruptedException
{
try
{
return (callModel(provider, request));
return (callModel(candidate, request));
}
catch(ModelProviderException e)
{
boolean anyStrict = request.getTools().stream().anyMatch(tool -> Boolean.TRUE.equals(tool.getStrict()));
if(!e.isStrictSchemaRejected() || !anyStrict)
if(!e.isStrictSchemaRejected() || !anyStrict || fallback.hasOutput())
{
throw (e);
}
LOG.warn("Assistant provider refused the strict tool schemas; retrying without strict tools", logPair("turnId", turn.getId()), logPair("model", model),
logPair("errorMessage", hipaaSafe(e).getProviderMessage()), logPair("requestId", hipaaSafe(e).getRequestId()));
logPair("errorMessage", hipaaSafe(e.getProviderMessage())), logPair("requestId", hipaaSafe(e.getRequestId())));
request.getTools().forEach(tool -> tool.setStrict(false));
return (callModel(provider, request));
return (callModel(candidate, request));
}
}
@@ -539,11 +532,19 @@ public class TurnEngine implements Runnable
/*******************************************************************************
** One streamed model call.
*******************************************************************************/
private ModelResponse callModel(AssistantModelProvider provider, ModelRequest request) throws QException, ModelProviderException, InterruptedException
private ModelResponse callModel(ProviderService.Active candidate, ModelRequest request) throws QException, ModelProviderException, InterruptedException
{
if(checkStop() != null)
{
return (new ModelResponse().withStopReason(ModelResponse.STOP_CANCELLED));
}
PolicyService.BudgetStatus configuredBudget = runtime.getPolicyService().checkConfiguredBudget(policy, userId, usage.billableTokens(), costMicros);
if(configuredBudget.hardExceeded())
{
throw (new BudgetService.Blocked(configuredBudget.message()));
}
classification.inherit(store.getConversation(userId, conversation.getId()));
BudgetService.Decision budgetDecision = runtime.getBudgetService().evaluate(userId, budgetOrigin,
classification.getClassification() == com.kingsrook.qbits.assistant.classification.DataClassification.PHI, false);
BudgetService.Decision budgetDecision = fallback.budget(budgetOrigin, candidate);
budgetDecision.notices().forEach(text -> notice("warning", text));
if(budgetDecision.stopped())
{
@@ -555,7 +556,7 @@ public class TurnEngine implements Runnable
active = selected;
model = active.model();
request.setModel(model);
request.getMessages().forEach(message -> message.withProviderState(null));
FallbackChain.neutralHistory(request);
request.getTools().forEach(tool -> tool.setStrict(false));
ModelCapabilities capabilities = active.provider().getCapabilities(model);
if(capabilities != null && Boolean.FALSE.equals(capabilities.getParallelToolCalls()))
@@ -567,9 +568,8 @@ public class TurnEngine implements Runnable
{
request.setMaxTokens(maxTokens);
}
capMicros = null;
}
provider = active.provider();
AssistantModelProvider provider = active.provider();
Long selectedCap = ProviderService.capMicros(active.row());
if(selectedCap != null && runtime.getProviderService().monthToDateCostMicros(active.providerId()) >= selectedCap)
{
@@ -599,8 +599,9 @@ public class TurnEngine implements Runnable
try
{
TurnDataRouting.beforeModelCall(runtime, classification, turn, active);
fallback.beforeCall(active);
modelCallCount++;
StreamHandle handle = provider.stream(request.copy(), listener);
StreamHandle handle = provider.stream(request.copy(), fallback.listener(listener));
running.setCurrentHandle(handle);
try
{
@@ -616,6 +617,11 @@ public class TurnEngine implements Runnable
catch(ModelProviderException e)
{
recordUsage(null, e, started);
if(e.getPartialUsage() != null)
{
usage.add(e.getPartialUsage());
costMicros += pricing().costMicros(e.getPartialUsage());
}
throw (e);
}
finally
@@ -982,6 +988,10 @@ public class TurnEngine implements Runnable
.withValue("errorCode", errorCode)
.withValue("cancelReason", running.getCancelReason())
.withValue("model", model)
.withValue("providerId", active == null ? null : active.providerId())
.withValue("routeDecision", turn.getRouteDecision())
.withValue("routeSource", turn.getRouteSource())
.withValue("fallbackCount", turn.getFallbackCount())
.withValue("endedAt", now)
.withValue("toolCallCount", toolCallCount)
.withValue("modelCallCount", modelCallCount)
@@ -39,6 +39,7 @@ public final class UsageRecorder
public static final String PURPOSE_TURN = "TURN";
public static final String PURPOSE_TEST = "TEST";
public static final String PURPOSE_ROUTE = "ROUTE";
public static final String STATUS_OK = "OK";
public static final String STATUS_ERROR = "ERROR";
@@ -17,6 +17,7 @@
package com.kingsrook.qbits.assistant.home;
import com.kingsrook.qbits.assistant.model.AssistantTurn;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.AssistantSecurity;
import com.kingsrook.qbits.assistant.engine.BudgetService;
@@ -73,6 +74,7 @@ public class AssistantHomeData
private final List<BudgetService.Snapshot> budgets = new ArrayList<>();
private final List<Totals> byProvider = new ArrayList<>();
private Long fallbackCount = 0L;
private Boolean hipaaMode;
private Integer hipaaRetentionDays;
private LocalDate today;
@@ -239,6 +241,16 @@ public class AssistantHomeData
// conversations started per day //
////////////////////////////////////////
Instant dayOneAt = firstDay.atStartOfDay().toInstant(ZoneOffset.UTC);
List<QRecord> routedTurns = runtime.getStore().queryRecords(AssistantTurn.TABLE_NAME, new QQueryFilter()
.withCriteria("startedAt", QCriteriaOperator.GREATER_THAN_OR_EQUALS, dayOneAt), max);
data.truncated |= routedTurns.size() >= max;
for(QRecord row : routedTurns)
{
if(row.getValueInteger("fallbackCount") != null)
{
data.fallbackCount += row.getValueInteger("fallbackCount");
}
}
List<QRecord> conversations = runtime.getStore().queryRecords(AssistantConversation.TABLE_NAME, new QQueryFilter()
.withCriteria("createDate", QCriteriaOperator.GREATER_THAN_OR_EQUALS, dayOneAt), max);
data.truncated |= conversations.size() >= max;
@@ -556,4 +568,14 @@ public class AssistantHomeData
{
return (byProvider);
}
/*******************************************************************************
** Provider switches in the displayed period, including recovered failures.
*******************************************************************************/
public Long getFallbackCount()
{
return (fallbackCount);
}
}
@@ -193,6 +193,11 @@ public final class AssistantHomeView
}
html.append("</div><div style=\"").append(CELLS).append("\">");
cell(html, "Type", nvl(p.typeLabel(), "Not set"));
cell(html, "Health", nvl(row.getHealthState(), "Unknown"));
if(row.getDownUntil() != null)
{
cell(html, "Down until", when(row.getDownUntil()));
}
cell(html, "Kind", AssistantRecordScreens.providerKindLabel(p.kind()));
cell(html, "Location", "LOCAL".equals(p.location()) ? "Local" : "Cloud");
if(Boolean.TRUE.equals(row.getHipaaEligible()))
@@ -423,6 +428,7 @@ public final class AssistantHomeView
private static void failures(StringBuilder html, AssistantHomeData data, String prefix)
{
open(html, "Recent failures");
html.append("<p>").append(esc(number(data.getFallbackCount()) + " fallbacks in the last " + AssistantHomeData.DAYS + " days.")).append("</p>");
if(data.getFailedCalls().isEmpty() && data.getFailedTools().isEmpty())
{
empty(html, "No failed model calls or tool calls in the last " + AssistantHomeData.DAYS + " days.");
@@ -92,6 +92,15 @@ public class AssistantProfileCustomizer implements TableCustomizerInterface
*******************************************************************************/
private static void checkProfile(QRecord record, QRecord merged)
{
String routingMode = merged.getValueString("routingMode");
if(routingMode == null)
{
record.setValue("routingMode", "OFF");
}
else if(!List.of("OFF", "HEURISTIC", "CLASSIFIER").contains(routingMode))
{
record.addError(new BadInputStatusMessage("Routing mode must be OFF, HEURISTIC or CLASSIFIER."));
}
String personality = merged.getValueString("personality");
if(personality != null && personality.length() > ProfileService.MAX_PERSONALITY_CHARS)
{
@@ -128,6 +128,7 @@ public final class AssistantRecordScreens
return (List.of(
QPossibleValueSource.newForTable(AssistantBudget.TABLE_NAME),
QPossibleValueSource.newForTable(AssistantModel.TABLE_NAME).withName("assistant_registry_model"),
QPossibleValueSource.newForTable(AssistantModel.TABLE_NAME).withName("assistant_budget_model"),
budgetEnum("scopeType", "INSTANCE", "USER", "ROLE", "PROVIDER", "MODEL"),
budgetEnum("period", "DAY", "MONTH"),
@@ -289,6 +290,7 @@ public final class AssistantRecordScreens
.withValueFormatAndFields(PVSValueFormatAndFields.LABEL_ONLY).withEnumValues(modelValues),
enumSource("assistant_provider_kind", kinds),
enumSource("assistant_provider_location", List.of(new QPossibleValue<>("LOCAL", "Local"), new QPossibleValue<>("CLOUD", "Cloud"))),
enumSource("assistant_routing_mode", List.of(new QPossibleValue<>("OFF", "Off"), new QPossibleValue<>("HEURISTIC", "Heuristic"), new QPossibleValue<>("CLASSIFIER", "Classifier"))),
enumSource("assistant_model_source", List.of(new QPossibleValue<>("DISCOVERED", "Discovered"), new QPossibleValue<>("MANUAL", "Manual"))),
QPossibleValueSource.newForTable(AssistantProvider.TABLE_NAME)));
}
@@ -365,6 +367,7 @@ public final class AssistantRecordScreens
section("connection", "Connection", "key", Tier.T2, "apiKeyEnvName", ProviderService.KEY_STATUS_FIELD, "baseUrl", "tlsCaBundle", "cloudProject"),
section("limits", "Limits", "speed", Tier.T2, "monthlyCostCapUsd", "maxOutputTokensPerRequest", "maxConcurrentRequests"),
section("registry", "Provider settings", "settings", Tier.T2, "kind", "location", "hipaaEligible", "region", "apiVersion", "connectTimeoutSeconds", "requestTimeoutSeconds"),
section("health", "Health", "network_check", Tier.T2, "healthState", "healthCheckedAt", "downUntil"),
section("discovery", "Last discovery", "memory", Tier.T2, "lastDiscoveryAt", "lastDiscoveryStatus", "lastDiscoveryError"),
section("lastTest", "Last connection test", "network_check", Tier.T2, "lastTestAt", "lastTestStatus", "lastTestLatencyMs", "lastTestModel", "lastTestError"));
}
@@ -406,7 +409,7 @@ public final class AssistantRecordScreens
micros(table, "costMicros");
counts(table, "inputTokens", "outputTokens", "cacheReadTokens", "cacheWriteTokens", "latencyMs");
table.getField("latencyMs").withLabel("Latency (ms)");
help(table, "purpose", "TURN: a step of an answer. TEST: a Test connection call.");
help(table, "purpose", "TURN: a step of an answer. ROUTE: a routing classifier call. TEST: a Test connection call.");
sections(table,
section("call", "Model Call", "receipt_long", Tier.T1, "id", "purpose", "status", "model", "providerId", "providerType", "userId", "conversationId", "turnId"),
section("result", "Result", "fact_check", Tier.T2, "latencyMs", "errorKind", "httpStatus", "requestId"),
@@ -467,11 +470,13 @@ public final class AssistantRecordScreens
*******************************************************************************/
private static void turn(QTableMetaData table)
{
link(table, "providerId", AssistantProvider.TABLE_NAME, "Final provider");
link(table, "userMessageId", AssistantMessage.TABLE_NAME, "User Message");
link(table, "assistantMessageId", AssistantMessage.TABLE_NAME, "Assistant Message");
micros(table, "costMicros");
counts(table, "inputTokens", "outputTokens", "cacheReadTokens", "cacheWriteTokens", "toolCallCount", "modelCallCount");
sections(table,
section("routing", "Routing", "route", Tier.T2, "routeDecision", "routeSource", "fallbackCount", "providerId"),
section("turn", "Turn", "sync_alt", Tier.T1, "id", "conversationId", "status", "model", "startedAt", "endedAt", "userMessageId", "assistantMessageId", "userId"),
section("outcome", "Outcome", "flag", Tier.T2, "stopReason", "cancelReason", "errorCode"),
section("usage", "Usage", "query_stats", Tier.T2, "inputTokens", "outputTokens", "cacheReadTokens", "cacheWriteTokens", "costMicros", "toolCallCount", "modelCallCount"),
@@ -520,6 +525,10 @@ public final class AssistantRecordScreens
help(table, "name", "An internal name for this profile.");
help(table, "displayName", "The assistant's name in the panel header and its messages, e.g. Ada.");
help(table, "icon", "A Material icon name for the panel, e.g. auto_awesome, support_agent, school.");
pvs(table, "routingMode", "assistant_routing_mode", "Routing mode");
pvs(table, "simpleModelId", "assistant_registry_model", "Simple model");
pvs(table, "complexModelId", "assistant_registry_model", "Complex model");
pvs(table, "classifierModelId", "assistant_registry_model", "Classifier model");
help(table, "personality", "Tone and style instructions, at most " + ProfileService.MAX_PERSONALITY_CHARS + " characters, e.g. \"Friendly and brief. Use the user's first name.\" Added to the assistant's instructions; it never changes its rules, tools or what the user may do.");
help(table, "greeting", "Shown when a new conversation opens.");
help(table, "starterPrompts", "Suggested first questions, one per line (at most " + ProfileService.MAX_STARTER_PROMPTS + ").");
@@ -528,6 +537,7 @@ public final class AssistantRecordScreens
help(table, "permissionName", "Optional: use this profile only for users who hold this permission (a role).");
sections(table,
section("profile", "Profile", "face", Tier.T1, "id", "name", "displayName", "icon", "enabled", "isDefault"),
section("routing", "Model routing", "route", Tier.T2, "routingMode", "simpleModelId", "complexModelId", "classifierModelId"),
section("voice", "Personality", "record_voice_over", Tier.T2, "personality", "greeting", "starterPrompts"),
section("scope", "Where it applies", "filter_alt", Tier.T2, "appName", "permissionName"));
}
@@ -878,6 +888,7 @@ public final class AssistantRecordScreens
{
link(table, "providerId", AssistantProvider.TABLE_NAME, "Provider");
pvs(table, "source", "assistant_model_source", "Source");
pvs(table, "fallbackModelId", "assistant_registry_model", "Fallback model");
table.getField("contextWindowTokens").withLabel("Context window (tokens)");
table.getField("maxOutputTokens").withLabel("Max output tokens");
table.getField("supportsReasoningEffort").withLabel("Supports reasoning effort");
@@ -885,7 +896,7 @@ public final class AssistantRecordScreens
table.getField("outputUsdPerMTok").withLabel("Output (USD per million tokens)");
table.getField("lastTestLatencyMs").withLabel("Last test latency (ms)");
sections(table,
section("model", "Model", "memory", Tier.T1, "id", "providerId", "modelId", "displayName", "enabled", "source", "discoveredAt"),
section("model", "Model", "memory", Tier.T1, "id", "providerId", "modelId", "displayName", "enabled", "source", "discoveredAt", "fallbackModelId"),
section("capabilities", "Capabilities", "settings", Tier.T2, "supportsTools", "supportsVision", "supportsStreaming", "supportsReasoningEffort", "contextWindowTokens", "maxOutputTokens"),
section("prices", "Prices", "payments", Tier.T2, "inputUsdPerMTok", "outputUsdPerMTok"),
section("lastTest", "Last connection test", "network_check", Tier.T2, "lastTestAt", "lastTestStatus", "lastTestLatencyMs", "lastTestError"));
@@ -17,6 +17,8 @@
package com.kingsrook.qbits.assistant.metadata;
import com.kingsrook.qbits.assistant.model.AssistantModel;
import com.kingsrook.qbits.assistant.engine.FallbackChain;
import java.io.Serializable;
import java.util.List;
import java.util.Map;
@@ -77,7 +79,7 @@ public class AssistantSettingsCustomizer implements TableCustomizerInterface
public List<QRecord> preInsertOrUpdate(AbstractActionInput input, List<QRecord> records, boolean isPreview, Optional<List<QRecord>> oldRecordList) throws QException
{
ProviderService service = service();
if(com.kingsrook.qbits.assistant.model.AssistantModel.TABLE_NAME.equals(tableName(input)))
if(AssistantModel.TABLE_NAME.equals(tableName(input)))
{
for(QRecord record : records)
{
@@ -93,6 +95,7 @@ public class AssistantSettingsCustomizer implements TableCustomizerInterface
}
}
}
validateFallbacks(records, oldRecordList);
return (records);
}
if(service == null || !AssistantProvider.TABLE_NAME.equals(tableName(input)))
@@ -204,6 +207,7 @@ public class AssistantSettingsCustomizer implements TableCustomizerInterface
{
String name = entry.getKey();
boolean system = name.equals("id") || name.equals("modifyDate") || name.startsWith("lastTest") || name.startsWith("lastDiscovery")
|| List.of("healthState", "healthCheckedAt", "downUntil").contains(name)
|| (name.equals("isDefault") && Boolean.FALSE.equals(entry.getValue()));
if(!system)
{
@@ -242,4 +246,54 @@ public class AssistantSettingsCustomizer implements TableCustomizerInterface
return (runtime == null ? null : runtime.getProviderService());
}
/*******************************************************************************
** Validate the resulting graph, including all changes in this save batch.
*******************************************************************************/
private void validateFallbacks(List<QRecord> records, Optional<List<QRecord>> oldRecords) throws QException
{
AssistantRuntime runtime = AssistantRuntime.of(QContext.getQInstance());
if(runtime == null)
{
return;
}
Map<Integer, AssistantModel> models = new java.util.LinkedHashMap<>();
for(AssistantModel model : runtime.getStore().listModels())
{
models.put(model.getId(), model);
}
Optional<Map<Serializable, QRecord>> old = oldRecordListToMap("id", oldRecords);
List<AssistantModel> proposed = new java.util.ArrayList<>();
for(QRecord record : records)
{
QRecord merged = new QRecord();
old.map(m -> m.get(record.getValue("id"))).ifPresent(o -> merged.setValues(new java.util.LinkedHashMap<>(o.getValues())));
record.getValues().forEach(merged::setValue);
AssistantModel model = new AssistantModel(merged);
proposed.add(model);
if(model.getId() != null)
{
models.put(model.getId(), model);
}
}
for(Integer i = 0; i < records.size(); i++)
{
String error = FallbackChain.validate(proposed.get(i), models);
if(error != null)
{
records.get(i).addError(new BadInputStatusMessage(error));
}
}
for(AssistantModel model : models.values())
{
String error = FallbackChain.validate(model, models, false);
if(error != null)
{
records.forEach(record -> record.addError(new BadInputStatusMessage(error)));
break;
}
}
}
}
@@ -32,6 +32,9 @@ public class AssistantModel extends QRecordEntity
{
public static final String TABLE_NAME = "assistant_model";
@QField()
private Integer fallbackModelId;
@QField(isEditable = false, isPrimaryKey = true)
private Integer id;
@@ -764,4 +767,35 @@ public class AssistantModel extends QRecordEntity
return (this);
}
/*******************************************************************************
** Get fallbackModelId.
*******************************************************************************/
public Integer getFallbackModelId()
{
return (this.fallbackModelId);
}
/*******************************************************************************
** Set fallbackModelId.
*******************************************************************************/
public void setFallbackModelId(Integer fallbackModelId)
{
this.fallbackModelId = fallbackModelId;
}
/*******************************************************************************
** With fallbackModelId.
*******************************************************************************/
public AssistantModel withFallbackModelId(Integer fallbackModelId)
{
this.fallbackModelId = fallbackModelId;
return (this);
}
}
@@ -32,6 +32,18 @@ public class AssistantProfile extends QRecordEntity
{
public static final String TABLE_NAME = "assistant_profile";
@QField(maxLength = 16)
private String routingMode = "OFF";
@QField()
private Integer simpleModelId;
@QField()
private Integer complexModelId;
@QField()
private Integer classifierModelId;
@QField(isEditable = false, isPrimaryKey = true)
private Integer id;
@@ -481,4 +493,128 @@ public class AssistantProfile extends QRecordEntity
return (this);
}
/*******************************************************************************
** Get routingMode.
*******************************************************************************/
public String getRoutingMode()
{
return (this.routingMode);
}
/*******************************************************************************
** Set routingMode.
*******************************************************************************/
public void setRoutingMode(String routingMode)
{
this.routingMode = routingMode;
}
/*******************************************************************************
** With routingMode.
*******************************************************************************/
public AssistantProfile withRoutingMode(String routingMode)
{
this.routingMode = routingMode;
return (this);
}
/*******************************************************************************
** Get simpleModelId.
*******************************************************************************/
public Integer getSimpleModelId()
{
return (this.simpleModelId);
}
/*******************************************************************************
** Set simpleModelId.
*******************************************************************************/
public void setSimpleModelId(Integer simpleModelId)
{
this.simpleModelId = simpleModelId;
}
/*******************************************************************************
** With simpleModelId.
*******************************************************************************/
public AssistantProfile withSimpleModelId(Integer simpleModelId)
{
this.simpleModelId = simpleModelId;
return (this);
}
/*******************************************************************************
** Get complexModelId.
*******************************************************************************/
public Integer getComplexModelId()
{
return (this.complexModelId);
}
/*******************************************************************************
** Set complexModelId.
*******************************************************************************/
public void setComplexModelId(Integer complexModelId)
{
this.complexModelId = complexModelId;
}
/*******************************************************************************
** With complexModelId.
*******************************************************************************/
public AssistantProfile withComplexModelId(Integer complexModelId)
{
this.complexModelId = complexModelId;
return (this);
}
/*******************************************************************************
** Get classifierModelId.
*******************************************************************************/
public Integer getClassifierModelId()
{
return (this.classifierModelId);
}
/*******************************************************************************
** Set classifierModelId.
*******************************************************************************/
public void setClassifierModelId(Integer classifierModelId)
{
this.classifierModelId = classifierModelId;
}
/*******************************************************************************
** With classifierModelId.
*******************************************************************************/
public AssistantProfile withClassifierModelId(Integer classifierModelId)
{
this.classifierModelId = classifierModelId;
return (this);
}
}
@@ -71,6 +71,15 @@ public class AssistantProvider extends QRecordEntity
public static final String TABLE_NAME = "assistant_provider";
@QField(isEditable = false, maxLength = 8)
private String healthState;
@QField(isEditable = false)
private Instant healthCheckedAt;
@QField(isEditable = false)
private Instant downUntil;
@QField(isEditable = false, isPrimaryKey = true)
private Integer id;
@@ -1116,4 +1125,97 @@ public class AssistantProvider extends QRecordEntity
this.cloudProject = cloudProject;
return (this);
}
/*******************************************************************************
** Get healthState.
*******************************************************************************/
public String getHealthState()
{
return (this.healthState);
}
/*******************************************************************************
** Set healthState.
*******************************************************************************/
public void setHealthState(String healthState)
{
this.healthState = healthState;
}
/*******************************************************************************
** With healthState.
*******************************************************************************/
public AssistantProvider withHealthState(String healthState)
{
this.healthState = healthState;
return (this);
}
/*******************************************************************************
** Get healthCheckedAt.
*******************************************************************************/
public Instant getHealthCheckedAt()
{
return (this.healthCheckedAt);
}
/*******************************************************************************
** Set healthCheckedAt.
*******************************************************************************/
public void setHealthCheckedAt(Instant healthCheckedAt)
{
this.healthCheckedAt = healthCheckedAt;
}
/*******************************************************************************
** With healthCheckedAt.
*******************************************************************************/
public AssistantProvider withHealthCheckedAt(Instant healthCheckedAt)
{
this.healthCheckedAt = healthCheckedAt;
return (this);
}
/*******************************************************************************
** Get downUntil.
*******************************************************************************/
public Instant getDownUntil()
{
return (this.downUntil);
}
/*******************************************************************************
** Set downUntil.
*******************************************************************************/
public void setDownUntil(Instant downUntil)
{
this.downUntil = downUntil;
}
/*******************************************************************************
** With downUntil.
*******************************************************************************/
public AssistantProvider withDownUntil(Instant downUntil)
{
this.downUntil = downUntil;
return (this);
}
}
@@ -44,6 +44,18 @@ public class AssistantTurn extends QRecordEntity
@QField(isEditable = false, maxLength = 500)
private String classificationTables;
@QField(isEditable = false, maxLength = 16)
private String routeDecision;
@QField(isEditable = false, maxLength = 16)
private String routeSource;
@QField(isEditable = false)
private Integer fallbackCount;
@QField(isEditable = false)
private Integer providerId;
@QField(isEditable = false, isPrimaryKey = true)
private Integer id;
@@ -875,4 +887,128 @@ public class AssistantTurn extends QRecordEntity
/*******************************************************************************
** Get routeDecision.
*******************************************************************************/
public String getRouteDecision()
{
return (this.routeDecision);
}
/*******************************************************************************
** Set routeDecision.
*******************************************************************************/
public void setRouteDecision(String routeDecision)
{
this.routeDecision = routeDecision;
}
/*******************************************************************************
** With routeDecision.
*******************************************************************************/
public AssistantTurn withRouteDecision(String routeDecision)
{
this.routeDecision = routeDecision;
return (this);
}
/*******************************************************************************
** Get routeSource.
*******************************************************************************/
public String getRouteSource()
{
return (this.routeSource);
}
/*******************************************************************************
** Set routeSource.
*******************************************************************************/
public void setRouteSource(String routeSource)
{
this.routeSource = routeSource;
}
/*******************************************************************************
** With routeSource.
*******************************************************************************/
public AssistantTurn withRouteSource(String routeSource)
{
this.routeSource = routeSource;
return (this);
}
/*******************************************************************************
** Get fallbackCount.
*******************************************************************************/
public Integer getFallbackCount()
{
return (this.fallbackCount);
}
/*******************************************************************************
** Set fallbackCount.
*******************************************************************************/
public void setFallbackCount(Integer fallbackCount)
{
this.fallbackCount = fallbackCount;
}
/*******************************************************************************
** With fallbackCount.
*******************************************************************************/
public AssistantTurn withFallbackCount(Integer fallbackCount)
{
this.fallbackCount = fallbackCount;
return (this);
}
/*******************************************************************************
** Get providerId.
*******************************************************************************/
public Integer getProviderId()
{
return (this.providerId);
}
/*******************************************************************************
** Set providerId.
*******************************************************************************/
public void setProviderId(Integer providerId)
{
this.providerId = providerId;
}
/*******************************************************************************
** With providerId.
*******************************************************************************/
public AssistantTurn withProviderId(Integer providerId)
{
this.providerId = providerId;
return (this);
}
}
@@ -36,6 +36,7 @@ public class ModelRequest
private ModelCapabilities capabilities;
private Boolean streaming = true;
private String model;
private String effort;
private Integer maxTokens;
@@ -52,6 +53,7 @@ public class ModelRequest
public ModelRequest copy()
{
return (new ModelRequest()
.withStreaming(streaming)
.withModel(model)
.withCapabilities(capabilities)
.withEffort(effort)
@@ -311,4 +313,25 @@ public class ModelRequest
return (this);
}
/*******************************************************************************
** Whether to request streaming transport (classifiers use one JSON reply).
*******************************************************************************/
public Boolean getStreaming()
{
return (streaming);
}
/*******************************************************************************
** Select streaming transport independently of model capability.
*******************************************************************************/
public ModelRequest withStreaming(Boolean streaming)
{
this.streaming = streaming;
return (this);
}
}
@@ -0,0 +1,784 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.engine;
import com.kingsrook.qbits.assistant.model.AssistantPendingAction;
import java.math.BigDecimal;
import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.time.ZoneOffset;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.classification.DataClassification;
import com.kingsrook.qbits.assistant.home.AssistantHomeData;
import com.kingsrook.qbits.assistant.home.AssistantHomeView;
import com.kingsrook.qbits.assistant.model.AssistantBudget;
import com.kingsrook.qbits.assistant.model.AssistantModel;
import com.kingsrook.qbits.assistant.model.AssistantProfile;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.model.AssistantTurn;
import com.kingsrook.qbits.assistant.model.AssistantUsage;
import com.kingsrook.qbits.assistant.persistence.AssistantStore;
import com.kingsrook.qbits.assistant.provider.AssistantModelProvider;
import com.kingsrook.qbits.assistant.provider.AssistantProviderType;
import com.kingsrook.qbits.assistant.provider.CompletableStreamHandle;
import com.kingsrook.qbits.assistant.provider.HealthResult;
import com.kingsrook.qbits.assistant.provider.ModelEventListener;
import com.kingsrook.qbits.assistant.provider.ModelPart;
import com.kingsrook.qbits.assistant.provider.ModelProviderException;
import com.kingsrook.qbits.assistant.provider.ModelRequest;
import com.kingsrook.qbits.assistant.provider.ModelResponse;
import com.kingsrook.qbits.assistant.provider.ProviderSettings;
import com.kingsrook.qbits.assistant.provider.StreamHandle;
import com.kingsrook.qbits.assistant.provider.fake.ScriptedModelProvider;
import com.kingsrook.qbits.assistant.testing.AssistantTestFixture;
import com.kingsrook.qbits.assistant.testing.BaseAssistantTest;
import com.kingsrook.qbits.assistant.testing.TurnRunner;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.CsvSource;
import org.junit.jupiter.params.provider.EnumSource;
import org.junit.jupiter.params.provider.ValueSource;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/*******************************************************************************
** Multi-provider dispatch through the fixture with persisted rows and fixed time.
*******************************************************************************/
class RoutingFallbackTest extends BaseAssistantTest
{
private static final String USER = "editor-routing";
private static final Instant NOW = Instant.now();
private final Map<String, ScriptedModelProvider> scripts = new LinkedHashMap<>();
private Integer probes = 0;
private Boolean healthy = true;
/*******************************************************************************
** Keyless adapter with distinct scripts for each endpoint.
*******************************************************************************/
@Override
protected void customizeConfig(AssistantQBitConfig value)
{
value.withProviderType(new RoutingType()).withAllowedBaseUrls(Set.of("https://a.invalid", "https://b.invalid", "https://c.invalid"));
}
/*******************************************************************************
** Fix time for budgets, routing and health.
*******************************************************************************/
@BeforeEach
void fixedClock()
{
config.getRuntime().close();
config.setRuntime(new AssistantRuntime(qInstance, config, Clock.fixed(NOW, ZoneOffset.UTC)));
}
/*******************************************************************************
** One adapter family with separate recording instances.
*******************************************************************************/
class RoutingType implements AssistantProviderType
{
/*******************************************************************************
** Type identifier.
*******************************************************************************/
public String getType()
{
return ("routing-test");
}
/*******************************************************************************
** Display label.
*******************************************************************************/
public String getLabel()
{
return ("Routing test");
}
/*******************************************************************************
** Registry supplies actual models.
*******************************************************************************/
public List<String> getModels()
{
return (List.of("a", "b", "c"));
}
/*******************************************************************************
** No secret needed.
*******************************************************************************/
public String getDefaultApiKeyEnvName()
{
return (null);
}
/*******************************************************************************
** No network traffic.
*******************************************************************************/
public String getDefaultBaseUrl()
{
return ("https://a.invalid");
}
/*******************************************************************************
** Resolve the endpoint's own script.
*******************************************************************************/
public AssistantModelProvider create(String key, String url)
{
return (scripts.get(url == null ? "https://a.invalid" : url));
}
/*******************************************************************************
** Count probes and assert the transport timeout.
*******************************************************************************/
public HealthResult checkHealth(ProviderSettings settings)
{
probes++;
assertThat(settings.getRequestTimeoutSeconds()).isEqualTo(2);
return (new HealthResult(healthy, 1L, "Test probe"));
}
}
/*******************************************************************************
** Admin setup; actual turns use the editor session.
*******************************************************************************/
private AssistantStore store()
{
QContext.init(qInstance, AssistantTestFixture.ownerSession("admin-routing", config));
return (config.getRuntime().getStore());
}
/*******************************************************************************
** One enabled registry model and endpoint.
*******************************************************************************/
private AssistantModel model(String name, String location, Boolean eligible) throws Exception
{
scripts.put("https://" + name + ".invalid", new ScriptedModelProvider(name));
AssistantProvider row = store().insert(new AssistantProvider().withName(name.toUpperCase()).withProviderType("routing-test")
.withBaseUrl("https://" + name + ".invalid").withModel(name).withEnabled(true).withIsDefault("a".equals(name)).withLocation(location).withHipaaEligible(eligible));
return (store().insert(new AssistantModel().withProviderId(row.getId()).withModelId(name).withEnabled(true)));
}
/*******************************************************************************
** Retrieve a recording endpoint.
*******************************************************************************/
private ScriptedModelProvider script(String name)
{
return (scripts.get("https://" + name + ".invalid"));
}
/*******************************************************************************
** Save through the real validation path.
*******************************************************************************/
private void link(AssistantModel from, AssistantModel to) throws Exception
{
store().updateFields(AssistantModel.TABLE_NAME, new QRecord().withValue("id", from.getId()).withValue("fallbackModelId", to.getId()));
}
/*******************************************************************************
** Default profile for this fixture.
*******************************************************************************/
private void profile(String mode, AssistantModel simple, AssistantModel complex, AssistantModel classifier) throws Exception
{
store().insert(new AssistantProfile().withName("routing").withDisplayName("Assistant").withIsDefault(true).withEnabled(true).withRoutingMode(mode)
.withSimpleModelId(simple.getId()).withComplexModelId(complex.getId()).withClassifierModelId(classifier == null ? null : classifier.getId()));
}
/*******************************************************************************
** Read all terminal fields.
*******************************************************************************/
private AssistantTurn turn(TurnRunner.Outcome outcome) throws Exception
{
return (store().getTurn(USER, outcome.turnId()));
}
/*******************************************************************************
** Each keyword and negative word-boundary case.
*******************************************************************************/
@ParameterizedTest
@CsvSource({"analyze sales,COMPLEX", "compare months,COMPLEX", "why,COMPLEX", "trend,COMPLEX", "explain this,COMPLEX", "summarize across stores,COMPLEX", "plan work,COMPLEX", "planet,SIMPLE", "hello,SIMPLE", "blog_post page,COMPLEX", "blog_post blog_post,SIMPLE", "blog_posts,SIMPLE"})
void heuristicWords(String question, String expected)
{
assertThat(ModelRouter.heuristic(question, Set.of("blog_post", "page"), false, 8, false)).isEqualTo(expected);
}
/*******************************************************************************
** Exact length/turn boundaries, pending writes, and tool exhaustion.
*******************************************************************************/
@Test
void heuristicBoundaries()
{
assertThat(ModelRouter.heuristic("x".repeat(600), Set.of(), false, 8, false)).isEqualTo("SIMPLE");
assertThat(ModelRouter.heuristic("x".repeat(601), Set.of(), false, 8, false)).isEqualTo("COMPLEX");
assertThat(ModelRouter.heuristic("hi", Set.of(), true, 0, false)).isEqualTo("COMPLEX");
assertThat(ModelRouter.heuristic("hi", Set.of(), false, 9, false)).isEqualTo("COMPLEX");
assertThat(ModelRouter.heuristic("hi", Set.of(), false, 0, true)).isEqualTo("COMPLEX");
assertThat(ModelRouter.heuristic(null, Set.of(), false, 0, false)).isEqualTo("SIMPLE");
}
/*******************************************************************************
** Exact classifier replies select a route; garbage uses the heuristic.
*******************************************************************************/
@ParameterizedTest
@CsvSource({"SIMPLE,a,CLASSIFIER", "COMPLEX,b,CLASSIFIER", "garbage,b,HEURISTIC"})
void classifierAnswers(String answer, String selected, String source) throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
AssistantModel c = model("c", "CLOUD", true);
store().updateFields(AssistantModel.TABLE_NAME, new QRecord().withValue("id", c.getId()).withValue("supportsTools", false));
profile("CLASSIFIER", a, b, c);
script("c").thenText(answer);
script(selected).thenText("Answer");
TurnRunner.Outcome result = runner.run(USER, "explain this");
assertThat(result.text()).isEqualTo("Answer");
assertThat(turn(result).getModel()).isEqualTo(selected);
assertThat(turn(result).getRouteSource()).isEqualTo(source);
ModelRequest request = script("c").getRequests().get(0);
assertThat(request.getMaxTokens()).isEqualTo(4);
assertThat(request.getStreaming()).isFalse();
assertThat(request.getTools()).isEmpty();
assertThat(request.getSystem()).hasSize(1);
assertThat(request.getSystem().get(0).getText()).isEqualTo(ModelRouter.CLASSIFIER_PROMPT);
assertThat(request.getMessages()).hasSize(1);
assertThat(request.getMessages().get(0).getParts()).hasSize(1);
assertThat(request.getMessages().get(0).getParts().get(0).getText()).isEqualTo("explain this");
List<AssistantUsage> rows = store().query(AssistantUsage.TABLE_NAME, new QQueryFilter(), AssistantUsage.class);
assertThat(rows).extracting(AssistantUsage::getPurpose).containsExactly("ROUTE", "TURN");
assertThat(turn(result).getInputTokens()).isEqualTo(2400L);
}
/*******************************************************************************
** Error categories leave the heuristic in charge and record ROUTE failures.
*******************************************************************************/
@ParameterizedTest
@EnumSource(value = ModelProviderException.Kind.class, names = {"TIMEOUT", "SERVER"})
void classifierErrors(ModelProviderException.Kind kind) throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
AssistantModel c = model("c", "CLOUD", true);
profile("CLASSIFIER", a, b, c);
script("c").thenError(kind, "unavailable");
script("b").thenText("Answer");
TurnRunner.Outcome result = runner.run(USER, "why");
assertThat(turn(result).getRouteSource()).isEqualTo("HEURISTIC");
assertThat(result.text()).isEqualTo("Answer");
assertThat(store().query(AssistantUsage.TABLE_NAME, new QQueryFilter(), AssistantUsage.class)).anyMatch(row -> "ROUTE".equals(row.getPurpose()) && "ERROR".equals(row.getStatus()));
}
/*******************************************************************************
** Real wait timeout cancels the classifier and opens its circuit.
*******************************************************************************/
@Test
void classifierDeadline() throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
AssistantModel c = model("c", "CLOUD", true);
profile("CLASSIFIER", a, b, c);
script("c").withDelay(Duration.ofSeconds(10)).thenText("SIMPLE");
script("b").thenText("Answer");
Long start = System.nanoTime();
TurnRunner.Outcome result = runner.run(USER, "why");
assertThat(Duration.ofNanos(System.nanoTime() - start)).isLessThan(Duration.ofSeconds(8));
assertThat(turn(result).getRouteSource()).isEqualTo("HEURISTIC");
assertThat(store().getProvider(c.getProviderId()).getHealthState()).isEqualTo("DOWN");
}
/*******************************************************************************
** PHI cannot reach an ineligible classifier.
*******************************************************************************/
@Test
void phiSkipsClassifier() throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
AssistantModel c = model("c", "CLOUD", false);
profile("CLASSIFIER", a, b, c);
config.withDataClassification("blog_post", DataClassification.PHI);
script("b").thenText("Answer");
TurnRunner.Outcome result = runner.run(USER, null, "why", Map.of("tableName", "blog_post"));
assertThat(script("c").getRequests()).isEmpty();
assertThat(turn(result).getRouteSource()).isEqualTo("HEURISTIC");
}
/*******************************************************************************
** OFF preserves row selection even for a complex question.
*******************************************************************************/
@Test
void offKeepsDefault() throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
profile("OFF", b, b, null);
script("a").thenText("Default");
TurnRunner.Outcome result = runner.run(USER, "why");
assertThat(result.text()).isEqualTo("Default");
assertThat(turn(result).getRouteDecision()).isEqualTo("OFF");
assertThat(turn(result).getRouteSource()).isNull();
}
/*******************************************************************************
** Every retryable category records both attempts and the final destination.
*******************************************************************************/
@ParameterizedTest
@EnumSource(value = ModelProviderException.Kind.class, names = {"NETWORK", "TIMEOUT", "OVERLOADED", "SERVER", "RATE_LIMITED"})
void fallbackBeforeOutput(ModelProviderException.Kind kind) throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
link(a, b);
script("a").thenError(kind, "unavailable");
script("b").thenText("Recovered");
TurnRunner.Outcome result = runner.run(USER, "hello");
assertThat(result.text()).isEqualTo("Recovered");
assertThat(turn(result).getFallbackCount()).isEqualTo(1);
assertThat(turn(result).getProviderId()).isEqualTo(b.getProviderId());
assertThat(turn(result).getModel()).isEqualTo("b");
assertThat(result.events(TurnEvent.NOTICE).toString()).contains("Provider A is unavailable; answered with B.");
List<AssistantUsage> rows = store().query(AssistantUsage.TABLE_NAME, new QQueryFilter(), AssistantUsage.class);
assertThat(rows).extracting(AssistantUsage::getStatus).containsExactly("ERROR", "OK");
assertThat(rows.get(0).getErrorKind()).isEqualTo(kind.name());
assertThat(turn(result).getInputTokens()).isEqualTo(2400L);
}
/*******************************************************************************
** Permanent errors cannot use another provider.
*******************************************************************************/
@ParameterizedTest
@EnumSource(value = ModelProviderException.Kind.class, names = {"INVALID_REQUEST", "AUTHENTICATION", "OTHER"})
void noFallbackForPermanentError(ModelProviderException.Kind kind) throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
link(a, b);
script("a").thenError(kind, "bad request");
TurnRunner.Outcome result = runner.run(USER, "hello");
assertThat(result.end()).containsEntry("status", "ERROR");
assertThat(script("b").getRequests()).isEmpty();
}
/*******************************************************************************
** Three models, two failures, one final answer.
*******************************************************************************/
@Test
void threeModelChain() throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
AssistantModel c = model("c", "CLOUD", true);
link(a, b);
link(b, c);
script("a").thenError(ModelProviderException.Kind.OVERLOADED, "busy");
script("b").thenError(ModelProviderException.Kind.SERVER, "busy");
script("c").thenText("Recovered");
TurnRunner.Outcome result = runner.run(USER, "hello");
assertThat(result.text()).isEqualTo("Recovered");
assertThat(turn(result).getFallbackCount()).isEqualTo(2);
}
/*******************************************************************************
** Real save validation rejects cycles and disabled destinations.
*******************************************************************************/
@Test
void rejectsInvalidChains() throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
link(a, b);
assertThatThrownBy(() -> link(b, a)).hasMessageContaining("cycle");
assertThatThrownBy(() -> link(a, a)).hasMessageContaining("cycle");
AssistantModel c = model("c", "CLOUD", true);
store().updateFields(AssistantModel.TABLE_NAME, new QRecord().withValue("id", c.getId()).withValue("enabled", false));
assertThatThrownBy(() -> link(a, c)).hasMessageContaining("enabled");
Map<Integer, AssistantModel> graph = new LinkedHashMap<>();
for(Integer i = 1; i <= 6; i++)
{
graph.put(i, new AssistantModel().withId(i).withEnabled(true).withFallbackModelId(i == 6 ? null : i + 1));
}
assertThat(FallbackChain.validate(graph.get(1), graph)).contains("four hops");
assertThat(FallbackChain.validate(graph.get(2), graph)).isNull();
}
/*******************************************************************************
** Skip unavailable intermediate endpoints but follow their next edge.
*******************************************************************************/
@ParameterizedTest
@ValueSource(strings = {"disabled", "down", "disabledModel"})
void skipsUnavailableCandidate(String state) throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
AssistantModel c = model("c", "CLOUD", true);
link(a, b);
link(b, c);
if("disabledModel".equals(state))
{
store().updateFields(AssistantModel.TABLE_NAME, new QRecord().withValue("id", b.getId()).withValue("enabled", false));
}
else
{
store().updateFields(AssistantProvider.TABLE_NAME, "disabled".equals(state)
? new QRecord().withValue("id", b.getProviderId()).withValue("enabled", false)
: new QRecord().withValue("id", b.getProviderId()).withValue("healthState", "DOWN").withValue("downUntil", NOW.plusSeconds(60)));
}
script("a").thenError(ModelProviderException.Kind.SERVER, "busy");
script("c").thenText("Recovered");
assertThat(runner.run(USER, "hello").text()).isEqualTo("Recovered");
assertThat(script("b").getRequests()).isEmpty();
}
/*******************************************************************************
** Explicit links never override PHI eligibility.
*******************************************************************************/
@Test
void phiNeverFallsToIneligible() throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", false);
link(a, b);
config.withDataClassification("blog_post", DataClassification.PHI);
script("a").thenError(ModelProviderException.Kind.OVERLOADED, "busy");
TurnRunner.Outcome result = runner.run(USER, null, "hello", Map.of("tableName", "blog_post"));
assertThat(result.end()).containsEntry("status", "ERROR");
assertThat(script("b").getRequests()).isEmpty();
}
/*******************************************************************************
** Automatic alternatives stay local, even when cloud is available.
*******************************************************************************/
@ParameterizedTest
@ValueSource(strings = {"LOCAL", "CLOUD"})
void implicitLocalOnly(String destination) throws Exception
{
model("a", "LOCAL", true);
model("b", destination, true);
script("a").thenError(ModelProviderException.Kind.OVERLOADED, "busy");
script("b").thenText("Recovered");
TurnRunner.Outcome result = runner.run(USER, "hello");
assertThat(result.end()).containsEntry("status", "LOCAL".equals(destination) ? "COMPLETE" : "ERROR");
assertThat(script("b").getRequests()).hasSize("LOCAL".equals(destination) ? 1 : 0);
}
/*******************************************************************************
** Destination STOP skips a candidate and continues to its successor.
*******************************************************************************/
@Test
void budgetSkipsCandidate() throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
AssistantModel c = model("c", "CLOUD", true);
link(a, b);
link(b, c);
store().insert(new AssistantBudget().withName("Stop B").withScopeType("MODEL").withScopeValue(String.valueOf(b.getId()))
.withPeriod("DAY").withMetric("TOKENS").withLimitValue(BigDecimal.ONE).withAtLimitAction("STOP"));
store().insert(new AssistantUsage().withUserId(USER).withProviderId(b.getProviderId()).withModel("b").withPurpose("TURN").withStatus("OK").withInputTokens(100L).withOutputTokens(0L).withCreateDate(NOW));
script("a").thenError(ModelProviderException.Kind.SERVER, "busy");
script("c").thenText("Recovered");
assertThat(runner.run(USER, "hello").text()).isEqualTo("Recovered");
assertThat(script("b").getRequests()).isEmpty();
}
/*******************************************************************************
** Circuit expiry, cached probes and successful recovery.
*******************************************************************************/
@Test
void healthCircuitAndProbeCache() throws Exception
{
AssistantModel a = model("a", "LOCAL", true);
ProviderService.Active active = config.getRuntime().getProviderService().resolveModel(a.getId(), false);
ProviderHealth health = new ProviderHealth(config.getRuntime());
health.beforeCall(active);
health.beforeCall(active);
assertThat(probes).isEqualTo(1);
health.failed(active, new ModelProviderException(ModelProviderException.Kind.NETWORK, null, "offline"));
assertThat(store().getProvider(a.getProviderId()).getDownUntil()).isEqualTo(NOW.plusSeconds(60));
assertThatThrownBy(() -> health.beforeCall(active)).isInstanceOf(ModelProviderException.class);
store().updateFields(AssistantProvider.TABLE_NAME, new QRecord().withValue("id", a.getProviderId()).withValue("downUntil", NOW).withValue("healthCheckedAt", NOW.minusSeconds(61)));
health.beforeCall(active);
assertThat(probes).isEqualTo(2);
assertThat(store().getProvider(a.getProviderId()).getHealthState()).isEqualTo("UP");
health.failed(active, new ModelProviderException(ModelProviderException.Kind.TIMEOUT, null, "offline"));
health.succeeded(active);
assertThat(store().getProvider(a.getProviderId()).getDownUntil()).isNull();
}
/*******************************************************************************
** Provider-private signatures from a successful tool round cannot reach B.
*******************************************************************************/
@Test
void midTurnSwitchRebuildsNeutralHistory() throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
link(a, b);
script("a").thenRespond(request -> new ModelResponse().withParts(new ArrayList<>(List.of(ModelPart.toolUse("tool-1", "count_records", input("table", "blog_post")))))
.withProviderState("private-signature").withStopReason(ModelResponse.STOP_TOOL_USE))
.thenError(ModelProviderException.Kind.OVERLOADED, "busy");
script("b").thenRespond(request ->
{
assertThat(request.getMessages()).allMatch(message -> message.getProviderState() == null);
assertThat(request.getMessages()).anyMatch(message -> message.getParts().stream().anyMatch(part -> part.getType() == ModelPart.Type.TOOL_RESULT));
return (new ModelResponse().withParts(new ArrayList<>(List.of(ModelPart.text("Recovered after tool")))).withStopReason(ModelResponse.STOP_END_TURN));
});
TurnRunner.Outcome result = runner.run(USER, "count posts");
assertThat(result.text()).isEqualTo("Recovered after tool");
assertThat(result.toolCalls()).hasSize(1);
assertThat(turn(result).getFallbackCount()).isEqualTo(1);
assertThat(script("a").getRequests().get(1).getMessages()).anyMatch(message -> "private-signature".equals(message.getProviderState()));
}
/*******************************************************************************
** Even buffered text or a tool start makes the failure terminal.
*******************************************************************************/
@ParameterizedTest
@ValueSource(strings = {"text", "tool"})
void partialOutputNeverFallsBack(String output) throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
link(a, b);
scripts.put("https://a.invalid", new ScriptedModelProvider("a")
{
/*******************************************************************************
** Emit a partial event, then fail the same call.
*******************************************************************************/
@Override
public StreamHandle stream(ModelRequest request, ModelEventListener listener)
{
if("text".equals(output))
{
listener.onTextDelta("partial");
}
else
{
listener.onToolUseStart("tool-1", "count_records");
}
CompletableStreamHandle handle = new CompletableStreamHandle();
handle.fail(new ModelProviderException(ModelProviderException.Kind.NETWORK, null, "offline"));
return (handle);
}
});
TurnRunner.Outcome result = runner.run(USER, "hello");
assertThat(result.end()).containsEntry("status", "ERROR");
assertThat(script("b").getRequests()).isEmpty();
assertThat(turn(result).getFallbackCount()).isZero();
assertThat(store().getProvider(a.getProviderId()).getHealthState()).isEqualTo("DOWN");
if("text".equals(output))
{
assertThat(result.text()).isEqualTo("partial");
}
}
/*******************************************************************************
** Failed local probes open the circuit without sending model requests.
*******************************************************************************/
@Test
void failedProbeAndHomeHealth() throws Exception
{
AssistantModel a = model("a", "LOCAL", true);
healthy = false;
TurnRunner.Outcome result = runner.run(USER, "hello");
assertThat(result.end()).containsEntry("status", "ERROR");
assertThat(script("a").getRequests()).isEmpty();
assertThat(store().getProvider(a.getProviderId()).getHealthState()).isEqualTo("DOWN");
String html = AssistantHomeView.render(AssistantHomeData.load(config.getRuntime()), "");
assertThat(html).contains("DOWN", "Down until", "0 fallbacks");
}
/*******************************************************************************
** Classifier tokens count before the answer call can dispatch.
*******************************************************************************/
@Test
void classifierConsumesBudget() throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
AssistantModel c = model("c", "CLOUD", true);
profile("CLASSIFIER", a, b, c);
store().insert(new AssistantBudget().withName("Small allowance").withScopeType("USER").withScopeValue(USER)
.withPeriod("DAY").withMetric("TOKENS").withLimitValue(BigDecimal.valueOf(100)).withAtLimitAction("STOP"));
script("c").thenText("SIMPLE");
TurnRunner.Outcome result = runner.run(USER, "hello");
assertThat(result.end()).containsEntry("status", "ERROR");
assertThat(script("a").getRequests()).isEmpty();
assertThat(store().query(AssistantUsage.TABLE_NAME, new QQueryFilter(), AssistantUsage.class)).extracting(AssistantUsage::getPurpose).containsExactly("ROUTE");
}
/*******************************************************************************
** Conversation context supplies tool exhaustion, pending writes and history.
*******************************************************************************/
@ParameterizedTest
@ValueSource(strings = {"tool_limit", "pending", "history"})
void heuristicConversationContext(String context) throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
profile("HEURISTIC", a, b, null);
script("a").thenText("Simple");
TurnRunner.Outcome first = runner.run(USER, "hello");
assertThat(turn(first).getRouteDecision()).isEqualTo("SIMPLE");
if("tool_limit".equals(context))
{
store().updateFields(AssistantTurn.TABLE_NAME, new QRecord().withValue("id", first.turnId()).withValue("stopReason", "tool_limit"));
}
else if("pending".equals(context))
{
store().insert(new AssistantPendingAction().withUserId(USER).withConversationId(first.conversationId()).withStatus("PENDING"));
}
else
{
for(Integer i = 0; i < 8; i++)
{
store().insert(new AssistantTurn().withUserId(USER).withConversationId(first.conversationId()).withStatus("COMPLETE"));
}
}
script("b").thenText("Complex");
TurnRunner.Outcome result = runner.run(USER, first.conversationId(), "hello", null);
assertThat(result.text()).isEqualTo("Complex");
assertThat(turn(result).getRouteDecision()).isEqualTo("COMPLEX");
}
/*******************************************************************************
** Without a profile, escalation supplies the complex model and effort.
*******************************************************************************/
@Test
void codeLevelEscalation() throws Exception
{
config.getModels().withEscalationEnabled(true).withEscalationModel("complex-default").withEscalationEffort("high");
provider.thenText("Complex");
TurnRunner.Outcome result = runner.run(USER, "why");
assertThat(result.text()).isEqualTo("Complex");
assertThat(provider.getRequests().get(0).getModel()).isEqualTo("complex-default");
assertThat(provider.getRequests().get(0).getEffort()).isEqualTo("high");
}
/*******************************************************************************
** Four edges are allowed; a fifth is rejected even when added at the tail.
*******************************************************************************/
@Test
void fourHopBoundary() throws Exception
{
AssistantModel a = model("a", "CLOUD", true);
AssistantModel b = model("b", "CLOUD", true);
AssistantModel c = model("c", "CLOUD", true);
AssistantModel d = store().insert(new AssistantModel().withProviderId(c.getProviderId()).withModelId("d").withEnabled(true));
AssistantModel e = store().insert(new AssistantModel().withProviderId(c.getProviderId()).withModelId("e").withEnabled(true));
AssistantModel f = store().insert(new AssistantModel().withProviderId(c.getProviderId()).withModelId("f").withEnabled(true));
link(a, b);
link(b, c);
link(c, d);
link(d, e);
assertThatThrownBy(() -> link(e, f)).hasMessageContaining("four hops");
script("a").thenError(ModelProviderException.Kind.OVERLOADED, "busy");
script("b").thenError(ModelProviderException.Kind.OVERLOADED, "busy");
script("c").thenError(ModelProviderException.Kind.OVERLOADED, "busy").thenError(ModelProviderException.Kind.OVERLOADED, "busy").thenText("Recovered");
TurnRunner.Outcome result = runner.run(USER, "hello");
assertThat(result.text()).isEqualTo("Recovered");
assertThat(turn(result).getModel()).isEqualTo("e");
assertThat(turn(result).getFallbackCount()).isEqualTo(4);
String html = AssistantHomeView.render(AssistantHomeData.load(config.getRuntime()), "");
assertThat(html).contains("4 fallbacks");
}
}