1import { Agent, type AgentEvent, type AgentTool } from "@mariozechner/pi-agent-core";
 2import { getEnvApiKey, type AssistantMessage } from "@mariozechner/pi-ai";
 3import type { RumiloConfig } from "../config/schema.js";
 4import { resolveModel } from "./model-resolver.js";
 5import { AgentError } from "../util/errors.js";
 6import { resolveConfigValue } from "../util/env.js";
 7
 8export interface AgentRunOptions {
 9  model: string;
10  systemPrompt: string;
11  tools: AgentTool[];
12  onEvent?: (event: AgentEvent) => void;
13  config: RumiloConfig;
14}
15
16export interface AgentRunResult {
17  message: string;
18  usage?: unknown;
19}
20
21/**
22 * Build a getApiKey callback for the Agent.
23 *
24 * Resolution order:
25 * 1. Custom model config — if a custom model for this provider defines an
26 *    `apiKey` field, resolve it via `resolveConfigValue` (supports env var
27 *    names, `$VAR` references, and `!shell` commands).
28 * 2. pi-ai’s built-in env-var lookup (`ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, etc.).
29 */
30export function buildGetApiKey(config: RumiloConfig): (provider: string) => string | undefined {
31  return (provider: string) => {
32    if (config.custom_models) {
33      for (const model of Object.values(config.custom_models)) {
34        if (model.provider === provider && model.api_key) {
35          return resolveConfigValue(model.api_key);
36        }
37      }
38    }
39
40    return getEnvApiKey(provider);
41  };
42}
43
44export async function runAgent(query: string, options: AgentRunOptions): Promise<AgentRunResult> {
45  const agent = new Agent({
46    initialState: {
47      systemPrompt: options.systemPrompt,
48      model: resolveModel(options.model, options.config),
49      tools: options.tools,
50    },
51    getApiKey: buildGetApiKey(options.config),
52  });
53
54  if (options.onEvent) {
55    agent.subscribe(options.onEvent);
56  }
57
58  await agent.prompt(query);
59
60  // Check for errors in agent state
61  if (agent.state.error) {
62    throw new AgentError(agent.state.error);
63  }
64
65  const last = agent.state.messages
66    .slice()
67    .reverse()
68    .find((msg): msg is AssistantMessage => msg.role === "assistant");
69
70  // Check if the last assistant message indicates an error
71  if (last?.stopReason === "error") {
72    throw new AgentError(last.errorMessage ?? "Agent stopped with an unknown error");
73  }
74
75  const text = last?.content
76    ?.filter((content): content is Extract<typeof content, { type: "text" }> => content.type === "text")
77    .map((content) => content.text)
78    .join("")
79    .trim();
80
81  if (text === undefined || text === "") {
82    throw new AgentError("Agent returned no text response");
83  }
84
85  return {
86    message: text,
87    usage: last?.usage,
88  };
89}