diff --git a/internal/conversation/session.go b/internal/conversation/session.go new file mode 100644 index 0000000..d07d837 --- /dev/null +++ b/internal/conversation/session.go @@ -0,0 +1,49 @@ +package conversation + +import "slices" + +// Role identifies a participant in a conversation. +type Role string + +const ( + RoleSystem Role = "system" + RoleUser Role = "user" + RoleAssistant Role = "assistant" +) + +// Message represents one entry in a conversation. +type Message struct { + Role Role + Content string +} + +// Session owns the ordered history for one conversation. +type Session struct { + messages []Message +} + +// NewSession constructs a session with a copy of the supplied initial context. +func NewSession(initial ...Message) *Session { + return &Session{messages: slices.Clone(initial)} +} + +// AppendUser appends a user message to the conversation. +func (session *Session) AppendUser(content string) { + session.messages = append(session.messages, Message{ + Role: RoleUser, + Content: content, + }) +} + +// AppendAssistant appends an assistant message to the conversation. +func (session *Session) AppendAssistant(content string) { + session.messages = append(session.messages, Message{ + Role: RoleAssistant, + Content: content, + }) +} + +// Messages returns a copy of the conversation history. +func (session *Session) Messages() []Message { + return slices.Clone(session.messages) +} diff --git a/internal/conversation/session_test.go b/internal/conversation/session_test.go new file mode 100644 index 0000000..830783d --- /dev/null +++ b/internal/conversation/session_test.go @@ -0,0 +1,253 @@ +package conversation + +import ( + "slices" + "testing" +) + +func TestNewSession(t *testing.T) { + tests := []struct { + name string + initial []Message + want []Message + }{ + { + name: "starts empty", + }, + { + name: "preserves initial context", + initial: []Message{ + { + Role: RoleSystem, + Content: "You are Ghost.", + }, + { + Role: RoleUser, + Content: "Establish uplink.", + }, + }, + want: []Message{ + { + Role: RoleSystem, + Content: "You are Ghost.", + }, + { + Role: RoleUser, + Content: "Establish uplink.", + }, + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + initial := slices.Clone(test.initial) + session := NewSession(initial...) + + if len(initial) > 0 { + initial[0] = Message{ + Role: RoleAssistant, + Content: "mutated outside the session", + } + } + + if got := session.Messages(); !slices.Equal(got, test.want) { + t.Errorf("Messages() = %#v, want %#v", got, test.want) + } + }) + } + + t.Run("sessions do not share history", func(t *testing.T) { + first := NewSession() + second := NewSession() + + first.AppendUser("first session") + + if got := second.Messages(); len(got) != 0 { + t.Errorf("second Messages() = %#v, want empty history", got) + } + }) +} + +func TestSessionAppendUser(t *testing.T) { + tests := []struct { + name string + initial []Message + content string + want []Message + }{ + { + name: "appends to empty history", + content: "Identify yourself.", + want: []Message{ + { + Role: RoleUser, + Content: "Identify yourself.", + }, + }, + }, + { + name: "appends after initial context", + initial: []Message{ + { + Role: RoleSystem, + Content: "You are Ghost.", + }, + }, + content: "Identify yourself.", + want: []Message{ + { + Role: RoleSystem, + Content: "You are Ghost.", + }, + { + Role: RoleUser, + Content: "Identify yourself.", + }, + }, + }, + { + name: "preserves content verbatim", + content: " signal with surrounding space ", + want: []Message{ + { + Role: RoleUser, + Content: " signal with surrounding space ", + }, + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + session := NewSession(test.initial...) + session.AppendUser(test.content) + + if got := session.Messages(); !slices.Equal(got, test.want) { + t.Errorf("Messages() = %#v, want %#v", got, test.want) + } + }) + } +} + +func TestSessionAppendAssistant(t *testing.T) { + tests := []struct { + name string + initial []Message + content string + want []Message + }{ + { + name: "appends to empty history", + content: "Ghost online.", + want: []Message{ + { + Role: RoleAssistant, + Content: "Ghost online.", + }, + }, + }, + { + name: "appends after user message", + initial: []Message{ + { + Role: RoleUser, + Content: "Identify yourself.", + }, + }, + content: "Ghost online.", + want: []Message{ + { + Role: RoleUser, + Content: "Identify yourself.", + }, + { + Role: RoleAssistant, + Content: "Ghost online.", + }, + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + session := NewSession(test.initial...) + session.AppendAssistant(test.content) + + if got := session.Messages(); !slices.Equal(got, test.want) { + t.Errorf("Messages() = %#v, want %#v", got, test.want) + } + }) + } +} + +func TestSessionMessages(t *testing.T) { + tests := []struct { + name string + initial []Message + build func(*Session) + want []Message + }{ + { + name: "returns empty history", + }, + { + name: "returns complete history in order", + initial: []Message{ + { + Role: RoleSystem, + Content: "You are Ghost.", + }, + }, + build: func(session *Session) { + session.AppendUser("First signal.") + session.AppendAssistant("First response.") + session.AppendUser("Second signal.") + }, + want: []Message{ + { + Role: RoleSystem, + Content: "You are Ghost.", + }, + { + Role: RoleUser, + Content: "First signal.", + }, + { + Role: RoleAssistant, + Content: "First response.", + }, + { + Role: RoleUser, + Content: "Second signal.", + }, + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + session := NewSession(test.initial...) + + if test.build != nil { + test.build(session) + } + + got := session.Messages() + if !slices.Equal(got, test.want) { + t.Fatalf("Messages() = %#v, want %#v", got, test.want) + } + + if len(got) > 0 { + got[0] = Message{ + Role: RoleAssistant, + Content: "mutated returned history", + } + } + + if gotAgain := session.Messages(); !slices.Equal(gotAgain, test.want) { + t.Errorf("Messages() after caller mutation = %#v, want %#v", gotAgain, test.want) + } + }) + } +} diff --git a/internal/ollama/client.go b/internal/ollama/client.go index a7da414..edfbd83 100644 --- a/internal/ollama/client.go +++ b/internal/ollama/client.go @@ -9,39 +9,32 @@ import ( "net/http" "net/url" "strings" + + "github.com/theantichris/ghost/v4/internal/conversation" ) const DefaultURL = "http://localhost:11434/api" const maxErrBodySize = 64 << 10 -// Role identifies a participant in an Ollama chat conversation. -type Role string - -const ( - RoleUser Role = "user" - RoleAssistant Role = "assistant" -) - -// Message represents one message in an Ollama chat conversation. -type Message struct { - Role Role `json:"role"` - Content string `json:"content"` -} - // Client communicates with the Ollama HTTP API. type Client struct { baseURL *url.URL httpClient *http.Client } +type chatMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + type chatRequest struct { - Model string `json:"model"` - Messages []Message `json:"messages"` - Stream bool `json:"stream"` + Model string `json:"model"` + Messages []chatMessage `json:"messages"` + Stream bool `json:"stream"` } type chatResponse struct { - Message Message `json:"message"` + Message chatMessage `json:"message"` } type errorResponse struct { @@ -70,14 +63,14 @@ func NewClient(baseURL string, httpClient *http.Client) (*Client, error) { } // Chat sends a non-streaming chat request to Ollama. -func (client *Client) Chat(ctx context.Context, model string, messages []Message) (Message, error) { +func (client *Client) Chat(ctx context.Context, model string, messages []conversation.Message) (conversation.Message, error) { payload, err := json.Marshal(chatRequest{ Model: model, - Messages: messages, + Messages: toChatMessages(messages), Stream: false, }) if err != nil { - return Message{}, fmt.Errorf("encode Ollama chat request: %w", err) + return conversation.Message{}, fmt.Errorf("encode Ollama chat request: %w", err) } endpoint := client.baseURL.JoinPath("chat") @@ -89,33 +82,50 @@ func (client *Client) Chat(ctx context.Context, model string, messages []Message bytes.NewReader(payload), ) if err != nil { - return Message{}, fmt.Errorf("create Ollama chat request: %w", err) + return conversation.Message{}, fmt.Errorf("create Ollama chat request: %w", err) } request.Header.Set("Content-Type", "application/json") response, err := client.httpClient.Do(request) if err != nil { - return Message{}, fmt.Errorf("send Ollama chat request: %w", err) + return conversation.Message{}, fmt.Errorf("send Ollama chat request: %w", err) } - defer response.Body.Close() + + defer func() { + _ = response.Body.Close() + }() if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { var apiError errorResponse decodeErr := json.NewDecoder(io.LimitReader(response.Body, maxErrBodySize)).Decode(&apiError) if decodeErr != nil || strings.TrimSpace(apiError.Error) == "" { - return Message{}, fmt.Errorf("Ollama chat request failed: %s", response.Status) + return conversation.Message{}, fmt.Errorf("ollama chat request failed: %s", response.Status) } - return Message{}, fmt.Errorf("Ollama chat request failed: %s: %s", response.Status, apiError.Error) + return conversation.Message{}, fmt.Errorf("ollama chat request failed: %s: %s", response.Status, apiError.Error) } var result chatResponse if err := json.NewDecoder(response.Body).Decode(&result); err != nil { - return Message{}, fmt.Errorf("decode Ollama chat response: %w", err) + return conversation.Message{}, fmt.Errorf("decode Ollama chat response: %w", err) } - return result.Message, nil + return fromChatMessage(result.Message), nil +} + +func toChatMessages(messages []conversation.Message) []chatMessage { + result := make([]chatMessage, len(messages)) + + for index, message := range messages { + result[index] = chatMessage{Role: string(message.Role), Content: message.Content} + } + + return result +} + +func fromChatMessage(message chatMessage) conversation.Message { + return conversation.Message{Role: conversation.Role(message.Role), Content: message.Content} } diff --git a/internal/ollama/client_test.go b/internal/ollama/client_test.go index 04cbb62..70fe748 100644 --- a/internal/ollama/client_test.go +++ b/internal/ollama/client_test.go @@ -7,6 +7,8 @@ import ( "net/http/httptest" "strings" "testing" + + "github.com/theantichris/ghost/v4/internal/conversation" ) func TestNewClient(t *testing.T) { @@ -55,22 +57,32 @@ func TestNewClient(t *testing.T) { } func TestClientChat(t *testing.T) { - messages := []Message{ + messages := []conversation.Message{ { - Role: "user", + Role: conversation.RoleUser, Content: "Identify yourself.", }, } - wantMessage := Message{ + wantRequestMessage := chatMessage{ + Role: "user", + Content: "Identify yourself.", + } + + responseMessage := chatMessage{ Role: "assistant", Content: "Ghost online.", } + wantMessage := conversation.Message{ + Role: conversation.RoleAssistant, + Content: "Ghost online.", + } + tests := []struct { name string handler http.HandlerFunc - wantMessage Message + wantMessage conversation.Message wantErrContains string }{ { @@ -106,13 +118,13 @@ func TestClientChat(t *testing.T) { t.Error("request stream = true, want false") } - if len(gotRequest.Messages) != 1 || gotRequest.Messages[0] != messages[0] { - t.Errorf("request messages = %#v, want %#v", gotRequest.Messages, messages) + if len(gotRequest.Messages) != 1 || gotRequest.Messages[0] != wantRequestMessage { + t.Errorf("request messages = %#v, want %#v", gotRequest.Messages, []chatMessage{wantRequestMessage}) } writer.Header().Set("Content-Type", "application/json") - if err := json.NewEncoder(writer).Encode(chatResponse{Message: wantMessage}); err != nil { + if err := json.NewEncoder(writer).Encode(chatResponse{Message: responseMessage}); err != nil { t.Errorf("encode response: %v", err) } }, diff --git a/main.go b/main.go index 1baefd8..02bc657 100644 --- a/main.go +++ b/main.go @@ -9,6 +9,7 @@ import ( "os" "strings" + "github.com/theantichris/ghost/v4/internal/conversation" "github.com/theantichris/ghost/v4/internal/ollama" ) @@ -17,7 +18,7 @@ var ( errPromptRequired = errors.New("prompt is required") ) -type chatFunc func(context.Context, string, []ollama.Message) (ollama.Message, error) +type chatFunc func(context.Context, string, []conversation.Message) (conversation.Message, error) func main() { client, err := ollama.NewClient(ollama.DefaultURL, nil) @@ -52,7 +53,10 @@ func run(ctx context.Context, args []string, output io.Writer, chat chatFunc) er return errPromptRequired } - response, err := chat(ctx, modelName, []ollama.Message{{Role: ollama.RoleUser, Content: prompt}}) + session := conversation.NewSession() + session.AppendUser(prompt) + + response, err := chat(ctx, modelName, session.Messages()) if err != nil { return fmt.Errorf("request Ollama chat response: %w", err) } diff --git a/main_test.go b/main_test.go index c8bb123..c653e2c 100644 --- a/main_test.go +++ b/main_test.go @@ -9,7 +9,7 @@ import ( "strings" "testing" - "github.com/theantichris/ghost/v4/internal/ollama" + "github.com/theantichris/ghost/v4/internal/conversation" ) type errorWriter struct { @@ -24,19 +24,19 @@ func TestRun(t *testing.T) { chatErr := errors.New("Ollama unavailable") writeErr := errors.New("output unavailable") - promptMessages := []ollama.Message{ - {Role: ollama.RoleUser, Content: "Identify yourself."}, + promptMessages := []conversation.Message{ + {Role: conversation.RoleUser, Content: "Identify yourself."}, } tests := []struct { name string args []string output io.Writer - response ollama.Message + response conversation.Message chatErr error wantChat bool wantModel string - wantMessages []ollama.Message + wantMessages []conversation.Message wantOutput string wantErr bool wantErrIs error @@ -46,7 +46,7 @@ func TestRun(t *testing.T) { name: "prints assistant response", args: []string{"--model", "qwen3:8b", "Identify", "yourself."}, output: &bytes.Buffer{}, - response: ollama.Message{Role: ollama.RoleAssistant, Content: "Ghost online."}, + response: conversation.Message{Role: conversation.RoleAssistant, Content: "Ghost online."}, wantChat: true, wantModel: "qwen3:8b", wantMessages: promptMessages, @@ -91,7 +91,7 @@ func TestRun(t *testing.T) { name: "wraps output error", args: []string{"--model", "qwen3:8b", "Identify yourself."}, output: errorWriter{err: writeErr}, - response: ollama.Message{Role: ollama.RoleAssistant, Content: "Ghost online."}, + response: conversation.Message{Role: conversation.RoleAssistant, Content: "Ghost online."}, wantChat: true, wantModel: "qwen3:8b", wantMessages: promptMessages, @@ -103,7 +103,7 @@ func TestRun(t *testing.T) { name: "does not duplicate final newline", args: []string{"--model", "qwen3:8b", "Identify yourself."}, output: &bytes.Buffer{}, - response: ollama.Message{Role: ollama.RoleAssistant, Content: "Ghost online.\n"}, + response: conversation.Message{Role: conversation.RoleAssistant, Content: "Ghost online.\n"}, wantChat: true, wantModel: "qwen3:8b", wantMessages: promptMessages, @@ -132,12 +132,12 @@ func TestRun(t *testing.T) { chatCalled := false var gotModel string - var gotMessages []ollama.Message + var gotMessages []conversation.Message - chat := func(_ context.Context, model string, messages []ollama.Message) (ollama.Message, error) { + chat := func(_ context.Context, model string, messages []conversation.Message) (conversation.Message, error) { chatCalled = true gotModel = model - gotMessages = append([]ollama.Message(nil), messages...) + gotMessages = append([]conversation.Message(nil), messages...) return test.response, test.chatErr }