backend/src/ai/engine.ts

import { ChatMessage, ChatRequest, SearchSummary } from "../types.js";
import { memoryStore } from "./memory.js";
import { summarizeExtractive, totalTokens } from "./context.js";
import { buildSystemPrompt, shouldSearch, trimHistory } from "./router.js";
import { searchWeb } from "../search/search.js";
import { streamChatCompletion } from "../utils/stream.js";

const MAX_CONTEXT_TOKENS = 2600;

export type EngineEvents = {
  onStatus: (message: string) => void;
  onToken: (token: string) => void;
  onDone: () => void;
  onError: (message: string) => void;
};

function buildSearchContext(summary: SearchSummary) {
  const lines: string[] = [];
  summary.results.forEach((r, idx) => {
    lines.push(`Source ${idx + 1}: ${r.title} (${r.url})`);
    lines.push(r.snippet);
  });
  summary.topContent.forEach((c, idx) => {
    lines.push(`Content ${idx + 1}: ${c.title}`);
    lines.push(c.content.slice(0, 1200));
  });
  return lines.join("\n");
}

export async function handleChat(request: ChatRequest, events: EngineEvents) {
  const session = memoryStore.get(request.sessionId);
  const userMessage: ChatMessage = {
    id: `msg_${Date.now()}`,
    role: "user",
    content: request.message,
    createdAt: Date.now()
  };

  session.messages.push(userMessage);

  const useSearch = shouldSearch(request.message, request.mode ?? "auto");
  let searchSummary: SearchSummary | null = null;

  if (useSearch) {
    events.onStatus("Searching web...");
    try {
      searchSummary = await searchWeb(request.message);
      events.onStatus("Summarizing sources...");
    } catch (err) {
      const msg = err instanceof Error ? err.message : "Search failed";
      events.onStatus("Search failed, continuing without web context.");
      events.onError(msg);
    }
  }

  if (totalTokens(session.messages) > MAX_CONTEXT_TOKENS) {
    const trimmed = trimHistory(session.messages, 12);
    session.summary = summarizeExtractive(trimmed, 6);
    session.messages = trimmed;
  }

  const searchContext = searchSummary ? buildSearchContext(searchSummary) : null;
  const systemPrompt = buildSystemPrompt(session.summary, searchContext);

  const payloadMessages = [
    { role: "system", content: systemPrompt },
    ...session.messages.map((m) => ({ role: m.role, content: m.content }))
  ];

  const assistantMessage: ChatMessage = {
    id: `msg_${Date.now()}_assistant`,
    role: "assistant",
    content: "",
    createdAt: Date.now()
  };

  try {
    await streamChatCompletion(payloadMessages, (token) => {
      assistantMessage.content += token;
      events.onToken(token);
    });
    session.messages.push(assistantMessage);
    memoryStore.set(request.sessionId, session);
    events.onDone();
  } catch (err) {
    const msg = err instanceof Error ? err.message : "LLM call failed";
    events.onError(msg);
  }
}