Rewrite test (#14)

ai_test.go

@@ -0,0 +1,132 @@
+package ai_test
+
+import (
+ "context"
+ "fmt"
+ "io"
+ "os"
+ "testing"
+ "time"
+
+ "github.com/sunshineplan/ai"
+ "github.com/sunshineplan/ai/chatgpt"
+ "github.com/sunshineplan/ai/gemini"
+)
+
+func init() {
+ if proxy := os.Getenv("AI_PROXY"); proxy != "" {
+ ai.SetProxy(proxy)
+ }
+}
+
+func testChat(ai ai.AI, prompt string) error {
+ ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+ defer cancel()
+ fmt.Println(prompt)
+ resp, err := ai.Chat(ctx, prompt)
+ if err != nil {
+ return err
+ }
+ fmt.Println(resp.Results())
+ fmt.Println("---")
+ return nil
+}
+
+func testChatStream(ai ai.AI, prompt string) error {
+ ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+ defer cancel()
+ fmt.Println(prompt)
+ stream, err := ai.ChatStream(ctx, prompt)
+ if err != nil {
+ return err
+ }
+ defer stream.Close()
+ for {
+ resp, err := stream.Next()
+ if err != nil {
+ if err == io.EOF {
+ break
+ }
+ return err
+ }
+ fmt.Println(resp.Results())
+ }
+ fmt.Println("---")
+ return nil
+}
+
+func testChatSession(ai ai.AI) error {
+ s := ai.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 {
+ return 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 {
+ return err
+ }
+ defer stream.Close()
+ for {
+ resp, err := stream.Next()
+ if err != nil {
+ if err == io.EOF {
+ break
+ }
+ return err
+ }
+ fmt.Println(resp.Results())
+ }
+ fmt.Println("---")
+ fmt.Println("History")
+ for _, i := range s.History() {
+ fmt.Println(i.Role, ":", i.Content)
+ }
+ fmt.Println("---")
+ return nil
+}
+
+func TestGemini(t *testing.T) {
+ apiKey := os.Getenv("GEMINI_API_KEY")
+ if apiKey == "" {
+ return
+ }
+ gemini, err := gemini.New(apiKey)
+ if err != nil {
+ t.Fatal(err)
+ }
+ defer gemini.Close()
+ if err := testChat(gemini, "Who are you?"); err != nil {
+ t.Error(err)
+ }
+ if err := testChatStream(gemini, "Who am I?"); err != nil {
+ t.Error(err)
+ }
+ if err := testChatSession(gemini); err != nil {
+ t.Error(err)
+ }
+}
+
+func TestChatGPT(t *testing.T) {
+ apiKey := os.Getenv("CHATGPT_API_KEY")
+ if apiKey == "" {
+ return
+ }
+ chatgpt := chatgpt.New(apiKey)
+ defer chatgpt.Close()
+ if err := testChat(chatgpt, "Who are you?"); err != nil {
+ t.Error(err)
+ }
+ if err := testChatStream(chatgpt, "Who am I?"); err != nil {
+ t.Error(err)
+ }
+ if err := testChatSession(chatgpt); err != nil {
+ t.Error(err)
+ }
+}