99 lines
2.6 KiB
Go
99 lines
2.6 KiB
Go
package gollama
|
|
|
|
import "context"
|
|
|
|
// SessionOption configures a conversation session.
|
|
type SessionOption func(*Session)
|
|
|
|
// Session maintains multi-turn chat history for a model.
|
|
type Session struct {
|
|
client *Client
|
|
model string
|
|
system string
|
|
opts OptionSet
|
|
messages []Message
|
|
maxHistory int // max user+assistant pairs; 0 = unlimited
|
|
}
|
|
|
|
// WithSessionSystem sets the session system prompt.
|
|
func WithSessionSystem(system string) SessionOption {
|
|
return func(s *Session) { s.system = system }
|
|
}
|
|
|
|
// WithSessionOptions sets default options for session sends.
|
|
func WithSessionOptions(opts OptionSet) SessionOption {
|
|
return func(s *Session) { s.opts = opts }
|
|
}
|
|
|
|
// WithMaxHistory limits stored turns (user+assistant pairs); oldest trimmed first.
|
|
func WithMaxHistory(pairs int) SessionOption {
|
|
return func(s *Session) { s.maxHistory = pairs }
|
|
}
|
|
|
|
// NewSession creates a stateful chat session for the given model.
|
|
func (c *Client) NewSession(model string, opts ...SessionOption) (*Session, error) {
|
|
if model == "" {
|
|
model = c.defaultModel
|
|
}
|
|
if model == "" {
|
|
return nil, errNoModel()
|
|
}
|
|
s := &Session{client: c, model: model}
|
|
for _, opt := range opts {
|
|
opt(s)
|
|
}
|
|
if s.system != "" {
|
|
s.messages = append(s.messages, System(s.system))
|
|
}
|
|
return s, nil
|
|
}
|
|
|
|
// Send appends a user message, calls the model, and stores the assistant reply.
|
|
func (s *Session) Send(ctx context.Context, message string, opts ...CallOption) (*Result, error) {
|
|
s.messages = append(s.messages, User(message))
|
|
callOpts := append([]CallOption{WithModel(s.model), WithOptions(s.opts)}, opts...)
|
|
result, err := s.client.ChatMessages(ctx, s.messages, callOpts...)
|
|
if err != nil {
|
|
// roll back user message on failure
|
|
s.messages = s.messages[:len(s.messages)-1]
|
|
return nil, err
|
|
}
|
|
s.messages = append(s.messages, Assistant(result.Text))
|
|
s.trimHistory()
|
|
return result, nil
|
|
}
|
|
|
|
// Messages returns a copy of the session history.
|
|
func (s *Session) Messages() []Message {
|
|
out := make([]Message, len(s.messages))
|
|
copy(out, s.messages)
|
|
return out
|
|
}
|
|
|
|
// Reset clears conversation history but keeps the system prompt.
|
|
func (s *Session) Reset() {
|
|
if s.system != "" {
|
|
s.messages = []Message{System(s.system)}
|
|
return
|
|
}
|
|
s.messages = nil
|
|
}
|
|
|
|
func (s *Session) trimHistory() {
|
|
if s.maxHistory <= 0 {
|
|
return
|
|
}
|
|
// count non-system messages as pairs
|
|
start := 0
|
|
if len(s.messages) > 0 && s.messages[0].Role == "system" {
|
|
start = 1
|
|
}
|
|
conv := s.messages[start:]
|
|
maxMsgs := s.maxHistory * 2
|
|
if len(conv) <= maxMsgs {
|
|
return
|
|
}
|
|
trim := len(conv) - maxMsgs
|
|
s.messages = append(s.messages[:start], conv[trim:]...)
|
|
}
|