318 lines
8.2 KiB
Go
318 lines
8.2 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
|
|
"github.com/edgedb/edgedb-go"
|
|
"github.com/flosch/pongo2"
|
|
"github.com/gofiber/fiber/v2"
|
|
)
|
|
|
|
type AnthropicChatCompletionRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []AnthropicMessage `json:"messages"`
|
|
MaxTokens int `json:"max_tokens"`
|
|
Temperature float64 `json:"temperature"`
|
|
}
|
|
|
|
type AnthropicMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
type AnthropicChatCompletionResponse struct {
|
|
ID string `json:"id"`
|
|
Content []AnthropicContentItem `json:"content"`
|
|
Model string `json:"model"`
|
|
StopReason string `json:"stop_reason"`
|
|
StopSequence string `json:"stop_sequence"`
|
|
Usage AnthropicUsage `json:"usage"`
|
|
}
|
|
|
|
type AnthropicContentItem struct {
|
|
Text string `json:"text"`
|
|
}
|
|
|
|
type AnthropicUsage struct {
|
|
InputTokens int32 `json:"input_tokens"`
|
|
OutputTokens int32 `json:"output_tokens"`
|
|
}
|
|
|
|
func init() {
|
|
var ModelInfosList = []ModelInfo{}
|
|
|
|
modelInfo := ModelInfo{
|
|
ID: "claude-3-haiku-20240307",
|
|
Name: "Haiku",
|
|
Icon: "anthropic",
|
|
MaxToken: 8192,
|
|
InputPrice: 0.50 / 1000000,
|
|
OutputPrice: 1.50 / 1000000,
|
|
}
|
|
ModelInfosList = append(ModelInfosList, modelInfo)
|
|
ModelsInfos = append(ModelsInfos, modelInfo)
|
|
|
|
modelInfo = ModelInfo{
|
|
ID: "claude-3-sonnet-20240229",
|
|
Name: "Sonnet",
|
|
Icon: "anthropic",
|
|
MaxToken: 8192,
|
|
InputPrice: 0.50 / 1000000,
|
|
OutputPrice: 1.50 / 1000000,
|
|
}
|
|
ModelInfosList = append(ModelInfosList, modelInfo)
|
|
ModelsInfos = append(ModelsInfos, modelInfo)
|
|
|
|
modelInfo = ModelInfo{
|
|
ID: "claude-3-opus-20240229",
|
|
Name: "Opus",
|
|
Icon: "anthropic",
|
|
MaxToken: 8192,
|
|
InputPrice: 0.50 / 1000000,
|
|
OutputPrice: 1.50 / 1000000,
|
|
}
|
|
ModelInfosList = append(ModelInfosList, modelInfo)
|
|
ModelsInfos = append(ModelsInfos, modelInfo)
|
|
|
|
companyInfo := CompanyInfo{
|
|
ID: "anthropic",
|
|
Name: "Anthropic",
|
|
ModelInfos: ModelInfosList,
|
|
}
|
|
CompanyInfos = append(CompanyInfos, companyInfo)
|
|
}
|
|
|
|
func addAnthropicMessage(modelID string, selected bool) edgedb.UUID {
|
|
Messages := getAllMessages()
|
|
|
|
chatCompletion, err := RequestAnthropic(modelID, Messages, 2048, 0.7) // TODO CHange parameters
|
|
if err != nil {
|
|
// Print error
|
|
fmt.Println("Error:", err)
|
|
} else if len(chatCompletion.Content) == 0 {
|
|
fmt.Println(chatCompletion)
|
|
fmt.Println("No response from Anthropic")
|
|
id := insertBotMessage("No response from Anthropic", selected, modelID)
|
|
return id
|
|
} else {
|
|
Content := chatCompletion.Content[0].Text
|
|
id := insertBotMessage(Content, selected, modelID)
|
|
return id
|
|
}
|
|
return edgedb.UUID{}
|
|
}
|
|
|
|
func EdgeMessages2AnthropicMessages(messages []Message) []AnthropicMessage {
|
|
AnthropicMessages := make([]AnthropicMessage, len(messages))
|
|
for i, msg := range messages {
|
|
var role string
|
|
switch msg.Role {
|
|
case "user":
|
|
role = "user"
|
|
case "bot":
|
|
role = "assistant"
|
|
default:
|
|
role = "system"
|
|
}
|
|
AnthropicMessages[i] = AnthropicMessage{
|
|
Role: role,
|
|
Content: msg.Content,
|
|
}
|
|
}
|
|
return AnthropicMessages
|
|
}
|
|
|
|
func TestAnthropicKey(apiKey string) bool {
|
|
url := "https://api.anthropic.com/v1/messages"
|
|
|
|
AnthropicMessages := []AnthropicMessage{
|
|
{
|
|
Role: "user",
|
|
Content: "Hello",
|
|
},
|
|
}
|
|
|
|
requestBody := AnthropicChatCompletionRequest{
|
|
Model: "claude-3-haiku-20240307",
|
|
Messages: AnthropicMessages,
|
|
MaxTokens: 10,
|
|
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("anthropic-version", "2023-06-01")
|
|
req.Header.Set("x-api-key", 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 AnthropicChatCompletionResponse
|
|
err = json.Unmarshal(body, &chatCompletionResponse)
|
|
return err == nil
|
|
}
|
|
|
|
func RequestAnthropic(model string, messages []Message, maxTokens int, temperature float64) (AnthropicChatCompletionResponse, error) {
|
|
var apiKey string
|
|
err := edgeClient.QuerySingle(edgeCtx, `
|
|
with
|
|
filtered_keys := (
|
|
select Key {
|
|
key
|
|
} filter .company = <str>$0
|
|
)
|
|
select filtered_keys.key limit 1
|
|
`, &apiKey, "anthropic")
|
|
if err != nil {
|
|
return AnthropicChatCompletionResponse{}, fmt.Errorf("error getting OpenAI API key: %w", err)
|
|
}
|
|
|
|
url := "https://api.anthropic.com/v1/messages"
|
|
|
|
AnthropicMessages := EdgeMessages2AnthropicMessages(messages)
|
|
|
|
requestBody := AnthropicChatCompletionRequest{
|
|
Model: model,
|
|
Messages: AnthropicMessages,
|
|
MaxTokens: maxTokens,
|
|
Temperature: temperature,
|
|
}
|
|
|
|
jsonBody, err := json.Marshal(requestBody)
|
|
if err != nil {
|
|
return AnthropicChatCompletionResponse{}, fmt.Errorf("error marshaling JSON: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
|
|
if err != nil {
|
|
return AnthropicChatCompletionResponse{}, fmt.Errorf("error creating request: %w", err)
|
|
}
|
|
|
|
req.Header.Set("content-Type", "application/json")
|
|
req.Header.Set("anthropic-version", "2023-06-01")
|
|
req.Header.Set("x-api-key", apiKey)
|
|
|
|
client := &http.Client{}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return AnthropicChatCompletionResponse{}, fmt.Errorf("error sending request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return AnthropicChatCompletionResponse{}, fmt.Errorf("error reading response body: %w", err)
|
|
}
|
|
|
|
var chatCompletionResponse AnthropicChatCompletionResponse
|
|
err = json.Unmarshal(body, &chatCompletionResponse)
|
|
if err != nil {
|
|
return AnthropicChatCompletionResponse{}, fmt.Errorf("error unmarshaling JSON: %w", err)
|
|
}
|
|
|
|
var usedModelInfo ModelInfo
|
|
for mi := range ModelsInfos {
|
|
if ModelsInfos[mi].ID == model {
|
|
usedModelInfo = ModelsInfos[mi]
|
|
}
|
|
}
|
|
var inputCost float32 = float32(chatCompletionResponse.Usage.InputTokens) * usedModelInfo.InputPrice
|
|
var outputCost float32 = float32(chatCompletionResponse.Usage.OutputTokens) * usedModelInfo.OutputPrice
|
|
addUsage(inputCost, outputCost, chatCompletionResponse.Usage.InputTokens, chatCompletionResponse.Usage.OutputTokens, model)
|
|
|
|
return chatCompletionResponse, nil
|
|
}
|
|
|
|
func addAnthropicKey(c *fiber.Ctx) error {
|
|
key := c.FormValue("key")
|
|
|
|
// Check if the key already exists
|
|
err := edgeClient.QuerySingle(edgeCtx, `
|
|
with
|
|
filtered_keys := (
|
|
select Key {
|
|
key
|
|
} filter .key = <str>$0 and .company = "anthropic"
|
|
)
|
|
select filtered_keys.key limit 1
|
|
`, &key, key)
|
|
if err == nil {
|
|
return c.SendString("")
|
|
}
|
|
|
|
if !TestAnthropicKey(key) {
|
|
fmt.Println("Invalid Anthropic API Key")
|
|
|
|
NextMessages := []NextMessage{}
|
|
nextMsg := NextMessage{
|
|
Icon: "bouvai2", // Assuming Icon is a field you want to include from Message
|
|
Content: "<br>" + markdownToHTML("Invalid Anthropic API Key"),
|
|
Hidden: false, // Assuming Hidden is a field you want to include from Message
|
|
Id: "0",
|
|
Name: "JADE",
|
|
}
|
|
NextMessages = append(NextMessages, nextMsg)
|
|
|
|
botOut, err := botTmpl.Execute(pongo2.Context{"Messages": NextMessages, "ConversationAreaId": 0})
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return c.SendString(botOut)
|
|
}
|
|
|
|
err = edgeClient.Execute(edgeCtx, `
|
|
UPDATE global currentUser.setting
|
|
SET {
|
|
keys += (
|
|
INSERT Key {
|
|
company := <str>$0,
|
|
key := <str>$1,
|
|
name := <str>$2,
|
|
}
|
|
)
|
|
}`, "anthropic", key, "Anthropic API Key")
|
|
if err != nil {
|
|
fmt.Println("Error in edgedb.QuerySingle: in addOpenaiKey")
|
|
fmt.Println(err)
|
|
}
|
|
|
|
NextMessages := []NextMessage{}
|
|
nextMsg := NextMessage{
|
|
Icon: "bouvai2", // Assuming Icon is a field you want to include from Message
|
|
Content: "<br>" + markdownToHTML("Key added successfully!"),
|
|
Hidden: false, // Assuming Hidden is a field you want to include from Message
|
|
Id: "0",
|
|
Name: "JADE",
|
|
}
|
|
NextMessages = append(NextMessages, nextMsg)
|
|
|
|
botOut, err := botTmpl.Execute(pongo2.Context{"Messages": NextMessages, "ConversationAreaId": 0})
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
return c.SendString(botOut)
|
|
}
|