From 3c765c94c1260b711c94654ec50af6c705d21b29 Mon Sep 17 00:00:00 2001 From: Maciej Jastrzebski Date: Wed, 11 Dec 2024 10:59:51 +0100 Subject: [PATCH 1/4] feat: token usage --- .../src/supervisor/finalizeWorkflow.ts | 15 ++++++++++--- packages/framework/src/supervisor/nextTick.ts | 21 ++++++++++++++----- packages/framework/src/telemetry.ts | 9 +++++--- packages/framework/src/types.ts | 6 ++++++ packages/framework/src/workflow.ts | 16 +++++++++----- 5 files changed, 51 insertions(+), 16 deletions(-) diff --git a/packages/framework/src/supervisor/finalizeWorkflow.ts b/packages/framework/src/supervisor/finalizeWorkflow.ts index dbfa4c3..20882f4 100644 --- a/packages/framework/src/supervisor/finalizeWorkflow.ts +++ b/packages/framework/src/supervisor/finalizeWorkflow.ts @@ -3,9 +3,14 @@ import { zodResponseFormat } from 'openai/helpers/zod' import { z } from 'zod' import { Provider } from '../models.js' -import { Message } from '../types.js' +import { Message, Usage } from '../types.js' -export async function finalizeWorkflow(provider: Provider, messages: Message[]): Promise { +export type FinalizeWorkflowResult = { + response: string + usage?: Usage +} + +export async function finalizeWorkflow(provider: Provider, messages: Message[]): Promise { const response = await provider.completions({ messages: [ { @@ -29,5 +34,9 @@ export async function finalizeWorkflow(provider: Provider, messages: Message[]): if (!result) { throw new Error('No parsed response received') } - return result.finalAnswer + + return { + response: result.finalAnswer, + usage: response.usage, + }; } diff --git a/packages/framework/src/supervisor/nextTick.ts b/packages/framework/src/supervisor/nextTick.ts index a84da55..fb9ec8c 100644 --- a/packages/framework/src/supervisor/nextTick.ts +++ b/packages/framework/src/supervisor/nextTick.ts @@ -1,4 +1,5 @@ -import { Workflow, WorkflowState } from '../workflow.js' +import type { Usage } from '../types.js' +import type { Workflow, WorkflowState } from '../workflow.js' import { finalizeWorkflow } from './finalizeWorkflow.js' import { nextTask } from './nextTask.js' import { runAgent } from './runAgent.js' @@ -12,11 +13,11 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis const { status, messages } = state /** - * When number of messages exceedes number of maximum iterations, we must force finish the workflow + * When number of messages exceeds number of maximum iterations, we must force finish the workflow * and produce best final answer */ if (messages.length > workflow.maxIterations) { - const content = await finalizeWorkflow(workflow.provider, messages) + const { response, usage } = await finalizeWorkflow(workflow.provider, messages) return { ...state, status: 'finished', @@ -24,7 +25,8 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis role: 'user', content, }), - } + usage: addUsage(state.usage, usage), + }; } /** @@ -77,7 +79,8 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis role: 'assistant', content: 'No agent found.', }), - } + usage: state.usage, + }; } /** @@ -138,3 +141,11 @@ export async function iterate(workflow: Workflow, state: WorkflowState) { workflow.snapshot({ prevState: state, nextState }) return nextState } + +function addUsage(prevUsage: Usage, usage: Usage | undefined) { + return { + prompt_tokens: prevUsage.prompt_tokens + (usage?.prompt_tokens ?? 0), + completion_tokens: prevUsage.completion_tokens + (usage?.completion_tokens ?? 0), + total_tokens: prevUsage.total_tokens + (usage?.total_tokens ?? 0), + } +} diff --git a/packages/framework/src/telemetry.ts b/packages/framework/src/telemetry.ts index a6522bb..de287bc 100644 --- a/packages/framework/src/telemetry.ts +++ b/packages/framework/src/telemetry.ts @@ -66,9 +66,12 @@ export const logger: Telemetry = ({ prevState, nextState }) => { break case 'finished': logMessage( - '🎉', - 'Workflow finished successfully!', - `Total messages: ${nextState.messages.length}` + "🎉", + "Workflow finished successfully!", + [ + `Total messages: ${nextState.messages.length}`, + `Total tokens: ${nextState.usage.total_tokens} (input: ${nextState.usage.prompt_tokens}, output: ${nextState.usage.completion_tokens})`, + ].join('\n') ) break case 'failed': diff --git a/packages/framework/src/types.ts b/packages/framework/src/types.ts index 6dbbf1b..112fb39 100644 --- a/packages/framework/src/types.ts +++ b/packages/framework/src/types.ts @@ -1,4 +1,5 @@ import type { ChatCompletionMessageParam } from 'openai/resources/chat/completions' +import type { CompletionUsage } from 'openai/resources/completions.mjs' /** * Utility type to get optional keys from T. @@ -22,3 +23,8 @@ export type RequiredOptionals = Required> */ export type Message = ChatCompletionMessageParam export type MessageContent = Message['content'] + +/** + * Usage type for completion + */ +export type Usage = CompletionUsage; diff --git a/packages/framework/src/workflow.ts b/packages/framework/src/workflow.ts index e1dc0f0..80e70e8 100644 --- a/packages/framework/src/workflow.ts +++ b/packages/framework/src/workflow.ts @@ -5,7 +5,7 @@ import s from 'dedent' import { Agent } from './agent.js' import { openai, Provider } from './models.js' import { noop, Telemetry } from './telemetry.js' -import { Message } from './types.js' +import { Message, Usage } from './types.js' type WorkflowOptions = { description: string @@ -35,9 +35,10 @@ export type Workflow = Required * Base workflow */ type BaseWorkflowState = { - id: string - messages: Message[] -} + id: string; + messages: Message[]; + usage: Usage; +}; /** * Different states workflow is in, in between execution from agents @@ -47,7 +48,7 @@ export type IdleWorkflowState = BaseWorkflowState & { } /** - * Supervisor selected the task, and is now pending assignement of an agent + * Supervisor selected the task, and is now pending assignment of an agent */ export type PendingWorkflowState = BaseWorkflowState & { status: 'pending' @@ -84,6 +85,11 @@ export const workflowState = (workflow: Workflow): IdleWorkflowState => { `, }, ], + usage: { + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + }, } } From 15d0a5fc1999c3eb7e545731b766d3dd092e246c Mon Sep 17 00:00:00 2001 From: Maciej Jastrzebski Date: Wed, 11 Dec 2024 11:43:51 +0100 Subject: [PATCH 2/4] tweaks --- packages/framework/src/supervisor/nextTask.ts | 13 +++++--- packages/framework/src/supervisor/nextTick.ts | 31 +++++++++++-------- packages/framework/src/supervisor/runAgent.ts | 21 ++++++++----- .../framework/src/supervisor/selectAgent.ts | 11 +++++-- packages/tools/src/usage.ts | 0 packages/tools/src/vision.ts | 6 ++++ 6 files changed, 55 insertions(+), 27 deletions(-) create mode 100644 packages/tools/src/usage.ts diff --git a/packages/framework/src/supervisor/nextTask.ts b/packages/framework/src/supervisor/nextTask.ts index 9705c32..d869034 100644 --- a/packages/framework/src/supervisor/nextTask.ts +++ b/packages/framework/src/supervisor/nextTask.ts @@ -3,9 +3,14 @@ import { zodResponseFormat } from 'openai/helpers/zod' import { z } from 'zod' import { Provider } from '../models.js' -import { Message } from '../types.js' +import { Message, Usage } from '../types.js' -export async function nextTask(provider: Provider, history: Message[]): Promise { +export type NextTaskResult = { + task: string | null + usage?: Usage +} + +export async function nextTask(provider: Provider, history: Message[]): Promise { const response = await provider.completions({ messages: [ { @@ -53,10 +58,10 @@ export async function nextTask(provider: Provider, history: Message[]): Promise< } if (!content.task) { - return null + return { task: null, usage: response.usage } } - return content.task + return { task: content.task, usage: response.usage } } catch (error) { throw new Error('Failed to determine next task') } diff --git a/packages/framework/src/supervisor/nextTick.ts b/packages/framework/src/supervisor/nextTick.ts index fb9ec8c..58815e4 100644 --- a/packages/framework/src/supervisor/nextTick.ts +++ b/packages/framework/src/supervisor/nextTick.ts @@ -23,10 +23,10 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis status: 'finished', messages: state.messages.concat({ role: 'user', - content, + content: response, }), - usage: addUsage(state.usage, usage), - }; + usage: combineUsage(state.usage, usage), + } } /** @@ -34,21 +34,23 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis */ if (status === 'idle') { const task = await nextTask(workflow.provider, messages) - if (task) { + if (task.task) { return { ...state, status: 'pending', agentRequest: [ { role: 'user', - content: task, + content: task.task, }, ], + usage: combineUsage(state.usage, task.usage), } } else { return { ...state, status: 'finished', + usage: combineUsage(state.usage, task.usage), } } } @@ -62,7 +64,8 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis ...state, status: 'assigned', agentStatus: 'idle', - agent: selectedAgent.role, + agent: selectedAgent.agent.role, + usage: combineUsage(state.usage, selectedAgent.usage), } } @@ -80,7 +83,7 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis content: 'No agent found.', }), usage: state.usage, - }; + } } /** @@ -104,18 +107,20 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis * * If further processing is required, we will carry `agentRequest` over to the next iteration. */ - const [agentResponse, status] = await runAgent(agent, state.messages, state.agentRequest) - if (status === 'complete') { + const agentResponse = await runAgent(agent, state.messages, state.agentRequest) + if (agentResponse.kind === 'complete') { return { ...state, status: 'idle', - messages: state.messages.concat(state.agentRequest[0], agentResponse), + messages: state.messages.concat(state.agentRequest[0], agentResponse.message), + usage: combineUsage(state.usage, agentResponse.usage), } } return { ...state, - agentStatus: status, - agentRequest: state.agentRequest.concat(agentResponse), + agentStatus: agentResponse.kind, + agentRequest: state.agentRequest.concat(agentResponse.message), + usage: combineUsage(state.usage, agentResponse.usage), } } @@ -142,7 +147,7 @@ export async function iterate(workflow: Workflow, state: WorkflowState) { return nextState } -function addUsage(prevUsage: Usage, usage: Usage | undefined) { +function combineUsage(prevUsage: Usage, usage: Usage | undefined) { return { prompt_tokens: prevUsage.prompt_tokens + (usage?.prompt_tokens ?? 0), completion_tokens: prevUsage.completion_tokens + (usage?.completion_tokens ?? 0), diff --git a/packages/framework/src/supervisor/runAgent.ts b/packages/framework/src/supervisor/runAgent.ts index 282de98..b6e976c 100644 --- a/packages/framework/src/supervisor/runAgent.ts +++ b/packages/framework/src/supervisor/runAgent.ts @@ -3,13 +3,19 @@ import { zodFunction, zodResponseFormat } from 'openai/helpers/zod' import { z } from 'zod' import { Agent } from '../agent.js' -import { Message } from '../types.js' +import { Message, Usage } from '../types.js' + +export type RunAgentResult = { + message: Message + kind: 'step' | 'complete' | 'tool' + usage?: Usage +} export async function runAgent( agent: Agent, agentContext: Message[], agentRequest: Message[] -): Promise<[Message, 'step' | 'complete' | 'tool']> { +): Promise { const tools = agent.tools ? Object.entries(agent.tools).map(([name, tool]) => zodFunction({ @@ -81,7 +87,7 @@ export async function runAgent( }) if (response.choices[0].message.tool_calls.length > 0) { - return [response.choices[0].message, 'tool'] + return { message: response.choices[0].message, kind: 'tool', usage: response.usage } } const result = response.choices[0].message.parsed @@ -93,11 +99,12 @@ export async function runAgent( throw new Error(result.response.reasoning) } - return [ - { + return { + message: { role: 'assistant', content: result.response.result, }, - result.response.kind, - ] + kind: result.response.kind, + usage: response.usage, + } } diff --git a/packages/framework/src/supervisor/selectAgent.ts b/packages/framework/src/supervisor/selectAgent.ts index c2631f4..1cb405a 100644 --- a/packages/framework/src/supervisor/selectAgent.ts +++ b/packages/framework/src/supervisor/selectAgent.ts @@ -4,13 +4,18 @@ import { z } from 'zod' import { Agent } from '../agent.js' import { Provider } from '../models.js' -import { Message } from '../types.js' +import { Message, Usage } from '../types.js' + +export type SelectAgentResult = { + agent: Agent + usage?: Usage +} export async function selectAgent( provider: Provider, agentRequest: Message[], agents: Agent[] -): Promise { +): Promise { const response = await provider.completions({ messages: [ { @@ -64,5 +69,5 @@ export async function selectAgent( throw new Error('Invalid agent') } - return agent + return { agent, usage: response.usage } } diff --git a/packages/tools/src/usage.ts b/packages/tools/src/usage.ts new file mode 100644 index 0000000..e69de29 diff --git a/packages/tools/src/vision.ts b/packages/tools/src/vision.ts index eb863fc..6fa777c 100644 --- a/packages/tools/src/vision.ts +++ b/packages/tools/src/vision.ts @@ -4,6 +4,7 @@ import path from 'node:path' import s from 'dedent' import { Provider } from 'fabrice-ai/models' import { tool } from 'fabrice-ai/tool' +import { CompletionUsage } from 'openai/resources/completions.mjs' import { zodResponseFormat } from 'openai/helpers/zod' import { z } from 'zod' @@ -12,6 +13,11 @@ const encodeImage = async (imagePath: string): Promise => { return `data:image/${path.extname(imagePath).toLowerCase().replace('.', '')};base64,${imageBuffer.toString('base64')}` } +export type CallOpenAIResult = { + text: string + usage?: CompletionUsage +} + async function callOpenAI( provider: Provider, prompt: string, From e36cabac7148cd22736fb606d8fc29ad98dded99 Mon Sep 17 00:00:00 2001 From: Maciej Jastrzebski Date: Wed, 11 Dec 2024 12:08:39 +0100 Subject: [PATCH 3/4] self code review --- packages/framework/src/supervisor/nextTick.ts | 34 +++++++++++-------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/packages/framework/src/supervisor/nextTick.ts b/packages/framework/src/supervisor/nextTick.ts index 58815e4..a814823 100644 --- a/packages/framework/src/supervisor/nextTick.ts +++ b/packages/framework/src/supervisor/nextTick.ts @@ -33,24 +33,24 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis * When workflow is idle, we must get next task to work on, or finish the workflow otherwise. */ if (status === 'idle') { - const task = await nextTask(workflow.provider, messages) - if (task.task) { + const { task, usage } = await nextTask(workflow.provider, messages) + if (task) { return { ...state, status: 'pending', agentRequest: [ { role: 'user', - content: task.task, + content: task, }, ], - usage: combineUsage(state.usage, task.usage), + usage: combineUsage(state.usage, usage), } } else { return { ...state, status: 'finished', - usage: combineUsage(state.usage, task.usage), + usage: combineUsage(state.usage, usage), } } } @@ -59,13 +59,17 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis * When workflow is pending, we must find best agent to work on it. */ if (status === 'pending') { - const selectedAgent = await selectAgent(workflow.provider, state.agentRequest, workflow.members) + const { agent, usage } = await selectAgent( + workflow.provider, + state.agentRequest, + workflow.members + ) return { ...state, status: 'assigned', agentStatus: 'idle', - agent: selectedAgent.agent.role, - usage: combineUsage(state.usage, selectedAgent.usage), + agent: agent.role, + usage: combineUsage(state.usage, usage), } } @@ -107,20 +111,20 @@ export async function nextTick(workflow: Workflow, state: WorkflowState): Promis * * If further processing is required, we will carry `agentRequest` over to the next iteration. */ - const agentResponse = await runAgent(agent, state.messages, state.agentRequest) - if (agentResponse.kind === 'complete') { + const { kind, message, usage } = await runAgent(agent, state.messages, state.agentRequest) + if (kind === 'complete') { return { ...state, status: 'idle', - messages: state.messages.concat(state.agentRequest[0], agentResponse.message), - usage: combineUsage(state.usage, agentResponse.usage), + messages: state.messages.concat(state.agentRequest[0], message), + usage: combineUsage(state.usage, usage), } } return { ...state, - agentStatus: agentResponse.kind, - agentRequest: state.agentRequest.concat(agentResponse.message), - usage: combineUsage(state.usage, agentResponse.usage), + agentStatus: kind, + agentRequest: state.agentRequest.concat(message), + usage: combineUsage(state.usage, usage), } } From 7fe1a69c8aeb9fdb1462dc8881086d24227367a3 Mon Sep 17 00:00:00 2001 From: Maciej Jastrzebski Date: Wed, 11 Dec 2024 12:14:44 +0100 Subject: [PATCH 4/4] =?UTF-8?q?=F0=9F=A7=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- packages/tools/src/usage.ts | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 packages/tools/src/usage.ts diff --git a/packages/tools/src/usage.ts b/packages/tools/src/usage.ts deleted file mode 100644 index e69de29..0000000