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

Make stream function configurable in agent

Browse source
This commit is contained in:
Mario Zechner 2025-09-15 20:31:53 +02:00
parent 1c9ab1ca24
commit 5f7a50deef
2 changed files with 246 additions and 140 deletions
Download patch file Download diff file Expand all files Collapse all files

8
packages/ai/src/agent/agent.ts
View file

@ -10,6 +10,7 @@ export function prompt(
context: AgentContext,
config: PromptConfig,
signal?: AbortSignal,
streamFn?: typeof streamSimple,
): EventStream<AgentEvent, AgentContext["messages"]> {
const stream = new EventStream<AgentEvent, AgentContext["messages"]>(
(event) => event.type === "agent_end",
@ -45,7 +46,7 @@ export function prompt(
firstTurn = false;
}
// Stream assistant response
const assistantMessage = await streamAssistantResponse(currentContext, config, signal, stream);
const assistantMessage = await streamAssistantResponse(currentContext, config, signal, stream, streamFn);
newMessages.push(assistantMessage);
// Check for tool calls
@ -74,6 +75,7 @@ async function streamAssistantResponse(
config: PromptConfig,
signal: AbortSignal | undefined,
stream: EventStream<AgentEvent, AgentContext["messages"]>,
streamFn?: typeof streamSimple,
): Promise<AssistantMessage> {
// Convert AgentContext to Context for streamSimple
// Use a copy of messages to avoid mutating the original context
@ -93,7 +95,9 @@ async function streamAssistantResponse(
tools: context.tools, // AgentTool extends Tool, so this works
};
const response = await streamSimple(config.model, processedContext, { ...config, signal });
// Use custom stream function if provided, otherwise use default streamSimple
const streamFunction = streamFn || streamSimple;
const response = await streamFunction(config.model, processedContext, { ...config, signal });
let partialMessage: AssistantMessage | null = null;
let addedPartial = false;