79 lines
1.9 KiB
Go
79 lines
1.9 KiB
Go
package gollama
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
|
|
"github.com/ollama/ollama/api"
|
|
)
|
|
|
|
// Complete sends a single prompt and returns the full response.
|
|
func (c *Client) Complete(ctx context.Context, prompt string, opts ...CallOption) (*Result, error) {
|
|
return c.generate(ctx, prompt, false, opts...)
|
|
}
|
|
|
|
// CompleteWithSystem sends a system prompt and user prompt.
|
|
func (c *Client) CompleteWithSystem(ctx context.Context, system, prompt string, opts ...CallOption) (*Result, error) {
|
|
opts = append([]CallOption{WithSystem(system)}, opts...)
|
|
return c.generate(ctx, prompt, false, opts...)
|
|
}
|
|
|
|
// CompleteRaw sends a prompt without template formatting.
|
|
func (c *Client) CompleteRaw(ctx context.Context, prompt string, opts ...CallOption) (*Result, error) {
|
|
return c.generate(ctx, prompt, true, opts...)
|
|
}
|
|
|
|
func (c *Client) generate(ctx context.Context, prompt string, raw bool, callOpts ...CallOption) (*Result, error) {
|
|
cfg, err := c.mergeCallConfig(callOpts...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
optMap, err := cfg.optionsMap()
|
|
if err != nil {
|
|
return nil, wrapError("complete", cfg.model, err)
|
|
}
|
|
|
|
stream := false
|
|
req := &api.GenerateRequest{
|
|
Model: cfg.model,
|
|
Prompt: prompt,
|
|
System: cfg.system,
|
|
Raw: raw,
|
|
Stream: &stream,
|
|
Options: optMap,
|
|
KeepAlive: cfg.keepAlive,
|
|
Think: cfg.think,
|
|
Format: cfg.format,
|
|
Truncate: cfg.truncate,
|
|
Shift: cfg.shift,
|
|
}
|
|
|
|
var text, thinking string
|
|
var final api.GenerateResponse
|
|
err = c.api.Generate(ctx, req, func(resp api.GenerateResponse) error {
|
|
text += resp.Response
|
|
thinking += resp.Thinking
|
|
if resp.Done {
|
|
final = resp
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, wrapError("complete", cfg.model, err)
|
|
}
|
|
if text == "" {
|
|
text = final.Response
|
|
}
|
|
if thinking == "" {
|
|
thinking = final.Thinking
|
|
}
|
|
|
|
return &Result{
|
|
Text: strings.TrimSpace(text),
|
|
Model: cfg.model,
|
|
Thinking: thinking,
|
|
Metrics: final.Metrics,
|
|
Raw: &final,
|
|
}, nil
|
|
}
|