82 lines
1.9 KiB
Go
82 lines
1.9 KiB
Go
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,
|
|
})
|
|
}))
|
|
}
|