Line data Source code
1 : import * as React from 'react';
2 :
3 : import type { AgentBridge } from './agent_bridge.js';
4 : import { AgentConnectionStatus } from './agent_client.js';
5 : import { newScriptConnection, type SessionScripts, type ScriptData } from '../session/scripts/session_scripts.js';
6 : import { ConnectorType } from '../connections/connector_info.js';
7 : import { buildCatalogNameIndex, collectMentionedIdentifiers } from './catalog_name_index.js';
8 : import { buildAgentPrompt, describeSqlDialect } from './agent_prompt.js';
9 : import { collectReferencedScripts } from './script_references.js';
10 : import { scriptDisplayName } from '../session/scripts/script_types.js';
11 : import { stringifyError } from '../platform/logger/logger.js';
12 : import { useLogger } from '../platform/logger/logger_provider.js';
13 :
14 1 : const LOG_CTX = 'agent_questions';
15 :
16 : /// A question that the agent answers or answered
17 : export interface AgentQuestion {
18 : /// The question in natural language
19 : question: string;
20 : /// The latest status line of the agent while it works, null once it finished
21 : status: string | null;
22 : /// The error of a failed request
23 : error: string | null;
24 : /// The final message of the agent
25 : note: string | null;
26 : /// The scripts that the agent wrote, in the order of their first edit
27 : scripts: number[];
28 : }
29 :
30 : /// Is the agent working on a question?
31 1 : export function isAgentAnswering(question: AgentQuestion): boolean {
32 1 : return question.status != null;
33 1 : }
34 :
35 : /// The session a question is answered in
36 : export interface AgentQuestionsArgs {
37 : /// The agent bridge, null outside of the provider
38 : agent: AgentBridge | null;
39 : /// The scripts of the session
40 : sessionScripts: SessionScripts;
41 : /// Runs a script that the agent wrote
42 : runScript: (next: SessionScripts, scriptData: ScriptData) => void;
43 : /// Follows the agent to a script, every time the agent edits or runs it
44 : followScript?: (scriptKey: number, target: AgentFollowTarget) => void;
45 : }
46 :
47 : /// Where the feed follows the agent to: the text of a script it edits or the results of a script it runs
48 : export type AgentFollowTarget = 'script' | 'result';
49 :
50 : /// The questions of a session that the agent answers
51 : export interface AgentQuestions {
52 : /// Can the session ask the agent?
53 : available: boolean;
54 : /// The running, failed and answered questions by question id, oldest first
55 : questions: ReadonlyMap<number, AgentQuestion>;
56 : /// Ask the agent a question, the agent picks the script it writes
57 : ask: (question: string) => Promise<void>;
58 : /// Stop the agent
59 : stop: (questionId: number) => void;
60 : /// Dismiss the error or the note of a question
61 : dismiss: (questionId: number) => void;
62 : }
63 :
64 1 : const NO_QUESTIONS: ReadonlyMap<number, AgentQuestion> = new Map();
65 :
66 : /// The status line of a request that waits for the agent
67 1 : const WAITING_STATUS = 'waiting for agent...';
68 :
69 : /// Answer questions through the agent.
70 : ///
71 : /// The agent writes its SQL into the scripts of the session through the dashql command, a new script or an existing
72 : /// one. The composer shows the question with its status lines. Once the agent finishes, a script it edited without
73 : /// running it runs.
74 1 : export function useAgentQuestions(args: AgentQuestionsArgs): AgentQuestions {
75 1 : const logger = useLogger();
76 1 : const latest = React.useRef(args);
77 1 : latest.current = args;
78 :
79 : // The bridge keeps the questions of the session, its turns finish while the user is in another session
80 1 : const sessionId = args.sessionScripts.sessionId;
81 1 : const store = args.agent?.questions ?? null;
82 1 : const subscribe = React.useCallback((listener: () => void) => store?.subscribe(listener) ?? (() => { }), [store]);
83 1 : const questions = React.useSyncExternalStore(subscribe, () => store?.questions(sessionId) ?? NO_QUESTIONS);
84 :
85 : // The feed runs and follows the scripts of the agent while it shows the session
86 1 : React.useEffect(() => store?.attach(sessionId, {
87 1 : sessionScripts: () => latest.current.sessionScripts,
88 1 : runScript: (next, scriptData) => latest.current.runScript(next, scriptData),
89 1 : followScript: (scriptKey, target) => latest.current.followScript?.(scriptKey, target),
90 1 : }), [store, sessionId]);
91 :
92 1 : const cancelTurn = React.useCallback((questionId: number, turnId: string) => {
93 1 : logger.info('Stopping the agent', { questionId: String(questionId), turnId }, LOG_CTX);
94 1 : latest.current.agent?.cancel(sessionId, turnId).catch(e => {
95 0 : logger.warn('Stopping the agent failed', { error: stringifyError(e) }, LOG_CTX);
96 1 : });
97 1 : }, [logger, sessionId]);
98 1 : const stop = React.useCallback((questionId: number) => {
99 1 : if (store == null) return;
100 : // A question without a turn id is stopped once the request returns it
101 1 : const turnId = store.turnId(sessionId, questionId);
102 1 : if (!store.finish(sessionId, questionId)) return;
103 1 : store.remove(sessionId, questionId);
104 1 : if (turnId != null) cancelTurn(questionId, turnId);
105 1 : }, [store, sessionId, cancelTurn]);
106 1 : const dismiss = React.useCallback((questionId: number) => store?.remove(sessionId, questionId), [store, sessionId]);
107 :
108 1 : const ask = React.useCallback(async (question: string) => {
109 1 : const { agent, sessionScripts } = latest.current;
110 1 : if (agent == null || !agent.hasAccess(sessionId)) {
111 0 : logger.error('Cannot ask a question without agent access', {}, LOG_CTX);
112 0 : return;
113 0 : }
114 1 : const store = agent.questions;
115 1 : const questionId = store.add(sessionId, { question, status: WAITING_STATUS, error: null, note: null, scripts: [] });
116 :
117 : // The question is answered on the connection that the next script of the session runs on
118 1 : const connection = newScriptConnection(sessionScripts, null);
119 1 : const prompt = buildAgentPrompt({
120 1 : question,
121 1 : dialect: describeSqlDialect(connection?.connectorInfo.connectorType ?? ConnectorType.HYPER),
122 1 : mentionedIdentifiers: connection == null ? [] : collectMentionedIdentifiers(buildCatalogNameIndex(connection.catalog), question),
123 1 : referencedScripts: collectReferencedScripts(sessionScripts.scriptOrder.map(scriptDisplayName), question),
124 1 : });
125 1 : logger.info('Asking the agent', { questionId: String(questionId), characters: String(question.length) }, LOG_CTX);
126 1 : const active = () => store.isWorking(sessionId, questionId);
127 1 : const update = (fn: (prev: AgentQuestion) => AgentQuestion | null) => store.update(sessionId, questionId, fn);
128 : // The scripts that the agent edited, and whether it ran them after the last edit
129 1 : const edits = new Map<number, boolean>();
130 1 : try {
131 1 : const turnId = await agent.request(sessionId, prompt, {
132 1 : status: (_state, text) => {
133 1 : if (active()) update(q => ({ ...q, status: text }));
134 1 : },
135 1 : edited: scriptKey => {
136 1 : if (!active()) return;
137 1 : edits.set(scriptKey, false);
138 1 : update(q => (q.scripts.includes(scriptKey) ? q : { ...q, scripts: [...q.scripts, scriptKey] }));
139 1 : store.feed(sessionId)?.followScript(scriptKey, 'script');
140 1 : },
141 1 : ran: scriptKey => {
142 1 : if (!active()) return;
143 1 : if (edits.has(scriptKey)) edits.set(scriptKey, true);
144 1 : store.feed(sessionId)?.followScript(scriptKey, 'result');
145 1 : },
146 1 : done: message => {
147 1 : if (!store.finish(sessionId, questionId)) return;
148 1 : const note = message.trim();
149 : // Without an answer and without scripts, nothing is left to show
150 1 : update(q => (note.length > 0 || q.scripts.length > 0 ? { ...q, status: null, note: note.length > 0 ? note : null } : null));
151 : // Run the first script that the agent wrote but did not run, while a feed shows the session
152 1 : const feed = store.feed(sessionId);
153 1 : if (feed == null) return;
154 1 : const current = feed.sessionScripts();
155 1 : const unrun = [...edits].find(([, ran]) => !ran);
156 1 : const scriptData = unrun != null ? current.scripts[unrun[0]] : undefined;
157 1 : if (scriptData == null) return;
158 1 : feed.followScript(scriptData.scriptKey, 'result');
159 1 : feed.runScript(current, scriptData);
160 1 : },
161 1 : failed: (error, cancelled) => {
162 1 : if (!store.finish(sessionId, questionId)) return;
163 0 : if (cancelled) store.remove(sessionId, questionId);
164 0 : else update(q => ({ ...q, status: null, error }));
165 1 : },
166 1 : });
167 1 : if (active()) store.setTurnId(sessionId, questionId, turnId);
168 1 : else cancelTurn(questionId, turnId);
169 1 : } catch (e: unknown) {
170 1 : if (!store.finish(sessionId, questionId)) return;
171 1 : logger.error('Asking the agent failed', { questionId: String(questionId), error: stringifyError(e) }, LOG_CTX);
172 1 : update(q => ({ ...q, status: null, error: stringifyError(e) }));
173 1 : }
174 1 : }, [logger, cancelTurn, sessionId]);
175 :
176 1 : const available = args.agent != null
177 1 : && args.agent.hasAccess(sessionId)
178 1 : && args.agent.status === AgentConnectionStatus.Connected;
179 1 : return { available, questions, ask, stop, dismiss };
180 1 : }
|
