Jade/RequestGooseai.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"github.com/edgedb/edgedb-go"
)

type GooseaiCompletionRequest struct {
	Model       string   `json:"model"`
	Prompt      []string `json:"prompt"`
	Temperature float64  `json:"temperature"`
}

type GooseaiCompletionResponse struct {
	ID      string          `json:"id"`
	Created int64           `json:"created"`
	Model   string          `json:"model"`
	Choices []GooseaiChoice `json:"choices"`
}

type GooseaiChoice struct {
	Text         string `json:"text"`
	FinishReason string `json:"finish_reason"`
	Index        int    `json:"index"`
}

func addGooseaiMessage(llm LLM, selected bool) edgedb.UUID {
	Messages := getAllSelectedMessages()

	chatCompletion, err := RequestGooseai(llm.Model.ModelID, Messages, float64(llm.Temperature))
	if err != nil {
		panic(err)
	} else if len(chatCompletion.Choices) == 0 {
		fmt.Println("No response from GooseAI")
		id := insertBotMessage("No response from GooseAI", selected, llm.ID)
		return id
	} else {
		Content := chatCompletion.Choices[0].Text
		id := insertBotMessage(Content, selected, llm.ID)
		return id
	}
}

func TestGooseaiKey(apiKey string) bool {
	url := "https://api.goose.ai/v1/engines/gpt-j-6b/completions"

	requestBody := GooseaiCompletionRequest{
		Model:       "gpt-j-6b",
		Prompt:      []string{"Hello, how are you?"},
		Temperature: 0,
	}

	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		return false
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
	if err != nil {
		return false
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+apiKey)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return false
	}

	// Print the response body
	fmt.Println(string(body))

	var chatCompletionResponse GooseaiCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		return false
	}
	if chatCompletionResponse.Choices[0].Text == "" {
		return false
	}
	return true
}

func RequestGooseai(model string, messages []Message, temperature float64) (GooseaiCompletionResponse, error) {
	var apiKey string
	err := edgeClient.QuerySingle(edgeCtx, `
		with
			filtered_keys := (
				select Key {
					key
				} filter .company.name = "gooseai" AND .<keys[is Setting].<setting[is User] = global currentUser
			)
		select filtered_keys.key limit 1
	`, &apiKey)
	if err != nil {
		return GooseaiCompletionResponse{}, fmt.Errorf("error getting GooseAI API key: %w", err)
	}

	url := "https://api.goose.ai/v1/engines/" + model + "/completions"

	requestBody := GooseaiCompletionRequest{
		Model:       model,
		Prompt:      []string{messages[len(messages)-1].Content},
		Temperature: temperature,
	}

	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		return GooseaiCompletionResponse{}, fmt.Errorf("error marshaling JSON: %w", err)
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
	if err != nil {
		return GooseaiCompletionResponse{}, fmt.Errorf("error creating request: %w", err)
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+apiKey)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return GooseaiCompletionResponse{}, fmt.Errorf("error sending request: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return GooseaiCompletionResponse{}, fmt.Errorf("error reading response body: %w", err)
	}

	var chatCompletionResponse GooseaiCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		return GooseaiCompletionResponse{}, fmt.Errorf("error unmarshaling JSON: %w", err)
	}

	addUsage(0, 0, 0, 0, model)

	return chatCompletionResponse, nil
}