Jade/RequestNim.go

package main

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

	"github.com/edgedb/edgedb-go"
	"github.com/gofiber/fiber/v2"
)

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

type NimChatCompletionResponse struct {
	ID      string      `json:"id"`
	Object  string      `json:"object"`
	Created int64       `json:"created"`
	Model   string      `json:"model"`
	Usage   NimUsage    `json:"usage"`
	Choices []NimChoice `json:"choices"`
}

type NimUsage struct {
	PromptTokens     int32 `json:"prompt_tokens"`
	CompletionTokens int32 `json:"completion_tokens"`
	TotalTokens      int32 `json:"total_tokens"`
}

type NimChoice struct {
	Message      RequestMessage `json:"message"`
	FinishReason string         `json:"finish_reason"`
	Index        int            `json:"index"`
}

func addNimMessage(c *fiber.Ctx, llm LLM, selected bool) edgedb.UUID {
	Messages := getAllSelectedMessages(c)

	chatCompletion, err := RequestNim(c, llm.Model.ModelID, Messages, float64(llm.Temperature), llm.Context, int(llm.MaxToken))
	if err != nil {
		fmt.Println("Error requesting NIM: ", err)
		id := insertBotMessage(c, "Error requesting NIM, model may not be available anymore. Better error message in development.", selected, llm.ID)
		return id
	} else if len(chatCompletion.Choices) == 0 {
		fmt.Println("No response from NIM")
		id := insertBotMessage(c, "No response from NIM", selected, llm.ID)
		return id
	} else {
		Content := chatCompletion.Choices[0].Message.Content
		id := insertBotMessage(c, Content, selected, llm.ID)
		return id
	}
}

func TestNimKey(apiKey string) bool {
	url := "https://integrate.api.nvidia.com/v1/chat/completions"
	//apiKey := "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"

	fmt.Println("Testing new Nvidia NIM key:", apiKey)

	// Convert messages to OpenAI format
	nimMessages := []RequestMessage{
		{
			Role:    "user",
			Content: "Hello",
		},
	}

	requestBody := NimChatCompletionRequest{
		Model:       "meta/llama3-8b-instruct",
		Messages:    nimMessages,
		Temperature: 0,
	}

	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		fmt.Println("Error when testing NIM key. Cant parse JSON request.")
		return false
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
	if err != nil {
		fmt.Println("Error when testing NIM key. Cant generate new request")
		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 {
		fmt.Println("Error when testing NIM key. Cant send request.")
		return false
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		fmt.Println("Error when testing NIM key. Cant read response.")
		return false
	}

	var chatCompletionResponse NimChatCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		fmt.Println(resp.Status)
		fmt.Println(resp.Body)

		fmt.Println("Error when testing NIM key. Cant unmarshal response.")
		return false
	}
	if chatCompletionResponse.Usage.CompletionTokens == 0 {
		fmt.Println(resp.Status)
		fmt.Println(resp.Body)

		fmt.Println("Error when testing NIM key. No completion token.")
		return false
	}

	Content := chatCompletionResponse.Choices[0].Message.Content
	fmt.Println(Content)

	return true
}

func RequestNim(c *fiber.Ctx, model string, messages []Message, temperature float64, context string, maxToken int) (NimChatCompletionResponse, error) {
	var apiKey string
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
		with
			filtered_keys := (
				select Key {
					key
				} filter .company.name = <str>$0 AND .<keys[is Setting].<setting[is User] = global currentUser
			)
		select filtered_keys.key limit 1
	`, &apiKey, "nim")
	if err != nil {
		return NimChatCompletionResponse{}, fmt.Errorf("error getting NIM API key: %w", err)
	}

	url := "https://integrate.api.nvidia.com/v1/chat/completions"

	requestBody := NimChatCompletionRequest{
		Model:       model,
		Messages:    Message2RequestMessage(messages, context),
		MaxTokens:   maxToken,
		Temperature: temperature,
	}

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

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
	if err != nil {
		return NimChatCompletionResponse{}, 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 NimChatCompletionResponse{}, fmt.Errorf("error sending request: %w", err)
	}
	defer resp.Body.Close()

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

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

	var usedModelInfo ModelInfo
	err = edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
		SELECT ModelInfo {
			inputPrice,
			outputPrice
		}
		FILTER .modelID = <str>$0
		LIMIT 1
	`, &usedModelInfo, model)
	if err != nil {
		return NimChatCompletionResponse{}, fmt.Errorf("error getting model info: %w", err)
	}

	var inputCost float32 = float32(chatCompletionResponse.Usage.PromptTokens) * usedModelInfo.InputPrice
	var outputCost float32 = float32(chatCompletionResponse.Usage.CompletionTokens) * usedModelInfo.OutputPrice
	addUsage(c, inputCost, outputCost, chatCompletionResponse.Usage.PromptTokens, chatCompletionResponse.Usage.CompletionTokens, model)

	return chatCompletionResponse, nil
}