agent.go

  1package ai
  2
  3import (
  4	"context"
  5)
  6
  7type StepResponse struct {
  8	Response
  9	// Messages generated during this step
 10	Messages []Message
 11}
 12
 13type StepCondition = func(steps []StepResponse) bool
 14
 15type PrepareStepFunctionOptions struct {
 16	Steps      []StepResponse
 17	StepNumber int
 18	Model      LanguageModel
 19	Messages   []Message
 20}
 21
 22type PrepareStepResult struct {
 23	SystemPrompt string
 24	Model        LanguageModel
 25	Messages     []Message
 26}
 27
 28type PrepareStepFunction = func(options PrepareStepFunctionOptions) PrepareStepResult
 29
 30type OnStepFinishedFunction = func(step StepResponse)
 31
 32type AgentSettings struct {
 33	Call
 34	Model LanguageModel
 35
 36	StopWhen       []StepCondition
 37	PrepareStep    PrepareStepFunction
 38	OnStepFinished OnStepFinishedFunction
 39}
 40
 41type Agent interface {
 42	Generate(context.Context, Call) (*Response, error)
 43	Stream(context.Context, Call) (StreamResponse, error)
 44}
 45
 46type agentOption = func(*AgentSettings)
 47
 48type agent struct {
 49	settings AgentSettings
 50}
 51
 52func NewAgent(model LanguageModel, opts ...agentOption) Agent {
 53	settings := AgentSettings{
 54		Model: model,
 55	}
 56	for _, o := range opts {
 57		o(&settings)
 58	}
 59	return &agent{
 60		settings: settings,
 61	}
 62}
 63
 64func mergeCall(agentOpts, opts Call) Call {
 65	if len(opts.Prompt) > 0 {
 66		agentOpts.Prompt = opts.Prompt
 67	}
 68	if opts.MaxOutputTokens != nil {
 69		agentOpts.MaxOutputTokens = opts.MaxOutputTokens
 70	}
 71	if opts.Temperature != nil {
 72		agentOpts.Temperature = opts.Temperature
 73	}
 74	if opts.TopP != nil {
 75		agentOpts.TopP = opts.TopP
 76	}
 77	if opts.TopK != nil {
 78		agentOpts.TopK = opts.TopK
 79	}
 80	if opts.PresencePenalty != nil {
 81		agentOpts.PresencePenalty = opts.PresencePenalty
 82	}
 83	if opts.FrequencyPenalty != nil {
 84		agentOpts.FrequencyPenalty = opts.FrequencyPenalty
 85	}
 86	if opts.Tools != nil {
 87		agentOpts.Tools = opts.Tools
 88	}
 89	if opts.Headers != nil {
 90		agentOpts.Headers = opts.Headers
 91	}
 92	if opts.ProviderOptions != nil {
 93		agentOpts.ProviderOptions = opts.ProviderOptions
 94	}
 95	return agentOpts
 96}
 97
 98// Generate implements Agent.
 99func (a *agent) Generate(ctx context.Context, opts Call) (*Response, error) {
100	// TODO: implement the agentic stuff
101	return a.settings.Model.Generate(ctx, mergeCall(a.settings.Call, opts))
102}
103
104// Stream implements Agent.
105func (a *agent) Stream(ctx context.Context, opts Call) (StreamResponse, error) {
106	// TODO: implement the agentic stuff
107	return a.settings.Model.Stream(ctx, mergeCall(a.settings.Call, opts))
108}