fix(oss): validate LLM fact output via FactRetrievalSchema before embedding (#4083)
This commit is contained in:
@@ -27,9 +27,12 @@ export class OllamaEmbedder implements Embedder {
|
|||||||
} catch (err) {
|
} catch (err) {
|
||||||
logger.error(`Error ensuring model exists: ${err}`);
|
logger.error(`Error ensuring model exists: ${err}`);
|
||||||
}
|
}
|
||||||
|
// Ollama's Go server requires prompt to be a string. Coerce defensively
|
||||||
|
// since callers may pass values parsed from untrusted LLM JSON output.
|
||||||
|
const prompt = typeof text === "string" ? text : JSON.stringify(text);
|
||||||
const response = await this.ollama.embeddings({
|
const response = await this.ollama.embeddings({
|
||||||
model: this.model,
|
model: this.model,
|
||||||
prompt: text,
|
prompt,
|
||||||
});
|
});
|
||||||
return response.embedding;
|
return response.embedding;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ import {
|
|||||||
HistoryManagerFactory,
|
HistoryManagerFactory,
|
||||||
} from "../utils/factory";
|
} from "../utils/factory";
|
||||||
import {
|
import {
|
||||||
|
FactRetrievalSchema,
|
||||||
getFactRetrievalMessages,
|
getFactRetrievalMessages,
|
||||||
getUpdateMemoryMessages,
|
getUpdateMemoryMessages,
|
||||||
parseMessages,
|
parseMessages,
|
||||||
@@ -261,7 +262,8 @@ export class Memory {
|
|||||||
const cleanResponse = removeCodeBlocks(response as string);
|
const cleanResponse = removeCodeBlocks(response as string);
|
||||||
let facts: string[] = [];
|
let facts: string[] = [];
|
||||||
try {
|
try {
|
||||||
facts = JSON.parse(cleanResponse).facts || [];
|
const parsed = FactRetrievalSchema.parse(JSON.parse(cleanResponse));
|
||||||
|
facts = parsed.facts;
|
||||||
} catch (e) {
|
} catch (e) {
|
||||||
console.error(
|
console.error(
|
||||||
"Failed to parse facts from LLM response:",
|
"Failed to parse facts from LLM response:",
|
||||||
|
|||||||
@@ -1,9 +1,18 @@
|
|||||||
import { z } from "zod";
|
import { z } from "zod";
|
||||||
|
|
||||||
|
// Accepts a string directly, or an object with a "fact" or "text" key
|
||||||
|
// (common malformed shapes from smaller LLMs like llama3.1:8b).
|
||||||
|
const factItem = z.union([
|
||||||
|
z.string(),
|
||||||
|
z.object({ fact: z.string() }).transform((o) => o.fact),
|
||||||
|
z.object({ text: z.string() }).transform((o) => o.text),
|
||||||
|
]);
|
||||||
|
|
||||||
// Define Zod schema for fact retrieval output
|
// Define Zod schema for fact retrieval output
|
||||||
export const FactRetrievalSchema = z.object({
|
export const FactRetrievalSchema = z.object({
|
||||||
facts: z
|
facts: z
|
||||||
.array(z.string())
|
.array(factItem)
|
||||||
|
.transform((arr) => arr.filter((s) => s.length > 0))
|
||||||
.describe("An array of distinct facts extracted from the conversation."),
|
.describe("An array of distinct facts extracted from the conversation."),
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user