Updated Ollama

Author: djthorpe

pkg/mistral/message.go

@@ -3,6 +3,7 @@ package mistral
 import (
 	"encoding/json"
 
+	// Packages
 	"github.com/mutablelogic/go-llm"
 	"github.com/mutablelogic/go-llm/pkg/tool"
 )
@@ -23,13 +24,11 @@ type Message struct {
 
 type RoleContent struct {
 	Role    string `json:"role,omitempty"`         // assistant, user, tool, system
+	Content any    `json:"content,omitempty"`      // string or array of text, reference, image_url
 	Id      string `json:"tool_call_id,omitempty"` // tool call - when role is tool
 	Name    string `json:"name,omitempty"`         // function name - when role is tool
-	Content any    `json:"content,omitempty"`      // string or array of text, reference, image_url
 }
 
-var _ llm.Completion = (*Message)(nil)
-
 // Completion Variation
 type Completion struct {
 	Index   uint64   `json:"index"`
@@ -38,6 +37,8 @@ type Completion struct {
 	Reason  string   `json:"finish_reason,omitempty"`
 }
 
+var _ llm.Completion = (*Message)(nil)
+
 type Content struct {
 	Type        string                       `json:"type,omitempty"` // text, reference, image_url
 	*Text       `json:"text,omitempty"`      // text content

pkg/ollama/chat.go

@@ -13,13 +13,13 @@ import (
 ///////////////////////////////////////////////////////////////////////////////
 // TYPES
 
-// Chat Response
+// Chat Completion Response
 type Response struct {
-	Model     string      `json:"model"`
-	CreatedAt time.Time   `json:"created_at"`
-	Message   MessageMeta `json:"message"`
-	Done      bool        `json:"done"`
-	Reason    string      `json:"done_reason,omitempty"`
+	Model     string    `json:"model"`
+	CreatedAt time.Time `json:"created_at"`
+	Done      bool      `json:"done"`
+	Reason    string    `json:"done_reason,omitempty"`
+	Message   `json:"message"`
 	Metrics
 }
 
@@ -33,6 +33,8 @@ type Metrics struct {
 	EvalDuration       time.Duration `json:"eval_duration,omitempty"`
 }
 
+var _ llm.Completion = (*Response)(nil)
+
 ///////////////////////////////////////////////////////////////////////////////
 // STRINGIFY
 
@@ -49,34 +51,36 @@ func (r Response) String() string {
 
 type reqChat struct {
 	Model     string                 `json:"model"`
-	Messages  []*MessageMeta         `json:"messages"`
-	Tools     []ToolFunction         `json:"tools,omitempty"`
+	Messages  []*Message             `json:"messages"`
+	Tools     []llm.Tool             `json:"tools,omitempty"`
 	Format    string                 `json:"format,omitempty"`
 	Options   map[string]interface{} `json:"options,omitempty"`
 	Stream    bool                   `json:"stream"`
 	KeepAlive *time.Duration         `json:"keep_alive,omitempty"`
 }
 
-func (ollama *Client) Chat(ctx context.Context, prompt llm.Context, opts ...llm.Opt) (*Response, error) {
+func (ollama *Client) Chat(ctx context.Context, context llm.Context, opts ...llm.Opt) (*Response, error) {
+	// Apply options
 	opt, err := llm.ApplyOpts(opts...)
 	if err != nil {
 		return nil, err
 	}
 
 	// Append the system prompt at the beginning
-	seq := make([]*MessageMeta, 0, len(prompt.(*session).seq)+1)
-	if system := opt.SystemPrompt(); system != "" {
-		seq = append(seq, &MessageMeta{
-			Role:    "system",
-			Content: opt.SystemPrompt(),
-		})
+	messages := make([]*Message, 0, len(context.(*session).seq)+1)
+	//if system := opt.SystemPrompt(); system != "" {
+	//	messages = append(messages, systemPrompt(system))
+	//}
+
+	// Always append the first message of each completion
+	for _, message := range context.(*session).seq {
+		messages = append(messages, message)
 	}
-	seq = append(seq, prompt.(*session).seq...)
 
 	// Request
 	req, err := client.NewJSONRequest(reqChat{
-		Model:     prompt.(*session).model.Name(),
-		Messages:  seq,
+		Model:     context.(*session).model.Name(),
+		Messages:  messages,
 		Tools:     optTools(ollama, opt),
 		Format:    optFormat(opt),
 		Options:   optOptions(opt),
@@ -88,52 +92,53 @@ func (ollama *Client) Chat(ctx context.Context, prompt llm.Context, opts ...llm.
 	}
 
 	//  Response
-	var response, delta Response
-	if err := ollama.DoWithContext(ctx, req, &delta, client.OptPath("chat"), client.OptJsonStreamCallback(func(v any) error {
-		if v, ok := v.(*Response); !ok || v == nil {
-			return llm.ErrConflict.Withf("Invalid stream response: %v", v)
-		} else {
-			response.Model = v.Model
-			response.CreatedAt = v.CreatedAt
-			response.Message.Role = v.Message.Role
-			response.Message.Content += v.Message.Content
-			if v.Done {
-				response.Done = v.Done
-				response.Metrics = v.Metrics
-				response.Reason = v.Reason
+	var response Response
+	reqopts := []client.RequestOpt{
+		client.OptPath("chat"),
+	}
+	if optStream(ollama, opt) {
+		reqopts = append(reqopts, client.OptJsonStreamCallback(func(v any) error {
+			if v, ok := v.(*Response); !ok || v == nil {
+				return llm.ErrConflict.Withf("Invalid stream response: %v", v)
+			} else if err := streamEvent(&response, v); err != nil {
+				return err
 			}
-		}
-
-		//Call the chat callback
-		if optStream(ollama, opt) {
 			if fn := opt.StreamFn(); fn != nil {
 				fn(&response)
 			}
-		}
-		return nil
-	})); err != nil {
-		return nil, err
+			return nil
+		}))
 	}
 
-	// We return the delta or the response
-	if optStream(ollama, opt) {
-		return &response, nil
-	} else {
-		return &delta, nil
+	// Response
+	if err := ollama.DoWithContext(ctx, req, &response, reqopts...); err != nil {
+		return nil, err
 	}
-}
 
-///////////////////////////////////////////////////////////////////////////////
-// INTERFACE - CONTEXT CONTENT
-
-func (response Response) Role() string {
-	return response.Message.Role
+	// Return success
+	return &response, nil
 }
 
-func (response Response) Text() string {
-	return response.Message.Content
-}
+///////////////////////////////////////////////////////////////////////////////
+// PRIVATE METHODS
 
-func (response Response) ToolCalls() []llm.ToolCall {
+func streamEvent(response, delta *Response) error {
+	if delta.Model != "" {
+		response.Model = delta.Model
+	}
+	if !delta.CreatedAt.IsZero() {
+		response.CreatedAt = delta.CreatedAt
+	}
+	if delta.Message.RoleContent.Role != "" {
+		response.Message.RoleContent.Role = delta.Message.RoleContent.Role
+	}
+	if delta.Message.RoleContent.Content != "" {
+		response.Message.RoleContent.Content += delta.Message.RoleContent.Content
+	}
+	if delta.Done {
+		response.Done = delta.Done
+		response.Metrics = delta.Metrics
+		response.Reason = delta.Reason
+	}
 	return nil
 }

pkg/ollama/chat_test.go renamed to pkg/ollama/chat_test.go_old

@@ -1,7 +1,10 @@
 package ollama_test
 
 import (
+	"flag"
+	"log"
 	"os"
+	"strconv"
 	"testing"
 
 	// Packages
@@ -10,23 +13,46 @@ import (
 	assert "github.com/stretchr/testify/assert"
 )
 
-func Test_client_001(t *testing.T) {
-	assert := assert.New(t)
-	client, err := ollama.New(GetEndpoint(t), opts.OptTrace(os.Stderr, true))
-	if assert.NoError(err) {
-		assert.NotNil(client)
-		t.Log(client)
+///////////////////////////////////////////////////////////////////////////////
+// TEST SET-UP
+
+var (
+	client *ollama.Client
+)
+
+func TestMain(m *testing.M) {
+	var verbose bool
+
+	// Verbose output
+	flag.Parse()
+	if f := flag.Lookup("test.v"); f != nil {
+		if v, err := strconv.ParseBool(f.Value.String()); err == nil {
+			verbose = v
+		}
 	}
+
+	// Endpoint
+	endpoint_url := os.Getenv("OLLAMA_URL")
+	if endpoint_url == "" {
+		log.Print("OLLAMA_URL not set")
+		os.Exit(0)
+	}
+
+	// Create client
+	var err error
+	client, err = ollama.New(endpoint_url, opts.OptTrace(os.Stderr, verbose))
+	if err != nil {
+		log.Println(err)
+		os.Exit(-1)
+	}
+	os.Exit(m.Run())
 }
 
 ///////////////////////////////////////////////////////////////////////////////
-// ENVIRONMENT
+// TESTS
 
-func GetEndpoint(t *testing.T) string {
-	key := os.Getenv("OLLAMA_URL")
-	if key == "" {
-		t.Skip("OLLAMA_URL not set, skipping tests")
-		t.SkipNow()
-	}
-	return key
+func Test_client_001(t *testing.T) {
+	assert := assert.New(t)
+	assert.NotNil(client)
+	t.Log(client)
 }

pkg/ollama/client_test.go

@@ -2,24 +2,17 @@ package ollama_test
 
 import (
 	"context"
-	"os"
 	"testing"
 
 	// Packages
-	opts "github.com/mutablelogic/go-client"
-	ollama "github.com/mutablelogic/go-llm/pkg/ollama"
+
 	assert "github.com/stretchr/testify/assert"
 )
 
 func Test_embed_001(t *testing.T) {
-	client, err := ollama.New(GetEndpoint(t), opts.OptTrace(os.Stderr, true))
-	if err != nil {
-		t.FailNow()
-	}
-
 	t.Run("Embedding", func(t *testing.T) {
 		assert := assert.New(t)
-		embedding, err := client.GenerateEmbedding(context.TODO(), "qwen:0.5b", []string{"world"})
+		embedding, err := client.GenerateEmbedding(context.TODO(), "qwen:0.5b", []string{"hello, world"})
 		if !assert.NoError(err) {
 			t.FailNow()
 		}

pkg/ollama/embedding_test.go

@@ -1,36 +1,80 @@
 package ollama
 
 import (
+	"fmt"
+
+	// Packages
 	llm "github.com/mutablelogic/go-llm"
+	tool "github.com/mutablelogic/go-llm/pkg/tool"
 )
 
 ///////////////////////////////////////////////////////////////////////////////
 // TYPES
 
-// Chat Message
-type MessageMeta struct {
-	Role         string     `json:"role"`
-	Content      string     `json:"content,omitempty"`
-	FunctionName string     `json:"name,omitempty"`       // Function name for a tool result
-	Images       []Data     `json:"images,omitempty"`     // Image attachments
-	ToolCalls    []ToolCall `json:"tool_calls,omitempty"` // Tool calls from the assistant
+// Message with text or object content
+type Message struct {
+	RoleContent
+	ToolCallArray `json:"tool_calls,omitempty"`
+}
+
+type RoleContent struct {
+	Role    string `json:"role,omitempty"`    // assistant, user, tool, system
+	Content string `json:"content,omitempty"` // string or array of text, reference, image_url
+	Images  []Data `json:"images,omitempty"`  // Image attachments
+	ToolResult
 }
 
+// A set of tool calls
+type ToolCallArray []ToolCall
+
 type ToolCall struct {
+	Type     string           `json:"type"` // function
 	Function ToolCallFunction `json:"function"`
 }
 
 type ToolCallFunction struct {
 	Index     int            `json:"index,omitempty"`
 	Name      string         `json:"name"`
-	Arguments map[string]any `json:"arguments"`
+	Arguments map[string]any `json:"arguments,omitempty"`
 }
 
 // Data represents the raw binary data of an image file.
 type Data []byte
 
-// ToolFunction
-type ToolFunction struct {
-	Type     string   `json:"type"` // function
-	Function llm.Tool `json:"function"`
+// ToolResult
+type ToolResult struct {
+	Name string `json:"name,omitempty"` // function name - when role is tool
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// PUBLIC METHODS - MESSAGE
+
+func (m Message) Num() int {
+	return 1
+}
+
+func (m Message) Role() string {
+	return m.RoleContent.Role
+}
+
+func (m Message) Text(index int) string {
+	if index != 0 {
+		return ""
+	}
+	return m.Content
+}
+
+func (m Message) ToolCalls(index int) []llm.ToolCall {
+	if index != 0 {
+		return nil
+	}
+
+	// Make the tool calls
+	calls := make([]llm.ToolCall, 0, len(m.ToolCallArray))
+	for _, call := range m.ToolCallArray {
+		calls = append(calls, tool.NewCall(fmt.Sprint(call.Function.Index), call.Function.Name, call.Function.Arguments))
+	}
+
+	// Return success
+	return calls
 }