1package fantasy
2
3import (
4 "cmp"
5 "context"
6 "encoding/json"
7 "errors"
8 "fmt"
9 "maps"
10 "slices"
11 "sync"
12)
13
14// StepResult represents the result of a single step in an agent execution.
15type StepResult struct {
16 Response
17 Messages []Message
18}
19
20// stepExecutionResult encapsulates the result of executing a step with stream processing.
21type stepExecutionResult struct {
22 StepResult StepResult
23 ShouldContinue bool
24}
25
26// StopCondition defines a function that determines when an agent should stop executing.
27type StopCondition = func(steps []StepResult) bool
28
29// StepCountIs returns a stop condition that stops after the specified number of steps.
30func StepCountIs(stepCount int) StopCondition {
31 return func(steps []StepResult) bool {
32 return len(steps) >= stepCount
33 }
34}
35
36// HasToolCall returns a stop condition that stops when the specified tool is called in the last step.
37func HasToolCall(toolName string) StopCondition {
38 return func(steps []StepResult) bool {
39 if len(steps) == 0 {
40 return false
41 }
42 lastStep := steps[len(steps)-1]
43 toolCalls := lastStep.Content.ToolCalls()
44 for _, toolCall := range toolCalls {
45 if toolCall.ToolName == toolName {
46 return true
47 }
48 }
49 return false
50 }
51}
52
53// HasContent returns a stop condition that stops when the specified content type appears in the last step.
54func HasContent(contentType ContentType) StopCondition {
55 return func(steps []StepResult) bool {
56 if len(steps) == 0 {
57 return false
58 }
59 lastStep := steps[len(steps)-1]
60 for _, content := range lastStep.Content {
61 if content.GetType() == contentType {
62 return true
63 }
64 }
65 return false
66 }
67}
68
69// FinishReasonIs returns a stop condition that stops when the specified finish reason occurs.
70func FinishReasonIs(reason FinishReason) StopCondition {
71 return func(steps []StepResult) bool {
72 if len(steps) == 0 {
73 return false
74 }
75 lastStep := steps[len(steps)-1]
76 return lastStep.FinishReason == reason
77 }
78}
79
80// MaxTokensUsed returns a stop condition that stops when total token usage exceeds the specified limit.
81func MaxTokensUsed(maxTokens int64) StopCondition {
82 return func(steps []StepResult) bool {
83 var totalTokens int64
84 for _, step := range steps {
85 totalTokens += step.Usage.TotalTokens
86 }
87 return totalTokens >= maxTokens
88 }
89}
90
91// PrepareStepFunctionOptions contains the options for preparing a step in an agent execution.
92type PrepareStepFunctionOptions struct {
93 Steps []StepResult
94 StepNumber int
95 Model LanguageModel
96 Messages []Message
97}
98
99// PrepareStepResult contains the result of preparing a step in an agent execution.
100type PrepareStepResult struct {
101 Model LanguageModel
102 Messages []Message
103 System *string
104 ToolChoice *ToolChoice
105 ActiveTools []string
106 DisableAllTools bool
107 Tools []AgentTool
108}
109
110// ToolCallRepairOptions contains the options for repairing a tool call.
111type ToolCallRepairOptions struct {
112 OriginalToolCall ToolCallContent
113 ValidationError error
114 AvailableTools []AgentTool
115 SystemPrompt string
116 Messages []Message
117}
118
119type (
120 // PrepareStepFunction defines a function that prepares a step in an agent execution.
121 PrepareStepFunction = func(ctx context.Context, options PrepareStepFunctionOptions) (context.Context, PrepareStepResult, error)
122
123 // OnStepFinishedFunction defines a function that is called when a step finishes.
124 OnStepFinishedFunction = func(step StepResult)
125
126 // RepairToolCallFunction defines a function that repairs a tool call.
127 RepairToolCallFunction = func(ctx context.Context, options ToolCallRepairOptions) (*ToolCallContent, error)
128)
129
130type agentSettings struct {
131 systemPrompt string
132 maxOutputTokens *int64
133 temperature *float64
134 topP *float64
135 topK *int64
136 presencePenalty *float64
137 frequencyPenalty *float64
138 headers map[string]string
139 providerOptions ProviderOptions
140
141 // TODO: add support for provider tools
142 tools []AgentTool
143 maxRetries *int
144
145 model LanguageModel
146
147 stopWhen []StopCondition
148 prepareStep PrepareStepFunction
149 repairToolCall RepairToolCallFunction
150 onRetry OnRetryCallback
151}
152
153// AgentCall represents a call to an agent.
154type AgentCall struct {
155 Prompt string `json:"prompt"`
156 Files []FilePart `json:"files"`
157 Messages []Message `json:"messages"`
158 MaxOutputTokens *int64
159 Temperature *float64 `json:"temperature"`
160 TopP *float64 `json:"top_p"`
161 TopK *int64 `json:"top_k"`
162 PresencePenalty *float64 `json:"presence_penalty"`
163 FrequencyPenalty *float64 `json:"frequency_penalty"`
164 ActiveTools []string `json:"active_tools"`
165 ProviderOptions ProviderOptions
166 OnRetry OnRetryCallback
167 MaxRetries *int
168
169 StopWhen []StopCondition
170 PrepareStep PrepareStepFunction
171 RepairToolCall RepairToolCallFunction
172}
173
174// Agent-level callbacks.
175type (
176 // OnAgentStartFunc is called when agent starts.
177 OnAgentStartFunc func()
178
179 // OnAgentFinishFunc is called when agent finishes.
180 OnAgentFinishFunc func(result *AgentResult) error
181
182 // OnStepStartFunc is called when a step starts.
183 OnStepStartFunc func(stepNumber int) error
184
185 // OnStepFinishFunc is called when a step finishes.
186 OnStepFinishFunc func(stepResult StepResult) error
187
188 // OnFinishFunc is called when entire agent completes.
189 OnFinishFunc func(result *AgentResult)
190
191 // OnErrorFunc is called when an error occurs.
192 OnErrorFunc func(error)
193)
194
195// Stream part callbacks - called for each corresponding stream part type.
196type (
197 // OnChunkFunc is called for each stream part (catch-all).
198 OnChunkFunc func(StreamPart) error
199
200 // OnWarningsFunc is called for warnings.
201 OnWarningsFunc func(warnings []CallWarning) error
202
203 // OnTextStartFunc is called when text starts.
204 OnTextStartFunc func(id string) error
205
206 // OnTextDeltaFunc is called for text deltas.
207 OnTextDeltaFunc func(id, text string) error
208
209 // OnTextEndFunc is called when text ends.
210 OnTextEndFunc func(id string) error
211
212 // OnReasoningStartFunc is called when reasoning starts.
213 OnReasoningStartFunc func(id string, reasoning ReasoningContent) error
214
215 // OnReasoningDeltaFunc is called for reasoning deltas.
216 OnReasoningDeltaFunc func(id, text string) error
217
218 // OnReasoningEndFunc is called when reasoning ends.
219 OnReasoningEndFunc func(id string, reasoning ReasoningContent) error
220
221 // OnToolInputStartFunc is called when tool input starts.
222 OnToolInputStartFunc func(id, toolName string) error
223
224 // OnToolInputDeltaFunc is called for tool input deltas.
225 OnToolInputDeltaFunc func(id, delta string) error
226
227 // OnToolInputEndFunc is called when tool input ends.
228 OnToolInputEndFunc func(id string) error
229
230 // OnToolCallFunc is called when tool call is complete.
231 OnToolCallFunc func(toolCall ToolCallContent) error
232
233 // OnToolResultFunc is called when tool execution completes.
234 OnToolResultFunc func(result ToolResultContent) error
235
236 // OnSourceFunc is called for source references.
237 OnSourceFunc func(source SourceContent) error
238
239 // OnStreamFinishFunc is called when stream finishes.
240 OnStreamFinishFunc func(usage Usage, finishReason FinishReason, providerMetadata ProviderMetadata) error
241)
242
243// AgentStreamCall represents a streaming call to an agent.
244type AgentStreamCall struct {
245 Prompt string `json:"prompt"`
246 Files []FilePart `json:"files"`
247 Messages []Message `json:"messages"`
248 MaxOutputTokens *int64
249 Temperature *float64 `json:"temperature"`
250 TopP *float64 `json:"top_p"`
251 TopK *int64 `json:"top_k"`
252 PresencePenalty *float64 `json:"presence_penalty"`
253 FrequencyPenalty *float64 `json:"frequency_penalty"`
254 ActiveTools []string `json:"active_tools"`
255 Headers map[string]string
256 ProviderOptions ProviderOptions
257 OnRetry OnRetryCallback
258 MaxRetries *int
259
260 StopWhen []StopCondition
261 PrepareStep PrepareStepFunction
262 RepairToolCall RepairToolCallFunction
263
264 // Agent-level callbacks
265 OnAgentStart OnAgentStartFunc // Called when agent starts
266 OnAgentFinish OnAgentFinishFunc // Called when agent finishes
267 OnStepStart OnStepStartFunc // Called when a step starts
268 OnStepFinish OnStepFinishFunc // Called when a step finishes
269 OnFinish OnFinishFunc // Called when entire agent completes
270 OnError OnErrorFunc // Called when an error occurs
271
272 // Stream part callbacks - called for each corresponding stream part type
273 OnChunk OnChunkFunc // Called for each stream part (catch-all)
274 OnWarnings OnWarningsFunc // Called for warnings
275 OnTextStart OnTextStartFunc // Called when text starts
276 OnTextDelta OnTextDeltaFunc // Called for text deltas
277 OnTextEnd OnTextEndFunc // Called when text ends
278 OnReasoningStart OnReasoningStartFunc // Called when reasoning starts
279 OnReasoningDelta OnReasoningDeltaFunc // Called for reasoning deltas
280 OnReasoningEnd OnReasoningEndFunc // Called when reasoning ends
281 OnToolInputStart OnToolInputStartFunc // Called when tool input starts
282 OnToolInputDelta OnToolInputDeltaFunc // Called for tool input deltas
283 OnToolInputEnd OnToolInputEndFunc // Called when tool input ends
284 OnToolCall OnToolCallFunc // Called when tool call is complete
285 OnToolResult OnToolResultFunc // Called when tool execution completes
286 OnSource OnSourceFunc // Called for source references
287 OnStreamFinish OnStreamFinishFunc // Called when stream finishes
288}
289
290// AgentResult represents the result of an agent execution.
291type AgentResult struct {
292 Steps []StepResult
293 // Final response
294 Response Response
295 TotalUsage Usage
296}
297
298// Agent represents an AI agent that can generate responses and stream responses.
299type Agent interface {
300 Generate(context.Context, AgentCall) (*AgentResult, error)
301 Stream(context.Context, AgentStreamCall) (*AgentResult, error)
302}
303
304// AgentOption defines a function that configures agent settings.
305type AgentOption = func(*agentSettings)
306
307type agent struct {
308 settings agentSettings
309}
310
311// NewAgent creates a new agent with the given language model and options.
312func NewAgent(model LanguageModel, opts ...AgentOption) Agent {
313 settings := agentSettings{
314 model: model,
315 }
316 for _, o := range opts {
317 o(&settings)
318 }
319 return &agent{
320 settings: settings,
321 }
322}
323
324func (a *agent) prepareCall(call AgentCall) AgentCall {
325 call.MaxOutputTokens = cmp.Or(call.MaxOutputTokens, a.settings.maxOutputTokens)
326 call.Temperature = cmp.Or(call.Temperature, a.settings.temperature)
327 call.TopP = cmp.Or(call.TopP, a.settings.topP)
328 call.TopK = cmp.Or(call.TopK, a.settings.topK)
329 call.PresencePenalty = cmp.Or(call.PresencePenalty, a.settings.presencePenalty)
330 call.FrequencyPenalty = cmp.Or(call.FrequencyPenalty, a.settings.frequencyPenalty)
331 call.MaxRetries = cmp.Or(call.MaxRetries, a.settings.maxRetries)
332
333 if len(call.StopWhen) == 0 && len(a.settings.stopWhen) > 0 {
334 call.StopWhen = a.settings.stopWhen
335 }
336 if call.PrepareStep == nil && a.settings.prepareStep != nil {
337 call.PrepareStep = a.settings.prepareStep
338 }
339 if call.RepairToolCall == nil && a.settings.repairToolCall != nil {
340 call.RepairToolCall = a.settings.repairToolCall
341 }
342 if call.OnRetry == nil && a.settings.onRetry != nil {
343 call.OnRetry = a.settings.onRetry
344 }
345
346 providerOptions := ProviderOptions{}
347 if a.settings.providerOptions != nil {
348 maps.Copy(providerOptions, a.settings.providerOptions)
349 }
350 if call.ProviderOptions != nil {
351 maps.Copy(providerOptions, call.ProviderOptions)
352 }
353 call.ProviderOptions = providerOptions
354
355 headers := map[string]string{}
356
357 if a.settings.headers != nil {
358 maps.Copy(headers, a.settings.headers)
359 }
360
361 return call
362}
363
364// Generate implements Agent.
365func (a *agent) Generate(ctx context.Context, opts AgentCall) (*AgentResult, error) {
366 opts = a.prepareCall(opts)
367 initialPrompt, err := a.createPrompt(a.settings.systemPrompt, opts.Prompt, opts.Messages, opts.Files...)
368 if err != nil {
369 return nil, err
370 }
371 var responseMessages []Message
372 var steps []StepResult
373
374 for {
375 stepInputMessages := append(initialPrompt, responseMessages...)
376 stepModel := a.settings.model
377 stepSystemPrompt := a.settings.systemPrompt
378 stepActiveTools := opts.ActiveTools
379 stepToolChoice := ToolChoiceAuto
380 disableAllTools := false
381 stepTools := a.settings.tools
382 if opts.PrepareStep != nil {
383 updatedCtx, prepared, err := opts.PrepareStep(ctx, PrepareStepFunctionOptions{
384 Model: stepModel,
385 Steps: steps,
386 StepNumber: len(steps),
387 Messages: stepInputMessages,
388 })
389 if err != nil {
390 return nil, err
391 }
392
393 ctx = updatedCtx
394
395 // Apply prepared step modifications
396 if prepared.Messages != nil {
397 stepInputMessages = prepared.Messages
398 }
399 if prepared.Model != nil {
400 stepModel = prepared.Model
401 }
402 if prepared.System != nil {
403 stepSystemPrompt = *prepared.System
404 }
405 if prepared.ToolChoice != nil {
406 stepToolChoice = *prepared.ToolChoice
407 }
408 if len(prepared.ActiveTools) > 0 {
409 stepActiveTools = prepared.ActiveTools
410 }
411 disableAllTools = prepared.DisableAllTools
412 if prepared.Tools != nil {
413 stepTools = prepared.Tools
414 }
415 }
416
417 // Recreate prompt with potentially modified system prompt
418 if stepSystemPrompt != a.settings.systemPrompt {
419 stepPrompt, err := a.createPrompt(stepSystemPrompt, opts.Prompt, opts.Messages, opts.Files...)
420 if err != nil {
421 return nil, err
422 }
423 // Replace system message part, keep the rest
424 if len(stepInputMessages) > 0 && len(stepPrompt) > 0 {
425 stepInputMessages[0] = stepPrompt[0] // Replace system message
426 }
427 }
428
429 preparedTools := a.prepareTools(stepTools, stepActiveTools, disableAllTools)
430
431 retryOptions := DefaultRetryOptions()
432 if opts.MaxRetries != nil {
433 retryOptions.MaxRetries = *opts.MaxRetries
434 }
435 retryOptions.OnRetry = opts.OnRetry
436 retry := RetryWithExponentialBackoffRespectingRetryHeaders[*Response](retryOptions)
437
438 result, err := retry(ctx, func() (*Response, error) {
439 return stepModel.Generate(ctx, Call{
440 Prompt: stepInputMessages,
441 MaxOutputTokens: opts.MaxOutputTokens,
442 Temperature: opts.Temperature,
443 TopP: opts.TopP,
444 TopK: opts.TopK,
445 PresencePenalty: opts.PresencePenalty,
446 FrequencyPenalty: opts.FrequencyPenalty,
447 Tools: preparedTools,
448 ToolChoice: &stepToolChoice,
449 ProviderOptions: opts.ProviderOptions,
450 })
451 })
452 if err != nil {
453 return nil, err
454 }
455
456 var stepToolCalls []ToolCallContent
457 for _, content := range result.Content {
458 if content.GetType() == ContentTypeToolCall {
459 toolCall, ok := AsContentType[ToolCallContent](content)
460 if !ok {
461 continue
462 }
463
464 // Validate and potentially repair the tool call
465 validatedToolCall := a.validateAndRepairToolCall(ctx, toolCall, stepTools, stepSystemPrompt, stepInputMessages, a.settings.repairToolCall)
466 stepToolCalls = append(stepToolCalls, validatedToolCall)
467 }
468 }
469
470 toolResults, err := a.executeTools(ctx, stepTools, stepToolCalls, nil)
471
472 // Build step content with validated tool calls and tool results
473 stepContent := []Content{}
474 toolCallIndex := 0
475 for _, content := range result.Content {
476 if content.GetType() == ContentTypeToolCall {
477 // Replace with validated tool call
478 if toolCallIndex < len(stepToolCalls) {
479 stepContent = append(stepContent, stepToolCalls[toolCallIndex])
480 toolCallIndex++
481 }
482 } else {
483 // Keep other content as-is
484 stepContent = append(stepContent, content)
485 }
486 }
487 // Add tool results
488 for _, result := range toolResults {
489 stepContent = append(stepContent, result)
490 }
491 currentStepMessages := toResponseMessages(stepContent)
492 responseMessages = append(responseMessages, currentStepMessages...)
493
494 stepResult := StepResult{
495 Response: Response{
496 Content: stepContent,
497 FinishReason: result.FinishReason,
498 Usage: result.Usage,
499 Warnings: result.Warnings,
500 ProviderMetadata: result.ProviderMetadata,
501 },
502 Messages: currentStepMessages,
503 }
504 steps = append(steps, stepResult)
505 shouldStop := isStopConditionMet(opts.StopWhen, steps)
506
507 if shouldStop || err != nil || len(stepToolCalls) == 0 || result.FinishReason != FinishReasonToolCalls {
508 break
509 }
510 }
511
512 totalUsage := Usage{}
513
514 for _, step := range steps {
515 usage := step.Usage
516 totalUsage.InputTokens += usage.InputTokens
517 totalUsage.OutputTokens += usage.OutputTokens
518 totalUsage.ReasoningTokens += usage.ReasoningTokens
519 totalUsage.CacheCreationTokens += usage.CacheCreationTokens
520 totalUsage.CacheReadTokens += usage.CacheReadTokens
521 totalUsage.TotalTokens += usage.TotalTokens
522 }
523
524 agentResult := &AgentResult{
525 Steps: steps,
526 Response: steps[len(steps)-1].Response,
527 TotalUsage: totalUsage,
528 }
529 return agentResult, nil
530}
531
532func isStopConditionMet(conditions []StopCondition, steps []StepResult) bool {
533 if len(conditions) == 0 {
534 return false
535 }
536
537 for _, condition := range conditions {
538 if condition(steps) {
539 return true
540 }
541 }
542 return false
543}
544
545func toResponseMessages(content []Content) []Message {
546 var assistantParts []MessagePart
547 var toolParts []MessagePart
548
549 for _, c := range content {
550 switch c.GetType() {
551 case ContentTypeText:
552 text, ok := AsContentType[TextContent](c)
553 if !ok {
554 continue
555 }
556 assistantParts = append(assistantParts, TextPart{
557 Text: text.Text,
558 ProviderOptions: ProviderOptions(text.ProviderMetadata),
559 })
560 case ContentTypeReasoning:
561 reasoning, ok := AsContentType[ReasoningContent](c)
562 if !ok {
563 continue
564 }
565 assistantParts = append(assistantParts, ReasoningPart{
566 Text: reasoning.Text,
567 ProviderOptions: ProviderOptions(reasoning.ProviderMetadata),
568 })
569 case ContentTypeToolCall:
570 toolCall, ok := AsContentType[ToolCallContent](c)
571 if !ok {
572 continue
573 }
574 assistantParts = append(assistantParts, ToolCallPart{
575 ToolCallID: toolCall.ToolCallID,
576 ToolName: toolCall.ToolName,
577 Input: toolCall.Input,
578 ProviderExecuted: toolCall.ProviderExecuted,
579 ProviderOptions: ProviderOptions(toolCall.ProviderMetadata),
580 })
581 case ContentTypeFile:
582 file, ok := AsContentType[FileContent](c)
583 if !ok {
584 continue
585 }
586 assistantParts = append(assistantParts, FilePart{
587 Data: file.Data,
588 MediaType: file.MediaType,
589 ProviderOptions: ProviderOptions(file.ProviderMetadata),
590 })
591 case ContentTypeSource:
592 // Sources are metadata about references used to generate the response.
593 // They don't need to be included in the conversation messages.
594 continue
595 case ContentTypeToolResult:
596 result, ok := AsContentType[ToolResultContent](c)
597 if !ok {
598 continue
599 }
600 toolParts = append(toolParts, ToolResultPart{
601 ToolCallID: result.ToolCallID,
602 Output: result.Result,
603 ProviderOptions: ProviderOptions(result.ProviderMetadata),
604 })
605 }
606 }
607
608 var messages []Message
609 if len(assistantParts) > 0 {
610 messages = append(messages, Message{
611 Role: MessageRoleAssistant,
612 Content: assistantParts,
613 })
614 }
615 if len(toolParts) > 0 {
616 messages = append(messages, Message{
617 Role: MessageRoleTool,
618 Content: toolParts,
619 })
620 }
621 return messages
622}
623
624func (a *agent) executeTools(ctx context.Context, allTools []AgentTool, toolCalls []ToolCallContent, toolResultCallback func(result ToolResultContent) error) ([]ToolResultContent, error) {
625 if len(toolCalls) == 0 {
626 return nil, nil
627 }
628
629 // Create a map for quick tool lookup
630 toolMap := make(map[string]AgentTool)
631 for _, tool := range allTools {
632 toolMap[tool.Info().Name] = tool
633 }
634
635 // Execute all tool calls sequentially in order
636 results := make([]ToolResultContent, 0, len(toolCalls))
637
638 for _, toolCall := range toolCalls {
639 result, isCriticalError := a.executeSingleTool(ctx, toolMap, toolCall, toolResultCallback)
640 results = append(results, result)
641 if isCriticalError {
642 if errorResult, ok := result.Result.(ToolResultOutputContentError); ok && errorResult.Error != nil {
643 return nil, errorResult.Error
644 }
645 }
646 }
647
648 return results, nil
649}
650
651// executeSingleTool executes a single tool and returns its result and a critical error flag.
652func (a *agent) executeSingleTool(ctx context.Context, toolMap map[string]AgentTool, toolCall ToolCallContent, toolResultCallback func(result ToolResultContent) error) (ToolResultContent, bool) {
653 result := ToolResultContent{
654 ToolCallID: toolCall.ToolCallID,
655 ToolName: toolCall.ToolName,
656 ProviderExecuted: false,
657 }
658
659 // Skip invalid tool calls - create error result (not critical)
660 if toolCall.Invalid {
661 result.Result = ToolResultOutputContentError{
662 Error: toolCall.ValidationError,
663 }
664 if toolResultCallback != nil {
665 _ = toolResultCallback(result)
666 }
667 return result, false
668 }
669
670 tool, exists := toolMap[toolCall.ToolName]
671 if !exists {
672 result.Result = ToolResultOutputContentError{
673 Error: errors.New("Error: Tool not found: " + toolCall.ToolName),
674 }
675 if toolResultCallback != nil {
676 _ = toolResultCallback(result)
677 }
678 return result, false
679 }
680
681 // Execute the tool
682 toolResult, err := tool.Run(ctx, ToolCall{
683 ID: toolCall.ToolCallID,
684 Name: toolCall.ToolName,
685 Input: toolCall.Input,
686 })
687 if err != nil {
688 result.Result = ToolResultOutputContentError{
689 Error: err,
690 }
691 result.ClientMetadata = toolResult.Metadata
692 if toolResultCallback != nil {
693 _ = toolResultCallback(result)
694 }
695 return result, true
696 }
697
698 result.ClientMetadata = toolResult.Metadata
699 if toolResult.IsError {
700 result.Result = ToolResultOutputContentError{
701 Error: errors.New(toolResult.Content),
702 }
703 } else if toolResult.Type == "image" || toolResult.Type == "media" {
704 result.Result = ToolResultOutputContentMedia{
705 Data: string(toolResult.Data),
706 MediaType: toolResult.MediaType,
707 Text: toolResult.Content,
708 }
709 } else {
710 result.Result = ToolResultOutputContentText{
711 Text: toolResult.Content,
712 }
713 }
714 if toolResultCallback != nil {
715 _ = toolResultCallback(result)
716 }
717 return result, false
718}
719
720// Stream implements Agent.
721func (a *agent) Stream(ctx context.Context, opts AgentStreamCall) (*AgentResult, error) {
722 // Convert AgentStreamCall to AgentCall for preparation
723 call := AgentCall{
724 Prompt: opts.Prompt,
725 Files: opts.Files,
726 Messages: opts.Messages,
727 MaxOutputTokens: opts.MaxOutputTokens,
728 Temperature: opts.Temperature,
729 TopP: opts.TopP,
730 TopK: opts.TopK,
731 PresencePenalty: opts.PresencePenalty,
732 FrequencyPenalty: opts.FrequencyPenalty,
733 ActiveTools: opts.ActiveTools,
734 ProviderOptions: opts.ProviderOptions,
735 MaxRetries: opts.MaxRetries,
736 OnRetry: opts.OnRetry,
737 StopWhen: opts.StopWhen,
738 PrepareStep: opts.PrepareStep,
739 RepairToolCall: opts.RepairToolCall,
740 }
741
742 call = a.prepareCall(call)
743
744 initialPrompt, err := a.createPrompt(a.settings.systemPrompt, call.Prompt, call.Messages, call.Files...)
745 if err != nil {
746 return nil, err
747 }
748
749 var responseMessages []Message
750 var steps []StepResult
751 var totalUsage Usage
752
753 // Start agent stream
754 if opts.OnAgentStart != nil {
755 opts.OnAgentStart()
756 }
757
758 for stepNumber := 0; ; stepNumber++ {
759 stepInputMessages := append(initialPrompt, responseMessages...)
760 stepModel := a.settings.model
761 stepSystemPrompt := a.settings.systemPrompt
762 stepActiveTools := call.ActiveTools
763 stepToolChoice := ToolChoiceAuto
764 disableAllTools := false
765 stepTools := a.settings.tools
766 // Apply step preparation if provided
767 if call.PrepareStep != nil {
768 updatedCtx, prepared, err := call.PrepareStep(ctx, PrepareStepFunctionOptions{
769 Model: stepModel,
770 Steps: steps,
771 StepNumber: stepNumber,
772 Messages: stepInputMessages,
773 })
774 if err != nil {
775 return nil, err
776 }
777
778 ctx = updatedCtx
779
780 if prepared.Messages != nil {
781 stepInputMessages = prepared.Messages
782 }
783 if prepared.Model != nil {
784 stepModel = prepared.Model
785 }
786 if prepared.System != nil {
787 stepSystemPrompt = *prepared.System
788 }
789 if prepared.ToolChoice != nil {
790 stepToolChoice = *prepared.ToolChoice
791 }
792 if len(prepared.ActiveTools) > 0 {
793 stepActiveTools = prepared.ActiveTools
794 }
795 disableAllTools = prepared.DisableAllTools
796 if prepared.Tools != nil {
797 stepTools = prepared.Tools
798 }
799 }
800
801 // Recreate prompt with potentially modified system prompt
802 if stepSystemPrompt != a.settings.systemPrompt {
803 stepPrompt, err := a.createPrompt(stepSystemPrompt, call.Prompt, call.Messages, call.Files...)
804 if err != nil {
805 return nil, err
806 }
807 if len(stepInputMessages) > 0 && len(stepPrompt) > 0 {
808 stepInputMessages[0] = stepPrompt[0]
809 }
810 }
811
812 preparedTools := a.prepareTools(stepTools, stepActiveTools, disableAllTools)
813
814 // Start step stream
815 if opts.OnStepStart != nil {
816 _ = opts.OnStepStart(stepNumber)
817 }
818
819 // Create streaming call
820 streamCall := Call{
821 Prompt: stepInputMessages,
822 MaxOutputTokens: call.MaxOutputTokens,
823 Temperature: call.Temperature,
824 TopP: call.TopP,
825 TopK: call.TopK,
826 PresencePenalty: call.PresencePenalty,
827 FrequencyPenalty: call.FrequencyPenalty,
828 Tools: preparedTools,
829 ToolChoice: &stepToolChoice,
830 ProviderOptions: call.ProviderOptions,
831 }
832
833 // Execute step with retry logic wrapping both stream creation and processing
834 retryOptions := DefaultRetryOptions()
835 if call.MaxRetries != nil {
836 retryOptions.MaxRetries = *call.MaxRetries
837 }
838 retryOptions.OnRetry = call.OnRetry
839 retry := RetryWithExponentialBackoffRespectingRetryHeaders[stepExecutionResult](retryOptions)
840
841 result, err := retry(ctx, func() (stepExecutionResult, error) {
842 // Create the stream
843 stream, err := stepModel.Stream(ctx, streamCall)
844 if err != nil {
845 return stepExecutionResult{}, err
846 }
847
848 // Process the stream
849 result, err := a.processStepStream(ctx, stream, opts, steps, stepTools)
850 if err != nil {
851 return stepExecutionResult{}, err
852 }
853
854 return result, nil
855 })
856 if err != nil {
857 if opts.OnError != nil {
858 opts.OnError(err)
859 }
860 return nil, err
861 }
862
863 steps = append(steps, result.StepResult)
864 totalUsage = addUsage(totalUsage, result.StepResult.Usage)
865
866 // Call step finished callback
867 if opts.OnStepFinish != nil {
868 _ = opts.OnStepFinish(result.StepResult)
869 }
870
871 // Add step messages to response messages
872 stepMessages := toResponseMessages(result.StepResult.Content)
873 responseMessages = append(responseMessages, stepMessages...)
874
875 // Check stop conditions
876 shouldStop := isStopConditionMet(call.StopWhen, steps)
877 if shouldStop || !result.ShouldContinue {
878 break
879 }
880 }
881
882 // Finish agent stream
883 agentResult := &AgentResult{
884 Steps: steps,
885 Response: steps[len(steps)-1].Response,
886 TotalUsage: totalUsage,
887 }
888
889 if opts.OnFinish != nil {
890 opts.OnFinish(agentResult)
891 }
892
893 if opts.OnAgentFinish != nil {
894 _ = opts.OnAgentFinish(agentResult)
895 }
896
897 return agentResult, nil
898}
899
900func (a *agent) prepareTools(tools []AgentTool, activeTools []string, disableAllTools bool) []Tool {
901 preparedTools := make([]Tool, 0, len(tools))
902
903 // If explicitly disabling all tools, return no tools
904 if disableAllTools {
905 return preparedTools
906 }
907
908 for _, tool := range tools {
909 // If activeTools has items, only include tools in the list
910 // If activeTools is empty, include all tools
911 if len(activeTools) > 0 && !slices.Contains(activeTools, tool.Info().Name) {
912 continue
913 }
914 info := tool.Info()
915 preparedTools = append(preparedTools, FunctionTool{
916 Name: info.Name,
917 Description: info.Description,
918 InputSchema: map[string]any{
919 "type": "object",
920 "properties": info.Parameters,
921 "required": info.Required,
922 },
923 ProviderOptions: tool.ProviderOptions(),
924 })
925 }
926 return preparedTools
927}
928
929// validateAndRepairToolCall validates a tool call and attempts repair if validation fails.
930func (a *agent) validateAndRepairToolCall(ctx context.Context, toolCall ToolCallContent, availableTools []AgentTool, systemPrompt string, messages []Message, repairFunc RepairToolCallFunction) ToolCallContent {
931 if err := a.validateToolCall(toolCall, availableTools); err == nil {
932 return toolCall
933 } else { //nolint: revive
934 if repairFunc != nil {
935 repairOptions := ToolCallRepairOptions{
936 OriginalToolCall: toolCall,
937 ValidationError: err,
938 AvailableTools: availableTools,
939 SystemPrompt: systemPrompt,
940 Messages: messages,
941 }
942
943 if repairedToolCall, repairErr := repairFunc(ctx, repairOptions); repairErr == nil && repairedToolCall != nil {
944 if validateErr := a.validateToolCall(*repairedToolCall, availableTools); validateErr == nil {
945 return *repairedToolCall
946 }
947 }
948 }
949
950 invalidToolCall := toolCall
951 invalidToolCall.Invalid = true
952 invalidToolCall.ValidationError = err
953 return invalidToolCall
954 }
955}
956
957// validateToolCall validates a tool call against available tools and their schemas.
958func (a *agent) validateToolCall(toolCall ToolCallContent, availableTools []AgentTool) error {
959 var tool AgentTool
960 for _, t := range availableTools {
961 if t.Info().Name == toolCall.ToolName {
962 tool = t
963 break
964 }
965 }
966
967 if tool == nil {
968 return fmt.Errorf("tool not found: %s", toolCall.ToolName)
969 }
970
971 // Validate JSON parsing
972 var input map[string]any
973 if err := json.Unmarshal([]byte(toolCall.Input), &input); err != nil {
974 return fmt.Errorf("invalid JSON input: %w", err)
975 }
976
977 // Basic schema validation (check required fields)
978 // TODO: more robust schema validation using JSON Schema or similar
979 toolInfo := tool.Info()
980 for _, required := range toolInfo.Required {
981 if _, exists := input[required]; !exists {
982 return fmt.Errorf("missing required parameter: %s", required)
983 }
984 }
985 return nil
986}
987
988func (a *agent) createPrompt(system, prompt string, messages []Message, files ...FilePart) (Prompt, error) {
989 if prompt == "" {
990 return nil, &Error{Title: "invalid argument", Message: "prompt can't be empty"}
991 }
992
993 var preparedPrompt Prompt
994
995 if system != "" {
996 preparedPrompt = append(preparedPrompt, NewSystemMessage(system))
997 }
998 preparedPrompt = append(preparedPrompt, messages...)
999 preparedPrompt = append(preparedPrompt, NewUserMessage(prompt, files...))
1000 return preparedPrompt, nil
1001}
1002
1003// WithSystemPrompt sets the system prompt for the agent.
1004func WithSystemPrompt(prompt string) AgentOption {
1005 return func(s *agentSettings) {
1006 s.systemPrompt = prompt
1007 }
1008}
1009
1010// WithMaxOutputTokens sets the maximum output tokens for the agent.
1011func WithMaxOutputTokens(tokens int64) AgentOption {
1012 return func(s *agentSettings) {
1013 s.maxOutputTokens = &tokens
1014 }
1015}
1016
1017// WithTemperature sets the temperature for the agent.
1018func WithTemperature(temp float64) AgentOption {
1019 return func(s *agentSettings) {
1020 s.temperature = &temp
1021 }
1022}
1023
1024// WithTopP sets the top-p value for the agent.
1025func WithTopP(topP float64) AgentOption {
1026 return func(s *agentSettings) {
1027 s.topP = &topP
1028 }
1029}
1030
1031// WithTopK sets the top-k value for the agent.
1032func WithTopK(topK int64) AgentOption {
1033 return func(s *agentSettings) {
1034 s.topK = &topK
1035 }
1036}
1037
1038// WithPresencePenalty sets the presence penalty for the agent.
1039func WithPresencePenalty(penalty float64) AgentOption {
1040 return func(s *agentSettings) {
1041 s.presencePenalty = &penalty
1042 }
1043}
1044
1045// WithFrequencyPenalty sets the frequency penalty for the agent.
1046func WithFrequencyPenalty(penalty float64) AgentOption {
1047 return func(s *agentSettings) {
1048 s.frequencyPenalty = &penalty
1049 }
1050}
1051
1052// WithTools sets the tools for the agent.
1053func WithTools(tools ...AgentTool) AgentOption {
1054 return func(s *agentSettings) {
1055 s.tools = append(s.tools, tools...)
1056 }
1057}
1058
1059// WithStopConditions sets the stop conditions for the agent.
1060func WithStopConditions(conditions ...StopCondition) AgentOption {
1061 return func(s *agentSettings) {
1062 s.stopWhen = append(s.stopWhen, conditions...)
1063 }
1064}
1065
1066// WithPrepareStep sets the prepare step function for the agent.
1067func WithPrepareStep(fn PrepareStepFunction) AgentOption {
1068 return func(s *agentSettings) {
1069 s.prepareStep = fn
1070 }
1071}
1072
1073// WithRepairToolCall sets the repair tool call function for the agent.
1074func WithRepairToolCall(fn RepairToolCallFunction) AgentOption {
1075 return func(s *agentSettings) {
1076 s.repairToolCall = fn
1077 }
1078}
1079
1080// WithMaxRetries sets the maximum number of retries for the agent.
1081func WithMaxRetries(maxRetries int) AgentOption {
1082 return func(s *agentSettings) {
1083 s.maxRetries = &maxRetries
1084 }
1085}
1086
1087// WithOnRetry sets the retry callback for the agent.
1088func WithOnRetry(callback OnRetryCallback) AgentOption {
1089 return func(s *agentSettings) {
1090 s.onRetry = callback
1091 }
1092}
1093
1094// processStepStream processes a single step's stream and returns the step result.
1095func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, opts AgentStreamCall, _ []StepResult, stepTools []AgentTool) (stepExecutionResult, error) {
1096 var stepContent []Content
1097 var stepToolCalls []ToolCallContent
1098 var stepUsage Usage
1099 stepFinishReason := FinishReasonUnknown
1100 var stepWarnings []CallWarning
1101 var stepProviderMetadata ProviderMetadata
1102
1103 activeToolCalls := make(map[string]*ToolCallContent)
1104 activeTextContent := make(map[string]string)
1105 type reasoningContent struct {
1106 content string
1107 options ProviderMetadata
1108 }
1109 activeReasoningContent := make(map[string]reasoningContent)
1110
1111 // Set up concurrent tool execution
1112 type toolExecutionRequest struct {
1113 toolCall ToolCallContent
1114 parallel bool
1115 }
1116 toolChan := make(chan toolExecutionRequest, 10)
1117 var toolExecutionWg sync.WaitGroup
1118 var toolStateMu sync.Mutex
1119 toolResults := make([]ToolResultContent, 0)
1120 var toolExecutionErr error
1121
1122 // Create a map for quick tool lookup
1123 toolMap := make(map[string]AgentTool)
1124 for _, tool := range stepTools {
1125 toolMap[tool.Info().Name] = tool
1126 }
1127
1128 // Semaphores for controlling parallelism
1129 parallelSem := make(chan struct{}, 5)
1130 var sequentialMu sync.Mutex
1131
1132 // Single coordinator goroutine that dispatches tools
1133 toolExecutionWg.Go(func() {
1134 for req := range toolChan {
1135 if req.parallel {
1136 parallelSem <- struct{}{}
1137 toolExecutionWg.Go(func() {
1138 defer func() { <-parallelSem }()
1139 result, isCriticalError := a.executeSingleTool(ctx, toolMap, req.toolCall, opts.OnToolResult)
1140 toolStateMu.Lock()
1141 toolResults = append(toolResults, result)
1142 if isCriticalError && toolExecutionErr == nil {
1143 if errorResult, ok := result.Result.(ToolResultOutputContentError); ok && errorResult.Error != nil {
1144 toolExecutionErr = errorResult.Error
1145 }
1146 }
1147 toolStateMu.Unlock()
1148 })
1149 } else {
1150 sequentialMu.Lock()
1151 result, isCriticalError := a.executeSingleTool(ctx, toolMap, req.toolCall, opts.OnToolResult)
1152 toolStateMu.Lock()
1153 toolResults = append(toolResults, result)
1154 if isCriticalError && toolExecutionErr == nil {
1155 if errorResult, ok := result.Result.(ToolResultOutputContentError); ok && errorResult.Error != nil {
1156 toolExecutionErr = errorResult.Error
1157 }
1158 }
1159 toolStateMu.Unlock()
1160 sequentialMu.Unlock()
1161 }
1162 }
1163 })
1164
1165 // Process stream parts
1166 for part := range stream {
1167 // Forward all parts to chunk callback
1168 if opts.OnChunk != nil {
1169 err := opts.OnChunk(part)
1170 if err != nil {
1171 return stepExecutionResult{}, err
1172 }
1173 }
1174
1175 switch part.Type {
1176 case StreamPartTypeWarnings:
1177 stepWarnings = part.Warnings
1178 if opts.OnWarnings != nil {
1179 err := opts.OnWarnings(part.Warnings)
1180 if err != nil {
1181 return stepExecutionResult{}, err
1182 }
1183 }
1184
1185 case StreamPartTypeTextStart:
1186 activeTextContent[part.ID] = ""
1187 if opts.OnTextStart != nil {
1188 err := opts.OnTextStart(part.ID)
1189 if err != nil {
1190 return stepExecutionResult{}, err
1191 }
1192 }
1193
1194 case StreamPartTypeTextDelta:
1195 if _, exists := activeTextContent[part.ID]; exists {
1196 activeTextContent[part.ID] += part.Delta
1197 }
1198 if opts.OnTextDelta != nil {
1199 err := opts.OnTextDelta(part.ID, part.Delta)
1200 if err != nil {
1201 return stepExecutionResult{}, err
1202 }
1203 }
1204
1205 case StreamPartTypeTextEnd:
1206 if text, exists := activeTextContent[part.ID]; exists {
1207 stepContent = append(stepContent, TextContent{
1208 Text: text,
1209 ProviderMetadata: part.ProviderMetadata,
1210 })
1211 delete(activeTextContent, part.ID)
1212 }
1213 if opts.OnTextEnd != nil {
1214 err := opts.OnTextEnd(part.ID)
1215 if err != nil {
1216 return stepExecutionResult{}, err
1217 }
1218 }
1219
1220 case StreamPartTypeReasoningStart:
1221 activeReasoningContent[part.ID] = reasoningContent{content: part.Delta, options: part.ProviderMetadata}
1222 if opts.OnReasoningStart != nil {
1223 content := ReasoningContent{
1224 Text: part.Delta,
1225 ProviderMetadata: part.ProviderMetadata,
1226 }
1227 err := opts.OnReasoningStart(part.ID, content)
1228 if err != nil {
1229 return stepExecutionResult{}, err
1230 }
1231 }
1232
1233 case StreamPartTypeReasoningDelta:
1234 if active, exists := activeReasoningContent[part.ID]; exists {
1235 active.content += part.Delta
1236 active.options = part.ProviderMetadata
1237 activeReasoningContent[part.ID] = active
1238 }
1239 if opts.OnReasoningDelta != nil {
1240 err := opts.OnReasoningDelta(part.ID, part.Delta)
1241 if err != nil {
1242 return stepExecutionResult{}, err
1243 }
1244 }
1245
1246 case StreamPartTypeReasoningEnd:
1247 if active, exists := activeReasoningContent[part.ID]; exists {
1248 if part.ProviderMetadata != nil {
1249 active.options = part.ProviderMetadata
1250 }
1251 content := ReasoningContent{
1252 Text: active.content,
1253 ProviderMetadata: active.options,
1254 }
1255 stepContent = append(stepContent, content)
1256 if opts.OnReasoningEnd != nil {
1257 err := opts.OnReasoningEnd(part.ID, content)
1258 if err != nil {
1259 return stepExecutionResult{}, err
1260 }
1261 }
1262 delete(activeReasoningContent, part.ID)
1263 }
1264
1265 case StreamPartTypeToolInputStart:
1266 activeToolCalls[part.ID] = &ToolCallContent{
1267 ToolCallID: part.ID,
1268 ToolName: part.ToolCallName,
1269 Input: "",
1270 ProviderExecuted: part.ProviderExecuted,
1271 }
1272 if opts.OnToolInputStart != nil {
1273 err := opts.OnToolInputStart(part.ID, part.ToolCallName)
1274 if err != nil {
1275 return stepExecutionResult{}, err
1276 }
1277 }
1278
1279 case StreamPartTypeToolInputDelta:
1280 if toolCall, exists := activeToolCalls[part.ID]; exists {
1281 toolCall.Input += part.Delta
1282 }
1283 if opts.OnToolInputDelta != nil {
1284 err := opts.OnToolInputDelta(part.ID, part.Delta)
1285 if err != nil {
1286 return stepExecutionResult{}, err
1287 }
1288 }
1289
1290 case StreamPartTypeToolInputEnd:
1291 if opts.OnToolInputEnd != nil {
1292 err := opts.OnToolInputEnd(part.ID)
1293 if err != nil {
1294 return stepExecutionResult{}, err
1295 }
1296 }
1297
1298 case StreamPartTypeToolCall:
1299 toolCall := ToolCallContent{
1300 ToolCallID: part.ID,
1301 ToolName: part.ToolCallName,
1302 Input: part.ToolCallInput,
1303 ProviderExecuted: part.ProviderExecuted,
1304 ProviderMetadata: part.ProviderMetadata,
1305 }
1306
1307 // Validate and potentially repair the tool call
1308 validatedToolCall := a.validateAndRepairToolCall(ctx, toolCall, stepTools, a.settings.systemPrompt, nil, opts.RepairToolCall)
1309 stepToolCalls = append(stepToolCalls, validatedToolCall)
1310 stepContent = append(stepContent, validatedToolCall)
1311
1312 if opts.OnToolCall != nil {
1313 err := opts.OnToolCall(validatedToolCall)
1314 if err != nil {
1315 return stepExecutionResult{}, err
1316 }
1317 }
1318
1319 // Determine if tool can run in parallel
1320 isParallel := false
1321 if tool, exists := toolMap[validatedToolCall.ToolName]; exists {
1322 isParallel = tool.Info().Parallel
1323 }
1324
1325 // Send tool call to execution channel
1326 toolChan <- toolExecutionRequest{toolCall: validatedToolCall, parallel: isParallel}
1327
1328 // Clean up active tool call
1329 delete(activeToolCalls, part.ID)
1330
1331 case StreamPartTypeSource:
1332 sourceContent := SourceContent{
1333 SourceType: part.SourceType,
1334 ID: part.ID,
1335 URL: part.URL,
1336 Title: part.Title,
1337 ProviderMetadata: part.ProviderMetadata,
1338 }
1339 stepContent = append(stepContent, sourceContent)
1340 if opts.OnSource != nil {
1341 err := opts.OnSource(sourceContent)
1342 if err != nil {
1343 return stepExecutionResult{}, err
1344 }
1345 }
1346
1347 case StreamPartTypeFinish:
1348 stepUsage = part.Usage
1349 stepFinishReason = part.FinishReason
1350 stepProviderMetadata = part.ProviderMetadata
1351 if opts.OnStreamFinish != nil {
1352 err := opts.OnStreamFinish(part.Usage, part.FinishReason, part.ProviderMetadata)
1353 if err != nil {
1354 return stepExecutionResult{}, err
1355 }
1356 }
1357
1358 case StreamPartTypeError:
1359 return stepExecutionResult{}, part.Error
1360 }
1361 }
1362
1363 // Close the tool execution channel and wait for all executions to complete
1364 close(toolChan)
1365 toolExecutionWg.Wait()
1366
1367 // Check for tool execution errors
1368 if toolExecutionErr != nil {
1369 return stepExecutionResult{}, toolExecutionErr
1370 }
1371
1372 // Add tool results to content if any
1373 if len(toolResults) > 0 {
1374 for _, result := range toolResults {
1375 stepContent = append(stepContent, result)
1376 }
1377 }
1378
1379 stepResult := StepResult{
1380 Response: Response{
1381 Content: stepContent,
1382 FinishReason: stepFinishReason,
1383 Usage: stepUsage,
1384 Warnings: stepWarnings,
1385 ProviderMetadata: stepProviderMetadata,
1386 },
1387 Messages: toResponseMessages(stepContent),
1388 }
1389
1390 // Determine if we should continue (has tool calls and not stopped)
1391 shouldContinue := len(stepToolCalls) > 0 && stepFinishReason == FinishReasonToolCalls
1392
1393 return stepExecutionResult{
1394 StepResult: stepResult,
1395 ShouldContinue: shouldContinue,
1396 }, nil
1397}
1398
1399func addUsage(a, b Usage) Usage {
1400 return Usage{
1401 InputTokens: a.InputTokens + b.InputTokens,
1402 OutputTokens: a.OutputTokens + b.OutputTokens,
1403 TotalTokens: a.TotalTokens + b.TotalTokens,
1404 ReasoningTokens: a.ReasoningTokens + b.ReasoningTokens,
1405 CacheCreationTokens: a.CacheCreationTokens + b.CacheCreationTokens,
1406 CacheReadTokens: a.CacheReadTokens + b.CacheReadTokens,
1407 }
1408}
1409
1410// WithHeaders sets the headers for the agent.
1411func WithHeaders(headers map[string]string) AgentOption {
1412 return func(s *agentSettings) {
1413 s.headers = headers
1414 }
1415}
1416
1417// WithProviderOptions sets the provider options for the agent.
1418func WithProviderOptions(providerOptions ProviderOptions) AgentOption {
1419 return func(s *agentSettings) {
1420 s.providerOptions = providerOptions
1421 }
1422}