ChatSession interface

sunshineplan · Mar 13, 2024 · 162a289 · 162a289
1 parent 38ac197
commit 162a289
Show file tree

Hide file tree

Showing 5 changed files with 149 additions and 28 deletions.
diff --git a/ai.go b/ai.go
@@ -16,7 +16,7 @@ type AI interface {
  Model
 
  Chatbot
- ChatSession() Chatbot
+ ChatSession() ChatSession
 
  Close() error
 }
@@ -33,6 +33,16 @@ type Chatbot interface {
  ChatStream(context.Context, ...string) (ChatStream, error)
 }
 
+type Message struct {
+ Content string
+ Role string
+}
+
+type ChatSession interface {
+ Chatbot
+ History() []Message
+}
+
 type ChatStream interface {
  Next() (ChatResponse, error)
  Close() error

diff --git a/chatgpt/chatgpt.go b/chatgpt/chatgpt.go
@@ -85,15 +85,15 @@ func (resp *ChatResponse[Response]) String() string {
 }
 
 func (ai *ChatGPT) createRequest(
- stream bool,
+ one bool,
  history []openai.ChatCompletionMessage,
  messages ...string,
 ) (req openai.ChatCompletionRequest) {
  req.Model = ai.model
  if ai.maxTokens != nil {
  req.MaxTokens = int(*ai.maxTokens)
  }
- if !stream && ai.count != nil {
+ if !one && ai.count != nil {
  req.N = int(*ai.count)
  }
  if ai.temperature != nil {
@@ -114,6 +114,7 @@ func (ai *ChatGPT) createRequest(
 
 func (chatgpt *ChatGPT) chat(
  ctx context.Context,
+ session bool,
  history []openai.ChatCompletionMessage,
  messages ...string,
 ) (resp openai.ChatCompletionResponse, err error) {
@@ -124,11 +125,11 @@ func (chatgpt *ChatGPT) chat(
  if err = chatgpt.wait(ctx); err != nil {
  return
  }
- return chatgpt.c.CreateChatCompletion(ctx, chatgpt.createRequest(false, history, messages...))
+ return chatgpt.c.CreateChatCompletion(ctx, chatgpt.createRequest(session, history, messages...))
 }
 
 func (ai *ChatGPT) Chat(ctx context.Context, messages ...string) (ai.ChatResponse, error) {
- resp, err := ai.chat(ctx, nil, messages...)
+ resp, err := ai.chat(ctx, false, nil, messages...)
  if err != nil {
  return nil, err
  }
@@ -148,7 +149,7 @@ func (stream *ChatStream) Next() (ai.ChatResponse, error) {
  if err != nil {
  if err == io.EOF {
  if stream.cs != nil {
- stream.cs.History = append(stream.cs.History, openai.ChatCompletionMessage{
+ stream.cs.history = append(stream.cs.history, openai.ChatCompletionMessage{
  Role: openai.ChatMessageRoleAssistant, Content: stream.merged})
  }
  }
@@ -189,31 +190,49 @@ func (ai *ChatGPT) ChatStream(ctx context.Context, messages ...string) (ai.ChatS
  return &ChatStream{stream, nil, ""}, nil
 }
 
-var _ ai.Chatbot = new(ChatSession)
+var _ ai.ChatSession = new(ChatSession)
 
 type ChatSession struct {
  ai *ChatGPT
- History []openai.ChatCompletionMessage
+ history []openai.ChatCompletionMessage
+}
+
+func addToHistory(history *[]openai.ChatCompletionMessage, role string, messages ...string) {
+ for _, i := range messages {
+ *history = append(
+ *history,
+ openai.ChatCompletionMessage{Role: role, Content: i},
+ )
+ }
 }
 
 func (session *ChatSession) Chat(ctx context.Context, messages ...string) (ai.ChatResponse, error) {
- resp, err := session.ai.chat(ctx, session.History, messages...)
+ resp, err := session.ai.chat(ctx, true, session.history, messages...)
  if err != nil {
  return nil, err
  }
- session.History = append(session.History, resp.Choices[0].Message)
+ addToHistory(&session.history, openai.ChatMessageRoleUser, messages...)
+ session.history = append(session.history, resp.Choices[0].Message)
  return &ChatResponse[openai.ChatCompletionResponse]{resp}, nil
 }
 
 func (session *ChatSession) ChatStream(ctx context.Context, messages ...string) (ai.ChatStream, error) {
- stream, err := session.ai.chatStream(ctx, session.History, messages...)
+ stream, err := session.ai.chatStream(ctx, session.history, messages...)
  if err != nil {
  return nil, err
  }
+ addToHistory(&session.history, openai.ChatMessageRoleUser, messages...)
  return &ChatStream{stream, session, ""}, nil
 }
 
-func (ai *ChatGPT) ChatSession() ai.Chatbot {
+func (session *ChatSession) History() (history []ai.Message) {
+ for _, i := range session.history {
+ history = append(history, ai.Message{Content: i.Content, Role: i.Role})
+ }
+ return
+}
+
+func (ai *ChatGPT) ChatSession() ai.ChatSession {
  return &ChatSession{ai: ai}
 }
 

diff --git a/chatgpt/chatgpt_test.go b/chatgpt/chatgpt_test.go
@@ -0,0 +1,78 @@
+package chatgpt
+
+import (
+ "context"
+ "fmt"
+ "io"
+ "os"
+ "testing"
+ "time"
+)
+
+func TestChatGPT(t *testing.T) {
+ apiKey := os.Getenv("CHATGPT_API_KEY")
+ if apiKey == "" {
+ return
+ }
+ chatgpt := New(apiKey)
+ defer chatgpt.Close()
+ ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+ defer cancel()
+ fmt.Println("Who are you?")
+ resp, err := chatgpt.Chat(ctx, "Who are you?")
+ if err != nil {
+ t.Fatal(err)
+ }
+ fmt.Println(resp.Results())
+ fmt.Println("---")
+ fmt.Println("Who am I?")
+ ctx, cancel = context.WithTimeout(context.Background(), time.Minute)
+ defer cancel()
+ stream, err := chatgpt.ChatStream(ctx, "Who am I?")
+ if err != nil {
+ t.Fatal(err)
+ }
+ defer stream.Close()
+ for {
+ resp, err := stream.Next()
+ if err != nil {
+ if err == io.EOF {
+ break
+ }
+ t.Fatal(err)
+ }
+ fmt.Println(resp.Results())
+ }
+ fmt.Println("---")
+ s := chatgpt.ChatSession()
+ ctx, cancel = context.WithTimeout(context.Background(), time.Minute)
+ defer cancel()
+ fmt.Println("Hello, I have 2 dogs in my house.")
+ resp, err = s.Chat(ctx, "Hello, I have 2 dogs in my house.")
+ if err != nil {
+ t.Fatal(err)
+ }
+ fmt.Println(resp.Results())
+ ctx, cancel = context.WithTimeout(context.Background(), time.Minute)
+ defer cancel()
+ fmt.Println("How many paws are in my house?")
+ stream, err = s.ChatStream(ctx, "How many paws are in my house?")
+ if err != nil {
+ t.Fatal(err)
+ }
+ defer stream.Close()
+ for {
+ resp, err := stream.Next()
+ if err != nil {
+ if err == io.EOF {
+ break
+ }
+ t.Fatal(err)
+ }
+ fmt.Println(resp.Results())
+ }
+ fmt.Println("---")
+ for _, i := range s.History() {
+ fmt.Println(i.Role, ":", i.Content)
+ }
+}
diff --git a/gemini/gemini.go b/gemini/gemini.go
@@ -59,6 +59,22 @@ func (ai *Gemini) SetMaxTokens(i int32) { ai.config.SetMaxOutputTokens(i) }
 func (ai *Gemini) SetTemperature(f float32) { ai.config.SetTemperature(f) }
 func (ai *Gemini) SetTopP(f float32) { ai.config.SetTopP(f) }
 
+func texts2parts(texts []string) (parts []genai.Part) {
+ for _, i := range texts {
+ parts = append(parts, genai.Text(i))
+ }
+ return
+}
+
+func parts2texts(parts []genai.Part) (texts []string) {
+ for _, i := range parts {
+ if text, ok := i.(genai.Text); ok {
+ texts = append(texts, string(text))
+ }
+ }
+ return
+}
+
 var _ ai.ChatResponse = new(ChatResponse)
 
 type ChatResponse struct {
@@ -68,13 +84,7 @@ type ChatResponse struct {
 func (resp *ChatResponse) Results() (res []string) {
  for _, i := range resp.Candidates {
  if i.Content != nil {
- var parts []string
- for _, part := range i.Content.Parts {
- if text, ok := part.(genai.Text); ok {
- parts = append(parts, string(text))
- }
- }
- res = append(res, strings.Join(parts, "\n"))
+ res = append(res, strings.Join(parts2texts(i.Content.Parts), "\n"))
  }
  }
  return
@@ -87,13 +97,6 @@ func (resp *ChatResponse) String() string {
  return ""
 }
 
-func texts2parts(texts []string) (parts []genai.Part) {
- for _, i := range texts {
- parts = append(parts, genai.Text(i))
- }
- return
-}
-
 func (ai *Gemini) Chat(ctx context.Context, parts ...string) (ai.ChatResponse, error) {
  if err := ai.wait(ctx); err != nil {
  return nil, err
@@ -137,7 +140,7 @@ func (ai *Gemini) ChatStream(ctx context.Context, parts ...string) (ai.ChatStrea
  return &ChatStream{ai.model.GenerateContentStream(ctx, texts2parts(parts)...)}, nil
 }
 
-var _ ai.Chatbot = new(ChatSession)
+var _ ai.ChatSession = new(ChatSession)
 
 type ChatSession struct {
  ai *Gemini
@@ -162,7 +165,14 @@ func (session *ChatSession) ChatStream(ctx context.Context, parts ...string) (ai
  return &ChatStream{session.cs.SendMessageStream(ctx, texts2parts(parts)...)}, nil
 }
 
-func (ai *Gemini) ChatSession() ai.Chatbot {
+func (session *ChatSession) History() (history []ai.Message) {
+ for _, i := range session.cs.History {
+ history = append(history, ai.Message{Content: strings.Join(parts2texts(i.Parts), "\n"), Role: i.Role})
+ }
+ return
+}
+
+func (ai *Gemini) ChatSession() ai.ChatSession {
  return &ChatSession{ai, ai.model.StartChat()}
 }
 

diff --git a/gemini/gemini_test.go b/gemini/gemini_test.go
@@ -79,4 +79,8 @@ func TestGemini(t *testing.T) {
  }
  fmt.Println(resp.Results())
  }
+ fmt.Println("---")
+ for _, i := range s.History() {
+ fmt.Println(i.Role, ":", i.Content)
+ }
 }