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

import { MAX_TOOL_ITERATIONS } 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 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[];
  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> };
}

/**
 * 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.";

/** Only reachable if the closing round comes back empty too. */
function exhaustedNotice(rounds: number): string {
  return `_Ran out of tool calls after ${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[] = [];
  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 = "";

  for (let iteration = 0; iteration < MAX_TOOL_ITERATIONS; 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 });

    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;
      }

      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,
      });

      const record: ToolCallRecord = {
        name,
        arguments: args,
        result: result.content,
        isError: result.isError,
      };
      toolCalls.push(record);
      options.events?.onToolCall?.(record);

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

  // Budget spent, so 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.
  options.events?.onStatus?.("Writing the answer");

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

  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(MAX_TOOL_ITERATIONS),
    thinking,
    toolCalls,
  };
}