chatgpt_test.go 1.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. package test
  2. import (
  3. "context"
  4. "fmt"
  5. "github.com/0xJacky/Nginx-UI/settings"
  6. "errors"
  7. "github.com/sashabaranov/go-openai"
  8. "github.com/uozi-tech/cosy/sandbox"
  9. "io"
  10. "os"
  11. "testing"
  12. )
  13. func TestChatGPT(t *testing.T) {
  14. sandbox.NewInstance("../../app.ini", "sqlite").
  15. Run(func(instance *sandbox.Instance) {
  16. c := openai.NewClient(settings.OpenAISettings.Token)
  17. ctx := context.Background()
  18. req := openai.ChatCompletionRequest{
  19. Model: openai.GPT3Dot5Turbo0301,
  20. Messages: []openai.ChatCompletionMessage{
  21. {
  22. Role: openai.ChatMessageRoleUser,
  23. Content: "帮我写一个 nginx 配置文件的示例",
  24. },
  25. },
  26. Stream: true,
  27. }
  28. stream, err := c.CreateChatCompletionStream(ctx, req)
  29. if err != nil {
  30. fmt.Printf("CompletionStream error: %v\n", err)
  31. return
  32. }
  33. defer stream.Close()
  34. for {
  35. response, err := stream.Recv()
  36. if errors.Is(err, io.EOF) {
  37. return
  38. }
  39. if err != nil {
  40. fmt.Printf("Stream error: %v\n", err)
  41. return
  42. }
  43. fmt.Printf("%v", response.Choices[0].Delta.Content)
  44. _ = os.Stdout.Sync()
  45. }
  46. })
  47. }