Updated

Author: djthorpe

agent.go

@@ -11,4 +11,8 @@ type Agent interface {
 
 	// Return the models
 	Models(context.Context) ([]Model, error)
+
+	// Return a model by name, or nil if not found.
+	// Panics on error.
+	Model(context.Context, string) Model
 }

pkg/agent/agent.go

@@ -105,6 +105,15 @@ func (a *Agent) Models(ctx context.Context) ([]llm.Model, error) {
 	return a.ListModels(ctx)
 }
 
+// Return a model
+func (a *Agent) Model(ctx context.Context, name string) llm.Model {
+	model, err := a.GetModel(ctx, name)
+	if err != nil {
+		panic(err)
+	}
+	return model
+}
+
 // Return the models from list of agents
 func (a *Agent) ListModels(ctx context.Context, names ...string) ([]llm.Model, error) {
 	var result error

pkg/anthropic/client.go

@@ -14,6 +14,7 @@ import (
 
 type Client struct {
 	*client.Client
+	cache map[string]llm.Model
 }
 
 var _ llm.Agent = (*Client)(nil)
@@ -41,7 +42,10 @@ func New(ApiKey string, opts ...client.ClientOpt) (*Client, error) {
 	}
 
 	// Return the client
-	return &Client{client}, nil
+	return &Client{
+		Client: client,
+		cache:  make(map[string]llm.Model),
+	}, nil
 }
 
 ///////////////////////////////////////////////////////////////////////////////

pkg/anthropic/model.go

@@ -34,7 +34,38 @@ type ModelMeta struct {
 
 // Agent interface
 func (anthropic *Client) Models(ctx context.Context) ([]llm.Model, error) {
-	return anthropic.ListModels(ctx)
+	// Cache models
+	if len(anthropic.cache) == 0 {
+		models, err := anthropic.ListModels(ctx)
+		if err != nil {
+			return nil, err
+		}
+		for _, model := range models {
+			name := model.Name()
+			anthropic.cache[name] = model
+		}
+	}
+
+	// Return models
+	result := make([]llm.Model, 0, len(anthropic.cache))
+	for _, model := range anthropic.cache {
+		result = append(result, model)
+	}
+	return result, nil
+}
+
+// Agent interface
+func (anthropic *Client) Model(ctx context.Context, model string) llm.Model {
+	// Cache models
+	if len(anthropic.cache) == 0 {
+		_, err := anthropic.Models(ctx)
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	// Return model
+	return anthropic.cache[model]
 }
 
 // Get a model by name

pkg/anthropic/session_test.go

@@ -83,9 +83,10 @@ func Test_session_002(t *testing.T) {
 			t.FailNow()
 		}
 
-		err := toolkit.Run(context.TODO(), session.ToolCalls()...)
+		result, err := toolkit.Run(context.TODO(), session.ToolCalls()...)
 		if !assert.NoError(err) {
 			t.FailNow()
 		}
+		assert.NotEmpty(result)
 	})
 }

pkg/ollama/chat_test.go

@@ -31,7 +31,7 @@ func Test_chat_001(t *testing.T) {
 
 	t.Run("ChatStream", func(t *testing.T) {
 		assert := assert.New(t)
-		response, err := client.Chat(context.TODO(), model.UserPrompt("why is the sky blue?"), llm.WithStream(func(stream llm.Context) {
+		response, err := client.Chat(context.TODO(), model.UserPrompt("why is the sky blue?"), llm.WithStream(func(stream llm.ContextContent) {
 			t.Log(stream)
 		}))
 		if !assert.NoError(err) {

pkg/ollama/model.go

@@ -91,6 +91,15 @@ func (ollama *Client) Models(ctx context.Context) ([]llm.Model, error) {
 	return ollama.ListModels(ctx)
 }
 
+// Agent interface
+func (ollama *Client) Model(ctx context.Context, name string) llm.Model {
+	model, err := ollama.GetModel(ctx, name)
+	if err != nil {
+		panic(err)
+	}
+	return model
+}
+
 // List models
 func (ollama *Client) ListModels(ctx context.Context) ([]llm.Model, error) {
 	type respListModel struct {

pkg/ollama/session_test.go

@@ -28,7 +28,7 @@ func Test_session_001(t *testing.T) {
 	// Session with a single user prompt - streaming
 	t.Run("stream", func(t *testing.T) {
 		assert := assert.New(t)
-		session := model.Context(llm.WithStream(func(stream llm.Context) {
+		session := model.Context(llm.WithStream(func(stream llm.ContextContent) {
 			t.Log("SESSION DELTA", stream)
 		}))
 		assert.NotNil(session)