Skip to content

Fix AI tool prompt alignment and diagnostics #644

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
codedogQBY wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix AI tool prompt alignment and diagnostics #644

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 23 additions & 1 deletion packages/core/src/ai/__tests__/system-prompt.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,29 @@ describe("buildSystemPrompt citations", () => {
expect(prompt).toContain("- addCitation");
expect(prompt).not.toContain("- getReadingProgress");
expect(prompt).not.toContain("Get overall reading progress");
expect(prompt).not.toContain("- ragSearch");
expect(prompt).not.toContain("ragSearch");
expect(prompt).not.toContain("ragContext");
expect(prompt).not.toContain("fallbackSearch");
expect(prompt).not.toContain("Semantic/keyword search across book content");
});

it("keeps workflow instructions aligned with library-only tools", () => {
const prompt = buildSystemPrompt({
book: makeBook(),
semanticContext: null,
enabledSkills: [],
isVectorized: true,
userLanguage: "en",
questionCategory: "library_request",
allowedToolNames: ["listBooks", "getReadingStats"],
});

expect(prompt).toContain("- listBooks");
expect(prompt).toContain("- getReadingStats");
expect(prompt).toContain("This turn does not expose book-content retrieval tools");
expect(prompt).not.toContain("ragSearch");
expect(prompt).not.toContain("fallbackSearch");
expect(prompt).not.toContain("addCitation");
expect(prompt).not.toContain("mindmap");
});
});
149 changes: 133 additions & 16 deletions packages/core/src/ai/agents/reading-agent.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,32 @@ const CHAPTER_TASK_RECURSION_LIMIT = 24;
const DEFAULT_TOOL_TIMEOUT_MS = 45_000;
const TOOL_EXECUTION_LIMIT = 12;
const REPEATED_TOOL_CALL_LIMIT = 2;
const TOOL_TIMEOUT_MS_BY_NAME: Record<string, number> = {
getSelection: 5_000,
getCurrentChapter: 5_000,
getReadingProgress: 5_000,
getSurroundingContext: 8_000,
getRecentHighlights: 8_000,
getAnnotations: 8_000,
addCitation: 20_000,
ragSearch: 30_000,
ragToc: 20_000,
ragContext: 30_000,
summarize: 35_000,
extractEntities: 35_000,
analyzeArguments: 35_000,
findQuotes: 35_000,
compareSections: 35_000,
fallbackSearch: 60_000,
fallbackToc: 45_000,
fallbackChapterContext: 60_000,
classifyBooks: 60_000,
tagBooks: 30_000,
manageBookTags: 30_000,
updateBookMetadata: 30_000,
manageBookGroups: 30_000,
mindmap: 10_000,
};

const CHAPTER_LOOKUP_STOP_TOOL_NAMES = new Set([
"resolveChapterReference",
Expand Down Expand Up @@ -553,13 +579,36 @@ function buildZodSchema(

switch (param.type) {
case "number":
fieldSchema = z.number().describe(param.description);
fieldSchema = z
.preprocess((value) => {
if (typeof value !== "string" || !value.trim()) return value;
const numberValue = Number(value);
return Number.isFinite(numberValue) ? numberValue : value;
}, z.number())
.describe(param.description);
break;
case "boolean":
fieldSchema = z.boolean().describe(param.description);
fieldSchema = z
.preprocess((value) => {
if (typeof value !== "string") return value;
const normalized = value.trim().toLowerCase();
if (normalized === "true") return true;
if (normalized === "false") return false;
return value;
}, z.boolean())
.describe(param.description);
break;
default:
fieldSchema = z.string().describe(param.description);
fieldSchema = /json/i.test(param.description)
? z
.preprocess((value) => {
if (Array.isArray(value) || (value && typeof value === "object")) {
return JSON.stringify(value);
}
return value;
}, z.string())
.describe(param.description)
: z.string().describe(param.description);
break;
}

Expand Down Expand Up @@ -592,16 +641,68 @@ function withToolTimeout<T>(promise: Promise<T>, timeoutMs: number, toolName: st
});
}

function getToolTimeoutMs(tool: ToolDefinition, defaultTimeoutMs: number): number {
return tool.timeoutMs ?? TOOL_TIMEOUT_MS_BY_NAME[tool.name] ?? defaultTimeoutMs;
}

function compactToolLogValue(value: unknown, maxLength = 240): unknown {
if (typeof value === "string") {
if (value.length <= maxLength) return value;
return `${value.slice(0, maxLength)}...(${value.length} chars)`;
}
if (Array.isArray(value)) {
return value.slice(0, 8).map((item) => compactToolLogValue(item, 120));
}
if (!value || typeof value !== "object") return value;

const result: Record<string, unknown> = {};
for (const [key, childValue] of Object.entries(value as Record<string, unknown>).slice(0, 12)) {
result[key] = compactToolLogValue(childValue, 120);
}
return result;
}

async function executeTool(
tool: ToolDefinition,
args: Record<string, unknown>,
timeoutMs: number,
): Promise<unknown> {
const startedAt = Date.now();
console.log(
"[ReadingAgent][tool-start]",
JSON.stringify({
name: tool.name,
timeoutMs,
args: compactToolLogValue(args),
}),
);
try {
return await withToolTimeout(Promise.resolve(tool.execute(args)), timeoutMs, tool.name);
const result = await withToolTimeout(Promise.resolve(tool.execute(args)), timeoutMs, tool.name);
console.log(
"[ReadingAgent][tool-end]",
JSON.stringify({
name: tool.name,
durationMs: Date.now() - startedAt,
ok: true,
result: compactToolLogValue(result),
}),
);
return result;
} catch (error) {
const message = error instanceof Error ? error.message : String(error);
const payload = {
name: tool.name,
durationMs: Date.now() - startedAt,
ok: false,
error: message,
};
if (/timed out/i.test(message)) {
console.warn("[ReadingAgent][tool-timeout]", JSON.stringify(payload));
} else {
console.warn("[ReadingAgent][tool-error]", JSON.stringify(payload));
}
return {
error: error instanceof Error ? error.message : String(error),
error: message,
};
}
}
Expand Down Expand Up @@ -947,7 +1048,7 @@ export async function* streamReadingAgent(
return JSON.stringify(cachedResult);
}

const result = await executeTool(tool, toolInput, toolTimeoutMs);
const result = await executeTool(tool, toolInput, getToolTimeoutMs(tool, toolTimeoutMs));
if (exactCacheKey) {
toolResultCache.set(exactCacheKey, result);
}
Expand Down Expand Up @@ -991,8 +1092,8 @@ export async function* streamReadingAgent(
);

// Track tool calls already emitted (from streaming chunks or on_chat_model_end)
// so we can deduplicate against on_tool_start events.
let pendingEarlyToolCalls = 0;
// so we can deduplicate against on_tool_start events without hiding unrelated calls.
const pendingEarlyToolStartNames: string[] = [];

// Accumulate tool_call_chunks from streaming to emit tool_call as early as possible.
// Key: chunk index, Value: { name accumulated so far, args accumulated so far }
Expand Down Expand Up @@ -1100,7 +1201,7 @@ export async function* streamReadingAgent(
// Emit as soon as we have a tool name (don't wait for full args)
if (entry.name && !entry.emitted) {
entry.emitted = true;
pendingEarlyToolCalls++;
pendingEarlyToolStartNames.push(entry.name);
yield {
type: "tool_call" as const,
name: entry.name,
Expand Down Expand Up @@ -1132,25 +1233,38 @@ export async function* streamReadingAgent(

if (output) {
if (Array.isArray(toolCalls)) {
const alreadyEmittedNames = [...pendingEarlyToolStartNames];
for (const tc of toolCalls) {
const toolName =
typeof tc?.name === "string"
? tc.name
: typeof tc?.function?.name === "string"
? tc.function.name
: "";
if (!toolName) continue;

// Check if already emitted from streaming chunks
if (pendingEarlyToolCalls > 0) {
// Already emitted — skip but don't decrement yet (that's for on_tool_start)
const alreadyEmittedIndex = alreadyEmittedNames.findIndex(
(name) => name === toolName,
);
if (alreadyEmittedIndex >= 0) {
alreadyEmittedNames.splice(alreadyEmittedIndex, 1);
continue;
}
let args: Record<string, unknown>;
const rawArgs = tc.args ?? tc.function?.arguments;
try {
args = (typeof tc.args === "string" ? JSON.parse(tc.args) : tc.args) as Record<
args = (typeof rawArgs === "string" ? JSON.parse(rawArgs) : rawArgs) as Record<
string,
unknown
>;
} catch {
args = {};
}
pendingEarlyToolCalls++;
pendingEarlyToolStartNames.push(toolName);
yield {
type: "tool_call" as const,
name: tc.name,
name: toolName,
args,
};
}
Expand All @@ -1161,8 +1275,11 @@ export async function* streamReadingAgent(
// Tool call started — skip if already emitted earlier
if (event.event === "on_tool_start") {
pendingToolCallNames.push(event.name);
if (pendingEarlyToolCalls > 0) {
pendingEarlyToolCalls--;
const pendingEarlyIndex = pendingEarlyToolStartNames.findIndex(
(name) => name === event.name,
);
if (pendingEarlyIndex >= 0) {
pendingEarlyToolStartNames.splice(pendingEarlyIndex, 1);
} else {
// Fallback: emit if not already emitted (e.g. non-OpenAI model)
yield {
Expand Down
Loading