RUN YOUR AGENT AS A SUBAGENT IN PI

const API = "https://api.example-stocks.com";

// Copies of the Oumi tool definitions. A name, parameter, or description that
// differs from training makes the model call tools that do not exist.
export const TOOLS = [
	{
		type: "function",
		function: {
			name: "search_company",
			description: "Find a company's stock ticker symbol by company name. Call this before get_stock_price.",
			parameters: {
				type: "object",
				properties: { name: { type: "string", description: "Company name, e.g. 'Apple'." } },
				required: ["name"],
			},
		},
	},
	{
		type: "function",
		function: {
			name: "get_stock_price",
			description: "Get the latest stock price for a ticker symbol returned by search_company.",
			parameters: {
				type: "object",
				properties: { ticker: { type: "string", description: "Stock ticker symbol, e.g. 'AAPL'." } },
				required: ["ticker"],
			},
		},
	},
];

export async function runTool(name: string, args: Record<string, string>, signal?: AbortSignal): Promise<string> {
	if (name === "search_company") return searchCompany(args.name, signal);
	if (name === "get_stock_price") return getStockPrice(args.ticker, signal);
	throw new Error(`Unknown tool: ${name}`);
}

async function searchCompany(name: string, signal?: AbortSignal): Promise<string> {
	const res = await fetch(`${API}/search?q=${encodeURIComponent(name)}`, { signal });
	if (res.status === 404) return `No company found matching "${name}".`;
	if (!res.ok) throw new Error(`Search returned HTTP ${res.status}`);

	const { ticker, name: found } = (await res.json()) as { ticker: string; name: string };
	return `Found ${found} (ticker: ${ticker}).`;
}

async function getStockPrice(ticker: string, signal?: AbortSignal): Promise<string> {
	const res = await fetch(`${API}/quote/${encodeURIComponent(ticker)}`, { signal });
	if (res.status === 404) return `No price found for "${ticker}".`;
	if (!res.ok) throw new Error(`Quote returned HTTP ${res.status}`);

	const { price, currency } = (await res.json()) as { price: number; currency: string };
	return `${ticker}: ${price} ${currency}.`;
}

import { readFileSync } from "node:fs";

import { runTool, TOOLS } from "./tools";

const ENDPOINT = "https://inference.oumi.ai/v1/chat/completions";
const MODEL = "projects/{project_id}/deployments/{deployment_id}";

// Stops a model that keeps calling tools. Set it above the number of tool
// calls that your longest training conversation makes.
const MAX_TOOL_STEPS = 8;

const SYSTEM_PROMPT = readFileSync(new URL("./system-prompt.md", import.meta.url), "utf8");

type Message = Record<string, unknown>;
type ToolCall = { id: string; type: "function"; function: { name: string; arguments: string } };

export class StockAgent {
	private history: Message[] = [{ role: "system", content: SYSTEM_PROMPT }];

	constructor(private readonly apiKey: string) {}

	// A turn that fails part way leaves tool calls without results, and your
	// deployment rejects a conversation like that. Removing the failed turn
	// keeps the next question working.
	async ask(question: string, onToolCall: (call: string) => void, signal?: AbortSignal): Promise<string> {
		const start = this.history.length;
		try {
			return await this.turn(question, onToolCall, signal);
		} catch (error) {
			this.history.length = start;
			throw error;
		}
	}

	private async turn(question: string, onToolCall: (call: string) => void, signal?: AbortSignal): Promise<string> {
		this.history.push({ role: "user", content: question });

		for (let step = 0; step < MAX_TOOL_STEPS; step++) {
			const reply = await this.complete(signal);
			const calls: ToolCall[] = reply.tool_calls ?? [];

			if (calls.length === 0) {
				this.history.push({ role: "assistant", content: reply.content });
				return reply.content;
			}

			this.history.push({ role: "assistant", content: reply.content ?? "", tool_calls: calls });
			for (const call of calls) {
				onToolCall(`${call.function.name}(${call.function.arguments})`);
				this.history.push({ role: "tool", tool_call_id: call.id, content: await runToolCall(call, signal) });
			}
		}

		throw new Error(`The stock agent made ${MAX_TOOL_STEPS} rounds of tool calls without an answer.`);
	}

	private async complete(signal?: AbortSignal): Promise<{ content: string; tool_calls?: ToolCall[] }> {
		const res = await fetch(ENDPOINT, {
			method: "POST",
			headers: { "Content-Type": "application/json", Authorization: `Bearer ${this.apiKey}` },
			body: JSON.stringify({ model: MODEL, messages: this.history, tools: TOOLS }),
			signal,
		});
		if (!res.ok) throw new Error(`Oumi returned HTTP ${res.status}: ${await res.text()}`);

		const body = (await res.json()) as { choices: { message: { content: string; tool_calls?: ToolCall[] } }[] };
		return body.choices[0].message;
	}
}

// A failed tool call goes back to the model as its result, as it does in pi,
// so the model can tell the user instead of inventing a value.
async function runToolCall(call: ToolCall, signal?: AbortSignal): Promise<string> {
	try {
		return await runTool(call.function.name, JSON.parse(call.function.arguments), signal);
	} catch (error) {
		return `Error: ${error instanceof Error ? error.message : String(error)}`;
	}
}

import { readFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";

import { Type } from "@earendil-works/pi-ai";
import { defineTool, type ExtensionAPI } from "@earendil-works/pi-coding-agent";

import { StockAgent } from "./agent";

const API_KEY_FILE = join(homedir(), ".config/oumi/api.key");

export default function (pi: ExtensionAPI) {
	// One subagent per pi session. It keeps its own history, so a follow-up
	// question can refer to a company that it already found.
	let agent: StockAgent | undefined;
	pi.on("session_start", () => {
		agent = undefined;
	});

	const askStockAgent = defineTool({
		name: "ask_stock_agent",
		label: "Stock Agent",
		description:
			"Ask the stock agent for the latest stock price of one or more companies. It finds the " +
			"ticker symbols itself and answers in prose. It remembers earlier questions in this " +
			"session, so a follow-up can refer to a company it already found. Set reset to true to " +
			"start a new conversation with it.",
		parameters: Type.Object({
			question: Type.String({ description: "The question, in plain English." }),
			reset: Type.Optional(Type.Boolean({ description: "Forget earlier questions before this one." })),
		}),

		async execute(_toolCallId, params, signal, onUpdate) {
			if (!agent || params.reset) agent = new StockAgent(readFileSync(API_KEY_FILE, "utf8").trim());

			// Shows each of the subagent's tool calls in pi while it works.
			const calls: string[] = [];
			const answer = await agent.ask(
				params.question,
				(call) => {
					calls.push(call);
					onUpdate?.({ content: [{ type: "text", text: calls.join("\n") }], details: { calls } });
				},
				signal,
			);

			return { content: [{ type: "text", text: answer }], details: { calls } };
		},
	});

	pi.registerTool(askStockAgent);
}