Added Nvidia NIM

EdgeDatabase.go

@@ -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"`

LLM.go

@@ -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
 		}

Request.go

@@ -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)
 			}
 

RequestNim.go

@@ -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
+}

TODO.md

@@ -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

dbschema/default.esdl

@@ -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;
         };

dbschema/migrations/00053-m1eqooh.edgeql

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

main.go

@@ -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,
 	}

views/partials/popover-settings.html

@@ -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"
@@ -182,4 +191,4 @@
         toggleIcon.classList.toggle('fa-chevron-up');
         toggleIcon.classList.toggle('fa-chevron-down');
     });
-</script>
+</script>