Agent Types

import type {
  Agent,
  AgentState,
  AgentMessage,
  AgentTool,
  AgentEvent,
  ThinkingLevel,
} from "@mariozechner/pi-agent-core";

interface AgentState {
  systemPrompt: string;
  model: Model<any>;
  thinkingLevel: ThinkingLevel;
  tools: AgentTool<any>[];
  messages: AgentMessage[];
  isStreaming: boolean;
  streamMessage: AgentMessage | null;
  pendingToolCalls: Set<string>;
  error?: string;
}

type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages];

declare module "@mariozechner/pi-agent-core" {
  interface CustomAgentMessages {
    artifact: ArtifactMessage;
    notification: NotificationMessage;
  }
}

interface AgentTool<TSchema extends TSchema = any> {
  name: string;
  description: string;
  parameters: TSchema;
  execute: (
    toolCallId: string,
    args: Static<TSchema>,
    signal?: AbortSignal,
    updateCallback?: AgentToolUpdateCallback
  ) => Promise<AgentToolResult<any>>;
}

interface AgentToolResult<T> {
  content: (TextContent | ImageContent)[];
  details: T;
}

type ThinkingLevel = "off" | "minimal" | "low" | "medium" | "high" | "xhigh";

interface AgentContext {
  systemPrompt: string;
  messages: AgentMessage[];
  tools: AgentTool<any>[];
}

interface AgentLoopConfig extends SimpleStreamOptions {
  model: Model<any>;
  convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
  transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
  getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
  getSteeringMessages?: () => Promise<AgentMessage[]>;
  getFollowUpMessages?: () => Promise<AgentMessage[]>;
}

type AgentEvent =
  | AgentStartEvent
  | AgentEndEvent
  | TurnStartEvent
  | TurnEndEvent
  | MessageStartEvent
  | MessageUpdateEvent
  | MessageEndEvent
  | ToolExecutionStartEvent
  | ToolExecutionUpdateEvent
  | ToolExecutionEndEvent;

interface AgentStartEvent {
  type: "agent_start";
}

interface AgentEndEvent {
  type: "agent_end";
  messages: AgentMessage[];
}

interface MessageStartEvent {
  type: "message_start";
  message: AgentMessage;
}

interface MessageUpdateEvent {
  type: "message_update";
  message: AgentMessage;
}

interface MessageEndEvent {
  type: "message_end";
  message: AgentMessage;
}

interface ToolExecutionStartEvent {
  type: "tool_execution_start";
  toolCallId: string;
  toolName: string;
  arguments: Record<string, any>;
}

interface ToolExecutionUpdateEvent {
  type: "tool_execution_update";
  toolCallId: string;
  delta: string;
}

interface ToolExecutionEndEvent {
  type: "tool_execution_end";
  toolCallId: string;
  result: AgentToolResult<any>;
}

interface TurnStartEvent {
  type: "turn_start";
}

interface TurnEndEvent {
  type: "turn_end";
}

type AgentToolUpdateCallback = (delta: string) => void;

const tool: AgentTool = {
  name: "search",
  description: "Search the web",
  parameters: Type.Object({ query: Type.String() }),
  execute: async (toolCallId, args, signal, updateCallback) => {
    updateCallback?.("Searching...");
    const results = await search(args.query);
    updateCallback?.("Done!");
    
    return {
      content: [{ type: "text", text: results }],
      details: { count: results.length },
    };
  },
};

type StreamFn = (
  model: Model<any>,
  context: Context,
  options?: SimpleStreamOptions
) => AssistantMessageEventStream | Promise<AssistantMessageEventStream>;

import type {
  Agent,
  AgentState,
  AgentMessage,
  AgentTool,
  AgentEvent,
} from "@mariozechner/pi-agent-core";
import { Agent } from "@mariozechner/pi-agent-core";
import { Type } from "@sinclair/typebox";

// Define custom tool
const searchTool: AgentTool = {
  name: "search",
  description: "Search the web",
  parameters: Type.Object({
    query: Type.String({ description: "Search query" }),
  }),
  execute: async (toolCallId, args, signal, updateCallback) => {
    updateCallback?.("Searching...");
    // ... implement search ...
    return {
      content: [{ type: "text", text: "Results..." }],
      details: { count: 10 },
    };
  },
};

// Create agent
const agent: Agent = new Agent({
  initialState: {
    systemPrompt: "You are helpful.",
    model: getModel("anthropic", "claude-4.5-sonnet-20250514"),
    thinkingLevel: "medium",
    tools: [searchTool],
    messages: [],
  },
});

// Send message
const stream = agent.prompt("Search for cats");

// Type-safe event handling
stream.on("tool_execution_start", (event: ToolExecutionStartEvent) => {
  console.log(`Executing: ${event.toolName}`);
});

// Await result
const messages: AgentMessage[] = await stream.result();