Skip to content

Commit

Permalink
add plugins
Browse files Browse the repository at this point in the history
  • Loading branch information
wladpaiva committed Oct 17, 2023
1 parent 41aea69 commit 84c2d7a
Show file tree
Hide file tree
Showing 9 changed files with 160 additions and 125 deletions.
15 changes: 2 additions & 13 deletions examples/0-template.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,7 @@
import {input} from '@inquirer/prompts'

import {AIbitat} from '../src'
import {terminal} from '../src/utils'

console.log('🚀 starting chat\n')
console.time('🚀 chat finished')
import {terminal} from '../src/plugins'

const aibitat = new AIbitat({
nodes: {
Expand All @@ -14,15 +11,7 @@ const aibitat = new AIbitat({
'🧑': {type: 'assistant'},
'🤖': {type: 'agent'},
},
})

aibitat.onMessage(terminal.print)
aibitat.onTerminate(() => console.timeEnd('🚀 chat finished'))

aibitat.onInterrupt(async node => {
const feedback = await terminal.askForFeedback(node)
await aibitat.continue(feedback)
})
}).use(terminal())

// Ask for the topic of the chat before starting the conversation
const topic = await input({
Expand Down
10 changes: 2 additions & 8 deletions examples/1-simple-math.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,7 @@
import {input} from '@inquirer/prompts'

import {AIbitat} from '../src'
import {terminal} from '../src/utils'

console.log('🚀 starting chat\n')
console.time('🚀 chat finished')
import {terminal} from '../src/plugins'

const aibitat = new AIbitat({
model: 'gpt-3.5-turbo',
Expand All @@ -19,10 +16,7 @@ const aibitat = new AIbitat({
},
'🤖': {type: 'agent'},
},
})

aibitat.onMessage(terminal.print)
aibitat.onTerminate(() => console.timeEnd('🚀 chat finished'))
}).use(terminal())

// Ask for the mathematical problem of the chat before starting the conversation
const math = await input({
Expand Down
14 changes: 2 additions & 12 deletions examples/2-group.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,7 @@
import {input} from '@inquirer/prompts'

import {AIbitat} from '../src'
import {terminal} from '../src/utils'

console.log('🚀 starting chat\n')
console.time('🚀 chat finished')
import {terminal} from '../src/plugins'

// Ask for the topic of the chat before starting the conversation
const topic = await input({
Expand Down Expand Up @@ -59,14 +56,7 @@ const aibitat = new AIbitat({
Provide an optimize post's meta description and title.`,
},
},
})

aibitat.onMessage(terminal.print)
aibitat.onTerminate(() => console.timeEnd('🚀 chat finished'))
aibitat.onInterrupt(async node => {
const feedback = await terminal.askForFeedback(node)
await aibitat.continue(feedback)
})
}).use(terminal())

await aibitat.start({
from: 'client',
Expand Down
15 changes: 2 additions & 13 deletions examples/3-basic-interaction.ts
Original file line number Diff line number Diff line change
@@ -1,8 +1,5 @@
import {AIbitat} from '../src'
import {terminal} from '../src/utils'

console.log('🚀 starting chat\n')
console.time('🚀 chat finished')
import {terminal} from '../src/plugins'

const aibitat = new AIbitat({
nodes: {
Expand All @@ -23,15 +20,7 @@ const aibitat = new AIbitat({
role: `You are a peer-reviewer and you do not solve math problems. Check the result from mathematician and then confirm. Just confirm, no talk.`,
},
},
})

aibitat.onMessage(terminal.print)
aibitat.onTerminate(() => console.timeEnd('🚀 chat finished'))

aibitat.onInterrupt(async node => {
const feedback = await terminal.askForFeedback(node)
await aibitat.continue(feedback)
})
}).use(terminal())

await aibitat.start({
from: 'client',
Expand Down
36 changes: 27 additions & 9 deletions src/aibitat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,14 @@ export type AIbitatProps = ProviderConfig & {
interrupt?: 'NEVER' | 'ALWAYS'
}

/**
* Plugin to use with the aibitat
*/
export type AIbitatPlugin = {
name: string
setup: (aibitat: AIbitat) => void
}

/**
* AIbitat is a class that manages the conversation between agents.
* It is designed to solve a task with LLM.
Expand Down Expand Up @@ -196,6 +204,14 @@ export class AIbitat {
return this._chats
}

/**
* Install a plugin.
*/
use(plugin: AIbitatPlugin) {
plugin.setup(this)
return this
}

/**
* Get the specific node configuration.
*
Expand Down Expand Up @@ -250,7 +266,7 @@ export class AIbitat {
* @param listener
* @returns
*/
public onTerminate(listener: (node: string) => void) {
public onTerminate(listener: (node: string, aibitat: AIbitat) => void) {
this.emitter.on('terminate', listener)
return this
}
Expand All @@ -261,7 +277,7 @@ export class AIbitat {
* @param node Last node to chat with
*/
private terminate(node: string) {
this.emitter.emit('terminate', node)
this.emitter.emit('terminate', node, this)
}

/**
Expand All @@ -270,7 +286,9 @@ export class AIbitat {
* @param listener
* @returns
*/
public onInterrupt(listener: (chat: {from: string; to: string}) => void) {
public onInterrupt(
listener: (chat: {from: string; to: string}, aibitat: AIbitat) => void,
) {
this.emitter.on('interrupt', listener)
return this
}
Expand All @@ -286,7 +304,7 @@ export class AIbitat {
...chat,
state: 'interrupt',
})
this.emitter.emit('interrupt', chat)
this.emitter.emit('interrupt', chat, this)
}

/**
Expand All @@ -296,7 +314,7 @@ export class AIbitat {
* @param listener
* @returns
*/
public onMessage(listener: (chat: ChatState) => void) {
public onMessage(listener: (chat: ChatState, aibitat: AIbitat) => void) {
this.emitter.on('message', listener)
return this
}
Expand All @@ -314,7 +332,7 @@ export class AIbitat {
}

this._chats.push(chat)
this.emitter.emit('message', chat)
this.emitter.emit('message', chat, this)
return this
}

Expand All @@ -324,7 +342,7 @@ export class AIbitat {
* @param listener
* @returns
*/
public onStart(listener: (chat: ChatState) => void) {
public onStart(listener: (chat: ChatState, aibitat: AIbitat) => void) {
this.emitter.on('start', listener)
return this
}
Expand All @@ -334,7 +352,7 @@ export class AIbitat {
*
* @param message The message to start the chat.
*/
async start(message: Chat) {
public async start(message: Chat) {
log(
`starting a chat from ${chalk.yellow(message.from)} to ${chalk.yellow(
message.to,
Expand All @@ -343,7 +361,7 @@ export class AIbitat {

// register the message in the chat history
this.newMessage(message)
this.emitter.emit('start', message)
this.emitter.emit('start', message, this)

// ask the node to reply
await this.chat({
Expand Down
1 change: 1 addition & 0 deletions src/plugins/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
export * from './terminal.ts'
124 changes: 124 additions & 0 deletions src/plugins/terminal.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
import {input} from '@inquirer/prompts'
import chalk from 'chalk'

import {AIbitatPlugin} from '..'

/**
* Print a message on the terminal
*
* @param message
* @param simulateStream
*/
export async function print(
message: {from: string; to: string; content?: string} & {
state: 'loading' | 'error' | 'success' | 'interrupt'
},
simulateStream: boolean = true,
) {
const replying = chalk.dim(`(to ${message.to})`)
const reference = `${chalk.magenta('✎')} ${chalk.bold(
message.from,
)} ${replying}:`

if (!simulateStream) {
console.log(reference)
console.log(message.content)
// Add an extra line after the message
console.log()
return
}

process.stdout.write(`${reference}\n`)

// Emulate streaming by breaking the cached response into chunks
const chunks = message.content?.split(' ') || []
const stream = new ReadableStream({
async start(controller) {
for (const chunk of chunks) {
const bytes = new TextEncoder().encode(chunk + ' ')
controller.enqueue(bytes)
await new Promise(r =>
setTimeout(
r,
// get a random number between 10ms and 50ms to simulate a random delay
Math.floor(Math.random() * 40) + 10,
),
)
}
controller.close()
},
})

// Stream the response to the chat
for await (const chunk of stream) {
process.stdout.write(new TextDecoder().decode(chunk))
}

// Add an extra line after the message
console.log()
console.log()
}

/**
* Ask for feedback to the user using the terminal
*
* @param node
* @returns
*/
export function askForFeedback(node: {from: string; to: string}) {
return input({
message: `Provide feedback to ${chalk.yellow(node.to)} as ${chalk.yellow(
node.from,
)}. Press enter to skip and use auto-reply, or type 'exit' to end the conversation: `,
})
}

/**
* Terminal plugin. It prints the messages on the terminal and asks for feedback
* while the conversation is running in the background.
*/
export function terminal({
simulateStream = true,
}: {
/**
* Simulate streaming by breaking the cached response into chunks.
* Helpful to make the conversation more realistic and faster.
* @default true
*/
simulateStream?: boolean
} = {}) {
return {
name: 'terminal',
setup(aibitat) {
let printing: Promise<void> | null = null

aibitat.onStart(() => {
console.log()
console.log('🚀 starting chat ...\n')
console.time('🚀 chat finished!')
printing = Promise.resolve()
})

aibitat.onMessage(async message => {
await printing
printing = print(message, simulateStream)
})

aibitat.onTerminate(() => console.timeEnd('🚀 chat finished'))

aibitat.onInterrupt(async (node, current) => {
await printing
const feedback = await askForFeedback(node)
// Add an extra line after the message
console.log()

if (feedback === 'exit') {
console.timeEnd('🚀 chat finished')
return process.exit(0)
}

await current.continue(feedback)
})
},
} as AIbitatPlugin
}
1 change: 0 additions & 1 deletion src/utils/index.ts

This file was deleted.

Loading

0 comments on commit 84c2d7a

Please sign in to comment.