Updated tool calling

Author: djthorpe

cmd/agent/chat.go

@@ -15,15 +15,16 @@ import (
 ////////////////////////////////////////////////////////////////////////////////
 // TYPES
 
-type GenerateCmd struct {
+type ChatCmd struct {
 	Model    string `arg:"" help:"Model name"`
 	NoStream bool   `flag:"nostream" help:"Disable streaming"`
+	System   string `flag:"system" help:"Set the system prompt"`
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
-func (cmd *GenerateCmd) Run(globals *Globals) error {
+func (cmd *ChatCmd) Run(globals *Globals) error {
 	return runagent(globals, func(ctx context.Context, client llm.Agent) error {
 		// Get the model
 		a, ok := client.(*agent.Agent)
@@ -35,8 +36,22 @@ func (cmd *GenerateCmd) Run(globals *Globals) error {
 			return err
 		}
 
+		// Set the options
+		opts := []llm.Opt{}
+		if !cmd.NoStream {
+			opts = append(opts, llm.WithStream(func(cc llm.ContextContent) {
+				fmt.Println(cc)
+			}))
+		}
+		if cmd.System != "" {
+			opts = append(opts, llm.WithSystemPrompt(cmd.System))
+		}
+		if globals.toolkit != nil {
+			opts = append(opts, llm.WithToolKit(globals.toolkit))
+		}
+
 		// Create a session
-		session := model.Context()
+		session := model.Context(opts...)
 
 		// Continue looping until end of input
 		for {
@@ -57,7 +72,16 @@ func (cmd *GenerateCmd) Run(globals *Globals) error {
 			if err := session.FromUser(ctx, input); err != nil {
 				return err
 			}
+
 			fmt.Println(session.Text())
+
+			// If there are tool calls, then do these
+			calls := session.ToolCalls()
+			if results, err := globals.toolkit.Run(ctx, calls...); err != nil {
+				return err
+			} else {
+				fmt.Println(results)
+			}
 		}
 	})
 }

cmd/agent/main.go

@@ -60,7 +60,7 @@ type CLI struct {
 
 	// Commands
 	Download DownloadModelCmd `cmd:"" help:"Download a model"`
-	Generate GenerateCmd      `cmd:"chat" help:"Start a chat session"`
+	Chat     ChatCmd          `cmd:"" help:"Start a chat session"`
 }
 
 ////////////////////////////////////////////////////////////////////////////////

context.go

@@ -5,8 +5,8 @@ import "context"
 //////////////////////////////////////////////////////////////////
 // TYPES
 
-// Context is fed to the agent to generate a response
-type Context interface {
+// ContextContent is the content of the last context message
+type ContextContent interface {
 	// Return the current session role, which can be system, assistant, user, tool, tool_result, ...
 	Role() string
 
@@ -15,6 +15,11 @@ type Context interface {
 
 	// Return the current session tool calls, or empty if no tool calls were made
 	ToolCalls() []ToolCall
+}
+
+// Context is fed to the agent to generate a response
+type Context interface {
+	ContextContent
 
 	// Generate a response from a user prompt (with attachments and
 	// other empheral options

opt.go

@@ -13,11 +13,12 @@ type Opt func(*Opts) error
 
 // set of options
 type Opts struct {
-	agents      map[string]Agent // Set of agents
-	toolkit     ToolKit          // Toolkit for tools
-	callback    func(Context)    // Streaming callback
-	attachments []*Attachment    // Attachments
-	options     map[string]any   // Additional options
+	agents      map[string]Agent     // Set of agents
+	toolkit     ToolKit              // Toolkit for tools
+	callback    func(ContextContent) // Streaming callback
+	attachments []*Attachment        // Attachments
+	system      string               // System prompt
+	options     map[string]any       // Additional options
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -45,10 +46,15 @@ func (o *Opts) ToolKit() ToolKit {
 }
 
 // Return the stream function
-func (o *Opts) StreamFn() func(Context) {
+func (o *Opts) StreamFn() func(ContextContent) {
 	return o.callback
 }
 
+// Return the system prompt
+func (o *Opts) SystemPrompt() string {
+	return o.system
+}
+
 // Return the array of registered agents
 func (o *Opts) Agents() []Agent {
 	result := make([]Agent, 0, len(o.agents))
@@ -102,6 +108,26 @@ func (o *Opts) GetBool(key string) bool {
 	return false
 }
 
+// Get an option value as an unsigned integer
+func (o *Opts) GetUint64(key string) uint64 {
+	if value, exists := o.options[key]; exists {
+		if v, ok := value.(uint64); ok {
+			return v
+		}
+	}
+	return 0
+}
+
+// Get an option value as a float64
+func (o *Opts) GetFloat64(key string) float64 {
+	if value, exists := o.options[key]; exists {
+		if v, ok := value.(float64); ok {
+			return v
+		}
+	}
+	return 0
+}
+
 // Get an option value as a duration
 func (o *Opts) GetDuration(key string) time.Duration {
 	if value, exists := o.options[key]; exists {
@@ -124,7 +150,7 @@ func WithToolKit(toolkit ToolKit) Opt {
 }
 
 // Set chat streaming function
-func WithStream(fn func(Context)) Opt {
+func WithStream(fn func(ContextContent)) Opt {
 	return func(o *Opts) error {
 		o.callback = fn
 		return nil
@@ -196,3 +222,11 @@ func WithTopK(v uint) Opt {
 		return nil
 	}
 }
+
+// Set system prompt
+func WithSystemPrompt(v string) Opt {
+	return func(o *Opts) error {
+		o.system = v
+		return nil
+	}
+}

pkg/agent/agent.go

@@ -56,7 +56,7 @@ func (m model) String() string {
 // PUBLIC METHODS
 
 // Return a list of tool names
-func (a *Agent) Tools() []string {
+func (a *Agent) ToolNames() []string {
 	if a.ToolKit() == nil {
 		return nil
 	}
@@ -67,11 +67,35 @@ func (a *Agent) Tools() []string {
 	return result
 }
 
+// Return a list of agent names
+func (a *Agent) AgentNames() []string {
+	var result []string
+	for _, a := range a.Agents() {
+		result = append(result, a.Name())
+	}
+	return result
+}
+
+// Return a list of agents
+func (a *Agent) AgentsWithName(name ...string) []llm.Agent {
+	all := a.Agents()
+	if len(name) == 0 {
+		return all
+	}
+	result := make([]llm.Agent, 0, len(name))
+	for _, a := range all {
+		if slices.Contains(name, a.Name()) {
+			result = append(result, a)
+		}
+	}
+	return result
+}
+
 // Return a comma-separated list of agent names
 func (a *Agent) Name() string {
 	var keys []string
-	for key := range a.Agents() {
-		keys = append(keys, key)
+	for _, agent := range a.Agents() {
+		keys = append(keys, agent.Name())
 	}
 	return strings.Join(keys, ",")
 }
@@ -82,22 +106,13 @@ func (a *Agent) Models(ctx context.Context) ([]llm.Model, error) {
 }
 
 // Return the models from list of agents
-func (a *Agent) ListModels(ctx context.Context, agents ...string) ([]llm.Model, error) {
+func (a *Agent) ListModels(ctx context.Context, names ...string) ([]llm.Model, error) {
 	var result error
 
-	// Ensure all agents are valid
+	// Gather models from agents
+	agents := a.AgentsWithName(names...)
+	models := make([]llm.Model, 0, len(agents)*10)
 	for _, agent := range agents {
-		if _, exists := a.agents[agent]; !exists {
-			result = errors.Join(result, llm.ErrNotFound.Withf("agent %q", agent))
-		}
-	}
-
-	// Gather models from all agents
-	models := make([]llm.Model, 0, 100)
-	for _, agent := range a.agents {
-		if len(agents) > 0 && !slices.Contains(agents, agent.Name()) {
-			continue
-		}
 		agentmodels, err := modelsForAgent(ctx, agent)
 		if err != nil {
 			result = errors.Join(result, err)
@@ -113,24 +128,12 @@ func (a *Agent) ListModels(ctx context.Context, agents ...string) ([]llm.Model,
 
 // Return a model by name. If no agents are specified, then all agents are considered.
 // If multiple agents are specified, then the first model found is returned.
-func (a *Agent) GetModel(ctx context.Context, name string, agents ...string) (llm.Model, error) {
-	if len(agents) == 0 {
-		for _, agent := range a.agents {
-			agents = append(agents, agent.Name())
-		}
-	}
-
-	// Ensure all agents are valid
+func (a *Agent) GetModel(ctx context.Context, name string, agentnames ...string) (llm.Model, error) {
 	var result error
-	for _, agent := range agents {
-		if _, exists := a.agents[agent]; !exists {
-			result = errors.Join(result, llm.ErrNotFound.Withf("agent %q", agent))
-		}
-	}
 
-	// Gather models from agents
+	agents := a.AgentsWithName(agentnames...)
 	for _, agent := range agents {
-		models, err := modelsForAgent(ctx, a.agents[agent], name)
+		models, err := modelsForAgent(ctx, agent, name)
 		if err != nil {
 			result = errors.Join(result, err)
 			continue

pkg/agent/opt.go

@@ -1,7 +1,6 @@
 package agent
 
 import (
-
 	// Packages
 	client "github.com/mutablelogic/go-client"
 	llm "github.com/mutablelogic/go-llm"
@@ -13,23 +12,23 @@ import (
 // PUBLIC METHODS
 
 func WithOllama(endpoint string, opts ...client.ClientOpt) llm.Opt {
-	return func(o any) error {
+	return func(o *llm.Opts) error {
 		client, err := ollama.New(endpoint, opts...)
 		if err != nil {
 			return err
 		} else {
-			return llm.WithAgent(client)
+			return llm.WithAgent(client)(o)
 		}
 	}
 }
 
 func WithAnthropic(key string, opts ...client.ClientOpt) llm.Opt {
-	return func(o any) error {
+	return func(o *llm.Opts) error {
 		client, err := anthropic.New(key, opts...)
 		if err != nil {
 			return err
 		} else {
-			return llm.WithAgent(client)
+			return llm.WithAgent(client)(o)
 		}
 	}
 }

pkg/anthropic/client.go

@@ -4,7 +4,6 @@ anthropic implements an API client for anthropic (https://docs.anthropic.com/en/
 package anthropic
 
 import (
-
 	// Packages
 	client "github.com/mutablelogic/go-client"
 	llm "github.com/mutablelogic/go-llm"