69 lines
2.8 KiB
TypeScript
69 lines
2.8 KiB
TypeScript
import { Agent } from "@mastra/core/agent";
|
|
import { z } from "zod";
|
|
import { defaultLanguageModel, resolveLanguageModel } from "@/mastra/model";
|
|
import {
|
|
eventKindSchema,
|
|
eventSubjectSchema,
|
|
evidenceDomainSchema,
|
|
relatedPersonSchema,
|
|
} from "../rectification-v4/contracts.ts";
|
|
import {
|
|
validatedModelAssistedEvidence,
|
|
type ExtractedLifeEventEvidence,
|
|
type ModelAssistedEventExtraction,
|
|
} from "../conversational-rectification/evidence-extractor.ts";
|
|
|
|
export const modelAssistedEventExtractionSchema = z.object({
|
|
sourceSpan: z.string().trim().min(1).max(4_000),
|
|
summary: z.string().trim().min(1).max(1_000),
|
|
domain: evidenceDomainSchema,
|
|
eventKind: eventKindSchema,
|
|
subject: eventSubjectSchema,
|
|
relatedPerson: relatedPersonSchema.nullable(),
|
|
dateText: z.string().trim().min(1).max(80).nullable(),
|
|
}).strict();
|
|
|
|
export type EventExtractorGenerator = (prompt: string) => Promise<Readonly<{ object: unknown }>>;
|
|
|
|
export async function extractEventWithModel(input: Readonly<{
|
|
rawText: string;
|
|
sourceTurnId: string;
|
|
asOfDate: string;
|
|
modelId?: string | null;
|
|
timeoutMs?: number;
|
|
generateExtraction?: EventExtractorGenerator;
|
|
}>): Promise<ExtractedLifeEventEvidence | null> {
|
|
const model = (input.modelId ? resolveLanguageModel(input.modelId) : null) ?? defaultLanguageModel();
|
|
if (!model && !input.generateExtraction) return null;
|
|
const agent = model ? new Agent({
|
|
id: `rectification-event-extractor-${model.id}`,
|
|
name: "Restricted Rectification Event Extractor",
|
|
model: model.model,
|
|
instructions: "Extract at most one explicitly stated dated life event. sourceSpan and dateText must be exact continuous substrings of the user text. Never infer or invent a date, normalized range, candidate time, score, id, or profile value. Return strict JSON only.",
|
|
}) : null;
|
|
const generate = input.generateExtraction ?? (async (prompt: string) => {
|
|
if (!agent) throw new Error("event_extractor_model_unavailable");
|
|
return agent.generate(prompt, {
|
|
abortSignal: AbortSignal.timeout(input.timeoutMs ?? 10_000),
|
|
structuredOutput: { schema: modelAssistedEventExtractionSchema, jsonPromptInjection: "inline" },
|
|
});
|
|
});
|
|
try {
|
|
const result = await generate(JSON.stringify({
|
|
task: "Extract one event that deterministic parsing could not classify. Use only literal text from userText.",
|
|
userText: input.rawText,
|
|
asOfDate: input.asOfDate,
|
|
allowedOutput: ["sourceSpan", "summary", "domain", "eventKind", "subject", "relatedPerson", "dateText"],
|
|
}));
|
|
const extraction = modelAssistedEventExtractionSchema.parse(result.object) as ModelAssistedEventExtraction;
|
|
return validatedModelAssistedEvidence({
|
|
rawText: input.rawText,
|
|
sourceTurnId: input.sourceTurnId,
|
|
asOfDate: input.asOfDate,
|
|
extraction,
|
|
});
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|