104 lines
2.6 KiB
Go
104 lines
2.6 KiB
Go
package gollama
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
|
|
"github.com/ollama/ollama/api"
|
|
)
|
|
|
|
// Message is a chat message with role and content.
|
|
type Message struct {
|
|
Role string
|
|
Content string
|
|
}
|
|
|
|
// User returns a user message.
|
|
func User(content string) Message { return Message{Role: "user", Content: content} }
|
|
|
|
// System returns a system message.
|
|
func System(content string) Message { return Message{Role: "system", Content: content} }
|
|
|
|
// Assistant returns an assistant message.
|
|
func Assistant(content string) Message { return Message{Role: "assistant", Content: content} }
|
|
|
|
func toAPIMessages(msgs []Message) []api.Message {
|
|
out := make([]api.Message, len(msgs))
|
|
for i, m := range msgs {
|
|
out[i] = api.Message{Role: m.Role, Content: m.Content}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func fromAPIMessage(m api.Message) Message {
|
|
return Message{Role: m.Role, Content: m.Content}
|
|
}
|
|
|
|
// Chat sends a single user message.
|
|
func (c *Client) Chat(ctx context.Context, message string, opts ...CallOption) (*Result, error) {
|
|
return c.ChatMessages(ctx, []Message{User(message)}, opts...)
|
|
}
|
|
|
|
// ChatWithSystem sends a system and user message.
|
|
func (c *Client) ChatWithSystem(ctx context.Context, system, message string, opts ...CallOption) (*Result, error) {
|
|
return c.ChatMessages(ctx, []Message{System(system), User(message)}, opts...)
|
|
}
|
|
|
|
// ChatMessages sends an explicit message list (history + new turns).
|
|
func (c *Client) ChatMessages(ctx context.Context, messages []Message, opts ...CallOption) (*Result, error) {
|
|
cfg, err := c.mergeCallConfig(opts...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
optMap, err := cfg.optionsMap()
|
|
if err != nil {
|
|
return nil, wrapError("chat", cfg.model, err)
|
|
}
|
|
|
|
msgs := toAPIMessages(messages)
|
|
if cfg.system != "" {
|
|
msgs = append([]api.Message{{Role: "system", Content: cfg.system}}, msgs...)
|
|
}
|
|
|
|
stream := false
|
|
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,
|
|
}
|
|
|
|
var text, thinking string
|
|
var final api.ChatResponse
|
|
err = c.api.Chat(ctx, req, func(resp api.ChatResponse) error {
|
|
text += resp.Message.Content
|
|
if resp.Message.Thinking != "" {
|
|
thinking += resp.Message.Thinking
|
|
}
|
|
if resp.Done {
|
|
final = resp
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, wrapError("chat", cfg.model, err)
|
|
}
|
|
|
|
if text == "" {
|
|
text = final.Message.Content
|
|
}
|
|
|
|
return &Result{
|
|
Text: strings.TrimSpace(text),
|
|
Model: cfg.model,
|
|
Thinking: thinking,
|
|
Metrics: final.Metrics,
|
|
Raw: &final,
|
|
}, nil
|
|
}
|