fix(tools): the tool-call service makes its own checks, for any caller #106

Merged
james.maes merged 1 commits from fix/tool-call-service-enforces-its-checks into develop 2026-11-12 15:56:46 +00:00
28 changed files with 1813 additions and 187 deletions

No files matched your search

+13
View File
@@ -1,5 +1,18 @@
# Changelog
## Unreleased - The tool-call service makes its own checks, for any caller
- `engine.ToolCallService` no longer trusts its caller for a check it can make itself. On every call, before the tool is reached, it checks that the request is whole (a session, a context made for that same session, a policy), that the person has assistant access, the kill switch (read fresh), the hard budget, and that the tool is offered to this person now under the request's policy. The chat's turn engine made these checks around the service; a caller with no chat turn would have skipped them. The service's javadoc lists what a caller still has to supply: rate limiting, how many calls a request may make and for how long, and classification.
- `execute` refuses a tool that is not offered. Before, only the turn engine's own list of allowed tools did, so a direct call ran `navigate` or `search_help` although the policy had switched them off (the kernel tools have no policy check of their own).
- A request without a policy is refused. `UiApiToolbox`, `KnowledgeTool` and `BulkLoadTools` treated a missing policy as "every tool is allowed"; they now refuse too.
- One call cannot mix two people: a request whose context was made for another session than its own is refused, for a call and for a confirmation. `ConfirmedToolCall` carries the confirming session.
- New `ToolCallService.executeAndLog`: runs the call and writes its `assistant_tool_call` row, redacting the input itself, for a refused call as for one that ran. The redaction and the row moved out of the turn engine (`TurnEngine.redactInput` is gone; what a presentation logs as its input is the tool's `inputForLog`). The turn engine still runs its calls in parallel under its own time limit and writes the rows in call order, so it uses the two halves of that path, `execute` and `log(request, outcome)`; an ArchUnit rule keeps every other class on `executeAndLog`. The chat's log rows are unchanged.
- Confirming a pending action (`POST /actions/{id}/confirm`) is refused while the kill switch is on (503 `KILL_SWITCH`) and when the policy no longer allows the tool (403). Both are checked before the action is claimed, so it stays pending. Before, a pending write could be confirmed while the kill switch was on.
- Tool families are registered explicitly. The classpath scan of `com.kingsrook.qbits.assistant.tools` is gone: any jar with a public `ToolFamily` under that package was registered and offered. The QBit's own families are the list `AssistantRuntime.ownToolFamilies()`; a host or another QBit registers one with the new `AssistantQBitConfig.withToolFamily(ToolFamily)` before the QBit is produced. The runtime logs the registered tool names once at startup (`Assistant tools registered`) and refuses to start on a duplicate name or when there is no tool. `ToolRegistry.discover` is replaced by `ToolRegistry.of(runtime, families)`.
- The tool list is the same bytes in every JVM. `ToolRegistry.definitions` now sorts the keys of every object of an input schema. A schema built with `Map.of` (`search_knowledge`, part of the bulk-load profile schema) iterated in an order that differs between JVMs and the provider mappings serialise in iteration order, which defeated prompt caching across pods and restarts. The order of `required` and other lists is unchanged. Visible effect: properties are listed alphabetically in every tool's schema, and the cached request prefix changes once with this version.
- Stricter than before, in the chat: a turn whose model call takes the person over a budget row's limit (`assistant_budget`, state STOPPED) now has the tool calls of that same response refused (`assistant:budget`) when they are the turn's first; before, they ran and the next model call stopped the turn. The policy's own hard limits already stopped the turn before its tools. Likewise a tool call that starts after the kill switch went on, and before the running turn's own poll noticed it, is refused (`assistant:killSwitch`) instead of run; the turn is cancelled either way.
- New `BudgetService.stoppedReason(userId, active)`: whether a budget has stopped a person, read without storing a budget event or raising an alert.
## Unreleased - One tool registry (no behaviour change); built against QQQ 4.1.0-RC.32
- The assistant builds against QQQ 4.1.0-RC.32 (was RC.25), the version applications pin. Nothing had to change for it; the chat's tool list is byte for byte the same under both versions.
+3 -1
View File
@@ -63,8 +63,10 @@ Always use the wrapper: Maven 3.10 cannot read the parent POM. The committed def
- **Data access:**
- The vector index stores no record content: only source/table/key/field/chunk, content hash and embedding. Retrieval re-reads every hit through the UI API as the asking user, then applies the existing get_record exposure/masking path.
- Only `persistence/` calls QueryAction, InsertAction and the other data actions (ArchUnit).
- Every tool is registered once in `tools/registry/ToolRegistry` (name, kind, checks, definition, execute) by a `ToolFamily` found under `tools/` at startup; `engine/KernelToolFamily` (navigate, search_help) is the one the runtime names. A new tool family is new files only.
- Every tool is registered once in `tools/registry/ToolRegistry` (name, kind, checks, definition, execute) by a registered `ToolFamily`: the QBit's own are the explicit list `AssistantRuntime.ownToolFamilies()`, a host's or another QBit's come from `AssistantQBitConfig.withToolFamily`. Nothing is found on the classpath. A new family of the QBit is new files and one entry in that list.
- A tool is looked up only in the registry. Outside `tools/`, only `engine/ToolCallService` runs a tool or writes an `assistant_tool_call` row (ArchUnit); the turn engine and the confirm route both go through it. A tool call runs as the signed-in person, never as the assistant.
- `ToolCallService` trusts no caller: on every call it checks itself that the request is whole (a session, a context made for that session, a policy), assistant access, the kill switch, the hard budget (per scope) and that the tool is offered to this person now, before the tool is reached; a confirmed write gets the first three and the policy's tool switch. A caller with no turn uses `executeAndLog`; only the turn engine may use `execute` and `log` separately (ArchUnit). What a caller still supplies (rate limiting, per-request call limits and timeouts, classification) is listed in the service's javadoc.
- A missing policy refuses; it never means "every tool".
- Data tools run only through `UiApiToolbox` → `UiApiGateway` (spec'd `/qqq/v1` operations with the user's `sessionUUID`); never core data actions for user data. Only `navigate` and `search_help` use `AiToolKernel.execute`.
- Data access (`DataAccessService`) and the policy only ever narrow what the UI API allows; a table whose AI exposure is NONE stays NONE.
- A write runs only after the owner confirms its pending action (unless the policy turns confirmation off), with the confirming request's session, once.
+1
View File
@@ -289,6 +289,7 @@ The nav lists the home and the main tables: **Providers**, **Conversations**, **
| `customRedactionPatterns` | none | Extra regular expressions the kernel redacts from tool results. |
| `killSwitchPollMillis` | 2000 | How often a running turn re-checks the kill switch while it waits on the model. |
| `providerTypes` | none | `withProviderType(...)`: the types the Providers screen offers (`AnthropicProviderType`). |
| `toolFamilies` | none | `withToolFamily(...)`: a tool family (`tools.registry.ToolFamily`) of the host or of another QBit. Its tools join the assistant's own in the tool registry when the QBit is produced. This is the only way a tool gets in from outside the QBit: a family class that is only on the classpath is never registered. The runtime refuses to start when two tools share a name, and logs the registered tool names once. |
| `allowedApiKeyEnvNames` | none | Extra environment variable names a provider row may name as its key. Each type's default is always allowed. |
| `allowedBaseUrls` | none | Base URL overrides a provider row may use. The type's own URL is always allowed. |
| `secretResolver` | `System::getenv` | How a key variable name becomes the key (a host secret store, or tests). |
@@ -38,6 +38,7 @@ import com.kingsrook.qbits.assistant.spi.AssistantModelProvider;
import com.kingsrook.qbits.assistant.spi.AssistantProviderType;
import com.kingsrook.qbits.assistant.spi.ModelPricing;
import com.kingsrook.qbits.assistant.spi.TrustedCaBundle;
import com.kingsrook.qbits.assistant.tools.registry.ToolFamily;
import com.kingsrook.qbits.assistant.uiapi.UiApiGateway;
import com.kingsrook.qqq.backend.core.actions.ai.AiExposure;
import com.kingsrook.qqq.backend.core.actions.ai.DataClassification;
@@ -166,6 +167,7 @@ public class AssistantQBitConfig implements QBitConfig
private transient MetaDataCustomizerInterface<QTableMetaData> tableMetaDataCustomizer;
private transient Map<String, AssistantModelProvider> modelProviders = new LinkedHashMap<>();
private transient Map<String, AssistantProviderType> providerTypes = new LinkedHashMap<>();
private transient List<ToolFamily> toolFamilies = new ArrayList<>();
private transient Function<String, String> secretResolver = System::getenv;
private transient Function<String, String> emailRecipientResolver;
private transient AssistantRuntime runtime;
@@ -1691,6 +1693,40 @@ public class AssistantQBitConfig implements QBitConfig
/*******************************************************************************
** Register a tool family of the host or of another QBit: its tools join
** the assistant's own in the tool registry when the QBit is produced. This
** is the only way a tool gets in from outside the QBit; a family class that
** is merely on the classpath is never registered.
**
** Call it before the QBit is produced. The runtime refuses to start when a
** registered tool has the name of another tool. A registered tool is
** offered and run like the assistant's own: as the signed-in person, after
** the checks of engine.ToolCallService, and it must apply its own checks to
** what it reads or writes (see tools.registry.ToolCheck).
*******************************************************************************/
public AssistantQBitConfig withToolFamily(ToolFamily toolFamily)
{
if(toolFamily != null)
{
toolFamilies.add(toolFamily);
}
return (this);
}
/*******************************************************************************
** The tool families a host or another QBit registered, in that order.
*******************************************************************************/
@JsonIgnore
public List<ToolFamily> getToolFamilies()
{
return (toolFamilies);
}
/*******************************************************************************
** The registered provider types, by type id.
*******************************************************************************/
@@ -21,6 +21,7 @@ import java.lang.ref.WeakReference;
import java.net.InetAddress;
import java.time.Clock;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
@@ -50,8 +51,13 @@ import com.kingsrook.qbits.assistant.rag.RagService;
import com.kingsrook.qbits.assistant.rag.VectorStore;
import com.kingsrook.qbits.assistant.scheduled.ScheduledTurn;
import com.kingsrook.qbits.assistant.tools.KnowledgeTool;
import com.kingsrook.qbits.assistant.tools.KnowledgeToolFamily;
import com.kingsrook.qbits.assistant.tools.PresentResultsToolFamily;
import com.kingsrook.qbits.assistant.tools.UiApiToolFamily;
import com.kingsrook.qbits.assistant.tools.UiApiToolbox;
import com.kingsrook.qbits.assistant.tools.bulkload.BulkLoadToolFamily;
import com.kingsrook.qbits.assistant.tools.bulkload.BulkLoadTools;
import com.kingsrook.qbits.assistant.tools.registry.ToolFamily;
import com.kingsrook.qbits.assistant.tools.registry.ToolRegistry;
import com.kingsrook.qbits.assistant.uiapi.CoreUiApiGateway;
import com.kingsrook.qbits.assistant.uiapi.UiApiCall;
@@ -194,11 +200,14 @@ public class AssistantRuntime
this.bulkLoadTools = new BulkLoadTools(qInstance, dataAccessService, config.getBulkLoadRunner());
this.attachmentService = new AttachmentService(config.getAttachments(), new UploadRefSource());
///////////////////////////////////////////////////////////////////////
// every tool, registered once: the families found under the tools //
// package, and the two record-free tools that stay on the kernel //
// every tool, registered once: the QBit's own families, named here, //
// and the ones a host or another QBit registered on the config. //
// Nothing is looked for on the classpath //
///////////////////////////////////////////////////////////////////////
this.toolRegistry = ToolRegistry.discover(this, List.of(new KernelToolFamily()));
this.toolCalls = new ToolCallService(toolRegistry, store);
List<ToolFamily> toolFamilies = new ArrayList<>(ownToolFamilies());
toolFamilies.addAll(config.getToolFamilies());
this.toolRegistry = ToolRegistry.of(this, toolFamilies);
this.toolCalls = new ToolCallService(this, toolRegistry, store);
this.contextBuilder = new ContextBuilder(config, kernel, toolRegistry).withBudgets(budgetService);
if(Boolean.TRUE.equals(config.getRegisterShutdownHook()))
{
@@ -328,6 +337,18 @@ public class AssistantRuntime
/*******************************************************************************
** The QBit's own tool families: the complete list. A new family of the
** QBit is added here; one of a host or another QBit is registered with
** AssistantQBitConfig.withToolFamily.
*******************************************************************************/
public static List<ToolFamily> ownToolFamilies()
{
return (List.of(new UiApiToolFamily(), new BulkLoadToolFamily(), new KnowledgeToolFamily(), new PresentResultsToolFamily(), new KernelToolFamily()));
}
/*******************************************************************************
** Every tool the assistant has; the only place a tool is looked up.
*******************************************************************************/
@@ -44,6 +44,7 @@ import com.kingsrook.qbits.assistant.engine.PolicyService;
import com.kingsrook.qbits.assistant.engine.ProfileService;
import com.kingsrook.qbits.assistant.engine.ProviderService;
import com.kingsrook.qbits.assistant.engine.RunningTurn;
import com.kingsrook.qbits.assistant.engine.ToolCallService;
import com.kingsrook.qbits.assistant.engine.TurnEngine;
import com.kingsrook.qbits.assistant.engine.TurnEvent;
import com.kingsrook.qbits.assistant.model.AssistantClientError;
@@ -1000,6 +1001,21 @@ public class AssistantApi
throw (new AssistantApiException(409, "Wait for the assistant to finish before confirming this action."));
}
///////////////////////////////////////////////////////////////////////
// the kill switch and the policy's switch for this tool refuse a //
// confirmation as they refuse a call: asked before the claim, so a //
// refused action stays pending //
///////////////////////////////////////////////////////////////////////
PolicyService.EffectivePolicy policy = runtime.getPolicyService().current(true);
UiApiToolbox.Context context = new UiApiToolbox.Context(session.getUuid(), session.getValue(QSession.VALUE_KEY_USER_TIMEZONE), runtime.getUiApiGateway(), policy.tools(), null, credentials, action.getTurnId() == null ? null : String.valueOf(action.getTurnId()));
AiToolResult refused = runtime.getToolCalls().confirmationRefusal(session, context, action.getToolCallId(), action.getToolName());
if(refused != null)
{
throw (ToolCallService.CHECK_KILL_SWITCH.equals(refused.getPermissionCheck())
? new AssistantApiException(503, "KILL_SWITCH", KILL_SWITCH_MESSAGE)
: new AssistantApiException(403, "NOT_ALLOWED", refused.getModelContent()));
}
///////////////////////////////////////////////////////////////////////
// claim it first, so a double click cannot run the write twice //
///////////////////////////////////////////////////////////////////////
@@ -1020,14 +1036,12 @@ public class AssistantApi
{
throw (new QException("The stored values of action " + id + " are not valid JSON", e));
}
PolicyService.EffectivePolicy policy = runtime.getPolicyService().current(true);
UiApiToolbox.Context context = new UiApiToolbox.Context(session.getUuid(), session.getValue(QSession.VALUE_KEY_USER_TIMEZONE), runtime.getUiApiGateway(), policy.tools(), null, credentials, action.getTurnId() == null ? null : String.valueOf(action.getTurnId()));
///////////////////////////////////////////////////////////////////////
// the registry's tool runs its confirmed write, as the user who //
// confirms (this request's session) //
///////////////////////////////////////////////////////////////////////
UiApiToolbox.Outcome outcome = runtime.getToolCalls().executeConfirmed(
new ConfirmedToolCall(context, action.getToolCallId(), action.getToolName(), action.getTableName(), action.getRecordId(), values, userId));
new ConfirmedToolCall(session, context, action.getToolCallId(), action.getToolName(), action.getTableName(), action.getRecordId(), values, userId));
AiToolResult result = outcome.getResult();
boolean ok = result.getOutcome() == AiToolOutcome.OK;
@@ -235,6 +235,30 @@ public class BudgetService
/*******************************************************************************
** What a person is told when a budget that applies to them (and to this
** provider and model, when there is one) has stopped them, or null. It only
** reads: no budget event is stored and no alert raised, so asking does not
** take the warning or the stop away from the model call that reports it.
*******************************************************************************/
public String stoppedReason(String userId, ProviderService.Active active) throws QException
{
for(AssistantBudget budget : store.budgets())
{
if(applies(budget, userId, active))
{
Snapshot snapshot = snapshot(budget, userId);
if("STOPPED".equals(snapshot.state()))
{
return (stoppedMessage(snapshot));
}
}
}
return (null);
}
/*******************************************************************************
** One of a person's own budgets as a sentence: what is used, of how much, and
** when the count starts again.
@@ -19,7 +19,11 @@ package com.kingsrook.qbits.assistant.engine;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.Map;
import java.util.Objects;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.AssistantSecurity;
import com.kingsrook.qbits.assistant.model.AssistantPendingAction;
import com.kingsrook.qbits.assistant.model.AssistantToolCall;
import com.kingsrook.qbits.assistant.persistence.AssistantStore;
@@ -28,43 +32,110 @@ import com.kingsrook.qbits.assistant.tools.registry.AssistantTool;
import com.kingsrook.qbits.assistant.tools.registry.ConfirmedToolCall;
import com.kingsrook.qbits.assistant.tools.registry.ToolCallOutcome;
import com.kingsrook.qbits.assistant.tools.registry.ToolCallRequest;
import com.kingsrook.qbits.assistant.tools.registry.ToolCallScope;
import com.kingsrook.qbits.assistant.tools.registry.ToolOffer;
import com.kingsrook.qbits.assistant.tools.registry.ToolRegistry;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolOutcome;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolResult;
import com.kingsrook.qqq.backend.core.actions.ai.output.AiRedactor;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.logging.QLogger;
import com.kingsrook.qqq.backend.core.model.metadata.QInstance;
import com.kingsrook.qqq.backend.core.model.session.QSession;
import com.kingsrook.qqq.backend.core.utils.JsonUtils;
import com.kingsrook.qqq.backend.core.utils.StringUtils;
import static com.kingsrook.qqq.backend.core.logging.LogUtils.logPair;
/*******************************************************************************
** "Run one tool call for this person under this policy, and log it", for
** every caller: the chat's turn engine and the confirm route today.
** "Run one tool call for this person, and log it", for every caller: the
** chat's turn engine and the confirm route today, a caller with no chat turn
** around it later. The service trusts no caller for a check it can make
** itself.
**
** The steps, which a caller takes in this order:
** 1. {@link #refusal}: a tool the person is not offered is refused before
** anything runs.
** 2. {@link #execute} (or {@link #executeConfirmed} for a write the person
** confirmed): the registry's tool runs the call with the person's own
** session, and applies its own checks (the policy's tool switch, the data
** access, what the person may do in the UI API).
** 3. {@link #log} (or {@link #logConfirmed}): one assistant_tool_call row,
** for a refused call as for one that ran.
** What the service enforces on every call, in this order, before the tool is
** reached (so a refusal happens before any data action):
** 1. the request is whole: a session, a context made for that same session,
** and a policy. A request without a policy is refused, never run
** unnarrowed.
** 2. assistant access: the assistant is enabled and the session's person
** holds assistant.hasAccess.
** 3. the kill switch, read fresh; a policy that cannot be read counts as on.
** 4. the hard budget (the policy's hard limits, and a budget that has stopped
** the person), when the first call of a scope arrives. It only reads: no
** budget event, no alert, nothing charged.
** 5. the offer: the tool is one this person is offered now under the
** request's policy (not switched off by the policy and, for a write or a
** bulk-load tool, with the writable table or the profiles it needs). A
** name the registry does not have is refused the same way.
** Then the registry's tool runs with the person's own session and applies its
** own checks (the policy's tool switch again, data access, field exposure,
** confirmation, what the person may do in the UI API).
**
** A confirmed write ({@link #executeConfirmed}) goes through 1 to 3 and the
** policy's switch for its tool. It spends no model budget, so the hard budget
** does not stop it.
**
** {@link #executeAndLog} is the one call a caller without a turn makes: it
** runs the call and writes its assistant_tool_call row, redacting the input
** itself, for a refused call as for one that ran. The turn engine runs its
** calls in parallel under its own time limit and writes the rows in call
** order afterwards, so it uses the two halves of that same path,
** {@link #execute} and {@link #log}; no other class may
** (ToolCallArchitectureTest).
**
** What a caller still has to supply, because only it knows its own requests:
** - rate limiting (the chat: RateLimiter, turns per minute and concurrent
** turns per person);
** - how many calls one request may make, how many at once and how long one
** may take (the chat: maxToolCallsPerTurn, maxParallelToolCalls,
** toolCallTimeoutSeconds and the turn's deadline), and the row of a call it
** stopped itself ({@link #log} with the result it gave instead);
** - classification: keeping the scope's data classification up to date from
** every result and persisting it (the chat: TurnClassification.observe and
** persist), and a truthful {@link ToolCallScope#hipaaRoute()};
** - the scope itself: one per request, never shared by two people, with the
** client channel its calls are logged under;
** - the session: the signed-in person's own, from the route's session setup,
** never one the caller made. Request credentials the context forwards to
** the UI API must be that same person's: the service compares the context's
** session id with the session, but it cannot check forwarded credentials;
** - the policy's tool settings in force for the request (the service refuses
** a missing policy, but does not compare the one it is given with the
** administrator's row).
**
** A call never runs as the assistant: the session is always the request's.
** The log row is written through the store, as the same person.
*******************************************************************************/
public class ToolCallService
{
private final ToolRegistry registry;
private final AssistantStore store;
private static final QLogger LOG = QLogger.getLogger(ToolCallService.class);
///////////////////////////////////////////////////////////////////////////
// the permission-check labels of the service's own refusals //
///////////////////////////////////////////////////////////////////////////
public static final String CHECK_REQUEST = "assistant:request";
public static final String CHECK_ACCESS = "assistant:access";
public static final String CHECK_KILL_SWITCH = "assistant:killSwitch";
public static final String CHECK_BUDGET = "assistant:budget";
public static final String CHECK_TOOL_SWITCH = "policy:tool";
private static final AiRedactor INPUT_REDACTOR = new AiRedactor(List.of());
private final AssistantRuntime runtime;
private final ToolRegistry registry;
private final AssistantStore store;
/*******************************************************************************
**
** The runtime is still being constructed here: its collaborators are asked
** for when a call arrives.
*******************************************************************************/
public ToolCallService(ToolRegistry registry, AssistantStore store)
public ToolCallService(AssistantRuntime runtime, ToolRegistry registry, AssistantStore store)
{
this.runtime = runtime;
this.registry = registry;
this.store = store;
}
@@ -72,39 +143,281 @@ public class ToolCallService
/*******************************************************************************
** One finished tool call, as the tool-call log keeps it.
**
** @param turnId the turn the call was made in
** @param conversationId the turn's conversation
** @param userId the person the call ran as
** @param toolCallId the model's id of the call
** @param toolName the tool called (as the model named it)
** @param clientChannel the channel label the call's requests carried
** @param inputJson the input, already redacted
** @param outcome the outcome's name
** @param summary the one-line summary shown in the panel
** @param result the tool's result
** @param extras the UI API operation's detail; null when the call
** had none (also when it was refused or timed out)
** What the service remembers for one scope: the offer its calls are checked
** against, and that the hard budget let the scope's first call through. It
** belongs to one session and one policy; a request with another is checked
** afresh.
*******************************************************************************/
public record Finished(Integer turnId, Integer conversationId, String userId, String toolCallId, String toolName, String clientChannel, String inputJson, String outcome,
String summary, AiToolResult result, UiApiToolbox.Outcome extras)
private static final class Admission
{
private QSession session;
private PolicyService.ToolPolicy policy;
private ToolOffer offer;
private boolean withinBudget;
/*******************************************************************************
** Bind a new admission to the request's person and policy.
*******************************************************************************/
private Admission bind(ToolCallRequest request)
{
session = request.session();
policy = request.context().policy();
offer = new ToolOffer(policy, session, request.context().turnId());
return (this);
}
}
/*******************************************************************************
** The person's session as the thread's QQQ context for as long as the
** service works; what was there before is put back.
*******************************************************************************/
private final class PersonsContext implements AutoCloseable
{
private final boolean entered;
private final QInstance priorInstance;
private final QSession priorSession;
/*******************************************************************************
**
*******************************************************************************/
private PersonsContext(QSession session)
{
priorInstance = QContext.getQInstance();
priorSession = QContext.getQSession();
entered = priorInstance == null || priorSession != session;
if(entered)
{
QContext.init(runtime.getQInstance(), session);
}
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public void close()
{
if(!entered)
{
return;
}
if(priorInstance != null && priorSession != null)
{
QContext.init(priorInstance, priorSession);
}
else
{
QContext.clear();
}
}
}
/*******************************************************************************
** Run one call as the request's person, after every check of the service.
** A refusal is a result, and nothing ran. The turn engine logs the outcome
** with {@link #log}; every other caller uses {@link #executeAndLog}.
*******************************************************************************/
public ToolCallOutcome execute(ToolCallRequest request)
{
AiToolResult incomplete = incomplete(request);
if(incomplete != null)
{
return (ToolCallOutcome.of(incomplete));
}
try(PersonsContext ignored = new PersonsContext(request.session()))
{
AiToolResult refused = refusal(request);
return (refused != null ? ToolCallOutcome.of(refused) : registry.find(request.toolName()).execute(request));
}
}
/*******************************************************************************
** Run one call and write its tool-call log row: the whole of a tool call
** for a caller with no turn. A refused call is logged like one that ran,
** and a tool that throws gives a fixed ERROR result and its row. When the
** row cannot be written the caller gets the exception, not the result: a
** result is never handed out unlogged. Only a call that names no person at
** all (which was refused) has no row.
*******************************************************************************/
public ToolCallOutcome executeAndLog(ToolCallRequest request) throws QException
{
ToolCallOutcome outcome;
try
{
outcome = execute(request);
}
catch(RuntimeException e)
{
LOG.warn("Assistant tool call threw", logPair("tool", request.toolName()), logPair("errorClass", e.getClass().getSimpleName()));
outcome = ToolCallOutcome.of(result(request).withOutcome(AiToolOutcome.ERROR).withSummary("Failed")
.withModelContent("The tool failed unexpectedly. Answer with what you have, or try another way."));
}
if(request == null || AssistantSecurity.userId(request.session()) == null)
{
LOG.warn("An assistant tool call that names no signed-in person was refused; it has no log row", logPair("tool", request == null ? null : request.toolName()));
return (outcome);
}
log(request, outcome);
return (outcome);
}
/*******************************************************************************
** Why a request cannot be checked at all, or null when it is whole: it
** names a person (a session), carries a context made for that same session,
** and a policy.
*******************************************************************************/
private static AiToolResult incomplete(ToolCallRequest request)
{
if(request == null || request.session() == null || request.context() == null)
{
return (denied(result(request), CHECK_REQUEST, "The tool call has no signed-in person."));
}
if(!Objects.equals(request.session().getUuid(), request.context().sessionUUID()))
{
return (denied(result(request), CHECK_REQUEST, "The tool call's session and its context are not the same person's."));
}
if(request.context().policy() == null)
{
return (denied(result(request), CHECK_REQUEST, "The tool call has no policy."));
}
return (null);
}
/*******************************************************************************
** The refusal of a whole request, or null when its tool may run: assistant
** access, the kill switch, the hard budget, then the offer.
*******************************************************************************/
private AiToolResult refusal(ToolCallRequest request)
{
AiToolResult result = result(request);
AiToolResult closed = accessOrKillSwitch(request.session(), result);
if(closed != null)
{
return (closed);
}
Admission admission = admission(request);
synchronized(admission)
{
if(!admission.withinBudget)
{
String reason = hardBudget(request.session());
if(reason != null)
{
return (denied(result, CHECK_BUDGET, reason));
}
admission.withinBudget = true;
}
AssistantTool tool = registry.find(request.toolName());
if(tool == null || tool.definition(admission.offer) == null)
{
return (notOffered(request.toolCallId(), request.toolName()));
}
}
return (null);
}
/*******************************************************************************
** Assistant access first (a person without it is told so even while the
** kill switch is on), then the kill switch, read fresh (a policy that
** cannot be read counts as on); null when both let the call through.
*******************************************************************************/
private AiToolResult accessOrKillSwitch(QSession session, AiToolResult result)
{
if(!runtime.getConfig().isEffectivelyEnabled())
{
return (denied(result, CHECK_ACCESS, "The assistant is not enabled."));
}
if(AssistantSecurity.userId(session) == null || !session.hasPermission(AssistantQBitConfig.ACCESS_PERMISSION))
{
return (denied(result, CHECK_ACCESS, "You do not have access to the assistant."));
}
if(runtime.getPolicyService().isKillSwitchOn())
{
return (denied(result, CHECK_KILL_SWITCH, "The assistant has been turned off by an administrator."));
}
return (null);
}
/*******************************************************************************
** What the person is told when a hard budget stops them, or null. It only
** reads. A budget that cannot be read stops the call.
*******************************************************************************/
private String hardBudget(QSession session)
{
try
{
String userId = AssistantSecurity.userId(session);
PolicyService.EffectivePolicy policy = runtime.getPolicyService().current(false);
PolicyService.BudgetStatus configured = runtime.getPolicyService().checkConfiguredBudget(policy, userId, 0, 0);
if(configured.hardExceeded())
{
return (configured.message());
}
return (runtime.getBudgetService().stoppedReason(userId, runtime.getProviderService().resolve(policy.mainModel())));
}
catch(QException | RuntimeException e)
{
LOG.warn("Could not read the assistant budgets; refusing the tool call (fail closed)", logPair("errorClass", e.getClass().getSimpleName()));
return ("The assistant's budgets could not be read. Try again in a moment.");
}
}
/*******************************************************************************
** The scope's admission when it was made for this request's session and
** policy; one of the request's own otherwise (also for a call without a
** scope), so one person's offer never answers for another.
*******************************************************************************/
private static Admission admission(ToolCallRequest request)
{
if(request.scope() == null)
{
return (new Admission().bind(request));
}
Admission admission = request.scope().state(Admission.class, Admission::new);
synchronized(admission)
{
if(admission.offer == null)
{
return (admission.bind(request));
}
boolean same = admission.session == request.session() && admission.policy.equals(request.context().policy());
return (same ? admission : new Admission().bind(request));
}
}
/*******************************************************************************
** The refusal of a call to a tool the person is not offered (a name the
** model made up, or a tool the policy or the host switched off), or null
** when the call may run.
** model made up, or a tool the policy or the host switched off).
*******************************************************************************/
public AiToolResult refusal(String toolCallId, String toolName, Set<String> offeredTools)
private static AiToolResult notOffered(String toolCallId, String toolName)
{
if(offeredTools.contains(toolName))
{
return (null);
}
return (new AiToolResult().withToolCallId(toolCallId).withToolName(toolName).withOutcome(AiToolOutcome.INVALID_INPUT)
.withSummary("Unknown tool").withModelContent("Unknown tool \"" + toolName + "\". Use only the tools provided."));
}
@@ -112,27 +425,89 @@ public class ToolCallService
/*******************************************************************************
** Run one call with the request's session. A name the registry does not
** have is refused like one that is not offered.
** An empty result for the request's call.
*******************************************************************************/
public ToolCallOutcome execute(ToolCallRequest request)
private static AiToolResult result(ToolCallRequest request)
{
AssistantTool tool = registry.find(request.toolName());
if(tool == null)
{
return (ToolCallOutcome.of(refusal(request.toolCallId(), request.toolName(), Set.of())));
}
return (tool.execute(request));
return (new AiToolResult().withToolCallId(request == null ? null : request.toolCallId()).withToolName(request == null ? null : request.toolName()));
}
/*******************************************************************************
** Run a write the person confirmed, with the confirming request's session.
** A name that proposes no write is refused.
** A refusal by one of the service's own checks.
*******************************************************************************/
private static AiToolResult denied(AiToolResult result, String check, String message)
{
return (result.withOutcome(AiToolOutcome.DENIED).withPermissionCheck(check).withSummary("Not allowed").withModelContent(message));
}
/*******************************************************************************
** Why a write the person confirmed may not run now, or null when it may:
** the confirming session and its context are the same person's and carry a
** policy, that person has assistant access, the kill switch is off, and the
** policy (the confirming request's and the administrator's current one)
** allows the tool. The confirm route asks before it claims the pending
** action, so a refused confirmation leaves the action pending;
** {@link #executeConfirmed} asks again itself.
*******************************************************************************/
public AiToolResult confirmationRefusal(QSession session, UiApiToolbox.Context context, String toolCallId, String toolName)
{
AiToolResult result = new AiToolResult().withToolCallId(toolCallId).withToolName(toolName);
if(session == null || context == null)
{
return (denied(result, CHECK_REQUEST, "The confirmation has no signed-in person."));
}
if(!Objects.equals(session.getUuid(), context.sessionUUID()))
{
return (denied(result, CHECK_REQUEST, "The confirmation's session and its context are not the same person's."));
}
if(context.policy() == null)
{
return (denied(result, CHECK_REQUEST, "The confirmation has no policy."));
}
try(PersonsContext ignored = new PersonsContext(session))
{
AiToolResult closed = accessOrKillSwitch(session, result);
if(closed != null)
{
return (closed);
}
if(!context.policy().allows(toolName) || !runtime.getPolicyService().current(false).tools().allows(toolName))
{
return (denied(result, CHECK_TOOL_SWITCH, "The administrator has not allowed the " + toolName + " tool."));
}
}
catch(QException e)
{
LOG.warn("Could not read the assistant policy; refusing the confirmation (fail closed)", logPair("errorClass", e.getClass().getSimpleName()));
return (denied(result, CHECK_KILL_SWITCH, "The assistant's policy could not be read. Try again in a moment."));
}
return (null);
}
/*******************************************************************************
** Run a write the person confirmed, with the confirming request's session,
** after {@link #confirmationRefusal}'s checks. A confirmation named for
** another person than its session's is refused; so is a name that proposes
** no write.
*******************************************************************************/
public UiApiToolbox.Outcome executeConfirmed(ConfirmedToolCall call)
{
AiToolResult refused = confirmationRefusal(call.session(), call.context(), call.toolCallId(), call.toolName());
if(refused == null && !Objects.equals(call.userId(), AssistantSecurity.userId(call.session())))
{
refused = denied(new AiToolResult().withToolCallId(call.toolCallId()).withToolName(call.toolName()), CHECK_REQUEST, "The confirmation's session and its user are not the same person.");
}
if(refused != null)
{
return (UiApiToolbox.Outcome.of(refused.withTableName(call.tableName()), null, null, null));
}
AssistantTool tool = registry.find(call.toolName());
if(tool == null)
{
@@ -144,35 +519,65 @@ public class ToolCallService
/*******************************************************************************
** The tool-call log row of a call made in a turn.
** The tool-call log row of a call: of {@link #execute}'s outcome, or of the
** result a caller gave instead of running or finishing the call (its own
** time limit, a stop). The service takes the person, the turn, the
** conversation and the channel from the request and its scope, and redacts
** the input itself; the caller passes nothing it has prepared.
*******************************************************************************/
public void log(Finished call) throws QException
public void log(ToolCallRequest request, ToolCallOutcome outcome) throws QException
{
AssistantTool tool = registry.find(call.toolName());
AiToolResult result = call.result();
UiApiToolbox.Outcome extras = call.extras();
store.insert(new AssistantToolCall()
.withTurnId(call.turnId())
.withConversationId(call.conversationId())
.withUserId(call.userId())
.withToolCallId(call.toolCallId())
.withToolName(call.toolName())
.withClientChannel(call.clientChannel())
.withInputJson(tool == null ? call.inputJson() : tool.forLog(call.inputJson()))
.withOutcome(call.outcome())
.withPermissionCheck(result.getPermissionCheck())
.withTableName(result.getTableName())
.withProcessName(result.getProcessName())
.withWidgetName(result.getWidgetName())
.withRowCount(result.getRowCount())
.withDurationMs(result.getDurationMs())
.withResultSummary(call.summary())
.withResultJson(tool == null ? result.getModelContent() : tool.forLog(result.getModelContent()))
.withRedactionsJson(result.getRedactions() == null || result.getRedactions().isEmpty() ? null : JsonUtils.toJson(result.getRedactions()))
.withUiApiOperation(extras == null ? (tool == null ? null : tool.operationWithoutDetail()) : extras.getOperation())
.withHttpStatus(extras == null ? null : extras.getHttpStatus())
.withAuditHref(extras == null ? null : extras.getAuditHref())
.withPendingActionId(extras == null ? null : extras.getPendingActionId()));
AssistantTool tool = registry.find(request.toolName());
AiToolResult result = outcome.result();
UiApiToolbox.Outcome extras = outcome.extras();
ToolCallScope scope = request.scope();
Map<String, Object> input = tool == null ? request.input() : tool.inputForLog(request.input(), result);
String inputJson = INPUT_REDACTOR.redact(JsonUtils.toJson(input), new HashSet<>());
try(PersonsContext ignored = new PersonsContext(request.session()))
{
store.insert(new AssistantToolCall()
.withTurnId(scope == null ? null : scope.turnId())
.withConversationId(scope == null ? null : scope.conversationId())
.withUserId(AssistantSecurity.userId(request.session()))
.withToolCallId(request.toolCallId())
.withToolName(request.toolName())
.withClientChannel(scope == null ? "assistant" : scope.clientChannel())
.withInputJson(tool == null ? inputJson : tool.forLog(inputJson))
.withOutcome(outcomeName(result))
.withPermissionCheck(result.getPermissionCheck())
.withTableName(result.getTableName())
.withProcessName(result.getProcessName())
.withWidgetName(result.getWidgetName())
.withRowCount(result.getRowCount())
.withDurationMs(result.getDurationMs())
.withResultSummary(summary(request.toolName(), result))
.withResultJson(tool == null ? result.getModelContent() : tool.forLog(result.getModelContent()))
.withRedactionsJson(result.getRedactions() == null || result.getRedactions().isEmpty() ? null : JsonUtils.toJson(result.getRedactions()))
.withUiApiOperation(extras == null ? (tool == null ? null : tool.operationWithoutDetail()) : extras.getOperation())
.withHttpStatus(extras == null ? null : extras.getHttpStatus())
.withAuditHref(extras == null ? null : extras.getAuditHref())
.withPendingActionId(extras == null ? null : extras.getPendingActionId()));
}
}
/*******************************************************************************
** The outcome's name as the log and the panel show it (none is an ERROR).
*******************************************************************************/
public static String outcomeName(AiToolResult result)
{
return (result.getOutcome() == null ? AiToolOutcome.ERROR.name() : result.getOutcome().name());
}
/*******************************************************************************
** The one-line summary of a result as the log and the panel show it.
*******************************************************************************/
public static String summary(String toolName, AiToolResult result)
{
return (StringUtils.hasContent(result.getSummary()) ? result.getSummary() : ToolLabels.done(toolName));
}
@@ -191,8 +596,8 @@ public class ToolCallService
.withUserId(userId)
.withToolCallId(action.getToolCallId())
.withToolName(action.getToolName())
.withInputJson(truncate(new AiRedactor(List.of()).redact(action.getParamsJson() == null ? "{}" : action.getParamsJson(), new HashSet<>()), 4000))
.withOutcome(result.getOutcome() == null ? AiToolOutcome.ERROR.name() : result.getOutcome().name())
.withInputJson(truncate(INPUT_REDACTOR.redact(action.getParamsJson() == null ? "{}" : action.getParamsJson(), new HashSet<>()), 4000))
.withOutcome(outcomeName(result))
.withPermissionCheck(result.getPermissionCheck())
.withTableName(action.getTableName())
.withRowCount(result.getRowCount())
@@ -24,7 +24,6 @@ import java.time.LocalDate;
import java.time.ZoneOffset;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
@@ -76,7 +75,6 @@ import com.kingsrook.qqq.backend.core.actions.ai.AiToolCall;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolOutcome;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolResult;
import com.kingsrook.qqq.backend.core.actions.ai.DataClassification;
import com.kingsrook.qqq.backend.core.actions.ai.output.AiRedactor;
import com.kingsrook.qqq.backend.core.context.CapturedContext;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.exceptions.QException;
@@ -118,7 +116,6 @@ public class TurnEngine implements Runnable
public static final String TIMEOUT_TEXT = "The assistant took too long to answer, so it was stopped. Ask again.";
public static final String RESTART_TEXT = "The server restarted while answering. Ask again.";
private static final AiRedactor INPUT_REDACTOR = new AiRedactor(List.of());
private static final int TEXT_FLUSH_CHARS = 48;
private static final long TEXT_FLUSH_MILLIS = 120;
@@ -366,7 +363,6 @@ public class TurnEngine implements Runnable
{
request.setParallelToolCalls(false);
}
Set<String> allowedTools = runtime.getContextBuilder().toolNames(toolPolicy);
boolean toolLimitReached = false;
while(true)
@@ -444,7 +440,7 @@ public class TurnEngine implements Runnable
break;
}
ToolBatch batch = executeTools(toolUses, allowedTools);
ToolBatch batch = executeTools(toolUses);
toolLimitReached = batch.limitReached;
for(ModelMessage message : request.getMessages())
@@ -1051,14 +1047,18 @@ public class TurnEngine implements Runnable
/*******************************************************************************
** Run a batch of tool calls through the kernel.
** Run a batch of tool calls through the tool-call service, which checks
** each one (a tool that is not offered is refused there) and runs it as the
** user. The limits of a turn are this engine's: how many calls, how many at
** once, how long each may take.
*******************************************************************************/
private ToolBatch executeTools(List<ModelPart> toolUses, Set<String> allowedTools) throws QException, InterruptedException
private ToolBatch executeTools(List<ModelPart> toolUses) throws QException, InterruptedException
{
int allowed = Math.max(0, config.getMaxToolCallsPerTurn() - toolCallCount);
boolean limitReached = false;
AiToolResult[] results = new AiToolResult[toolUses.size()];
UiApiToolbox.Outcome[] outcomes = new UiApiToolbox.Outcome[toolUses.size()];
ToolCallRequest[] requests = new ToolCallRequest[toolUses.size()];
List<Future<?>> futures = new ArrayList<>();
List<Integer> futureIndexes = new ArrayList<>();
AtomicLongArray startedNanos = new AtomicLongArray(toolUses.size());
@@ -1079,13 +1079,7 @@ public class TurnEngine implements Runnable
phase("tool", use.getToolName());
emit(TurnEvent.TOOL_START, map("messageId", assistantMessage.getId(), "toolCallId", use.getToolUseId(), "name", use.getToolName(), "summary", ToolLabels.running(use.getToolName())));
AiToolResult refusal = runtime.getToolCalls().refusal(use.getToolUseId(), use.getToolName(), allowedTools);
if(refusal != null)
{
results[i] = refusal;
continue;
}
requests[i] = new ToolCallRequest(session, toolContext(), use.getToolUseId(), use.getToolName(), use.getInput(), toolScope);
final int index = i;
AiToolCall call = new AiToolCall().withId(use.getToolUseId()).withName(use.getToolName()).withInput(use.getInput());
futureIndexes.add(index);
@@ -1107,10 +1101,11 @@ public class TurnEngine implements Runnable
if(!running.isCancelled())
{
///////////////////////////////////////////////////////////////
// the registry's tool runs the call, as the user: with the //
// user's own session, policy and pending-action store //
// the service checks the call, then the registry's tool //
// runs it, as the user: with the user's own session, policy //
// and pending-action store //
///////////////////////////////////////////////////////////////
ToolCallOutcome done = runtime.getToolCalls().execute(new ToolCallRequest(session, toolContext(), call.getId(), call.getName(), call.getInput(), toolScope));
ToolCallOutcome done = runtime.getToolCalls().execute(requests[index]);
outcomes[index] = done.extras();
results[index] = done.result();
}
@@ -1208,7 +1203,7 @@ public class TurnEngine implements Runnable
timedOut[i] = !stopped;
}
classification.observe(use, result);
recordToolResult(use, result, extras);
recordToolResult(use, requests[i], result, extras);
if(timedOut[i])
{
notice("error", "A step took too long and was stopped (" + ToolLabels.running(use.getToolName()) + ").");
@@ -1245,10 +1240,10 @@ public class TurnEngine implements Runnable
** checked first (RenderShapes): one that cannot be shown is left out and an
** error notice says so.
*******************************************************************************/
private void recordToolResult(ModelPart use, AiToolResult result, UiApiToolbox.Outcome extras)
private void recordToolResult(ModelPart use, ToolCallRequest request, AiToolResult result, UiApiToolbox.Outcome extras)
{
String outcome = result.getOutcome() == null ? AiToolOutcome.ERROR.name() : result.getOutcome().name();
String summary = StringUtils.hasContent(result.getSummary()) ? result.getSummary() : ToolLabels.done(use.getToolName());
String outcome = ToolCallService.outcomeName(result);
String summary = ToolCallService.summary(use.getToolName(), result);
List<Map<String, Object>> refs = new ArrayList<>();
for(AiRecordRef ref : result.getRecordRefs() == null ? List.<AiRecordRef>of() : result.getRecordRefs())
@@ -1307,9 +1302,11 @@ public class TurnEngine implements Runnable
try
{
String inputJson = redactInput(PresentResultsTool.NAME.equals(use.getToolName()) ? (presentation ? result.getRender() : Map.of()) : use.getInput());
runtime.getToolCalls().log(new ToolCallService.Finished(turn.getId(), conversation.getId(), userId, use.getToolUseId(), use.getToolName(),
scheduledTask == null ? "assistant/" + turn.getId() : "assistant/scheduled/" + scheduledTask.getId(), inputJson, outcome, summary, result, extras));
//////////////////////////////////////////////////////////////////////
// the service writes the row: it redacts the input itself and //
// takes the person, the turn and the channel from the request //
//////////////////////////////////////////////////////////////////////
runtime.getToolCalls().log(request, new ToolCallOutcome(result, extras));
}
catch(QException e)
{
@@ -1363,6 +1360,18 @@ public class TurnEngine implements Runnable
/*******************************************************************************
** The channel this turn's requests carry: the turn's, or the scheduled
** task's.
*******************************************************************************/
@Override
public String clientChannel()
{
return (scheduledTask == null ? "assistant/" + turn.getId() : "assistant/scheduled/" + scheduledTask.getId());
}
/*******************************************************************************
** The X-QQQ-AI-Route the bulk-load tools send.
*******************************************************************************/
@@ -1431,17 +1440,6 @@ public class TurnEngine implements Runnable
/*******************************************************************************
** The tool input as JSON with secrets redacted (for the tool-call log).
*******************************************************************************/
private String redactInput(Map<String, Object> input)
{
String json = JsonUtils.toJson(input);
return (INPUT_REDACTOR.redact(json, new HashSet<>()));
}
/*******************************************************************************
** Only same-origin application links are passed on.
*******************************************************************************/
@@ -133,7 +133,7 @@ public class KnowledgeTool
Long started = System.nanoTime();
try
{
if(context.policy() != null && !context.policy().allows(NAME))
if(context.policy() == null || !context.policy().allows(NAME))
{
return (result.withOutcome(AiToolOutcome.DENIED).withSummary("Not allowed").withModelContent("The policy does not allow search_knowledge."));
}
@@ -18,6 +18,7 @@ package com.kingsrook.qbits.assistant.tools;
import java.util.List;
import java.util.Map;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.spi.ModelToolDefinition;
import com.kingsrook.qbits.assistant.tools.registry.AssistantTool;
@@ -28,6 +29,8 @@ import com.kingsrook.qbits.assistant.tools.registry.ToolFamily;
import com.kingsrook.qbits.assistant.tools.registry.ToolKind;
import com.kingsrook.qbits.assistant.tools.registry.ToolOffer;
import com.kingsrook.qbits.assistant.tools.registry.ToolRegistry;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolOutcome;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolResult;
/*******************************************************************************
@@ -127,5 +130,18 @@ public class PresentResultsToolFamily implements ToolFamily
PresentResultsTool presenter = request.scope().state(PresentResultsTool.class, () -> new PresentResultsTool(runtime.getQInstance()));
return (ToolCallOutcome.of(presenter.execute(request.context(), request.toolCallId(), request.input())));
}
/*******************************************************************************
** The log keeps the presentation that was shown (checked and cut to what
** the person may see), never the model's raw input; nothing for a call
** that showed nothing.
*******************************************************************************/
@Override
public Map<String, Object> inputForLog(Map<String, Object> input, AiToolResult result)
{
return (AiToolOutcome.OK.equals(result.getOutcome()) && result.getRender() != null ? result.getRender() : Map.of());
}
}
}
@@ -486,8 +486,11 @@ public class UiApiToolbox
throw (AiToolException.invalidInput("Unknown tool \"" + toolName + "\". Use only the tools provided."));
}
outcome.operation = operation[0] + " " + operation[1];
if(context.policy() != null && !context.policy().allows(toolName))
if(context.policy() == null || !context.policy().allows(toolName))
{
///////////////////////////////////////////////////////////////
// no policy is not "every tool": a call without one refuses //
///////////////////////////////////////////////////////////////
deny(result, "policy:tool", "The administrator has not allowed the " + toolName + " tool.");
return (outcome);
}
@@ -203,7 +203,7 @@ public class BulkLoadTools
long started = System.nanoTime();
try
{
if(context.policy() != null && !context.policy().allows(toolName))
if(context.policy() == null || !context.policy().allows(toolName))
{
deny(result, "policy:tool", "The administrator has not allowed the " + toolName + " tool.");
return (UiApiToolbox.Outcome.of(result, null, null, null));
@@ -18,8 +18,10 @@ package com.kingsrook.qbits.assistant.tools.registry;
import java.util.List;
import java.util.Map;
import com.kingsrook.qbits.assistant.spi.ModelToolDefinition;
import com.kingsrook.qbits.assistant.tools.UiApiToolbox;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolResult;
/*******************************************************************************
@@ -84,6 +86,18 @@ public interface AssistantTool
/*******************************************************************************
** What the tool-call log keeps as a call's input, before the service
** redacts it: the model's input, unless the tool's input is not what
** should be kept (a presentation keeps what was shown instead).
*******************************************************************************/
default Map<String, Object> inputForLog(Map<String, Object> input, AiToolResult result)
{
return (input);
}
/*******************************************************************************
** The input or result JSON as the tool-call log may keep it (a tool whose
** JSON can carry more than the log should hold reduces it here).
@@ -19,20 +19,23 @@ package com.kingsrook.qbits.assistant.tools.registry;
import java.util.Map;
import com.kingsrook.qbits.assistant.tools.UiApiToolbox;
import com.kingsrook.qqq.backend.core.model.session.QSession;
/*******************************************************************************
** A proposed write the person has confirmed: the pending action's tool call,
** to be run once with the confirming request's own session (the context).
**
** @param context the confirming person's session, gateway and policy
** @param session the confirming person's own session
** @param context that same session for the UI API, with the gateway and
** the policy's tool settings
** @param toolCallId the proposing tool call's id
** @param toolName the tool that proposed the write
** @param tableName the pending action's table
** @param recordId the pending action's record, if it has one
** @param values the stored values of the proposal
** @param userId the confirming person
** @param userId the confirming person (the session's)
*******************************************************************************/
public record ConfirmedToolCall(UiApiToolbox.Context context, String toolCallId, String toolName, String tableName, String recordId, Map<String, Object> values, String userId)
public record ConfirmedToolCall(QSession session, UiApiToolbox.Context context, String toolCallId, String toolName, String tableName, String recordId, Map<String, Object> values, String userId)
{
}
@@ -26,11 +26,14 @@ import com.kingsrook.qqq.backend.core.model.session.QSession;
/*******************************************************************************
** One tool call, for one person.
** One tool call, for one person. The service that runs it refuses a request
** whose context was not made for its session (two identities in one call)
** or carries no policy.
**
** @param session the signed-in person's own session
** @param context that session for the UI API, with the gateway, the
** policy's tool settings and where a proposed write is kept
** @param context that same session for the UI API (its sessionUUID is the
** session's), with the gateway, the policy's tool settings
** and where a proposed write is kept
** @param toolCallId the model's id of the call
** @param toolName the tool called
** @param input the model's input
@@ -55,6 +55,17 @@ public interface ToolCallScope
/*******************************************************************************
** The client channel label the scope's calls are logged under: assistant,
** with the turn when there is one.
*******************************************************************************/
default String clientChannel()
{
return (turnId() == null ? "assistant" : "assistant/" + turnId());
}
/*******************************************************************************
** State a tool keeps for as long as the scope lives (in chat: one turn),
** created on first use and then shared by the scope's calls.
@@ -24,10 +24,16 @@ import com.kingsrook.qbits.assistant.AssistantRuntime;
/*******************************************************************************
** A group of tools that share their code (the UI API data tools, the
** bulk-load tools, ...). A family is how tools get into the
** {@link ToolRegistry}: a public top-level class under
** {@link ToolRegistry#FAMILY_PACKAGE} that implements this interface and has
** a public no-argument constructor is found when the runtime starts. Adding
** a tool family is therefore adding files; no existing class names it.
** {@link ToolRegistry}, and only a registered family does: the QBit's own
** are named by the runtime (AssistantRuntime.ownToolFamilies), and a host or
** another QBit registers one with AssistantQBitConfig.withToolFamily before
** the QBit is produced. A family class that is only on the classpath
** registers nothing.
**
** A registered family's tools run inside the assistant with the signed-in
** person's session, so registering one is trusting its code: every call
** still goes through the checks of engine.ToolCallService first, and the
** tool applies its own checks after them (see {@link ToolCheck}).
**
** tools() runs while the runtime is still being constructed: keep the
** runtime and ask it for collaborators when a tool is offered or called,
@@ -17,20 +17,18 @@
package com.kingsrook.qbits.assistant.tools.registry;
import java.io.IOException;
import java.lang.reflect.Modifier;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.spi.ModelToolDefinition;
import com.kingsrook.qqq.backend.core.utils.ClassPathUtils;
import com.kingsrook.qqq.backend.core.logging.QLogger;
import com.kingsrook.qqq.backend.core.utils.StringUtils;
import static com.kingsrook.qqq.backend.core.logging.LogUtils.logPair;
/*******************************************************************************
@@ -38,22 +36,29 @@ import com.kingsrook.qqq.backend.core.utils.StringUtils;
** looked up: the chat's tool list, a tool call and a confirmed write all go
** through it, so they cannot disagree about which tools exist.
**
** Tools come from {@link ToolFamily} classes found under
** {@link #FAMILY_PACKAGE} when the runtime starts.
** Tools come only from the {@link ToolFamily} instances the runtime is
** given: the QBit's own, which the runtime names, and those a host or
** another QBit registered with AssistantQBitConfig.withToolFamily. Nothing
** is looked for on the classpath: a class that merely exists is not a tool.
*******************************************************************************/
public class ToolRegistry
{
public static final String FAMILY_PACKAGE = "com.kingsrook.qbits.assistant.tools";
private static final QLogger LOG = QLogger.getLogger(ToolRegistry.class);
private final Map<String, AssistantTool> tools;
/*******************************************************************************
** A registry of exactly these tools. Two tools may not share a name.
** A registry of exactly these tools. Two tools may not share a name, and a
** registry without a tool is an error, never an assistant without tools.
*******************************************************************************/
public ToolRegistry(Collection<AssistantTool> tools)
{
if(tools.isEmpty())
{
throw (new IllegalStateException("The assistant has no tools: no tool family was registered, or none has a tool"));
}
Map<String, AssistantTool> byName = new TreeMap<>();
for(AssistantTool tool : tools)
{
@@ -72,41 +77,21 @@ public class ToolRegistry
/*******************************************************************************
** The registry of a runtime: the tools of every family found under
** {@link #FAMILY_PACKAGE}, and of the families given (those that may not
** live in that package). Finding no family is an error, never an assistant
** without tools.
** The registry of a runtime: the tools of exactly the families given, and
** one log line naming them. Starting fails when two tools share a name or
** there is no tool at all.
*******************************************************************************/
public static ToolRegistry discover(AssistantRuntime runtime, List<ToolFamily> given)
public static ToolRegistry of(AssistantRuntime runtime, List<ToolFamily> families)
{
List<ToolFamily> families = new ArrayList<>();
try
{
for(Class<?> type : ClassPathUtils.getClassesInPackage(FAMILY_PACKAGE))
{
if(ToolFamily.class.isAssignableFrom(type) && !type.isInterface() && !Modifier.isAbstract(type.getModifiers()))
{
families.add((ToolFamily) type.getDeclaredConstructor().newInstance());
}
}
}
catch(IOException | ReflectiveOperationException e)
{
throw (new IllegalStateException("The assistant's tool families could not be loaded from " + FAMILY_PACKAGE, e));
}
if(families.isEmpty())
{
throw (new IllegalStateException("No assistant tool family was found under " + FAMILY_PACKAGE));
}
families.sort(Comparator.comparing(family -> family.getClass().getName()));
families.addAll(given);
List<AssistantTool> tools = new ArrayList<>();
for(ToolFamily family : families)
{
tools.addAll(family.tools(runtime));
}
return (new ToolRegistry(tools));
ToolRegistry registry = new ToolRegistry(tools);
LOG.info("Assistant tools registered", logPair("count", registry.tools.size()), logPair("tools", String.join(",", registry.tools.keySet())),
logPair("families", families.stream().map(family -> family.getClass().getName()).toList()));
return (registry);
}
@@ -133,7 +118,11 @@ public class ToolRegistry
/*******************************************************************************
** The definitions of the tools this offer's person is offered now, sorted by
** name (a stable cache prefix).
** name, each input schema with the keys of every object sorted (a stable
** cache prefix). A schema built with Map.of iterates in an order that
** differs from one JVM to the next, and the provider mappings serialise a
** schema in the order it iterates; sorted here, every pod and every restart
** sends the same bytes, whichever family built the schema.
*******************************************************************************/
public List<ModelToolDefinition> definitions(ToolOffer offer)
{
@@ -143,7 +132,8 @@ public class ToolRegistry
ModelToolDefinition definition = tool.definition(offer);
if(definition != null)
{
definitions.add(definition);
definitions.add(new ModelToolDefinition().withName(definition.getName()).withDescription(definition.getDescription()).withStrict(definition.getStrict())
.withInputSchema(definition.getInputSchema() == null ? null : sortedObject(definition.getInputSchema())));
}
}
return (definitions);
@@ -151,6 +141,39 @@ public class ToolRegistry
/*******************************************************************************
** A copy of a JSON object with its keys sorted, at every depth.
*******************************************************************************/
private static Map<String, Object> sortedObject(Map<?, ?> object)
{
Map<String, Object> sorted = new TreeMap<>();
object.forEach((key, value) -> sorted.put(String.valueOf(key), sortedValue(value)));
return (sorted);
}
/*******************************************************************************
** A JSON value with the keys of every object in it sorted. A list keeps its
** order, which is part of what it says.
*******************************************************************************/
private static Object sortedValue(Object value)
{
if(value instanceof Map<?, ?> object)
{
return (sortedObject(object));
}
if(value instanceof List<?> list)
{
List<Object> copy = new ArrayList<>();
list.forEach(item -> copy.add(sortedValue(item)));
return (copy);
}
return (value);
}
/*******************************************************************************
** Strict only when the host names the tool in strictToolNames: the provider
** compiles every strict schema of a request into one grammar, which it
@@ -25,6 +25,7 @@ import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import com.kingsrook.qbits.assistant.engine.PolicyService;
import com.kingsrook.qbits.assistant.metadata.AssistantChannelWall;
import com.kingsrook.qbits.assistant.testing.AssistantTestFixture;
import com.kingsrook.qbits.assistant.testing.BaseAssistantTest;
@@ -111,7 +112,7 @@ class AssistantChannelWallTest extends BaseAssistantTest
assertThat(uiMetadata.status()).isEqualTo(200);
QContext.init(qInstance, AssistantTestFixture.ownerSession(AssistantTestFixture.ADMIN, config));
UiApiToolbox.Outcome listed = config.getRuntime().getToolbox().execute(new UiApiToolbox.Context(AssistantTestFixture.ADMIN, null,
config.getRuntime().getUiApiGateway(), null, null, null, "42"), "list", UiApiToolbox.LIST_TABLES, Map.of());
config.getRuntime().getUiApiGateway(), PolicyService.ToolPolicy.defaults(), null, null, "42"), "list", UiApiToolbox.LIST_TABLES, Map.of());
for(String table : qInstance.getTables().keySet().stream().filter(name -> name.startsWith("assistant")).toList())
{
assertThat(((Map<?, ?>) metadata.json().get("tables")).containsKey(table)).isFalse();
@@ -0,0 +1,659 @@
/*
* 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.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Supplier;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.access.DataAccessMode;
import com.kingsrook.qbits.assistant.api.AssistantApi;
import com.kingsrook.qbits.assistant.api.AssistantApiException;
import com.kingsrook.qbits.assistant.model.AssistantDataAccess;
import com.kingsrook.qbits.assistant.model.AssistantPendingAction;
import com.kingsrook.qbits.assistant.model.AssistantPolicy;
import com.kingsrook.qbits.assistant.model.AssistantToolCall;
import com.kingsrook.qbits.assistant.spi.ModelToolDefinition;
import com.kingsrook.qbits.assistant.testing.AssistantTestFixture;
import com.kingsrook.qbits.assistant.testing.BaseAssistantTest;
import com.kingsrook.qbits.assistant.testing.TurnRunner;
import com.kingsrook.qbits.assistant.testing.UiApiTestServer;
import com.kingsrook.qbits.assistant.tools.UiApiToolbox;
import com.kingsrook.qbits.assistant.tools.bulkload.BulkLoadTools;
import com.kingsrook.qbits.assistant.tools.registry.AssistantTool;
import com.kingsrook.qbits.assistant.tools.registry.ConfirmedToolCall;
import com.kingsrook.qbits.assistant.tools.registry.ToolCallOutcome;
import com.kingsrook.qbits.assistant.tools.registry.ToolCallRequest;
import com.kingsrook.qbits.assistant.tools.registry.ToolCallScope;
import com.kingsrook.qbits.assistant.tools.registry.ToolCheck;
import com.kingsrook.qbits.assistant.tools.registry.ToolFamily;
import com.kingsrook.qbits.assistant.tools.registry.ToolKind;
import com.kingsrook.qbits.assistant.tools.registry.ToolOffer;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolOutcome;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolResult;
import com.kingsrook.qqq.backend.core.actions.ai.DataClassification;
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.session.QSession;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import static com.kingsrook.qbits.assistant.testing.AssistantTestFixture.ADMIN;
import static com.kingsrook.qbits.assistant.testing.AssistantTestFixture.ANALYST;
import static com.kingsrook.qbits.assistant.testing.AssistantTestFixture.EDITOR;
import static com.kingsrook.qbits.assistant.testing.AssistantTestFixture.LIMITED;
import static com.kingsrook.qbits.assistant.testing.AssistantTestFixture.NO_ACCESS;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/*******************************************************************************
** The tool-call service is safe for a caller that has no chat turn around
** it: every check a call needs is made by the service itself, before the
** tool is reached, and the call is logged by the service with its input
** redacted. Each test calls the service directly, the way a later caller
** will, with nothing of the turn engine in between; the UI API test server
** sees every request the assistant makes, so "nothing ran" is what it did
** not see.
*******************************************************************************/
class ToolCallServiceTest extends BaseAssistantTest
{
/////////////////////////////////////////////////////////////////////////
// an obviously made-up number in the shape the input redactor removes //
/////////////////////////////////////////////////////////////////////////
private static final String FAKE_SSN = "123-45-6789";
private static final String THROWING_TOOL = "throwing_tool";
/*******************************************************************************
** A registered family whose tool always throws.
*******************************************************************************/
@Override
protected void customizeConfig(AssistantQBitConfig config)
{
config.withToolFamily(new ThrowingFamily());
}
/*******************************************************************************
**
*******************************************************************************/
@BeforeEach
void forgetEarlierRequests()
{
UiApiTestServer.SCHEDULED_REQUESTS.clear();
}
/*******************************************************************************
** Fix 1: the service refuses a tool the person is not offered now, whoever
** calls: a tool the policy switched off (the kernel's navigate has no policy
** check of its own, and ran before), and a write tool while no table is
** writable. The refusal is the one of a name without a tool, and nothing is
** requested. A tool the same policy allows runs.
*******************************************************************************/
@Test
void aToolThatIsNotOfferedIsRefusedByTheServiceItself()
{
PolicyService.ToolPolicy onlyListTables = new PolicyService.ToolPolicy(true, Set.of(UiApiToolbox.LIST_TABLES), 50, true);
assertUnknownTool(service().execute(request(EDITOR, "navigate", input("target", "blog_post"), onlyListTables)), "navigate");
assertUnknownTool(service().execute(request(EDITOR, UiApiToolbox.COUNT_RECORDS, input("table", "blog_post", "filter", filter("AND")), onlyListTables)), UiApiToolbox.COUNT_RECORDS);
assertUnknownTool(service().execute(request(ADMIN, UiApiToolbox.CREATE_RECORD, input("table", "blog_post", "values", Map.of("title", "Never proposed")), PolicyService.ToolPolicy.defaults())),
UiApiToolbox.CREATE_RECORD);
assertThat(UiApiTestServer.SCHEDULED_REQUESTS).as("nothing was requested").isEmpty();
ToolCallOutcome allowed = service().execute(request(EDITOR, UiApiToolbox.LIST_TABLES, Map.of(), onlyListTables));
assertThat(allowed.result().getOutcome()).isEqualTo(AiToolOutcome.OK);
assertThat(UiApiTestServer.SCHEDULED_REQUESTS).isNotEmpty().allSatisfy(seen -> assertThat(seen).containsEntry("session", EDITOR));
}
/*******************************************************************************
** Fix 2: no policy is not "every tool". The service refuses a request
** without one, and the three tools that read the switch themselves refuse
** too when they are handed a context without a policy.
*******************************************************************************/
@Test
void aCallWithoutAPolicyIsRefused()
{
ToolCallOutcome outcome = service().execute(request(EDITOR, UiApiToolbox.LIST_TABLES, Map.of(), null));
assertRefused(outcome, ToolCallService.CHECK_REQUEST);
assertThat(outcome.result().getModelContent()).contains("no policy");
UiApiToolbox.Context noPolicy = context(EDITOR, null);
AssistantRuntime runtime = config.getRuntime();
QContext.init(qInstance, AssistantTestFixture.ownerSession(EDITOR, config));
AiToolResult data = runtime.getToolbox().execute(noPolicy, "call-1", UiApiToolbox.LIST_TABLES, Map.of()).getResult();
assertThat(data.getOutcome()).isEqualTo(AiToolOutcome.DENIED);
assertThat(data.getPermissionCheck()).isEqualTo("policy:tool");
AiToolResult knowledge = runtime.getKnowledgeTool().execute(noPolicy, "call-2", Map.of("query", "posts"), DataClassification.PUBLIC, null, null);
assertThat(knowledge.getOutcome()).isEqualTo(AiToolOutcome.DENIED);
AiToolResult bulkLoad = runtime.getBulkLoadTools().execute(noPolicy, "call-3", BulkLoadTools.LIST, Map.of("tableName", "blog_post"), false).getResult();
assertThat(bulkLoad.getOutcome()).isEqualTo(AiToolOutcome.DENIED);
assertThat(bulkLoad.getPermissionCheck()).isEqualTo("policy:tool");
assertThat(UiApiTestServer.SCHEDULED_REQUESTS).as("nothing was requested").isEmpty();
}
/*******************************************************************************
** Fix 3: one call cannot mix two people. A session with a context made for
** another session is refused, for a call and for a confirmation, and so is
** a confirmation named for another person than its session's.
*******************************************************************************/
@Test
void aCallWhoseSessionAndContextAreTwoPeopleIsRefused()
{
QSession analyst = AssistantTestFixture.ownerSession(ANALYST, config);
ToolCallOutcome mixed = service().execute(new ToolCallRequest(analyst, context(EDITOR, PolicyService.ToolPolicy.defaults()), "call-1", UiApiToolbox.COUNT_RECORDS,
input("table", "blog_post", "filter", filter("AND")), new Scope()));
assertRefused(mixed, ToolCallService.CHECK_REQUEST);
assertThat(mixed.result().getModelContent()).contains("not the same person");
UiApiToolbox.Outcome mixedConfirmation = service().executeConfirmed(new ConfirmedToolCall(analyst, context(EDITOR, PolicyService.ToolPolicy.defaults()), "call-2",
UiApiToolbox.CREATE_RECORD, "blog_post", null, Map.of("title", "Never created"), ANALYST));
assertThat(mixedConfirmation.getResult().getOutcome()).isEqualTo(AiToolOutcome.DENIED);
assertThat(mixedConfirmation.getResult().getPermissionCheck()).isEqualTo(ToolCallService.CHECK_REQUEST);
UiApiToolbox.Outcome namedForAnother = service().executeConfirmed(new ConfirmedToolCall(analyst, context(ANALYST, PolicyService.ToolPolicy.defaults()), "call-3",
UiApiToolbox.CREATE_RECORD, "blog_post", null, Map.of("title", "Never created"), EDITOR));
assertThat(namedForAnother.getResult().getOutcome()).isEqualTo(AiToolOutcome.DENIED);
assertThat(namedForAnother.getResult().getPermissionCheck()).isEqualTo(ToolCallService.CHECK_REQUEST);
assertThat(UiApiTestServer.SCHEDULED_REQUESTS).as("nothing was requested").isEmpty();
}
/*******************************************************************************
** Fix 4: executeAndLog runs the call and writes its row itself, with the
** input redacted, for a call that ran, for one that was refused and for one
** whose tool threw. The caller prepares nothing.
*******************************************************************************/
@Test
void executeAndLogWritesTheRowWithRedactedInputAlsoForARefusedCall() throws Exception
{
Map<String, Object> input = input("table", "blog_post", "filter", filter("AND", "title", "EQUALS", List.of(FAKE_SSN)));
ToolCallOutcome ran = service().executeAndLog(request(EDITOR, UiApiToolbox.COUNT_RECORDS, input, PolicyService.ToolPolicy.defaults()));
assertThat(ran.result().getOutcome()).isEqualTo(AiToolOutcome.OK);
PolicyService.ToolPolicy onlyListTables = new PolicyService.ToolPolicy(true, Set.of(UiApiToolbox.LIST_TABLES), 50, true);
ToolCallOutcome refused = service().executeAndLog(request(EDITOR, UiApiToolbox.COUNT_RECORDS, input, onlyListTables));
assertUnknownTool(refused, UiApiToolbox.COUNT_RECORDS);
ToolCallOutcome threw = service().executeAndLog(request(EDITOR, THROWING_TOOL, input("note", FAKE_SSN), PolicyService.ToolPolicy.defaults()));
assertThat(threw.result().getOutcome()).isEqualTo(AiToolOutcome.ERROR);
assertThat(threw.result().getModelContent()).doesNotContain("boom");
List<AssistantToolCall> rows = everyToolCall();
assertThat(rows).hasSize(3).allSatisfy(row ->
{
assertThat(row.getUserId()).isEqualTo(EDITOR);
assertThat(row.getToolCallId()).isEqualTo("call-1");
assertThat(row.getClientChannel()).isEqualTo("assistant");
assertThat(row.getInputJson()).doesNotContain(FAKE_SSN);
});
AssistantToolCall ranRow = rows.get(0);
assertThat(ranRow.getToolName()).isEqualTo(UiApiToolbox.COUNT_RECORDS);
assertThat(ranRow.getOutcome()).isEqualTo("OK");
assertThat(ranRow.getInputJson()).contains("blog_post").contains("title");
assertThat(ranRow.getUiApiOperation()).isEqualTo("POST /table/{tableName}/count");
assertThat(ranRow.getHttpStatus()).isEqualTo(200);
AssistantToolCall refusedRow = rows.get(1);
assertThat(refusedRow.getToolName()).isEqualTo(UiApiToolbox.COUNT_RECORDS);
assertThat(refusedRow.getOutcome()).isEqualTo("INVALID_INPUT");
assertThat(refusedRow.getResultSummary()).isEqualTo("Unknown tool");
assertThat(refusedRow.getInputJson()).contains("blog_post");
assertThat(refusedRow.getHttpStatus()).isNull();
AssistantToolCall threwRow = rows.get(2);
assertThat(threwRow.getToolName()).isEqualTo(THROWING_TOOL);
assertThat(threwRow.getOutcome()).isEqualTo("ERROR");
assertThat(threwRow.getResultSummary()).isEqualTo("Failed");
assertThat(requestsFor("blog_post")).as("only the call that was offered reached the UI API").hasSize(1);
}
/*******************************************************************************
** Fix 4: a call that names no person at all is refused and, having nobody
** to belong to, leaves no row.
*******************************************************************************/
@Test
void aCallThatNamesNoPersonIsRefusedAndHasNoRow() throws Exception
{
ToolCallOutcome nobody = service().executeAndLog(new ToolCallRequest(null, context(EDITOR, PolicyService.ToolPolicy.defaults()), "call-1", UiApiToolbox.LIST_TABLES, Map.of(), new Scope()));
assertRefused(nobody, ToolCallService.CHECK_REQUEST);
assertThat(service().executeAndLog(null).result().getOutcome()).isEqualTo(AiToolOutcome.DENIED);
assertThat(everyToolCall()).isEmpty();
assertThat(UiApiTestServer.SCHEDULED_REQUESTS).isEmpty();
}
/*******************************************************************************
** Fix 6: a write the person confirms while the kill switch is on is refused
** before the pending action is claimed, so it stays pending and can be
** confirmed once the switch is off. The service refuses it too when it is
** asked directly.
*******************************************************************************/
@Test
void aConfirmationWhileTheKillSwitchIsOnIsRefusedAndTheActionStaysPending() throws Exception
{
Integer actionId = proposeAPost();
policy(new AssistantPolicy().withName("default").withKillSwitch(true));
UiApiTestServer.SCHEDULED_REQUESTS.clear();
QSession admin = AssistantTestFixture.sessionFor(ADMIN);
AssistantApi api = new AssistantApi(config.getRuntime(), admin);
QContext.init(qInstance, admin);
assertThatThrownBy(() -> api.confirmAction(actionId)).isInstanceOfSatisfying(AssistantApiException.class, e ->
{
assertThat(e.getStatus()).isEqualTo(503);
assertThat(e.getCode()).isEqualTo("KILL_SWITCH");
});
QContext.init(qInstance, admin);
assertThat(api.getAction(actionId)).containsEntry("status", AssistantPendingAction.STATUS_PENDING);
assertThat(requestsFor("blog_post")).as("nothing was written").isEmpty();
assertThat(everyToolCall()).as("only the proposal's row").hasSize(1);
UiApiToolbox.Outcome direct = service().executeConfirmed(new ConfirmedToolCall(admin, context(ADMIN, PolicyService.ToolPolicy.defaults()), "call-1", UiApiToolbox.CREATE_RECORD,
"blog_post", null, Map.of("title", "Never created"), ADMIN));
assertThat(direct.getResult().getOutcome()).isEqualTo(AiToolOutcome.DENIED);
assertThat(direct.getResult().getPermissionCheck()).isEqualTo(ToolCallService.CHECK_KILL_SWITCH);
assertThat(requestsFor("blog_post")).as("nothing was written").isEmpty();
}
/*******************************************************************************
** Fix 6: so is a confirmation of a tool the policy has switched off since
** the write was proposed.
*******************************************************************************/
@Test
void aConfirmationOfAToolThePolicySwitchedOffIsRefusedAndTheActionStaysPending() throws Exception
{
Integer actionId = proposeAPost();
policy(new AssistantPolicy().withName("default").withAllowedTools(UiApiToolbox.LIST_TABLES));
UiApiTestServer.SCHEDULED_REQUESTS.clear();
QSession admin = AssistantTestFixture.sessionFor(ADMIN);
AssistantApi api = new AssistantApi(config.getRuntime(), admin);
QContext.init(qInstance, admin);
assertThatThrownBy(() -> api.confirmAction(actionId)).isInstanceOfSatisfying(AssistantApiException.class, e -> assertThat(e.getStatus()).isEqualTo(403));
QContext.init(qInstance, admin);
assertThat(api.getAction(actionId)).containsEntry("status", AssistantPendingAction.STATUS_PENDING);
assertThat(requestsFor("blog_post")).as("nothing was written").isEmpty();
UiApiToolbox.Outcome direct = service().executeConfirmed(new ConfirmedToolCall(admin, context(ADMIN, PolicyService.ToolPolicy.defaults()), "call-1", UiApiToolbox.CREATE_RECORD,
"blog_post", null, Map.of("title", "Never created"), ADMIN));
assertThat(direct.getResult().getOutcome()).isEqualTo(AiToolOutcome.DENIED);
assertThat(direct.getResult().getPermissionCheck()).isEqualTo(ToolCallService.CHECK_TOOL_SWITCH);
}
/*******************************************************************************
** Fix 9: with no chat turn around the call, the service itself requires
** assistant access, and checks it before the kill switch: a person without
** access is told so even while the switch is on.
*******************************************************************************/
@Test
void assistantAccessIsRequiredByTheService() throws Exception
{
assertRefused(service().execute(request(NO_ACCESS, UiApiToolbox.LIST_TABLES, Map.of(), PolicyService.ToolPolicy.defaults())), ToolCallService.CHECK_ACCESS);
policy(new AssistantPolicy().withName("default").withKillSwitch(true));
assertRefused(service().execute(request(NO_ACCESS, UiApiToolbox.LIST_TABLES, Map.of(), PolicyService.ToolPolicy.defaults())), ToolCallService.CHECK_ACCESS);
UiApiToolbox.Outcome confirmation = service().executeConfirmed(new ConfirmedToolCall(AssistantTestFixture.sessionFor(NO_ACCESS), context(NO_ACCESS, PolicyService.ToolPolicy.defaults()),
"call-1", UiApiToolbox.CREATE_RECORD, "blog_post", null, Map.of("title", "Never created"), NO_ACCESS));
assertThat(confirmation.getResult().getPermissionCheck()).isEqualTo(ToolCallService.CHECK_ACCESS);
assertThat(UiApiTestServer.SCHEDULED_REQUESTS).as("nothing was requested").isEmpty();
}
/*******************************************************************************
** Fix 9: the kill switch stops a call that no turn admitted.
*******************************************************************************/
@Test
void theKillSwitchIsCheckedByTheService() throws Exception
{
assertThat(service().execute(request(EDITOR, UiApiToolbox.LIST_TABLES, Map.of(), PolicyService.ToolPolicy.defaults())).result().getOutcome()).isEqualTo(AiToolOutcome.OK);
UiApiTestServer.SCHEDULED_REQUESTS.clear();
policy(new AssistantPolicy().withName("default").withKillSwitch(true));
ToolCallOutcome outcome = service().execute(request(EDITOR, UiApiToolbox.LIST_TABLES, Map.of(), PolicyService.ToolPolicy.defaults()));
assertRefused(outcome, ToolCallService.CHECK_KILL_SWITCH);
assertThat(UiApiTestServer.SCHEDULED_REQUESTS).as("nothing was requested").isEmpty();
}
/*******************************************************************************
** Fix 9: the hard budget stops a call that no turn admitted, for the person
** who is over it and for nobody else; a scope that was let through is not
** asked again (the check is made when a scope's first call arrives).
*******************************************************************************/
@Test
void theHardBudgetIsCheckedByTheService() throws Exception
{
provider.thenText("Hello.");
runner.run(EDITOR, "Hello");
UiApiTestServer.SCHEDULED_REQUESTS.clear();
ToolCallRequest admitted = request(EDITOR, UiApiToolbox.LIST_TABLES, Map.of(), PolicyService.ToolPolicy.defaults());
assertThat(service().execute(admitted).result().getOutcome()).isEqualTo(AiToolOutcome.OK);
UiApiTestServer.SCHEDULED_REQUESTS.clear();
config.getBudgets().withUserDailyTokensHard(1L);
ToolCallOutcome outcome = service().execute(request(EDITOR, UiApiToolbox.LIST_TABLES, Map.of(), PolicyService.ToolPolicy.defaults()));
assertRefused(outcome, ToolCallService.CHECK_BUDGET);
assertThat(outcome.result().getModelContent()).contains("daily assistant usage limit");
assertThat(UiApiTestServer.SCHEDULED_REQUESTS).as("nothing was requested").isEmpty();
assertThat(service().execute(request(LIMITED, UiApiToolbox.LIST_TABLES, Map.of(), PolicyService.ToolPolicy.defaults())).result().getPermissionCheck())
.as("a person who has used nothing").isNotEqualTo(ToolCallService.CHECK_BUDGET);
assertThat(service().execute(admitted).result().getOutcome()).as("the scope admitted before the limit").isEqualTo(AiToolOutcome.OK);
}
/*******************************************************************************
**
*******************************************************************************/
private ToolCallService service()
{
return (config.getRuntime().getToolCalls());
}
/*******************************************************************************
** A call of a person's, in a scope of its own, with no turn.
*******************************************************************************/
private ToolCallRequest request(String person, String toolName, Map<String, Object> input, PolicyService.ToolPolicy policy)
{
return (new ToolCallRequest(AssistantTestFixture.ownerSession(person, config), context(person, policy), "call-1", toolName, input, new Scope()));
}
/*******************************************************************************
** A person's tool context, outside a turn.
*******************************************************************************/
private UiApiToolbox.Context context(String person, PolicyService.ToolPolicy policy)
{
return (new UiApiToolbox.Context(person, "UTC", config.getRuntime().getUiApiGateway(), policy, null));
}
/*******************************************************************************
** A refusal by one of the service's own checks.
*******************************************************************************/
private static void assertRefused(ToolCallOutcome outcome, String check)
{
assertThat(outcome.extras()).isNull();
assertThat(outcome.result().getOutcome()).isEqualTo(AiToolOutcome.DENIED);
assertThat(outcome.result().getPermissionCheck()).isEqualTo(check);
}
/*******************************************************************************
** The refusal of a tool that is not offered: the one of an unknown name.
*******************************************************************************/
private static void assertUnknownTool(ToolCallOutcome outcome, String toolName)
{
assertThat(outcome.extras()).as(toolName).isNull();
assertThat(outcome.result().getOutcome()).as(toolName).isEqualTo(AiToolOutcome.INVALID_INPUT);
assertThat(outcome.result().getSummary()).as(toolName).isEqualTo("Unknown tool");
assertThat(outcome.result().getModelContent()).isEqualTo("Unknown tool \"" + toolName + "\". Use only the tools provided.");
}
/*******************************************************************************
** A write the administrator's turn proposed: the pending action's id.
*******************************************************************************/
private Integer proposeAPost() throws Exception
{
QContext.init(qInstance, AssistantTestFixture.sessionFor(ADMIN));
try
{
config.getRuntime().getStore().insert(new AssistantDataAccess().withObjectType(AssistantDataAccess.TYPE_TABLE).withObjectName("blog_post").withMode(DataAccessMode.READ_WRITE.name()));
config.getRuntime().getDataAccessService().invalidate();
}
finally
{
QContext.clear();
}
provider.thenToolUse(UiApiToolbox.CREATE_RECORD, input("table", "blog_post", "values", Map.of("title", "Proposed in a turn", "status", "DRAFT"))).thenText("Please confirm.");
TurnRunner.Outcome outcome = runner.run(ADMIN, "Create a draft post");
assertThat(outcome.toolCalls()).hasSize(1);
assertThat(outcome.toolCalls().get(0).getPendingActionId()).isNotNull();
return (outcome.toolCalls().get(0).getPendingActionId());
}
/*******************************************************************************
** The requests the UI API server saw for a table.
*******************************************************************************/
private static List<Map<String, String>> requestsFor(String tableName)
{
return (UiApiTestServer.SCHEDULED_REQUESTS.stream().filter(request -> request.get("path").contains("/table/" + tableName)).toList());
}
/*******************************************************************************
** Every tool-call log row there is, whoever it belongs to, oldest first.
*******************************************************************************/
private List<AssistantToolCall> everyToolCall() throws Exception
{
QContext.init(qInstance, AssistantTestFixture.sessionFor(ADMIN));
try
{
return (config.getRuntime().getStore().query(AssistantToolCall.TABLE_NAME, new QQueryFilter(), AssistantToolCall.class).stream()
.sorted((a, b) -> a.getId().compareTo(b.getId())).toList());
}
finally
{
QContext.clear();
}
}
/*******************************************************************************
** The administrator's policy row.
*******************************************************************************/
private void policy(AssistantPolicy policy) throws Exception
{
QContext.init(qInstance, AssistantTestFixture.sessionFor(ADMIN));
try
{
config.getRuntime().getStore().insert(policy);
config.getRuntime().getPolicyService().invalidate();
}
finally
{
QContext.clear();
}
}
/*******************************************************************************
** A scope with no turn: what a caller outside the chat gives the service.
*******************************************************************************/
private static final class Scope implements ToolCallScope
{
private final Map<Class<?>, Object> state = new HashMap<>();
/*******************************************************************************
**
*******************************************************************************/
@Override
public DataClassification classification()
{
return (DataClassification.PUBLIC);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public Integer conversationId()
{
return (null);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public Integer turnId()
{
return (null);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public boolean hipaaRoute()
{
return (false);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public <T> T state(Class<T> type, Supplier<T> create)
{
return (type.cast(state.computeIfAbsent(type, key -> create.get())));
}
}
/*******************************************************************************
** A registered family with one tool, which throws.
*******************************************************************************/
private static final class ThrowingFamily implements ToolFamily
{
/*******************************************************************************
**
*******************************************************************************/
@Override
public List<AssistantTool> tools(AssistantRuntime runtime)
{
return (List.of(new AssistantTool()
{
/*******************************************************************************
**
*******************************************************************************/
@Override
public String name()
{
return (THROWING_TOOL);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public ToolKind kind()
{
return (ToolKind.READ);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public List<ToolCheck> checks()
{
return (List.of(ToolCheck.TOOL_SWITCH));
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public ModelToolDefinition definition(ToolOffer offer)
{
return (offer.allows(THROWING_TOOL) ? new ModelToolDefinition().withName(THROWING_TOOL).withDescription("Throws.").withStrict(false).withInputSchema(Map.of("type", "object")) : null);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public ToolCallOutcome execute(ToolCallRequest request)
{
throw (new IllegalStateException("boom"));
}
}));
}
}
}
@@ -386,7 +386,7 @@ class AnalysisSaveTest extends BaseAssistantTest
QContext.init(qInstance, AssistantTestFixture.ownerSession(user, config));
try
{
return (config.getRuntime().getToolbox().execute(new UiApiToolbox.Context(user, null, config.getRuntime().getUiApiGateway(), policy,
return (config.getRuntime().getToolbox().execute(new UiApiToolbox.Context(user, null, config.getRuntime().getUiApiGateway(), policy == null ? PolicyService.ToolPolicy.defaults() : policy,
write -> config.getRuntime().getStore().insert(new AssistantPendingAction().withUserId(user).withToolCallId(write.toolCallId()).withToolName(write.toolName()).withTableName(write.tableName())
.withRecordId(write.recordId()).withUiApiOperation(write.operation()).withParamsJson(com.kingsrook.qqq.backend.core.utils.JsonUtils.toJson(write.values()))
.withSummary(write.summary()).withStatus(AssistantPendingAction.STATUS_PENDING)
@@ -286,7 +286,7 @@ class AnalysisToolsTest extends BaseAssistantTest
QContext.init(qInstance, AssistantTestFixture.ownerSession(user, config));
try
{
return (config.getRuntime().getToolbox().execute(new UiApiToolbox.Context(user, null, config.getRuntime().getUiApiGateway(), policy, null), "call-1", AnalysisTools.ANALYZE_RECORDS, in));
return (config.getRuntime().getToolbox().execute(new UiApiToolbox.Context(user, null, config.getRuntime().getUiApiGateway(), policy == null ? PolicyService.ToolPolicy.defaults() : policy, null), "call-1", AnalysisTools.ANALYZE_RECORDS, in));
}
finally
{
@@ -0,0 +1,116 @@
/*
* 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.tools;
import java.util.List;
import java.util.Map;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.spi.ModelToolDefinition;
import com.kingsrook.qbits.assistant.tools.registry.AssistantTool;
import com.kingsrook.qbits.assistant.tools.registry.ToolCallOutcome;
import com.kingsrook.qbits.assistant.tools.registry.ToolCallRequest;
import com.kingsrook.qbits.assistant.tools.registry.ToolCheck;
import com.kingsrook.qbits.assistant.tools.registry.ToolFamily;
import com.kingsrook.qbits.assistant.tools.registry.ToolKind;
import com.kingsrook.qbits.assistant.tools.registry.ToolOffer;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolOutcome;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolResult;
/*******************************************************************************
** A tool family as another jar could bring one: a public class with a public
** no-argument constructor under the assistant's tools package. Being on the
** classpath registers nothing (ToolRegistryTest); a host that wants it
** registers it with AssistantQBitConfig.withToolFamily.
*******************************************************************************/
public class ClasspathOnlyToolFamily implements ToolFamily
{
public static final String NAME = "classpath_only_tool";
/*******************************************************************************
**
*******************************************************************************/
@Override
public List<AssistantTool> tools(AssistantRuntime runtime)
{
return (List.of(new Tool()));
}
/*******************************************************************************
** A tool that reads nothing and says it ran.
*******************************************************************************/
private static final class Tool implements AssistantTool
{
/*******************************************************************************
**
*******************************************************************************/
@Override
public String name()
{
return (NAME);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public ToolKind kind()
{
return (ToolKind.READ);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public List<ToolCheck> checks()
{
return (List.of(ToolCheck.TOOL_SWITCH));
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public ModelToolDefinition definition(ToolOffer offer)
{
return (offer.allows(NAME) ? new ModelToolDefinition().withName(NAME).withDescription("Says that it ran.").withStrict(false).withInputSchema(Map.of("type", "object")) : null);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public ToolCallOutcome execute(ToolCallRequest request)
{
return (ToolCallOutcome.of(new AiToolResult().withToolCallId(request.toolCallId()).withToolName(NAME).withOutcome(AiToolOutcome.OK).withSummary("Ran").withModelContent("ran")));
}
}
}
@@ -184,7 +184,7 @@ class UiApiToolsSecurityTest extends BaseAssistantTest
QContext.init(qInstance, AssistantTestFixture.ownerSession(user, config));
try
{
return (config.getRuntime().getToolbox().execute(new UiApiToolbox.Context(user, null, config.getRuntime().getUiApiGateway(), policy,
return (config.getRuntime().getToolbox().execute(new UiApiToolbox.Context(user, null, config.getRuntime().getUiApiGateway(), policy == null ? PolicyService.ToolPolicy.defaults() : policy,
write -> config.getRuntime().getStore().insert(new AssistantPendingAction().withUserId(user).withToolName(write.toolName()).withTableName(write.tableName())
.withRecordId(write.recordId()).withUiApiOperation(write.operation()).withParamsJson(com.kingsrook.qqq.backend.core.utils.JsonUtils.toJson(write.values()))
.withSummary(write.summary()).withStatus(AssistantPendingAction.STATUS_PENDING)
@@ -211,7 +211,7 @@ class BulkLoadToolsTest extends BaseAssistantTest
BulkLoadTools older = new BulkLoadTools(qInstance, config.getRuntime().getDataAccessService(), null);
FakeBulkLoadServer oldApi = new FakeBulkLoadServer(false);
QContext.init(qInstance, AssistantTestFixture.ownerSession(AssistantTestFixture.ADMIN, config));
UiApiToolbox.Outcome outcome = older.execute(new UiApiToolbox.Context("s", null, oldApi, null, null), "c", BulkLoadTools.SUGGEST,
UiApiToolbox.Outcome outcome = older.execute(new UiApiToolbox.Context("s", null, oldApi, PolicyService.ToolPolicy.defaults(), null), "c", BulkLoadTools.SUGGEST,
Map.of("tableName", "blog_post", "uploadRef", "up_1"), false);
assertThat(outcome.getResult().getOutcome()).isEqualTo(AiToolOutcome.ERROR);
assertThat(outcome.getResult().getModelContent()).contains("does not offer");
@@ -18,6 +18,7 @@ package com.kingsrook.qbits.assistant.tools.registry;
import com.kingsrook.qbits.assistant.engine.ToolCallService;
import com.kingsrook.qbits.assistant.engine.TurnEngine;
import com.kingsrook.qbits.assistant.model.AssistantToolCall;
import com.tngtech.archunit.core.domain.JavaClasses;
import com.tngtech.archunit.core.importer.ClassFileImporter;
@@ -33,7 +34,8 @@ import static com.tngtech.archunit.lang.syntax.ArchRuleDefinition.noClasses;
** the tools package, so it may touch neither the store, the kernel nor a QQQ
** action):
** - outside the tools package, only ToolCallService runs a tool;
** - only ToolCallService writes a tool-call log row.
** - only ToolCallService writes a tool-call log row;
** - only the turn engine runs a call and logs it in two steps.
*******************************************************************************/
class ToolCallArchitectureTest
{
@@ -85,4 +87,21 @@ class ToolCallArchitectureTest
.check(classes);
}
/*******************************************************************************
** A caller cannot forget the log row. Only the turn engine uses the two
** halves of a call (execute in its workers, log in call order afterwards);
** every other caller has executeAndLog, which does both.
*******************************************************************************/
@Test
void onlyTheTurnEngineRunsACallAndLogsItInTwoSteps()
{
noClasses().that().resideInAPackage(BASE).and().doNotHaveFullyQualifiedName(ToolCallService.class.getName()).and().doNotHaveFullyQualifiedName(TurnEngine.class.getName())
.should().callMethod(ToolCallService.class, "execute", ToolCallRequest.class)
.orShould().callMethod(ToolCallService.class, "log", ToolCallRequest.class, ToolCallOutcome.class)
.because("a tool call that is run is logged: callers other than the turn engine use ToolCallService.executeAndLog")
.check(classes);
}
}
@@ -18,21 +18,29 @@ package com.kingsrook.qbits.assistant.tools.registry;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.engine.KernelToolFamily;
import com.kingsrook.qbits.assistant.engine.PolicyService;
import com.kingsrook.qbits.assistant.engine.ToolCallService;
import com.kingsrook.qbits.assistant.provider.fake.ScriptedModelProvider;
import com.kingsrook.qbits.assistant.spi.ModelToolDefinition;
import com.kingsrook.qbits.assistant.testing.AssistantTestFixture;
import com.kingsrook.qbits.assistant.testing.BaseAssistantTest;
import com.kingsrook.qbits.assistant.tools.AnalysisTools;
import com.kingsrook.qbits.assistant.tools.ClasspathOnlyToolFamily;
import com.kingsrook.qbits.assistant.tools.KnowledgeTool;
import com.kingsrook.qbits.assistant.tools.PresentResultsTool;
import com.kingsrook.qbits.assistant.tools.UiApiToolbox;
import com.kingsrook.qbits.assistant.tools.bulkload.BulkLoadTools;
import com.kingsrook.qqq.backend.core.actions.ai.AiToolOutcome;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.model.metadata.QInstance;
import org.junit.jupiter.api.Test;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -101,17 +109,161 @@ class ToolRegistryTest extends BaseAssistantTest
/*******************************************************************************
** The families under the tools package are found without being named; the
** kernel's two tools are the ones the runtime gives.
** The registry holds the tools of exactly the families it is given: the
** QBit's own, which the runtime names, are every tool the code has.
*******************************************************************************/
@Test
void familiesUnderTheToolsPackageAreFound()
void onlyTheFamiliesGivenAreRegistered()
{
ToolRegistry withoutKernel = ToolRegistry.discover(config.getRuntime(), List.of());
assertThat(withoutKernel.tools()).extracting(AssistantTool::name).doesNotContain("navigate", "search_help").hasSize(KINDS.size() - 2);
ToolRegistry own = ToolRegistry.of(config.getRuntime(), AssistantRuntime.ownToolFamilies());
assertThat(own.tools()).extracting(AssistantTool::name).containsExactlyElementsOf(KINDS.keySet());
ToolRegistry complete = ToolRegistry.discover(config.getRuntime(), List.of(new KernelToolFamily()));
assertThat(complete.tools()).extracting(AssistantTool::name).containsExactlyElementsOf(KINDS.keySet());
ToolRegistry kernelOnly = ToolRegistry.of(config.getRuntime(), List.of(new KernelToolFamily()));
assertThat(kernelOnly.tools()).extracting(AssistantTool::name).containsExactly("navigate", "search_help");
}
/*******************************************************************************
** A family that is only on the classpath (ClasspathOnlyToolFamily is what a
** scan of the tools package used to pick up: public, a no-argument
** constructor, under that package) is not registered, its tool is not
** offered, and a call to it is refused like any name without a tool.
*******************************************************************************/
@Test
void aFamilyThatIsOnlyOnTheClasspathIsNotRegisteredOrOffered()
{
assertThat(new ClasspathOnlyToolFamily().tools(config.getRuntime())).extracting(AssistantTool::name).containsExactly(ClasspathOnlyToolFamily.NAME);
assertThat(config.getRuntime().getToolRegistry().find(ClasspathOnlyToolFamily.NAME)).isNull();
assertThat(config.getRuntime().getToolRegistry().tools()).extracting(AssistantTool::name).containsExactlyElementsOf(KINDS.keySet());
assertThat(offeredToTheEditor(qInstance, config)).doesNotContain(ClasspathOnlyToolFamily.NAME);
ToolCallOutcome outcome = config.getRuntime().getToolCalls().execute(request(config, ClasspathOnlyToolFamily.NAME));
assertThat(outcome.result().getOutcome()).isEqualTo(AiToolOutcome.INVALID_INPUT);
assertThat(outcome.result().getSummary()).isEqualTo("Unknown tool");
}
/*******************************************************************************
** The one way in from outside the QBit: a host (or another QBit) registers
** the family on the config. Its tool is then in the registry beside the
** assistant's own, offered, and run through the service.
*******************************************************************************/
@Test
void aHostRegistersAFamilyOnTheConfig() throws Exception
{
AssistantQBitConfig hostConfig = new AssistantQBitConfig().withProviderName(ScriptedModelProvider.DEFAULT_NAME).withModelProvider(new ScriptedModelProvider())
.withToolFamily(new ClasspathOnlyToolFamily()).withToolFamily(null);
QInstance host = AssistantTestFixture.build(hostConfig);
try
{
assertThat(hostConfig.getToolFamilies()).hasSize(1);
assertThat(hostConfig.getRuntime().getToolRegistry().tools()).extracting(AssistantTool::name).contains(ClasspathOnlyToolFamily.NAME).containsAll(KINDS.keySet()).hasSize(KINDS.size() + 1);
assertThat(offeredToTheEditor(host, hostConfig)).contains(ClasspathOnlyToolFamily.NAME);
ToolCallOutcome outcome = hostConfig.getRuntime().getToolCalls().execute(request(hostConfig, ClasspathOnlyToolFamily.NAME));
assertThat(outcome.result().getOutcome()).isEqualTo(AiToolOutcome.OK);
assertThat(outcome.result().getModelContent()).isEqualTo("ran");
}
finally
{
hostConfig.getRuntime().close();
}
}
/*******************************************************************************
** Two JVMs give the provider the same bytes: a schema built with Map.of
** iterates in a different order in each, so the registry sorts the keys of
** every object of a schema. Here the same schema is built twice with every
** object's keys in opposite orders (as two JVMs would hold it), and the
** tool list serialised as it iterates, with no sorting by the serialiser.
*******************************************************************************/
@Test
void theToolListIsTheSameBytesWhateverOrderASchemaWasBuiltIn() throws Exception
{
String forward = serialised(new ToolRegistry(List.of(new WithSchema(schema(false)))));
String backward = serialised(new ToolRegistry(List.of(new WithSchema(schema(true)))));
assertThat(backward).isEqualTo(forward);
assertThat(forward).isEqualTo("[{\"name\":\"with_schema\",\"description\":\"A tool with a schema.\",\"inputSchema\":{"
+ "\"properties\":{\"alpha\":{\"description\":\"The alpha.\",\"type\":\"string\"},"
+ "\"beta\":{\"items\":{\"description\":\"One beta.\",\"type\":\"string\"},\"type\":\"array\"},"
+ "\"gamma\":{\"anyOf\":[{\"description\":\"A number.\",\"type\":\"integer\"},{\"description\":\"A word.\",\"type\":\"string\"}]}},"
+ "\"required\":[\"gamma\",\"alpha\"],\"type\":\"object\"},\"strict\":false}]");
}
/*******************************************************************************
** The definitions the registry offers, as JSON in the order they iterate.
*******************************************************************************/
private static String serialised(ToolRegistry registry) throws Exception
{
return (new ObjectMapper().writeValueAsString(registry.definitions(new ToolOffer(PolicyService.ToolPolicy.defaults(), null, null))));
}
/*******************************************************************************
** One schema, with the keys of every object put in one order or the other.
** The lists (required, anyOf) keep their order: it is part of what they say.
*******************************************************************************/
private static Map<String, Object> schema(boolean reversed)
{
Map<String, Object> properties = object(reversed,
"alpha", object(reversed, "type", "string", "description", "The alpha."),
"beta", object(reversed, "type", "array", "items", object(reversed, "type", "string", "description", "One beta.")),
"gamma", object(reversed, "anyOf", List.of(object(reversed, "type", "integer", "description", "A number."), object(reversed, "type", "string", "description", "A word."))));
return (object(reversed, "type", "object", "properties", properties, "required", List.of("gamma", "alpha")));
}
/*******************************************************************************
** A map that iterates its keys in the order given, or in the opposite one.
*******************************************************************************/
private static Map<String, Object> object(boolean reversed, Object... keyValues)
{
Map<String, Object> map = new LinkedHashMap<>();
for(int i = 0; i < keyValues.length; i += 2)
{
int at = reversed ? keyValues.length - 2 - i : i;
map.put((String) keyValues[at], keyValues[at + 1]);
}
return (map);
}
/*******************************************************************************
** The names of the tools the chat offers the editor.
*******************************************************************************/
private static List<String> offeredToTheEditor(QInstance instance, AssistantQBitConfig config)
{
QContext.init(instance, AssistantTestFixture.ownerSession(AssistantTestFixture.EDITOR, config));
try
{
return (config.getRuntime().getContextBuilder().tools().stream().map(ModelToolDefinition::getName).toList());
}
finally
{
QContext.clear();
}
}
/*******************************************************************************
** A call of the editor's to a tool, outside a turn.
*******************************************************************************/
private static ToolCallRequest request(AssistantQBitConfig config, String toolName)
{
UiApiToolbox.Context context = new UiApiToolbox.Context(AssistantTestFixture.EDITOR, "UTC", config.getRuntime().getUiApiGateway(), PolicyService.ToolPolicy.defaults(), null);
return (new ToolCallRequest(AssistantTestFixture.ownerSession(AssistantTestFixture.EDITOR, config), context, "call-1", toolName, Map.of(), null));
}
@@ -128,6 +280,17 @@ class ToolRegistryTest extends BaseAssistantTest
assertThatThrownBy(() -> new ToolRegistry(twice)).isInstanceOf(IllegalStateException.class).hasMessageContaining(twice.get(0).name());
assertThatThrownBy(() -> new ToolRegistry(List.of(new Unnamed()))).isInstanceOf(IllegalStateException.class).hasMessageContaining("name");
///////////////////////////////////////////////////////////////////////
// no tool at all, and a registered family whose tool has the name //
// of one of the assistant's own //
///////////////////////////////////////////////////////////////////////
assertThatThrownBy(() -> new ToolRegistry(List.of())).isInstanceOf(IllegalStateException.class).hasMessageContaining("no tools");
assertThatThrownBy(() -> ToolRegistry.of(config.getRuntime(), List.of())).isInstanceOf(IllegalStateException.class).hasMessageContaining("no tools");
List<ToolFamily> clash = new ArrayList<>(AssistantRuntime.ownToolFamilies());
clash.add(new KernelToolFamily());
assertThatThrownBy(() -> ToolRegistry.of(config.getRuntime(), clash)).isInstanceOf(IllegalStateException.class).hasMessageContaining("navigate");
}
@@ -162,7 +325,7 @@ class ToolRegistryTest extends BaseAssistantTest
for(String name : List.of(UiApiToolbox.LIST_TABLES, AnalysisTools.ANALYZE_RECORDS, BulkLoadTools.LIST, KnowledgeTool.NAME, PresentResultsTool.NAME, "navigate", "no_such_tool"))
{
UiApiToolbox.Outcome before = config.getRuntime().getToolbox().executeConfirmedWrite(context(), "call-1", name, "blog_post", "1", Map.of("title", "x"));
UiApiToolbox.Outcome now = config.getRuntime().getToolCalls().executeConfirmed(new ConfirmedToolCall(context(), "call-1", name, "blog_post", "1", Map.of("title", "x"), AssistantTestFixture.EDITOR));
UiApiToolbox.Outcome now = config.getRuntime().getToolCalls().executeConfirmed(new ConfirmedToolCall(AssistantTestFixture.sessionFor(AssistantTestFixture.EDITOR), context(), "call-1", name, "blog_post", "1", Map.of("title", "x"), AssistantTestFixture.EDITOR));
assertThat(now.getResult().getOutcome()).as(name).isEqualTo(before.getResult().getOutcome()).isEqualTo(AiToolOutcome.INVALID_INPUT);
assertThat(now.getResult().getSummary()).as(name).isEqualTo(before.getResult().getSummary());
@@ -214,6 +377,81 @@ class ToolRegistryTest extends BaseAssistantTest
/*******************************************************************************
** A tool that is always offered, with the schema it was given.
*******************************************************************************/
private static final class WithSchema implements AssistantTool
{
private final Map<String, Object> schema;
/*******************************************************************************
**
*******************************************************************************/
private WithSchema(Map<String, Object> schema)
{
this.schema = schema;
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public String name()
{
return ("with_schema");
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public ToolKind kind()
{
return (ToolKind.READ);
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public List<ToolCheck> checks()
{
return (List.of(ToolCheck.TOOL_SWITCH));
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public ModelToolDefinition definition(ToolOffer offer)
{
return (new ModelToolDefinition().withName(name()).withDescription("A tool with a schema.").withStrict(false).withInputSchema(schema));
}
/*******************************************************************************
**
*******************************************************************************/
@Override
public ToolCallOutcome execute(ToolCallRequest request)
{
return (null);
}
}
/*******************************************************************************
** A tool that says nothing about itself.
*******************************************************************************/