ask command is now mistral!

lib/replicate.go

@@ -6,10 +6,51 @@ import (
 	"errors"
 	"io"
 	"net/http"
+	"strings"
 
 	"github.com/replicate/replicate-go"
 )
 
+func ReplicateTextGeneration(prompt string) (string, error) {
+	client, clientError := replicate.NewClient(replicate.WithTokenFromEnv())
+	if clientError != nil {
+		return "", 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(), "mistralai/mixtral-8x7b-instruct-v0.1:7b3212fbaf88310cfef07a061ce94224e82efc8403c26fc67e8f6c065de51f21", input, &webhook)
+
+	if predictionError != nil {
+		return "", predictionError
+	}
+
+	test, ok := prediction.([]interface{})
+
+	if !ok {
+		return "", errors.New("there was an error generating the image based on this prompt... this usually happens when the generated image violates safety requirements")
+	}
+
+	strs := make([]string, len(test))
+	for i, v := range test {
+		str, ok := v.(string)
+		if !ok {
+			return "", errors.New("element is not a string")
+		}
+		strs[i] = str
+	}
+
+	result := strings.Join(strs, "")
+
+	return result, nil
+}
+
 func ReplicateImageGeneration(prompt string) (*bytes.Buffer, error) {
 	client, clientError := replicate.NewClient(replicate.WithTokenFromEnv())
 	if clientError != nil {