113 lines
2.1 KiB
Go
113 lines
2.1 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"gitea.stevedudenhoeffer.com/steve/answer/pkg/toolbox"
|
|
|
|
"github.com/tmc/langchaingo/llms"
|
|
)
|
|
|
|
type Context struct {
|
|
context.Context
|
|
|
|
Messages []llms.MessageContent
|
|
Agent *Agent
|
|
}
|
|
|
|
func From(ctx context.Context, a *Agent) *Context {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
|
|
return &Context{
|
|
Context: ctx,
|
|
Agent: a,
|
|
}
|
|
}
|
|
|
|
func (c *Context) WithMessage(m llms.MessageContent) *Context {
|
|
c.Messages = append(c.Messages, m)
|
|
return c
|
|
}
|
|
|
|
func (c *Context) WithMessages(m ...llms.MessageContent) *Context {
|
|
c.Messages = append(c.Messages, m...)
|
|
return c
|
|
}
|
|
|
|
func (c *Context) WithToolResults(r ...toolbox.ToolResult) *Context {
|
|
msg := llms.MessageContent{
|
|
Role: llms.ChatMessageTypeTool,
|
|
}
|
|
|
|
for _, v := range r {
|
|
res := v.Result
|
|
if v.Error != nil {
|
|
res = "error executing: " + v.Error.Error()
|
|
}
|
|
msg.Parts = append(msg.Parts, llms.ToolCallResponse{
|
|
ToolCallID: v.ID,
|
|
Name: v.Name,
|
|
Content: res,
|
|
})
|
|
}
|
|
|
|
return c.WithMessage(msg)
|
|
}
|
|
|
|
func (c *Context) WithAgent(a *Agent) *Context {
|
|
return &Context{
|
|
Context: c.Context,
|
|
Agent: a,
|
|
}
|
|
}
|
|
|
|
func (c *Context) WithCancel() (*Context, func()) {
|
|
ctx, cancel := context.WithCancel(c.Context)
|
|
return &Context{
|
|
Context: ctx,
|
|
Agent: c.Agent,
|
|
}, cancel
|
|
}
|
|
|
|
func (c *Context) WithDeadline(deadline time.Time) (*Context, func()) {
|
|
ctx, cancel := context.WithDeadline(c.Context, deadline)
|
|
return &Context{
|
|
Context: ctx,
|
|
Agent: c.Agent,
|
|
}, cancel
|
|
}
|
|
|
|
func (c *Context) WithTimeout(timeout time.Duration) (*Context, func()) {
|
|
ctx, cancel := context.WithTimeout(c.Context, timeout)
|
|
return &Context{
|
|
Context: ctx,
|
|
Agent: c.Agent,
|
|
}, cancel
|
|
}
|
|
|
|
func (c *Context) WithValue(key, value interface{}) *Context {
|
|
return &Context{
|
|
Context: context.WithValue(c.Context, key, value),
|
|
Agent: c.Agent,
|
|
}
|
|
}
|
|
|
|
func (c *Context) Done() <-chan struct{} {
|
|
return c.Context.Done()
|
|
}
|
|
|
|
func (c *Context) Err() error {
|
|
return c.Context.Err()
|
|
}
|
|
|
|
func (c *Context) Value(key interface{}) interface{} {
|
|
return c.Context.Value(key)
|
|
}
|
|
|
|
func (c *Context) Deadline() (time.Time, bool) {
|
|
return c.Context.Deadline()
|
|
}
|