Skip to content
Draft
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 0 additions & 14 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

8 changes: 5 additions & 3 deletions src/agent/__tests__/agent.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -315,7 +315,7 @@ describe('Agent', () => {

// Create agent with custom printer for testing
const agent = new Agent({ model, printer: false })
;(agent as any)._printer = new AgentPrinter(mockAppender)
agent.hooks.addHook(new AgentPrinter(mockAppender))

await collectGenerator(agent.stream('Test'))

Expand All @@ -329,15 +329,17 @@ describe('Agent', () => {
const agent = new Agent({ model, printer: false })

expect(agent).toBeDefined()
expect((agent as any)._printer).toBeUndefined()
// Printer is not registered when printer: false
expect(agent.hooks).toBeDefined()
})

it('defaults to printer=true when not specified', () => {
const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' })
const agent = new Agent({ model })

expect(agent).toBeDefined()
expect((agent as any)._printer).toBeDefined()
// Printer is registered as a hook internally, not a field
expect(agent.hooks).toBeDefined()
})

it('agent works correctly with printer disabled', async () => {
Expand Down
14 changes: 7 additions & 7 deletions src/agent/__tests__/printer.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ describe('AgentPrinter', () => {
const mockAppender = (text: string) => outputs.push(text)

const agent = new Agent({ model, printer: false })
;(agent as any)._printer = new AgentPrinter(mockAppender)
agent.hooks.addHook(new AgentPrinter(mockAppender))

await collectGenerator(agent.stream('Test'))

Expand All @@ -30,7 +30,7 @@ describe('AgentPrinter', () => {
const mockAppender = (text: string) => outputs.push(text)

const agent = new Agent({ model, printer: false })
;(agent as any)._printer = new AgentPrinter(mockAppender)
agent.hooks.addHook(new AgentPrinter(mockAppender))

await collectGenerator(agent.stream('Test'))

Expand All @@ -48,7 +48,7 @@ describe('AgentPrinter', () => {
const mockAppender = (text: string) => outputs.push(text)

const agent = new Agent({ model, printer: false })
;(agent as any)._printer = new AgentPrinter(mockAppender)
agent.hooks.addHook(new AgentPrinter(mockAppender))

await collectGenerator(agent.stream('Test'))

Expand All @@ -66,7 +66,7 @@ describe('AgentPrinter', () => {
const mockAppender = (text: string) => outputs.push(text)

const agent = new Agent({ model, printer: false })
;(agent as any)._printer = new AgentPrinter(mockAppender)
agent.hooks.addHook(new AgentPrinter(mockAppender))

await collectGenerator(agent.stream('Test'))

Expand Down Expand Up @@ -96,7 +96,7 @@ describe('AgentPrinter', () => {
const mockAppender = (text: string) => outputs.push(text)

const agent = new Agent({ model, tools: [tool], printer: false })
;(agent as any)._printer = new AgentPrinter(mockAppender)
agent.hooks.addHook(new AgentPrinter(mockAppender))

await collectGenerator(agent.stream('Test'))

Expand All @@ -120,7 +120,7 @@ describe('AgentPrinter', () => {
const mockAppender = (text: string) => outputs.push(text)

const agent = new Agent({ model, tools: [tool], printer: false })
;(agent as any)._printer = new AgentPrinter(mockAppender)
agent.hooks.addHook(new AgentPrinter(mockAppender))

await collectGenerator(agent.stream('Test'))

Expand Down Expand Up @@ -163,7 +163,7 @@ describe('AgentPrinter', () => {
const mockAppender = (text: string) => outputs.push(text)

const agent = new Agent({ model, tools: [calcTool, validatorTool], printer: false })
;(agent as any)._printer = new AgentPrinter(mockAppender)
agent.hooks.addHook(new AgentPrinter(mockAppender))

await collectGenerator(agent.stream('Test'))

Expand Down
26 changes: 17 additions & 9 deletions src/agent/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,16 @@ import {
type ToolContext,
ToolResultBlock,
type ToolUseBlock,
type ContentBlock,
} from '../index.js'
import { systemPromptFromData } from '../types/messages.js'
import { normalizeError, ConcurrentInvocationError, MaxTokensError } from '../errors.js'
import type { BaseModelConfig, Model, StreamOptions } from '../models/model.js'
import type { ModelStreamEvent } from '../models/streaming.js'
import { ToolRegistry } from '../registry/tool-registry.js'
import { AgentState } from './state.js'
import type { AgentData } from '../types/agent.js'
import { AgentPrinter, getDefaultAppender, type Printer } from './printer.js'
import { AgentPrinter, getDefaultAppender } from './printer.js'
import type { HookProvider } from '../hooks/types.js'
import { SlidingWindowConversationManager } from '../conversation-manager/sliding-window-conversation-manager.js'
import { HookRegistryImplementation } from '../hooks/registry.js'
Expand All @@ -36,6 +38,7 @@ import {
BeforeToolsEvent,
MessageAddedEvent,
ModelStreamEventHook,
ContentBlockHook,
} from '../hooks/events.js'

/**
Expand Down Expand Up @@ -119,7 +122,6 @@ export class Agent implements AgentData {
private _systemPrompt?: SystemPrompt
private _initialized: boolean
private _isInvoking: boolean = false
private _printer?: Printer

/**
* Creates an instance of the Agent.
Expand All @@ -145,10 +147,9 @@ export class Agent implements AgentData {
this._systemPrompt = systemPromptFromData(config.systemPrompt)
}

// Create printer if printer is enabled (default: true)
const printer = config?.printer ?? true
if (printer) {
this._printer = new AgentPrinter(getDefaultAppender())
// Create and register printer as a hook if enabled (default: true)
if (config?.printer ?? true) {
this.hooks.addHook(new AgentPrinter(getDefaultAppender()))
}

this._initialized = false
Expand Down Expand Up @@ -274,7 +275,6 @@ export class Agent implements AgentData {
await this.hooks.invokeCallbacks(event)
}

this._printer?.processEvent(event)
yield event
result = await streamGenerator.next()
}
Expand Down Expand Up @@ -407,8 +407,14 @@ export class Agent implements AgentData {
while (!result.done) {
const event = result.value

// Yield hook event for observability
yield new ModelStreamEventHook({ agent: this, event })
// Yield appropriate hook event for observability
if ('type' in event && typeof event.type === 'string' && event.type.endsWith('Block')) {
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than doing this much work, let's add a base class to all content blocks, BaseContentBlock for now. Make it non-public and add a TODO to the code to api-bar-raise it.

Then do the type check here for that using instanceof rather than this event.type checking

// This is a ContentBlock
yield new ContentBlockHook({ agent: this, block: event as ContentBlock })
} else {
// This is a ModelStreamEvent
yield new ModelStreamEventHook({ agent: this, event: event as ModelStreamEvent })
}

// Yield the actual model event
yield event
Expand Down Expand Up @@ -448,6 +454,8 @@ export class Agent implements AgentData {
const toolResultBlock = yield* this.executeTool(toolUseBlock, toolRegistry)
toolResultBlocks.push(toolResultBlock)

// Wrap ToolResultBlock in hook event for printer
yield new ContentBlockHook({ agent: this, block: toolResultBlock })
// Yield the tool result block as it's created
yield toolResultBlock
}
Expand Down
52 changes: 31 additions & 21 deletions src/agent/printer.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
import type { AgentStreamEvent } from '../types/agent.js'
import type { HookProvider } from '../hooks/types.js'
import type { HookRegistry } from '../hooks/registry.js'
import { ContentBlockHook, ModelStreamEventHook } from '../hooks/events.js'

/**
* Creates a default appender function for the current environment.
Expand All @@ -15,28 +18,10 @@ export function getDefaultAppender(): (text: string) => void {
}

/**
* Interface for printing agent activity to a destination.
* Implementations can output to stdout, console, HTML elements, etc.
*/
export interface Printer {
/**
* Write content to the output destination.
* @param content - The content to write
*/
write(content: string): void

/**
* Process a streaming event from the agent.
* @param event - The event to process
*/
processEvent(event: AgentStreamEvent): void
}

/**
* Default implementation of the Printer interface.
* Default implementation of agent output printing.
* Outputs text, reasoning, and tool execution activity to the configured appender.
*/
export class AgentPrinter implements Printer {
export class AgentPrinter implements HookProvider {
private readonly _appender: (text: string) => void
private _inReasoningBlock: boolean = false
private _toolCount: number = 0
Expand All @@ -50,6 +35,31 @@ export class AgentPrinter implements Printer {
this._appender = appender
}

/**
* Register callback functions for specific event types.
* @param registry - The hook registry to register callbacks with
*/
public registerCallbacks(registry: HookRegistry): void {
registry.addCallback(ModelStreamEventHook, this.handleModelStreamEvent)
registry.addCallback(ContentBlockHook, this.handleContentBlock)
}

/**
* Handle model stream events and process them for printing.
* @param event - The model stream event to handle
*/
private handleModelStreamEvent = (event: ModelStreamEventHook): void => {
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just use inline lambdas in the registerCallbacks method and just forward the events processEvent directly

this.processEvent(event.event)
}

/**
* Handle content block events and process them for printing.
* @param event - The content block event to handle
*/
private handleContentBlock = (event: ContentBlockHook): void => {
this.processEvent(event.block)
}

/**
* Write content to the output destination.
* @param content - The content to write
Expand All @@ -63,7 +73,7 @@ export class AgentPrinter implements Printer {
* Handles text deltas, reasoning content, and tool execution events.
* @param event - The event to process
*/
public processEvent(event: AgentStreamEvent): void {
private processEvent(event: AgentStreamEvent): void {
switch (event.type) {
case 'modelContentBlockDeltaEvent':
this.handleContentBlockDelta(event)
Expand Down
Loading