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