Added Nvidia NIM

This commit is contained in:
Adrien Bouvais 2024-08-03 16:22:33 +02:00
parent c3d702a09a
commit 183bec7203
11 changed files with 248 additions and 7 deletions

View File

@ -77,7 +77,7 @@ type LLM struct {
ID edgedb.UUID `edgedb:"id"`
Name string `edgedb:"name"`
Context string `edgedb:"context"`
MaxToken int32 `edgedb:"max_token"`
MaxToken int32 `edgedb:"max_tokens"`
Temperature float32 `edgedb:"temperature"`
Model ModelInfo `edgedb:"modelInfo"`
Endpoint CustomEndpoint `edgedb:"custom_endpoint"`

6
LLM.go
View File

@ -65,7 +65,7 @@ func deleteLLMtoDelete(c *fiber.Ctx) {
maxToken, err := strconv.Atoi(maxTokenStr)
if err != nil {
maxToken = 0
maxToken = 1024
}
fmt.Println("Adding LLM with maxtoken:", maxToken)
@ -80,7 +80,7 @@ func deleteLLMtoDelete(c *fiber.Ctx) {
context := <str>$1,
temperature := <float32>$2,
position := countLLM + 1,
max_token := <int32>$6,
max_tokens := <int32>$6,
modelInfo := (INSERT ModelInfo {
name := <str>$0,
modelID := <str>$5,
@ -108,7 +108,7 @@ func deleteLLMtoDelete(c *fiber.Ctx) {
context := <str>$1,
temperature := <float32>$2,
position := countLLM + 1,
max_token := <int32>$4,
max_tokens := <int32>$4,
modelInfo := (SELECT ModelInfo FILTER .modelID = <str>$3 LIMIT 1),
user := global currentUser
}

View File

@ -55,6 +55,7 @@ func GeneratePlaceholderHTML(c *fiber.Ctx, message string, selectedLLMIds []stri
name,
context,
temperature,
max_tokens,
custom_endpoint : {
id,
endpoint,
@ -164,6 +165,8 @@ func GenerateMultipleMessagesHandler(c *fiber.Ctx) error {
addMessageFunc = addPerplexityMessage
case "fireworks":
addMessageFunc = addFireworkMessage
case "nim":
addMessageFunc = addNimMessage
}
var messageID edgedb.UUID
@ -193,7 +196,7 @@ func GenerateMultipleMessagesHandler(c *fiber.Ctx) error {
FILTER .id = <uuid>$0;
`, &message, messageID)
if err != nil {
fmt.Println("Error getting message")
fmt.Println("Error getting message for the placeholder. The function addProviderMessage seem to not return any message ID.")
panic(err)
}

207
RequestNim.go Normal file
View File

@ -0,0 +1,207 @@
package main
import (
"fmt"
"io"
"net/http"
"encoding/json"
"bytes"
"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
}

11
TODO.md Normal file
View File

@ -0,0 +1,11 @@
# Bugs
[ ] The SSE event that sometime fails
[ ] 2 icons ?
[ ] On first response, code block width are too long
[X] Change Terms of service to say that I use one cookie
[ ] Change the lastSelectedLLMs, 2 users can't use the same time...
# Features
[ ] Errors messages to know what happend
[ ] Add Deepseek API
[ ] Add Nvidia NIM

View File

@ -107,7 +107,7 @@ module default {
required user: User {
on target delete delete source;
};
required max_token: int32;
required max_tokens: int32;
custom_endpoint: CustomEndpoint {
on source delete delete target;
};

View File

@ -0,0 +1,9 @@
CREATE MIGRATION m1eqooh5xjbysafocihnromqeyrleo57w6txsr36tu73rkju2eyfcq
ONTO m1ad7psqkud35f6dmlabih4aqrlwrcoq2vlnmzfavf3dy6avyxsw2q
{
ALTER TYPE default::LLM {
ALTER PROPERTY max_token {
RENAME TO max_tokens;
};
};
};

View File

@ -191,6 +191,7 @@ func addKeys(c *fiber.Ctx) error {
"groq": c.FormValue("groq_key"),
"gooseai": c.FormValue("goose_key"),
"google": c.FormValue("google_key"),
"nim": c.FormValue("nim_key"),
"perplexity": c.FormValue("perplexity_key"),
"fireworks": c.FormValue("fireworks_key"),
}
@ -202,6 +203,7 @@ func addKeys(c *fiber.Ctx) error {
"groq": TestGroqKey,
"gooseai": TestGooseaiKey,
"google": TestGoogleKey,
"nim": TestNimKey,
"perplexity": TestPerplexityKey,
"fireworks": TestFireworkKey,
}

BIN
static/icons/nim.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 171 KiB

BIN
static/icons/old/nim.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 38 KiB

View File

@ -60,6 +60,15 @@
</span>
</p>
</div>
<div class="field has-addons is-hidden" id="nim-field">
<p class="control has-icons-left is-expanded">
<input class="input is-small {% if NimExists %}is-success{% endif %}" type="text"
placeholder="NIM API key" name="nim_key" autocomplete="off">
<span class="icon is-small is-left">
<i class="fas fa-lock"></i>
</span>
</p>
</div>
<div class="field has-addons is-hidden" id="perplexity-field">
<p class="control has-icons-left is-expanded">
<input class="input is-small {% if PerplexityExists %}is-success{% endif %}" type="text"