First AI generated version

session.go

@@ -0,0 +1,98 @@
+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:]...)
+}