Jade/RequestAnthropic.go

package main

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

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

type AnthropicChatCompletionRequest struct {
	Model       string           `json:"model"`
	Messages    []RequestMessage `json:"messages"`
	MaxTokens   int              `json:"max_tokens"`
	Temperature float64          `json:"temperature"`
	Context     string           `json:"system"`
}

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 addAnthropicMessage(llm LLM, selected bool) edgedb.UUID {
	Messages := getAllSelectedMessages()

	chatCompletion, err := RequestAnthropic(llm.Model.ModelID, Messages, int(llm.Model.MaxToken), float64(llm.Temperature), llm.Context)
	if err != nil {
		fmt.Print("Error fetching user profile")
		panic(err)
	} else if len(chatCompletion.Content) == 0 {
		fmt.Println("No response from Anthropic")
		id := insertBotMessage("No response from Anthropic", selected, llm.ID)
		return id
	} else {
		id := insertBotMessage(chatCompletion.Content[0].Text, selected, llm.ID)
		return id
	}
}

func TestAnthropicKey(apiKey string) bool {
	url := "https://api.anthropic.com/v1/messages"

	AnthropicMessages := []RequestMessage{
		{
			Role:    "user",
			Content: "Hello",
		},
	}

	requestBody := AnthropicChatCompletionRequest{
		Model:       "claude-3-haiku-20240307",
		Messages:    AnthropicMessages,
		MaxTokens:   10,
		Temperature: 0,
		Context:     "",
	}

	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)
	if err != nil {
		return false
	}
	if chatCompletionResponse.Content == nil {
		return false
	}
	return true
}

func RequestAnthropic(model string, messages []Message, maxTokens int, temperature float64, context string) (AnthropicChatCompletionResponse, error) {
	var apiKey struct {
		Key string `edgedb:"key"`
	}
	err := edgeClient.QuerySingle(edgeCtx, `
		SELECT Key {
			key
		}
		FILTER .<keys[is Setting].<setting[is User] = global currentUser and .company.name = "anthropic"
		LIMIT 1
	`, &apiKey, "anthropic")
	if err != nil {
		return AnthropicChatCompletionResponse{}, fmt.Errorf("error getting Anthropic API key: %w", err)
	}

	url := "https://api.anthropic.com/v1/messages"

	requestBody := AnthropicChatCompletionRequest{
		Model:       model,
		Messages:    Message2RequestMessage(messages, ""),
		MaxTokens:   maxTokens,
		Temperature: temperature,
		Context:     context,
	}

	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("x-api-key", apiKey.Key)
	req.Header.Set("content-Type", "application/json")
	req.Header.Set("anthropic-version", "2023-06-01")

	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
	err = edgeClient.QuerySingle(edgeCtx, `
		SELECT ModelInfo {
			inputPrice,
			outputPrice
		}
		FILTER .modelID = <str>$0
		LIMIT 1
	`, &usedModelInfo, model)
	if err != nil {
		return AnthropicChatCompletionResponse{}, fmt.Errorf("error getting model info: %w", err)
	}

	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
}