Cleaned up more business logic into the lib package

lib/openai.go

@@ -0,0 +1,78 @@
+package lib
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"os"
+
+	"github.com/joho/godotenv"
+	"github.com/sashabaranov/go-openai"
+)
+
+var client *openai.Client
+
+func init() {
+	godotenv.Load(".env")
+	apiKey := os.Getenv("OPENAI_API_KEY")
+	if apiKey == "" {
+		fmt.Println("OPENAI_API_KEY environment variable not set")
+		os.Exit(1)
+	}
+	client = openai.NewClient(apiKey)
+}
+
+func OpenAITextGeneration(prompt string) (string, error) {
+	resp, err := client.CreateChatCompletion(
+		context.Background(),
+		openai.ChatCompletionRequest{
+			Model: openai.GPT4TurboPreview,
+			Messages: []openai.ChatCompletionMessage{
+				{
+					Role:    openai.ChatMessageRoleUser,
+					Content: prompt,
+				},
+			},
+		},
+	)
+
+	if err != nil {
+		fmt.Printf("ChatCompletion error: %v\n", err)
+		return "", err
+	}
+
+	return resp.Choices[0].Message.Content, nil
+}
+
+func OpenAIImageGeneration(prompt string) (*bytes.Buffer, error) {
+	// Send the generation request to DALL·E 3
+	resp, err := client.CreateImage(context.Background(), openai.ImageRequest{
+		Prompt: prompt,
+		Model:  "dall-e-3",
+		Size:   "1024x1024",
+	})
+	if err != nil {
+		log.Printf("Image creation error: %v\n", err)
+		return nil, fmt.Errorf("failed to generate image")
+	}
+
+	imageRes, err := http.Get(resp.Data[0].URL)
+
+	if err != nil {
+		return nil, err
+	}
+
+	defer imageRes.Body.Close()
+
+	imageBytes, err := io.ReadAll(imageRes.Body)
+
+	if err != nil {
+		return nil, err
+	}
+
+	imageFile := bytes.NewBuffer(imageBytes)
+	return imageFile, nil
+}

lib/replicate.go

@@ -0,0 +1,59 @@
+package lib
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"io"
+	"net/http"
+
+	"github.com/replicate/replicate-go"
+)
+
+func ReplicateImageGeneration(prompt string) (*bytes.Buffer, error) {
+	client, clientError := replicate.NewClient(replicate.WithTokenFromEnv())
+	if clientError != nil {
+		return nil, clientError
+	}
+
+	input := replicate.PredictionInput{
+		"prompt": prompt,
+	}
+	webhook := replicate.Webhook{
+		URL:    "https://example.com/webhook",
+		Events: []replicate.WebhookEventType{"start", "completed"},
+	}
+
+	prediction, predictionError := client.Run(context.Background(), "stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b", input, &webhook)
+
+	if predictionError != nil {
+		return nil, predictionError
+	}
+
+	test, ok := prediction.([]interface{})
+
+	if !ok {
+		return nil, errors.New("there was an error generating the image based on this prompt... this usually happens when the generated image violates safety requirements")
+	}
+
+	imgUrl, ok := test[0].(string)
+
+	if !ok {
+		return nil, errors.New("there was an error generating the image based on this prompt... this usually happens when the generated image violates safety requirements")
+	}
+
+	imageRes, imageGetErr := http.Get(imgUrl)
+	if imageGetErr != nil {
+		return nil, imageGetErr
+	}
+
+	defer imageRes.Body.Close()
+
+	imageBytes, imgReadErr := io.ReadAll(imageRes.Body)
+	if imgReadErr != nil {
+		return nil, imgReadErr
+	}
+
+	imageFile := bytes.NewBuffer(imageBytes)
+	return imageFile, nil
+}