package agent import ( "context" "fmt" "io" "log/slog" "github.com/davecgh/go-spew/spew" gollm "gitea.stevedudenhoeffer.com/steve/go-llm" ) type Agent struct { // ToolBox is the toolbox to use for the agent. ToolBox *gollm.ToolBox // Model is the model to use for the agent. Model gollm.ChatCompletion // OnNewFunction is a callback that, if non-nil, will be called when a new function is called by the LLM. // The "answer" and "no_answer" functions are not included in this callback. // Return an error to stop the function from being called. OnNewFunction func(ctx context.Context, funcName string, question string, parameter string) (any, error) // OnFunctionFinished is a callback that, if non-nil, will be called when a function has finished executing. The // function name is passed in, as well as the question, the parameter, all similar to OnNewFunction. The result of // the function is also passed in, as well as any error that occurred. Finally, the result passed from the // OnNewFunction that preceded this function is passed in as well. OnFunctionFinished func(ctx context.Context, funcName string, question string, parameter string, result string, err error, newFunctionResult any) error req gollm.Request } func NewAgent(req gollm.Request) *Agent { return &Agent{req: req} } type Response struct { Text string Sources []string } func deferClose(cl io.Closer) { if cl != nil { _ = cl.Close() } } func (a *Agent) AddConversation(in gollm.Input) { a.req.Conversation = append(a.req.Conversation, in) } func (a *Agent) AddMessage(msg gollm.Message) { slog.Info("adding message", "message", msg) a.req.Messages = append(a.req.Messages, msg) } // Execute will execute the current request with the given messages. The messages will be appended to the current // request, but they will _not_ be saved into the embedded request. However, the embedded request will be // generated with the on the results from the ChatComplete call. func (a *Agent) Execute(ctx context.Context, msgs ...gollm.Message) error { ctx, cancel := context.WithCancel(ctx) defer cancel() req := a.req slog.Info("executing", "request", req, "messages", msgs) for _, c := range req.Conversation { slog.Info("conversation", "message", c) } req.Messages = append(req.Messages, msgs...) for _, m := range req.Messages { slog.Info("messages", "message", m) } req.Toolbox = a.ToolBox fmt.Println("req:") spew.Dump(req) res, err := a.Model.ChatComplete(ctx, req) fmt.Println("res:") spew.Dump(res) if err != nil { return err } if len(res.Choices) == 0 { return nil } choice := res.Choices[0] var callsOutput = make(chan gollm.ToolCallResponse, len(choice.Calls)) fnCall := func(call gollm.ToolCall) gollm.ToolCallResponse { str, err := a.ToolBox.Execute(gollm.NewContext(ctx, a.req, &choice, &call), call) if err != nil { return gollm.ToolCallResponse{ ID: call.ID, Error: err, } } return gollm.ToolCallResponse{ ID: call.ID, Result: str, } } for _, call := range choice.Calls { go func(call gollm.ToolCall) { var arg any var err error if a.OnNewFunction != nil { arg, err = a.OnNewFunction(ctx, call.FunctionCall.Name, choice.Content, call.FunctionCall.Arguments) if err != nil { callsOutput <- gollm.ToolCallResponse{ ID: call.ID, Error: err, } return } } callRes := fnCall(call) if a.OnFunctionFinished != nil { err = a.OnFunctionFinished(ctx, call.FunctionCall.Name, choice.Content, call.FunctionCall.Arguments, callRes.Result, callRes.Error, arg) if err != nil { callsOutput <- gollm.ToolCallResponse{ ID: call.ID, Error: err, } return } } callsOutput <- callRes }(call) } var answers []gollm.ToolCallResponse for i := 0; i < len(choice.Calls); i++ { result := <-callsOutput answers = append(answers, result) } close(callsOutput) slog.Info("generating new request", "answers", answers, "choice", choice) a.req = gollm.NewContext(ctx, a.req, &choice, nil).ToNewRequest(answers...) return nil }