First AI generated version

stream.go

@@ -0,0 +1,81 @@
+package gollama
+
+import (
+	"context"
+
+	"github.com/ollama/ollama/api"
+)
+
+// CompleteStream streams completion tokens through fn.
+func (c *Client) CompleteStream(ctx context.Context, prompt string, fn func(Chunk) error, opts ...CallOption) error {
+	cfg, err := c.mergeCallConfig(opts...)
+	if err != nil {
+		return err
+	}
+	optMap, err := cfg.optionsMap()
+	if err != nil {
+		return wrapError("complete_stream", cfg.model, err)
+	}
+
+	stream := true
+	req := &api.GenerateRequest{
+		Model:     cfg.model,
+		Prompt:    prompt,
+		System:    cfg.system,
+		Stream:    &stream,
+		Options:   optMap,
+		KeepAlive: cfg.keepAlive,
+		Think:     cfg.think,
+		Format:    cfg.format,
+		Truncate:  cfg.truncate,
+		Shift:     cfg.shift,
+	}
+
+	return wrapError("complete_stream", cfg.model, c.api.Generate(ctx, req, func(resp api.GenerateResponse) error {
+		return fn(Chunk{
+			Text:     resp.Response,
+			Thinking: resp.Thinking,
+			Done:     resp.Done,
+			Metrics:  resp.Metrics,
+		})
+	}))
+}
+
+// ChatStream streams chat response tokens through fn.
+func (c *Client) ChatStream(ctx context.Context, messages []Message, fn func(Chunk) error, opts ...CallOption) error {
+	cfg, err := c.mergeCallConfig(opts...)
+	if err != nil {
+		return err
+	}
+	optMap, err := cfg.optionsMap()
+	if err != nil {
+		return wrapError("chat_stream", cfg.model, err)
+	}
+
+	msgs := toAPIMessages(messages)
+	if cfg.system != "" {
+		msgs = append([]api.Message{{Role: "system", Content: cfg.system}}, msgs...)
+	}
+
+	stream := true
+	req := &api.ChatRequest{
+		Model:     cfg.model,
+		Messages:  msgs,
+		Stream:    &stream,
+		Options:   optMap,
+		KeepAlive: cfg.keepAlive,
+		Think:     cfg.think,
+		Format:    cfg.format,
+		Truncate:  cfg.truncate,
+		Shift:     cfg.shift,
+	}
+
+	return wrapError("chat_stream", cfg.model, c.api.Chat(ctx, req, func(resp api.ChatResponse) error {
+		return fn(Chunk{
+			Text:     resp.Message.Content,
+			Thinking: resp.Message.Thinking,
+			Done:     resp.Done,
+			Metrics:  resp.Metrics,
+		})
+	}))
+}