Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
9 changes: 9 additions & 0 deletions src/cli/groups/ask.ts
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,15 @@ export function addAskCommands(program: Command): void {
return;
}

// An ungrounded run is not an answer, so it is not printed as one:
// no model attribution line, and the next step is the command that
// fills the index rather than a follow-up question.
if (!out.grounded) {
console.log(`\n${out.answer}\n`);
suggest({ command: "sec index", why: "build the index this question needs" });
return;
}

console.log(`\n${out.answer}\n`);
for (const reference of out.references) {
console.log(` [${reference.index}] ${reference.title}`);
Expand Down
140 changes: 140 additions & 0 deletions src/task/kb/AskTask.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,140 @@
/**
* @license
* Copyright 2026 Steven Roussey <sroussey@gmail.com>
* SPDX-License-Identifier: Apache-2.0
*/

import { beforeEach, describe, expect, it, vi } from "vitest";

/**
* `AiChatWithKbTask` is the seam: everything in `AskTask` that is checkable
* without a model is on this side of it — whether an answer with no retrieved
* text is printed as an answer, and what the scope flags do.
*/
const runMock = vi.fn();
const chunkCountMock = vi.fn(async () => 0);

vi.mock("workglow", async () => {
const actual = await vi.importActual<Record<string, unknown>>("workglow");
return {
...actual,
AiChatWithKbTask: class {
run(input: unknown) {
return runMock(input);
}
},
};
});

vi.mock("../../kb/secKnowledgeBase", () => ({
SEC_KB_ID: "sec",
getSecKnowledgeBase: async () => ({ chunkCount: chunkCountMock }),
}));

vi.mock("../../config/models", () => ({
secGenerationModel: () => ({ modelId: "onnx:test-model", reason: "test" }),
}));

const { AskTask } = await import("./AskTask");

const REFERENCE = {
index: 1,
title: "10-K · 2024-11-01 · 0000320193-24-000123",
url: "https://example.invalid/filing",
snippet: "Revenue increased…",
score: 0.71,
};

describe("AskTask grounding", () => {
beforeEach(() => {
runMock.mockReset();
chunkCountMock.mockReset();
chunkCountMock.mockResolvedValue(0);
});

it("does not present a model's memory as an answer when nothing was retrieved", async () => {
// The defect this guards: eleven sentences of confident, unsourced
// financial prose about a company whose filings the database does not hold,
// returned with exit 0 and `references: []`.
runMock.mockResolvedValue({
text: "Apple's revenue reached approximately $383 billion in fiscal year 2023.",
references: [],
});

const out = await new AskTask().run({ question: "What is Apple's revenue?" } as never);

expect(out.grounded).toBe(false);
expect(out.answer).not.toContain("383");
expect(out.answer).toContain("sec index");
expect(out.references).toEqual([]);
});

it("names the index build when nothing is indexed at all", async () => {
runMock.mockResolvedValue({ text: "The CEO of Tesla is Elon Musk.", references: [] });
chunkCountMock.mockResolvedValue(0);

const out = await new AskTask().run({ question: "Who is the CEO of Tesla?" } as never);

expect(out.answer).toContain("Nothing is indexed");
expect(out.answer).toContain("sec update documents");
expect(out.answer).not.toContain("Elon Musk");
});

it("distinguishes an index that holds chunks but matched none of them", async () => {
// Not the same failure, and not the same fix: the index is built, so
// telling the operator to build it would be wrong.
runMock.mockResolvedValue({ text: "made up", references: [] });
chunkCountMock.mockResolvedValue(4210);

const out = await new AskTask().run({ question: "anything" } as never);

expect(out.answer).toContain("4210 chunk(s) indexed");
expect(out.answer).not.toContain("Nothing is indexed");
});

it("passes the answer through, marked grounded, when filing text was cited", async () => {
runMock.mockResolvedValue({ text: "Revenue was $383bn.", references: [REFERENCE] });

const out = await new AskTask().run({ question: "revenue?" } as never);

expect(out.grounded).toBe(true);
expect(out.answer).toBe("Revenue was $383bn.");
expect(out.references).toHaveLength(1);
expect(out.references[0]).toMatchObject({ index: 1, score: 0.71 });
});

it("states its own score floor rather than inheriting the library's", async () => {
runMock.mockResolvedValue({ text: "x", references: [REFERENCE] });

await new AskTask().run({ question: "q" } as never);

expect(runMock.mock.calls[0]![0]).toMatchObject({ minScore: 0.3, maxIterations: 1 });
});

it("carries every scope flag into the prompt it sends", async () => {
runMock.mockResolvedValue({ text: "x", references: [REFERENCE] });

await new AskTask().run({
question: "what happened?",
cik: 320193,
form: "10-K",
since: "2024-01-01",
accession: "0000320193-24-000123",
} as never);

const prompt = String((runMock.mock.calls[0]![0] as { prompt: string }).prompt);
expect(prompt).toContain("what happened?");
expect(prompt).toContain("CIK 320193");
expect(prompt).toContain("form 10-K");
expect(prompt).toContain("filed on or after 2024-01-01");
expect(prompt).toContain("accession 0000320193-24-000123");
});

it("sends the bare question when no scope is given", async () => {
runMock.mockResolvedValue({ text: "x", references: [REFERENCE] });

await new AskTask().run({ question: "just this" } as never);

expect((runMock.mock.calls[0]![0] as { prompt: string }).prompt).toBe("just this");
});
});
75 changes: 67 additions & 8 deletions src/task/kb/AskTask.ts
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,47 @@ export interface AskTaskOutput {
/** The model that answered, and why it was the one. */
readonly modelId: string;
readonly modelReason: string;
/**
* Whether filing text was actually retrieved and cited.
*
* A field rather than something a caller infers from `references.length`: a
* consumer should not have to work out that an answer is fiction from an
* array being empty, and `--json` has no other way to tell.
*/
readonly grounded: boolean;
}

/**
* Score floor for a chunk to count as a match, stated here rather than
* inherited.
*
* `AiChatWithKbTask` defaults it to the same 0.3, but the floor is what decides
* whether an answer is grounded, so `sec ask` names its own instead of moving
* whenever the library's default does.
*/
const ASK_MIN_SCORE = 0.3;

/**
* What to say when nothing was retrieved.
*
* The model is not asked and its text is not printed. A 350M local model is the
* default path — deliberately, so `ask` works with no API key — and it will
* answer a question about Apple's revenue from memory however firmly the system
* prompt tells it not to. An instruction is a request; this is the guard.
*/
function ungroundedAnswer(indexedChunks: number): string {
if (indexedChunks === 0) {
return (
"Nothing is indexed, so there is no filing text to answer from.\n\n" +
"Run `sec index` to build the index — or `sec update documents` first, if no " +
"filings have been converted yet."
);
}
return (
`Nothing in the index matched this question closely enough to answer from ` +
`(${indexedChunks} chunk(s) indexed, none scoring at or above ${ASK_MIN_SCORE}).\n\n` +
"Try rewording it, widen the scope flags, or run `sec index` to cover more filings."
);
}

/**
Expand Down Expand Up @@ -89,14 +130,15 @@ export class AskTask extends Task<TaskPorts<AskTaskInput>, TaskPorts<AskTaskOutp
references: Type.Array(Type.Unknown()),
modelId: Type.String(),
modelReason: Type.String(),
grounded: Type.Boolean(),
});
}

async execute(input: TaskPorts<AskTaskInput>): Promise<TaskPorts<AskTaskOutput>> {
// Resolved before the KB is touched, so a machine with no usable model says
// so before it spends time embedding a query.
const model = secGenerationModel();
await getSecKnowledgeBase();
const kb = await getSecKnowledgeBase();
if (!globalServiceRegistry.has(HUMAN_CONNECTOR)) {
globalServiceRegistry.registerInstance(HUMAN_CONNECTOR, ONE_SHOT_CONNECTOR);
}
Expand All @@ -114,6 +156,7 @@ export class AskTask extends Task<TaskPorts<AskTaskInput>, TaskPorts<AskTaskOutp
// one-shot `sec ask` has not registered and should not need to.
maxIterations: 1,
topKPerKb: input.topK ?? 8,
minScore: ASK_MIN_SCORE,
prompt: scope === undefined ? input.question : `${input.question}\n\n(${scope})`,
system:
"You answer questions about SEC filings using only the retrieved excerpts. " +
Expand All @@ -131,17 +174,33 @@ export class AskTask extends Task<TaskPorts<AskTaskInput>, TaskPorts<AskTaskOutp
}[];
};

const references = (result.references ?? []).map((reference) => ({
index: reference.index,
title: reference.title,
url: reference.url,
snippet: reference.snippet,
score: reference.score,
}));

// Retrieved nothing means the model answered from memory, and the answer is
// about SEC filings the database does not contain. Suppressed here rather
// than in the renderer, so `--json` consumers get the same refusal.
if (references.length === 0) {
return {
answer: ungroundedAnswer(await kb.chunkCount()),
references: [],
modelId: model.modelId,
modelReason: model.reason,
grounded: false,
};
}

return {
answer: result.text,
references: (result.references ?? []).map((reference) => ({
index: reference.index,
title: reference.title,
url: reference.url,
snippet: reference.snippet,
score: reference.score,
})),
references,
modelId: model.modelId,
modelReason: model.reason,
grounded: true,
};
}
}
Expand Down