diff --git a/Bouncer/src/background/local-model.ts b/Bouncer/src/background/local-model.ts index e4d4a9a..ee3f776 100644 --- a/Bouncer/src/background/local-model.ts +++ b/Bouncer/src/background/local-model.ts @@ -26,6 +26,19 @@ const DOWNLOAD_RETRY_DELAY_MS = 2000; // Keys that belong on the ModelRecord (appConfig), not chatOpts. const MODEL_RECORD_KEYS = new Set(['model', 'model_lib', 'model_type']); +/** JSON schema for structured classification output from local models. */ +const CLASSIFICATION_RESPONSE_FORMAT = { + type: 'json_object' as const, + schema: JSON.stringify({ + type: 'object', + properties: { + reasoning: { type: 'string' }, + match: { type: ['string', 'null'] }, + }, + required: ['reasoning', 'match'], + }), +}; + // ==================== Pure helpers ==================== // Build both the appConfig (ModelRecord for CreateMLCEngine) and chatOpts @@ -63,6 +76,22 @@ export function parseLocalModelResponse(rawResponse: string | null): { shouldHid return { shouldHide: false, reasoning: 'Empty model response — model returned no output' }; } + // Try JSON structured output first + try { + const parsed = JSON.parse(rawResponse); + if (typeof parsed === 'object' && parsed !== null && 'reasoning' in parsed) { + const match = parsed.match; + const shouldHide = typeof match === 'string' && match.length > 0 && match !== 'null'; + const reasoning = shouldHide + ? `${parsed.reasoning} (Matched: ${match})` + : String(parsed.reasoning); + return { shouldHide, reasoning }; + } + } catch { + // Not JSON — fall through to regex parsing + } + + // Fallback: freeform text parsing (backward compatibility) let reasoning = rawResponse; let shouldHide = false; @@ -324,9 +353,10 @@ export class LocalEngine { async generate( messages: ChatMessage[], maxTokens: number, - { priority = 0, temperature, onStart }: { priority?: number; temperature?: number; onStart?: () => void } = {} + { priority = 0, temperature, onStart, responseFormat }: { priority?: number; temperature?: number; onStart?: () => void; responseFormat?: Record } = {} ): Promise { const requestOpts: Record = { messages, max_tokens: maxTokens }; + if (responseFormat) requestOpts.response_format = responseFormat; if (temperature !== undefined) requestOpts.temperature = temperature; const request = buildInferenceRequest(this._modelConfig || ({} as Record), requestOpts); @@ -674,7 +704,7 @@ export async function callLocalInference( let rawResponse: string; try { - rawResponse = await localEngine.generate(messages, 40, { priority, onStart }); + rawResponse = await localEngine.generate(messages, 40, { priority, onStart, responseFormat: CLASSIFICATION_RESPONSE_FORMAT }); } catch (imgError) { if ((imgError as Error).message === 'Inference preempted') throw imgError; if (useImages) { @@ -684,7 +714,7 @@ export async function callLocalInference( { role: "system", content: LOCAL_SYSTEM_PROMPT }, { role: "user", content: textOnlyContent } ]; - rawResponse = await localEngine.generate(textMessages, 40, { priority, onStart }); + rawResponse = await localEngine.generate(textMessages, 40, { priority, onStart, responseFormat: CLASSIFICATION_RESPONSE_FORMAT }); } else { throw imgError; } diff --git a/Bouncer/src/shared/prompts.ts b/Bouncer/src/shared/prompts.ts index d563b6b..99b8113 100644 --- a/Bouncer/src/shared/prompts.ts +++ b/Bouncer/src/shared/prompts.ts @@ -3,25 +3,9 @@ import type { ChatMessage, EvaluationPostData } from '../types'; // System prompt for local models processing one post at a time -export const LOCAL_SYSTEM_PROMPT = `You filter posts. Write 10-15 words identifying what the post is about, then state if it matches a filter category. +export const LOCAL_SYSTEM_PROMPT = `You filter posts. Classify whether a post matches any of the given filter categories. -Example outputs (NOTE: you may not be filtering on these categories!): - - -sports -The Lakers won last night against the Bucks! -Post about NBA basketball game results, which is sports content. Matches sports. - - - -politics -I love cooking dinner each night with my husband -Post about making dinner at home, which is Food/lifestyle content, not politics content. No match. - - -You will be provided with a post () and a list of filter categories (). -Assess whether the topic of the post relates to any of the topics in the filter categories list. -Your reasoning must be AT MOST 15 words, and MUST end with a statement of "Matches " or "No match". +Respond with JSON: {"reasoning": "<10-15 words about what the post is about>", "match": ""} Be precise in your judgment; only match posts that clearly and directly relate to the filter categories.`; diff --git a/Bouncer/tests/background/local-model.test.ts b/Bouncer/tests/background/local-model.test.ts index 3a332cb..ed99106 100644 --- a/Bouncer/tests/background/local-model.test.ts +++ b/Bouncer/tests/background/local-model.test.ts @@ -677,6 +677,47 @@ describe('parseLocalModelResponse', () => { }); }); +describe('parseLocalModelResponse — structured JSON output', () => { + it('parses a JSON match response', () => { + const raw = '{"reasoning":"NBA game results, sports content","match":"sports"}'; + const result = parseLocalModelResponse(raw); + expect(result.shouldHide).toBe(true); + expect(result.reasoning).toContain('sports'); + }); + + it('parses a JSON no-match response', () => { + const raw = '{"reasoning":"Cooking dinner at home, lifestyle content","match":null}'; + const result = parseLocalModelResponse(raw); + expect(result.shouldHide).toBe(false); + expect(result.reasoning).toContain('Cooking'); + }); + + it('parses JSON with match as empty string as no-match', () => { + const raw = '{"reasoning":"General lifestyle post","match":""}'; + const result = parseLocalModelResponse(raw); + expect(result.shouldHide).toBe(false); + }); + + it('falls back to regex parsing for non-JSON freeform output', () => { + const raw = 'Post about basketball game results. Matches sports.'; + const result = parseLocalModelResponse(raw); + expect(result.shouldHide).toBe(true); + expect(result.reasoning).toContain('sports'); + }); + + it('falls back to regex for malformed JSON', () => { + const raw = '{"reasoning": broken json Matches politics.'; + const result = parseLocalModelResponse(raw); + expect(result.shouldHide).toBe(true); + }); + + it('treats match:"null" string as no-match', () => { + const raw = '{"reasoning":"Not related","match":"null"}'; + const result = parseLocalModelResponse(raw); + expect(result.shouldHide).toBe(false); + }); +}); + // ==================== LocalEngine.generate / preempt / ensureLoaded / teardown ==================== describe('LocalEngine generate + preempt + lifecycle', () => {