1import { Agent, ProviderTransport, type AgentEvent } from "@mariozechner/pi-agent";
  2import { getApiKey, type AgentTool, 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 { expandEnvVars } 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 that checks custom model headers first,
 23 * then falls back to pi-ai's built-in env-var lookup.
 24 *
 25 * Custom models may specify `Authorization: "Bearer $SOME_ENV_VAR"` in their
 26 * headers config. We extract the bearer token and expand env var references
 27 * so the value reaches the OpenAI-compatible SDK as a real API key.
 28 */
 29export function buildGetApiKey(config: RumiloConfig): (provider: string) => string | undefined {
 30  return (provider: string) => {
 31    if (config.custom_models) {
 32      for (const model of Object.values(config.custom_models)) {
 33        if (model.provider === provider && model.headers) {
 34          const authHeader = model.headers["Authorization"] ?? model.headers["authorization"];
 35          if (authHeader) {
 36            const expanded = expandEnvVars(authHeader);
 37            const match = expanded.match(/^Bearer\s+(.+)$/i);
 38            if (match) {
 39              return match[1];
 40            }
 41            return expanded;
 42          }
 43        }
 44      }
 45    }
 46
 47    // Fall back to pi-ai's built-in env var resolution
 48    // (handles anthropic → ANTHROPIC_API_KEY, openai → OPENAI_API_KEY, etc.)
 49    return getApiKey(provider);
 50  };
 51}
 52
 53export async function runAgent(query: string, options: AgentRunOptions): Promise<AgentRunResult> {
 54  const agent = new Agent({
 55    initialState: {
 56      systemPrompt: options.systemPrompt,
 57      model: resolveModel(options.model, options.config),
 58      tools: options.tools,
 59    },
 60    transport: new ProviderTransport({
 61      getApiKey: buildGetApiKey(options.config),
 62    }),
 63  });
 64
 65  if (options.onEvent) {
 66    agent.subscribe(options.onEvent);
 67  }
 68
 69  await agent.prompt(query);
 70
 71  // Check for errors in agent state
 72  if (agent.state.error) {
 73    throw new AgentError(agent.state.error);
 74  }
 75
 76  const last = agent.state.messages
 77    .slice()
 78    .reverse()
 79    .find((msg): msg is AssistantMessage => msg.role === "assistant");
 80
 81  // Check if the last assistant message indicates an error
 82  if (last?.stopReason === "error") {
 83    throw new AgentError(last.errorMessage ?? "Agent stopped with an unknown error");
 84  }
 85
 86  const text = last?.content
 87    ?.filter((content): content is Extract<typeof content, { type: "text" }> => content.type === "text")
 88    .map((content) => content.text)
 89    .join("")
 90    .trim();
 91
 92  if (text === undefined || text === "") {
 93    throw new AgentError("Agent returned no text response");
 94  }
 95
 96  return {
 97    message: text,
 98    usage: last?.usage,
 99  };
100}