164 lines
4.1 KiB
Go
164 lines
4.1 KiB
Go
// It work but I disable it because it is not chat API
|
|
// It is text completion, not chat completion. But they will soon release API for chat
|
|
// So I leave it here for now
|
|
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"`
|
|
MaxToken int32 `json:"max_tokens"`
|
|
}
|
|
|
|
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 {
|
|
fmt.Print("Error fetching user profile")
|
|
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,
|
|
MaxToken: 10,
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
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"
|
|
|
|
var prompt string
|
|
for _, message := range messages {
|
|
prompt += message.Content
|
|
}
|
|
|
|
requestBody := GooseaiCompletionRequest{
|
|
Model: model,
|
|
Prompt: []string{prompt},
|
|
Temperature: temperature,
|
|
MaxToken: 300,
|
|
}
|
|
|
|
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
|
|
}
|