434 lines
16 KiB
Markdown
434 lines
16 KiB
Markdown
# Pi Agent Architecture - Top-Down Overview
|
|
|
|
## Executive Summary
|
|
|
|
The Pi Agent is a **stateful, event-driven agent framework** built in TypeScript. It provides:
|
|
|
|
1. **Core Agent** - Low-level agent loop with message/tool streaming
|
|
2. **Agent Harness** - High-level session management with persistence, branching, and compaction
|
|
|
|
Both layers follow the **same core pattern**: stream LLM response → execute tools → emit events → repeat.
|
|
|
|
---
|
|
|
|
## Architecture Layers
|
|
|
|
```
|
|
┌─────────────────────────────────────────────────────────────────────────────────────┐
|
|
│ APPLICATION LAYER │
|
|
│ • Creates Agent/AgentHarness instances │
|
|
│ • Subscribes to events for UI updates │
|
|
│ • Provides tools and model configuration │
|
|
└─────────────────────────────────────────────────────────────────────────────────────┘
|
|
│
|
|
┌─────────────────────────────┼─────────────────────────────┐
|
|
│ │ │
|
|
▼ ▼ ▼
|
|
┌────────────────────────┐ ┌────────────────────────────────┐ ┌─────────────────┐
|
|
│ Agent (Core) │ │ Agent Harness (High-Level) │ │ Agent-Loop │
|
|
│ │ │ │ │ │
|
|
│ • State management │ │ • Session persistence │ │ • Turn │
|
|
│ • Event streaming │ │ • Branching/compaction │ │ • Tool exec │
|
|
│ • Steering/follow-up │ │ • Skills/templates │ │ • Message │
|
|
│ queues │ │ • Tool context binding │ │ streaming │
|
|
│ • Hook system │ │ • State snapshots │ │ │
|
|
└────────────────────────┘ └────────────────────────────────┘ └─────────────────┘
|
|
│ │
|
|
▼ ▼
|
|
┌──────────────────────────┐ ┌─────────────────┐
|
|
│ LLM Provider API │ │ Session Repo │
|
|
│ (via @earendil-works) │ │ (JSONL/ │
|
|
└──────────────────────────┘ │ Memory) │
|
|
└─────────────────┘
|
|
```
|
|
|
|
---
|
|
|
|
## Core Concepts
|
|
|
|
### 1. AgentMessage
|
|
|
|
```typescript
|
|
type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages]
|
|
```
|
|
|
|
The unified message type that combines:
|
|
- **LLM messages**: `user`, `assistant`, `toolResult` (from pi-ai)
|
|
- **Custom messages**: Application-specific types (via declaration merging)
|
|
|
|
### 2. AgentEvent
|
|
|
|
```typescript
|
|
type AgentEvent =
|
|
| { type: "agent_start" }
|
|
| { type: "agent_end"; messages: AgentMessage[] }
|
|
| { type: "turn_start" }
|
|
| { type: "turn_end"; message: AgentMessage; toolResults: ToolResultMessage[] }
|
|
| { type: "message_start"; message: AgentMessage }
|
|
| { type: "message_update"; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent }
|
|
| { type: "message_end"; message: AgentMessage }
|
|
| { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any }
|
|
| { type: "tool_execution_update"; toolCallId: string; toolName: string; args: any; partialResult: any }
|
|
| { type: "tool_execution_end"; toolCallId: string; toolName: string; result: any; isError: boolean }
|
|
```
|
|
|
|
**Event Flow per Turn:**
|
|
```
|
|
turn_start
|
|
message_start (user prompt)
|
|
message_end
|
|
message_start (assistant streaming)
|
|
message_update (multiple - as chunks arrive)
|
|
message_end
|
|
tool_execution_start (if tool calls present)
|
|
tool_execution_update (if tool streams partial results)
|
|
tool_execution_end
|
|
turn_end
|
|
```
|
|
|
|
### 3. AgentTool
|
|
|
|
```typescript
|
|
interface AgentTool<TParameters extends TSchema, TDetails> {
|
|
name: string;
|
|
label: string;
|
|
description: string;
|
|
parameters: TSchema;
|
|
execute(
|
|
toolCallId: string,
|
|
params: Static<TParameters>,
|
|
signal?: AbortSignal,
|
|
onUpdate?: AgentToolUpdateCallback<TDetails>
|
|
): Promise<AgentToolResult<TDetails>>;
|
|
}
|
|
```
|
|
|
|
---
|
|
|
|
## Process Flow Diagrams
|
|
|
|
### Prompt Flow (High-Level)
|
|
|
|
```
|
|
User Input
|
|
│
|
|
▼
|
|
Agent.prompt("Hello")
|
|
│
|
|
├─► normalizePromptInput() → AgentMessage[]
|
|
│
|
|
├─► runWithLifecycle()
|
|
│ ├─► Set isStreaming=true
|
|
│ └─► Create abort controller
|
|
│
|
|
▼
|
|
runAgentLoop()
|
|
│
|
|
├─► Emit: agent_start
|
|
├─► Emit: turn_start
|
|
├─► Emit: message_start/end (prompts)
|
|
│
|
|
▼
|
|
runLoop() - Main Loop
|
|
│
|
|
├─► Check steering queue (drain if any)
|
|
├─► Check follow-up queue (skip if not first turn)
|
|
│
|
|
▼
|
|
streamAssistantResponse()
|
|
│
|
|
├─► transformContext() [optional]
|
|
├─► convertToLlm() → Message[]
|
|
├─► Build Context {systemPrompt, messages, tools}
|
|
├─► Resolve API key
|
|
├─► Call streamFn(model, context, options)
|
|
│
|
|
▼
|
|
Assistant Message Stream
|
|
│
|
|
├─► message_start (assistant)
|
|
├─► message_update (text chunks)
|
|
├─► message_update (toolCall blocks)
|
|
├─► message_end
|
|
│
|
|
▼
|
|
executeToolCalls()
|
|
│
|
|
├─► Check if sequential/parallel execution
|
|
├─► For each tool call:
|
|
│ ├─► prepareToolCall()
|
|
│ │ ├─► Find tool by name
|
|
│ │ ├─► Validate arguments
|
|
│ │ └─► beforeToolCall() hook
|
|
│ │
|
|
│ ├─► executePreparedToolCall()
|
|
│ │ └─► tool.execute() with onUpdate callback
|
|
│ │
|
|
│ └─► finalizeExecutedToolCall()
|
|
│ └─► afterToolCall() hook
|
|
│
|
|
├─► Emit: tool_execution_start/update/end
|
|
└─► Emit: message_start/end (toolResult)
|
|
│
|
|
▼
|
|
turn_end
|
|
│
|
|
├─► Check prepareNextTurn hook
|
|
├─► Check shouldStopAfterTurn hook
|
|
├─► Drain steering queue
|
|
└─► Drain follow-up queue
|
|
│
|
|
├─► If steering/follow-up exists → repeat loop
|
|
└─► If no more messages → agent_end
|
|
```
|
|
|
|
### Tool Execution Flow (Detailed)
|
|
|
|
```
|
|
Tool Call from LLM
|
|
│
|
|
▼
|
|
prepareToolCall()
|
|
│
|
|
├─► Find tool in currentContext.tools
|
|
│ └─► If not found → immediate error
|
|
│
|
|
├─► prepareToolCallArguments() [optional]
|
|
│
|
|
├─► validateToolArguments()
|
|
│ └─► If invalid → immediate error
|
|
│
|
|
└─► beforeToolCall() hook
|
|
├─► Return {block: true} → error
|
|
└─► Continue
|
|
│
|
|
▼
|
|
executePreparedToolCall()
|
|
│
|
|
├─► Call tool.execute() with onUpdate callback
|
|
│ └─► tool calls onUpdate(partialResult) during execution
|
|
│
|
|
├─► onUpdate() → emit tool_execution_update
|
|
└─► Return {result, isError}
|
|
│
|
|
▼
|
|
finalizeExecutedToolCall()
|
|
│
|
|
└─► afterToolCall() hook
|
|
├─► Override content/details/usage/terminate
|
|
└─► Return {toolCall, result, isError}
|
|
│
|
|
▼
|
|
emitToolExecutionEnd()
|
|
│
|
|
└─► Emit: tool_execution_end
|
|
│
|
|
▼
|
|
createToolResultMessage()
|
|
│
|
|
└─► Create ToolResultMessage with:
|
|
├─► toolCallId
|
|
├─► toolName
|
|
├─► content
|
|
├─► details
|
|
├─► usage
|
|
└─► isError
|
|
│
|
|
▼
|
|
emitToolResultMessage()
|
|
│
|
|
├─► Emit: message_start
|
|
└─► Emit: message_end
|
|
```
|
|
|
|
### Session Persistence Flow
|
|
|
|
```
|
|
AgentHarness.handleAgentEvent()
|
|
│
|
|
├─► message_end → session.appendMessage()
|
|
│ └─► Storage: write entry to JSONL file
|
|
│
|
|
├─► turn_end → flushPendingSessionWrites()
|
|
│ ├─► Write all pending entries
|
|
│ ├─► Emit: save_point
|
|
│ └─► session.getStorage().setLeafId()
|
|
│
|
|
└─► agent_end → flushPendingSessionWrites()
|
|
├─► Write leaf entry pointing to last message
|
|
└─► Emit: settled
|
|
│
|
|
▼
|
|
Session Tree Structure:
|
|
root
|
|
├─► message (user prompt #1)
|
|
├─► message (assistant #1)
|
|
├─► tool_result (result #1)
|
|
├─► turn_end
|
|
├─► message (user prompt #2)
|
|
├─► message (assistant #2)
|
|
├─► compaction (summary of history)
|
|
├─► message (assistant continues)
|
|
└─► leaf → points to current head
|
|
```
|
|
|
|
---
|
|
|
|
## Hook System
|
|
|
|
### Agent-Level Hooks (agent-loop.ts)
|
|
|
|
```typescript
|
|
interface AgentLoopConfig {
|
|
// Message transformation
|
|
convertToLlm: (messages: AgentMessage[]) => Message[]
|
|
transformContext?: (messages: AgentMessage[]) => AgentMessage[]
|
|
|
|
// Lifecycle hooks
|
|
beforeToolCall?: (context: BeforeToolCallContext) => BeforeToolCallResult
|
|
afterToolCall?: (context: AfterToolCallContext) => AfterToolCallResult
|
|
shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean
|
|
prepareNextTurn?: (context: PrepareNextTurnContext) => AgentLoopTurnUpdate
|
|
|
|
// Queue draining
|
|
getSteeringMessages?: () => AgentMessage[]
|
|
getFollowUpMessages?: () => AgentMessage[]
|
|
}
|
|
```
|
|
|
|
### Harness-Level Hooks (agent-harness.ts)
|
|
|
|
```typescript
|
|
// Hook types in AgentHarnessEventResultMap:
|
|
type HookName =
|
|
| "before_agent_start"
|
|
| "context"
|
|
| "tool_call"
|
|
| "tool_result"
|
|
| "session_before_compact"
|
|
| "session_before_tree"
|
|
| "before_provider_request"
|
|
| "before_provider_payload"
|
|
```
|
|
|
|
**Hook Execution Order per Turn:**
|
|
```
|
|
1. before_agent_start (harness)
|
|
2. context (harness) → transformContext
|
|
3. streamAssistantResponse
|
|
├─► Before provider request (harness)
|
|
├─► convertToLlm (agent)
|
|
└─► LLM call
|
|
4. For each tool call:
|
|
├─► tool_call (harness) → beforeToolCall
|
|
├─► Execute tool
|
|
└─► tool_result (harness) → afterToolCall
|
|
5. turn_end
|
|
6. shouldStopAfterTurn (agent)
|
|
7. prepareNextTurn (agent)
|
|
8. Drain steering/follow-up queues
|
|
```
|
|
|
|
---
|
|
|
|
## Data Flow Summary
|
|
|
|
```
|
|
┌────────────────────────────────────────────────────────────────────────────────┐
|
|
│ AGENT LIFECYCLE - DATA FLOW │
|
|
├────────────────────────────────────────────────────────────────────────────────┤
|
|
│ 1. INPUT │
|
|
│ • prompt("Hello") → normalizePromptInput() │
|
|
│ → AgentMessage[] │
|
|
│ 2. INITIATE │
|
|
│ • createMutableAgentState() │
|
|
│ • runWithLifecycle() │
|
|
│ 3. LOOP CONTROL │
|
|
│ • runLoop() │
|
|
│ ├─► Steering queue? → drain and inject │
|
|
│ └─► Follow-up queue? (after first turn) │
|
|
│ 4. LLM STREAMING │
|
|
│ • transformContext() [optional] │
|
|
│ • convertToLlm() │
|
|
│ • streamFn() │
|
|
│ → AssistantMessage stream (text + toolCalls) │
|
|
│ 5. TOOL EXECUTION │
|
|
│ • executeToolCalls() │
|
|
│ ├─► prepareToolCall() │
|
|
│ │ ├─► beforeToolCall() hook │
|
|
│ │ └─► Validate args │
|
|
│ ├─► executePreparedToolCall() │
|
|
│ │ └─► tool.execute() │
|
|
│ └─► finalizeExecutedToolCall() │
|
|
│ └─► afterToolCall() hook │
|
|
│ 6. UPDATE STATE │
|
|
│ • Push assistant message to state.messages │
|
|
│ • Push toolResult messages to state.messages │
|
|
│ 7. TERMINATION CHECK │
|
|
│ • shouldStopAfterTurn? → exit │
|
|
│ • prepareNextTurn? → update context/model │
|
|
│ • Drain steering/follow-up → continue │
|
|
│ 8. FINISH │
|
|
│ • emit agent_end │
|
|
│ • finishRun() → reset isStreaming │
|
|
└────────────────────────────────────────────────────────────────────────────────┘
|
|
```
|
|
|
|
---
|
|
|
|
## Key Design Patterns
|
|
|
|
### 1. Event-Driven Architecture
|
|
|
|
- All external communication via `AgentEvent` stream
|
|
- Hooks can be async and are awaited in order
|
|
- Abort signal propagated through all operations
|
|
|
|
### 2. State Isolation
|
|
|
|
- `AgentState` is read-only externally
|
|
- `AgentHarness` snapshots state per turn
|
|
- Context transforms return new arrays (immutability)
|
|
|
|
### 3. Layered Abstraction
|
|
|
|
```
|
|
Low-level (agent-loop.ts)
|
|
• Pure async iteration
|
|
• No session management
|
|
• No tool context binding
|
|
|
|
High-level (agent-harness.ts)
|
|
• Session persistence
|
|
• Branching/compaction
|
|
• Hook system for customization
|
|
```
|
|
|
|
### 4. Extensibility Points
|
|
|
|
- **Custom messages**: Extend `CustomAgentMessages` interface
|
|
- **Custom hooks**: Add handlers via `subscribe()`/`on()`
|
|
- **Tool context**: Pass `toolContext` to harness constructor
|
|
- **Storage**: Implement `SessionStorage` interface
|
|
|
|
---
|
|
|
|
## Learning Path
|
|
|
|
1. **Start with types.ts** - Understand `AgentMessage`, `AgentEvent`, `AgentTool`
|
|
2. **Read agent-loop.ts** - See how messages flow through the loop
|
|
3. **Study agent.ts** - See how Agent wraps the loop with state management
|
|
4. **Read agent-harness.ts** - See how session management hooks into the loop
|
|
5. **Explore session/* files** - Understand persistence and branching
|
|
6. **Study tools/* files** - See concrete tool implementations
|
|
|
|
---
|
|
|
|
## Next Steps
|
|
|
|
See individual markdown files in this folder for:
|
|
- `AGENT-LOOP-DETAILED.md` - Deep dive into the agent loop
|
|
- `HOOK-SYSTEM.md` - Complete hook documentation
|
|
- `SESSION-ARCHITECTURE.md` - Session persistence details
|
|
- `TOOL-EXECUTION.md` - Tool execution mechanics
|