ex-brain 0.1.1 → 0.2.0

package/src/ai/timeline-extractor.ts

@@ -1,5 +1,7 @@
 import type { ResolvedLLM } from "../settings";
 import type { TimelineEntry } from "../types";
+import { callLLM, resolveApiKey, isLLMConfigured } from "./llm-client";
+import { jsonrepair } from "jsonrepair";
 
 // ---------------------------------------------------------------------------
 // Types
@@ -50,14 +52,14 @@ export async function extractTimelineEvents(
   input: TimelineExtractionInput,
   llm: ResolvedLLM,
 ): Promise<TimelineExtractionResult> {
-  const apiKey = resolveApiKey(llm);
-  if (!apiKey) {
+  if (!isLLMConfigured(llm)) {
     // Fallback: regex-based extraction
     return fallbackExtract(input);
   }
 
   const prompt = buildExtractionPrompt(input);
-  const resp = await callLLM(llm, prompt, 2048);
+  const systemPrompt = "You are a timeline extraction assistant. Extract events from unstructured text. Always output valid JSON array. Be concise and factual.";
+  const resp = await callLLM(llm, prompt, 2048, systemPrompt);
   
   if (!resp) {
     return fallbackExtract(input);
@@ -94,7 +96,8 @@ export async function extractTimelineFromRelation(
   }
 
   const prompt = buildRelationTimelinePrompt(relation, defaultDate);
-  const resp = await callLLM(llm, prompt, 512);
+  const systemPrompt = "You are a timeline extraction assistant. Extract events from relationships. Always output valid JSON array.";
+  const resp = await callLLM(llm, prompt, 512, systemPrompt);
   
   if (!resp) return null;
 
@@ -186,50 +189,6 @@ Examples:
 /no_think`;
 }
 
-// ---------------------------------------------------------------------------
-// LLM Call
-// ---------------------------------------------------------------------------
-
-async function callLLM(llm: ResolvedLLM, prompt: string, maxTokens: number): Promise<string> {
-  const apiKey = resolveApiKey(llm);
-  if (!apiKey) return "";
-
-  const body = {
-    model: llm.model,
-    messages: [
-      { role: "system", content: "You are a timeline extraction assistant. Extract events from unstructured text. Always output valid JSON array. Be concise and factual." },
-      { role: "user", content: prompt },
-    ],
-    temperature: 0.1,
-    max_tokens: maxTokens,
-    enable_thinking: false,
-  };
-
-  try {
-    const resp = await fetch(
-      llm.baseURL.endsWith("/") ? llm.baseURL + "chat/completions" : llm.baseURL + "/chat/completions",
-      {
-        method: "POST",
-        headers: { "Content-Type": "application/json", Authorization: `Bearer ${apiKey}` },
-        body: JSON.stringify(body),
-      },
-    );
-
-    if (!resp.ok) {
-      const text = await resp.text();
-      console.warn(`[timeline-extractor] LLM call failed (${resp.status}): ${text.slice(0, 200)}`);
-      return "";
-    }
-
-    const data = await resp.json();
-    return data.choices?.[0]?.message?.content?.trim() ?? "";
-  } catch (error) {
-    const msg = error instanceof Error ? error.message : String(error);
-    console.warn(`[timeline-extractor] LLM call error: ${msg}`);
-    return "";
-  }
-}
-
 // ---------------------------------------------------------------------------
 // Response Parsing
 // ---------------------------------------------------------------------------
@@ -239,7 +198,9 @@ function parseExtractionResponse(resp: string, pageSlug: string): TimelineEntry[
   if (!match) return [];
 
   try {
-    const parsed = JSON.parse(match[0]) as unknown[];
+    // Use jsonrepair to fix common LLM JSON issues
+    const repaired = jsonrepair(match[0]);
+    const parsed = JSON.parse(repaired) as unknown[];
     const entries: TimelineEntry[] = [];
     
     for (const e of parsed) {
@@ -249,12 +210,18 @@ function parseExtractionResponse(resp: string, pageSlug: string): TimelineEntry[
       const date = normalizeDate(String(entry.date ?? ""));
       if (!date) continue;
       
+      // Get importance from the response, default to 3
+      const importance = typeof entry.importance === "number" 
+        ? Math.max(1, Math.min(5, Math.round(entry.importance))) 
+        : 3;
+      
       entries.push({
         pageSlug,
         date,
         source: "extracted",
         summary: String(entry.summary ?? "").slice(0, 120),
         detail: String(entry.detail ?? ""),
+        importance,
       });
     }
     
@@ -344,15 +311,19 @@ function normalizeDate(raw: string, defaultDate?: string): string {
   const chineseMatch = trimmed.match(/(\d{4})年(\d{1,2})月(\d{1,2})日/);
   if (chineseMatch) {
     const [, year, month, day] = chineseMatch;
-    return `${year}-${month.padStart(2, "0")}-${day.padStart(2, "0")}`;
+    if (year && month && day) {
+      return `${year}-${month.padStart(2, "0")}-${day.padStart(2, "0")}`;
+    }
   }
 
   // Chinese format without year: 1月15日
   const chineseNoYearMatch = trimmed.match(/(\d{1,2})月(\d{1,2})日/);
   if (chineseNoYearMatch && defaultDate) {
     const [, month, day] = chineseNoYearMatch;
-    const year = defaultDate.slice(0, 4);
-    return `${year}-${month.padStart(2, "0")}-${day.padStart(2, "0")}`;
+    if (month && day) {
+      const year = defaultDate.slice(0, 4);
+      return `${year}-${month.padStart(2, "0")}-${day.padStart(2, "0")}`;
+    }
   }
 
   // English month names
@@ -374,10 +345,12 @@ function normalizeDate(raw: string, defaultDate?: string): string {
   const englishMatch = trimmed.match(/(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\.?\s+(\d{1,2})(?:st|nd|rd|th)?(?:,?\s+(\d{4}))?/i);
   if (englishMatch) {
     const [, monthName, day, year] = englishMatch;
-    const month = monthMap[monthName.toLowerCase().slice(0, 3)];
-    if (month) {
-      const finalYear = year || (defaultDate ? defaultDate.slice(0, 4) : new Date().getFullYear().toString());
-      return `${finalYear}-${month}-${day.padStart(2, "0")}`;
+    if (monthName && day) {
+      const month = monthMap[monthName.toLowerCase().slice(0, 3)];
+      if (month) {
+        const finalYear = year || (defaultDate ? defaultDate.slice(0, 4) : new Date().getFullYear().toString());
+        return `${finalYear}-${month}-${day.padStart(2, "0")}`;
+      }
     }
   }
 
@@ -416,12 +389,6 @@ function normalizeDate(raw: string, defaultDate?: string): string {
 // Helpers
 // ---------------------------------------------------------------------------
 
-function resolveApiKey(llm: ResolvedLLM): string {
-  if (llm.apiKey) return llm.apiKey;
-  if (llm.apiKeyEnv) return process.env[llm.apiKeyEnv] ?? "";
-  return "";
-}
-
 function deduplicateEntries(entries: TimelineEntry[]): TimelineEntry[] {
   const seen = new Map<string, TimelineEntry>();