lam-agent-vicinae/src/lib/agent.ts

import { createSourceList } from "./citations";
import { DEFAULT_MAX_TOOL_ROUNDS, MAX_REPEATED_TOOL_CALLS } from "./constants";
import { streamChat } from "./ollama";
import { executeTool, toolLabel, type ToolContext } from "./tools";
import type { ApiMessage, ToolCall, ToolCallRecord, ToolDefinition } from "./types";

/**
 * The tool-calling loop.
 *
 * Kept free of any React or Vicinae imports so the whole thing can be driven
 * from a plain unit test with a stubbed `fetch`.
 */

export interface AgentEvents {
  onContentDelta?: (delta: string) => void;
  onThinkingDelta?: (delta: string) => void;
  onToolCall?: (record: ToolCallRecord) => void;
  /**
   * A tool is about to run, with how many times this turn has already made the
   * identical call, so the UI can tell progress from a rut.
   */
  onToolStart?: (name: string, label: string, count: number) => void;
  /** A short line for the UI's activity indicator, e.g. "Searching the web". */
  onStatus?: (status: string) => void;
}

export interface AgentOptions {
  host: string;
  apiKey: string;
  model: string;
  /** Omit or pass "off" to leave reasoning disabled. */
  think?: string;
  messages: ApiMessage[];
  tools: ToolDefinition[];
  /**
   * Rounds of tool calls this turn may spend before it has to answer.
   *
   * Omitted means {@link DEFAULT_MAX_TOOL_ROUNDS}, and `null` means no cap at
   * all. Unlimited is spelled as an explicit `null` rather than as an omitted
   * option because the opposite arrangement makes "no cap" the default for a
   * function whose whole job is to stop looping, and a caller that forgets to
   * pass a preference should not be the one that removes the guard.
   */
  maxToolRounds?: number | null;
  toolContext: Omit<ToolContext, "signal" | "now"> & { now?: Date };
  signal?: AbortSignal;
  events?: AgentEvents;
}

export interface AgentResult {
  content: string;
  thinking: string;
  toolCalls: ToolCallRecord[];
}

/**
 * Picks the authoritative tool-call list out of a stream.
 *
 * Ollama sends tool calls complete in the final streamed message rather than
 * dribbling them out across chunks the way OpenAI does, so the last non-empty
 * array wins. Exported because the assumption is worth pinning down in a test:
 * if Ollama ever starts streaming partial calls, this is the one place to fix.
 */
export function pickToolCalls(
  previous: ToolCall[] | undefined,
  incoming: ToolCall[] | undefined,
): ToolCall[] | undefined {
  return incoming && incoming.length > 0 ? incoming : previous;
}

function normalizeToolCall(call: ToolCall): { name: string; arguments: Record<string, unknown> } {
  const name = typeof call.function?.name === "string" ? call.function.name : "";
  const rawArgs = call.function?.arguments;
  const args = rawArgs && typeof rawArgs === "object" && !Array.isArray(rawArgs) ? rawArgs : {};
  return { name, arguments: args as Record<string, unknown> };
}

/**
 * A stable key for "the same call".
 *
 * Keys are sorted so that a model which writes the same arguments in a different
 * order still counts as repeating itself, which is the shape the rut usually
 * takes. Values are compared as JSON, so two searches for the same query count as
 * one call, which is the point: the second one will not return anything the first
 * did not.
 */
function callSignature(name: string, args: Record<string, unknown>): string {
  const sorted = Object.keys(args)
    .sort()
    .map((key) => [key, args[key]] as const);

  return `${name}(${JSON.stringify(sorted)})`;
}

/** Why the tool loop stopped short: the budget ran out, or the model got stuck. */
type StopReason = { kind: "budget"; rounds: number } | { kind: "repeat"; name: string; count: number };

/** Sent to the model in place of a tool result, so the protocol stays valid and it learns why. */
function repeatedNotice(name: string, count: number): string {
  return `Refused: this turn has now asked for \`${name}\` with identical arguments ${count} times, so running it again will not return anything new. No more tool calls are available. Answer with what you have already gathered.`;
}

/**
 * The last word, when the tool budget runs out.
 *
 * The loop used to end by throwing away every preamble the model had written and
 * returning a notice about itself, so a long research question produced no answer
 * at all. One closing round, with the tools taken away, turns that into an answer
 * built from whatever was already gathered, with the gaps named. The instruction
 * asks for the gaps because an answer that knows what it is missing is worth more
 * than a confident half of one.
 */
const CLOSING_INSTRUCTION =
  "You have reached the limit on tool calls for this turn, so no more tools are available. Answer now using what you have already gathered. If something is missing, say plainly which parts you could not confirm.";

/**
 * The closing instruction when the loop stopped because the model was repeating itself.
 *
 * Worth saying outright rather than reusing the budget wording: the model has no
 * way to know it was caught going in circles, and naming the call is what turns
 * "I ran out of tools" into something it can act on.
 */
function closingInstruction(reason: StopReason): string {
  if (reason.kind === "budget") return CLOSING_INSTRUCTION;

  return `This turn asked for \`${reason.name}\` with identical arguments ${reason.count} times, so no more tools are available. Answer now using what you have already gathered. If something is missing, say plainly which parts you could not confirm.`;
}

/** Only reachable if the closing round comes back empty too. */
function exhaustedNotice(reason: StopReason): string {
  if (reason.kind === "repeat") {
    return `_Stopped because \`${reason.name}\` kept repeating the same call, and could not write an answer from what it had. Try asking something narrower._`;
  }

  return `_Ran out of tool calls after ${reason.rounds} rounds and could not write an answer. Try asking something narrower._`;
}

export async function runAgent(options: AgentOptions): Promise<AgentResult> {
  const conversation: ApiMessage[] = [...options.messages];
  const toolCalls: ToolCallRecord[] = [];
  /**
   * The turn's sources, so one number means one page for the whole turn.
   *
   * Created here rather than passed in: it is scoped to a single run of this loop,
   * and every call the loop makes shares it, which is what keeps three searches in
   * one answer from each starting again at 1.
   */
  const sources = createSourceList();
  /** `Infinity` needs no special case: the loop simply never reaches its own end. */
  const rounds =
    options.maxToolRounds === null ? Infinity : (options.maxToolRounds ?? DEFAULT_MAX_TOOL_ROUNDS);
  /** How many times this turn has asked for each identical call. */
  const repeats = new Map<string, number>();
  let thinking = "";
  /**
   * The most recent preamble the model wrote.
   *
   * Not normally shown, because a preamble is what it says before finding out
   * anything. Kept anyway, so that a turn which runs out of budget still has
   * something to fall back on if the closing round produces nothing.
   */
  let lastPreamble = "";
  let stopped: StopReason | undefined;

  for (let iteration = 0; iteration < rounds; iteration++) {
    let content = "";
    let streamedThinking = "";
    let pending: ToolCall[] | undefined;

    for await (const chunk of streamChat({
      host: options.host,
      apiKey: options.apiKey,
      model: options.model,
      messages: conversation,
      tools: options.tools.length > 0 ? options.tools : undefined,
      think: options.think,
      signal: options.signal,
    })) {
      if (chunk.thinking) {
        streamedThinking += chunk.thinking;
        options.events?.onThinkingDelta?.(chunk.thinking);
      }
      if (chunk.content) {
        content += chunk.content;
        options.events?.onContentDelta?.(chunk.content);
      }
      pending = pickToolCalls(pending, chunk.toolCalls);
    }

    thinking += streamedThinking;

    // No tool calls means this was the model's answer, rather than something it
    // said on the way to finding out.
    if (!pending || pending.length === 0) {
      return { content, thinking, toolCalls };
    }

    if (content.trim()) lastPreamble = content;

    conversation.push({ role: "assistant", content, tool_calls: pending });

    let repeated: { name: string; count: number } | undefined;

    for (const rawCall of pending) {
      const { name, arguments: args } = normalizeToolCall(rawCall);

      if (!name) {
        conversation.push({
          role: "tool",
          content: "Malformed tool call: missing function name.",
          tool_name: "unknown",
        });
        toolCalls.push({ name: "unknown", arguments: {}, result: "Malformed tool call.", isError: true });
        continue;
      }

      const signature = callSignature(name, args);
      const count = (repeats.get(signature) ?? 0) + 1;
      repeats.set(signature, count);

      // Every tool call in a round needs a result, so a refused one is answered
      // with the refusal rather than skipped: dropping it would leave the next
      // request carrying a tool call the protocol says must be answered.
      if (count >= MAX_REPEATED_TOOL_CALLS) {
        repeated = { name, count };

        const refusal = repeatedNotice(name, count);
        const record: ToolCallRecord = {
          name,
          arguments: args,
          result: refusal,
          isError: true,
          repeated: true,
        };

        conversation.push({ role: "tool", content: refusal, tool_name: name });
        toolCalls.push(record);
        options.events?.onToolCall?.(record);
        continue;
      }

      options.events?.onToolStart?.(name, toolLabel(name), count);
      options.events?.onStatus?.(toolLabel(name));

      const result = await executeTool(name, args, {
        host: options.toolContext.host,
        apiKey: options.toolContext.apiKey,
        place: options.toolContext.place,
        now: options.toolContext.now,
        signal: options.signal,
        sources,
        // The policy and the ledger travel with the turn, so a second write to a path
        // the user already approved does not ask again, and a new question starts
        // asking from the beginning.
        files: options.toolContext.files,
      });

      const record: ToolCallRecord = {
        name,
        arguments: args,
        result: result.content,
        isError: result.isError,
        // Only what this call added. The turn's whole list is the flattening of
        // its records, so the running total is deliberately not repeated here.
        sources: result.sources,
      };
      toolCalls.push(record);
      options.events?.onToolCall?.(record);

      conversation.push({ role: "tool", content: result.content, tool_name: name });
    }

    // The rest of the round still ran, because every call in it needs an answer,
    // but the turn is over: the model is going in circles and one more round
    // would only add one more circle.
    if (repeated) {
      stopped = { kind: "repeat", name: repeated.name, count: repeated.count };
      break;
    }
  }

  // The loop is done, either because the budget is spent or because the model got
  // stuck. Either way the turn gets one closing round with the tools taken away.
  // Ending on an answer drawn from what was gathered beats ending on a notice
  // about the loop itself, and the model is the only one that can say what it
  // did and did not manage to confirm.
  const reason: StopReason = stopped ?? { kind: "budget", rounds };

  options.events?.onStatus?.("Writing the answer");

  conversation.push({ role: "user", content: closingInstruction(reason) });

  let closing = "";

  for await (const chunk of streamChat({
    host: options.host,
    apiKey: options.apiKey,
    model: options.model,
    messages: conversation,
    // Deliberately no tools: the answer has to be written from what is already here.
    tools: undefined,
    think: options.think,
    signal: options.signal,
  })) {
    if (chunk.thinking) {
      thinking += chunk.thinking;
      options.events?.onThinkingDelta?.(chunk.thinking);
    }
    if (chunk.content) {
      closing += chunk.content;
      options.events?.onContentDelta?.(chunk.content);
    }
  }

  return {
    content: closing.trim() || lastPreamble.trim() || exhaustedNotice(reason),
    thinking,
    toolCalls,
  };
}