ollama tool streaming hack

adrianliechti · Oct 8, 2024 · 2c0e018 · 2c0e018
1 parent 45217e9
commit 2c0e018
Show file tree

Hide file tree

Showing 7 changed files with 79 additions and 30 deletions.
diff --git a/config/config_chain.go b/config/config_chain.go
@@ -172,10 +172,6 @@ func agentChain(cfg chainConfig, context chainContext) (chain.Provider, error) {
 		options = append(options, agent.WithMessages(context.Messages...))
 	}
 
-	if cfg.Temperature != nil {
-		options = append(options, agent.WithTemperature(*cfg.Temperature))
-	}
-
 	return agent.New(options...)
 }
 

diff --git a/pkg/chain/agent/chain.go b/pkg/chain/agent/chain.go
@@ -20,8 +20,6 @@ type Chain struct {
 
 	tools    []tool.Tool
 	messages []provider.Message
-
-	temperature *float32
 }
 
 type Option func(*Chain)
@@ -58,12 +56,6 @@ func WithTools(tool ...tool.Tool) Option {
 	}
 }
 
-func WithTemperature(temperature float32) Option {
-	return func(c *Chain) {
-		c.temperature = &temperature
-	}
-}
-
 func (c *Chain) Complete(ctx context.Context, messages []provider.Message, options *provider.CompleteOptions) (*provider.Completion, error) {
 	if options == nil {
 		options = new(provider.CompleteOptions)
@@ -104,14 +96,14 @@ func (c *Chain) Complete(ctx context.Context, messages []provider.Message, optio
 
 	for {
 		inputOptions := &provider.CompleteOptions{
-			Temperature: options.Temperature,
-			Tools:       to.Values(inputTools),
-
 			Stream: options.Stream,
-		}
 
-		if inputOptions.Temperature == nil {
-			inputOptions.Temperature = c.temperature
+			Tools: to.Values(inputTools),
+
+			MaxTokens:   options.MaxTokens,
+			Temperature: options.Temperature,
+
+			Format: options.Format,
 		}
 
 		completion, err := c.completer.Complete(ctx, input, inputOptions)

diff --git a/pkg/chain/reasoning/chain.go b/pkg/chain/reasoning/chain.go
@@ -79,6 +79,7 @@ func (c *Chain) Complete(ctx context.Context, messages []provider.Message, optio
 	})
 
 	inputOptions := &provider.CompleteOptions{
+		MaxTokens:   options.MaxTokens,
 		Temperature: options.Temperature,
 
 		Format: provider.CompletionFormatJSON,

diff --git a/pkg/provider/custom/completer.go b/pkg/provider/custom/completer.go
@@ -51,7 +51,7 @@ func NewCompleter(url string, options ...Option) (*Completer, error) {
 
 func (c *Completer) Complete(ctx context.Context, messages []provider.Message, options *provider.CompleteOptions) (*provider.Completion, error) {
 	if options == nil {
-		options = &provider.CompleteOptions{}
+		options = new(provider.CompleteOptions)
 	}
 
 	req := &CompletionRequest{

diff --git a/pkg/provider/ollama/completer.go b/pkg/provider/ollama/completer.go
@@ -1,12 +1,18 @@
 package ollama
 
 import (
+	"context"
 	"strings"
 
+	"github.com/adrianliechti/llama/pkg/provider"
 	"github.com/adrianliechti/llama/pkg/provider/openai"
 )
 
-type Completer = openai.Completer
+var _ provider.Completer = (*Completer)(nil)
+
+type Completer struct {
+	completer *openai.Completer
+}
 
 func NewCompleter(url, model string, options ...Option) (*Completer, error) {
 	if url == "" {
@@ -16,11 +22,61 @@ func NewCompleter(url, model string, options ...Option) (*Completer, error) {
 	url = strings.TrimRight(url, "/")
 	url = strings.TrimSuffix(url, "/v1")
 
-	c := &Config{}
+	cfg := &Config{}
 
 	for _, option := range options {
-		option(c)
+		option(cfg)
+	}
+
+	opts := []openai.Option{}
+
+	if cfg.client != nil {
+		opts = append(opts, openai.WithClient(cfg.client))
+	}
+
+	completer, err := openai.NewCompleter(url+"/v1", model, opts...)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &Completer{
+		completer: completer,
+	}, nil
+}
+
+func (c *Completer) Complete(ctx context.Context, messages []provider.Message, options *provider.CompleteOptions) (*provider.Completion, error) {
+	if options == nil {
+		options = new(provider.CompleteOptions)
+	}
+
+	inputOptions := &provider.CompleteOptions{
+		Stream: options.Stream,
+
+		Stop:  options.Stop,
+		Tools: options.Tools,
+
+		MaxTokens:   options.MaxTokens,
+		Temperature: options.Temperature,
+
+		Format: options.Format,
+	}
+
+	if len(options.Tools) > 0 {
+		inputOptions.Stream = nil
+	}
+
+	result, err := c.completer.Complete(ctx, messages, inputOptions)
+
+	if err != nil {
+		return nil, err
+	}
+
+	if inputOptions.Stream == nil && options.Stream != nil {
+		if err = options.Stream(ctx, *result); err != nil {
+			return nil, err
+		}
 	}
 
-	return openai.NewCompleter(url+"/v1", model, c.options...)
+	return result, nil
 }
diff --git a/pkg/provider/ollama/config.go b/pkg/provider/ollama/config.go
@@ -2,18 +2,16 @@ package ollama
 
 import (
 	"net/http"
-
-	"github.com/adrianliechti/llama/pkg/provider/openai"
 )
 
 type Config struct {
-	options []openai.Option
+	client *http.Client
 }
 
 type Option func(*Config)
 
 func WithClient(client *http.Client) Option {
 	return func(c *Config) {
-		c.options = append(c.options, openai.WithClient(client))
+		c.client = client
 	}
 }
diff --git a/pkg/provider/ollama/embedder.go b/pkg/provider/ollama/embedder.go
@@ -16,11 +16,17 @@ func NewEmbedder(url, model string, options ...Option) (*Embedder, error) {
 	url = strings.TrimRight(url, "/")
 	url = strings.TrimSuffix(url, "/v1")
 
-	c := &Config{}
+	cfg := &Config{}
 
 	for _, option := range options {
-		option(c)
+		option(cfg)
 	}
 
-	return openai.NewEmbedder(url+"/v1", model, c.options...)
+	opts := []openai.Option{}
+
+	if cfg.client != nil {
+		opts = append(opts, openai.WithClient(cfg.client))
+	}
+
+	return openai.NewEmbedder(url+"/v1", model, opts...)
 }