getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-15 23:01:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: add thinkingBudgets option to customize token budgets

Browse source
This commit is contained in:
Melih Mucuk 2026-01-07 15:13:26 +03:00
parent 10e651f99b
commit 0f27eae77e
8 changed files with 75 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

23
packages/agent/src/agent.ts
View file

@ -10,6 +10,7 @@ import {
type Model,
streamSimple,
type TextContent,
type ThinkingBudgets,
} from "@mariozechner/pi-ai";
import { agentLoop, agentLoopContinue } from "./agent-loop.js";
import type {
@ -71,6 +72,11 @@ export interface AgentOptions {
* Useful for expiring tokens (e.g., GitHub Copilot OAuth).
*/
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
/**
* Custom token budgets for thinking levels (token-based providers only).
*/
thinkingBudgets?: ThinkingBudgets;
}
export class Agent {
@ -99,6 +105,7 @@ export class Agent {
public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
private runningPrompt?: Promise<void>;
private resolveRunningPrompt?: () => void;
private _thinkingBudgets?: ThinkingBudgets;
constructor(opts: AgentOptions = {}) {
this._state = { ...this._state, ...opts.initialState };
@ -109,6 +116,7 @@ export class Agent {
this.streamFn = opts.streamFn || streamSimple;
this._sessionId = opts.sessionId;
this.getApiKey = opts.getApiKey;
this._thinkingBudgets = opts.thinkingBudgets;
}
/**
@ -126,6 +134,20 @@ export class Agent {
this._sessionId = value;
}
/**
* Get the current thinking budgets.
*/
get thinkingBudgets(): ThinkingBudgets | undefined {
return this._thinkingBudgets;
}
/**
* Set custom thinking budgets for token-based providers.
*/
set thinkingBudgets(value: ThinkingBudgets | undefined) {
this._thinkingBudgets = value;
}
get state(): AgentState {
return this._state;
}
@ -310,6 +332,7 @@ export class Agent {
model,
reasoning,
sessionId: this._sessionId,
thinkingBudgets: this._thinkingBudgets,
convertToLlm: this.convertToLlm,
transformContext: this.transformContext,
getApiKey: this.getApiKey,