From 3eef03b303fd43bf93fb6385b65a0cbb7841e025 Mon Sep 17 00:00:00 2001
From: Enrico Ros <enrico.ros@gmail.com>
Date: Mon, 22 Apr 2024 02:52:19 -0700
Subject: [PATCH] Groq: update models

---
 src/modules/llms/server/openai/models.data.ts | 54 +++++++++++++++----
 .../llms/server/openai/openai.router.ts       |  5 +-
 2 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/src/modules/llms/server/openai/models.data.ts b/src/modules/llms/server/openai/models.data.ts
index 5d2c8d80a..ffe9e66bf 100644
--- a/src/modules/llms/server/openai/models.data.ts
+++ b/src/modules/llms/server/openai/models.data.ts
@@ -842,23 +842,48 @@ export function perplexityAIModelSort(a: ModelDescriptionSchema, b: ModelDescrip
   return b.label.localeCompare(a.label);
 }
 
-// Groq
+
+// Groq - https://console.groq.com/docs/models
 
 const _knownGroqModels: ManualMappings = [
-  // {
-  //   id: 'lama2-70b-4096',
-  //   label: 'Llama 2 70B Chat',
-  //   description: 'Llama 2 is a collection of pretrained and fine-tuned generative text models.',
-  //   contextWindow: 4096,
-  //   interfaces: [LLM_IF_OAI_Chat],
-  // },
+  {
+    isLatest: true,
+    idPrefix: 'llama3-70b-8192',
+    label: 'Llama 3 · 70B',
+    description: 'LLaMA3 70b developed by Meta with a context window of 8,192 tokens.',
+    contextWindow: 8192,
+    interfaces: [LLM_IF_OAI_Chat],
+  },
+  {
+    // isLatest: true,
+    idPrefix: 'llama3-8b-8192',
+    label: 'Llama 3 · 8B',
+    description: 'LLaMA3 8b developed by Meta with a context window of 8,192 tokens.',
+    contextWindow: 8192,
+    interfaces: [LLM_IF_OAI_Chat],
+  },
+  {
+    idPrefix: 'llama2-70b-4096',
+    label: 'Llama 2 · 70B',
+    description: 'LLaMA2 70b developed by Meta with a context window of 4,096 tokens.',
+    contextWindow: 4096,
+    interfaces: [LLM_IF_OAI_Chat],
+    hidden: true,
+  },
   {
     idPrefix: 'mixtral-8x7b-32768',
-    label: 'Mixtral 8x7B Instruct v0.1',
-    description: 'The Mixtral-8x7B Large Language Model (LLM) is a pretrained generative Sparse Mixture of Experts.',
+    label: 'Mixtral 8x7B',
+    description: 'Mixtral 8x7b developed by Mistral with a context window of 32,768 tokens.',
     contextWindow: 32768,
     interfaces: [LLM_IF_OAI_Chat],
   },
+  {
+    idPrefix: 'gemma-7b-it',
+    label: 'Gemma 1.1 · 7B Instruct',
+    description: 'Gemma 7b developed by Google with a context window of 8,192 tokens.',
+    contextWindow: 8192,
+    interfaces: [LLM_IF_OAI_Chat],
+  },
 ];
 
 export function groqModelToModelDescription(_model: unknown): ModelDescriptionSchema {
@@ -873,6 +898,15 @@ export function groqModelToModelDescription(_model: unknown): ModelDescriptionSc
   });
 }
 
+export function groqModelSortFn(a: ModelDescriptionSchema, b: ModelDescriptionSchema): number {
+  // sort as per their order in the known models
+  const aIndex = _knownGroqModels.findIndex(base => a.id.startsWith(base.idPrefix));
+  const bIndex = _knownGroqModels.findIndex(base => b.id.startsWith(base.idPrefix));
+  if (aIndex !== -1 && bIndex !== -1)
+    return aIndex - bIndex;
+  return a.id.localeCompare(b.id);
+}
+
 
 // Helpers
 
diff --git a/src/modules/llms/server/openai/openai.router.ts b/src/modules/llms/server/openai/openai.router.ts
index e896cd08c..83bcf9dd5 100644
--- a/src/modules/llms/server/openai/openai.router.ts
+++ b/src/modules/llms/server/openai/openai.router.ts
@@ -11,7 +11,7 @@ import { Brand } from '~/common/app.config';
 import { fixupHost } from '~/common/util/urlUtils';
 
 import { OpenAIWire, WireOpenAICreateImageOutput, wireOpenAICreateImageOutputSchema, WireOpenAICreateImageRequest } from './openai.wiretypes';
-import { azureModelToModelDescription, groqModelToModelDescription, lmStudioModelToModelDescription, localAIModelToModelDescription, mistralModelsSort, mistralModelToModelDescription, oobaboogaModelToModelDescription, openAIModelFilter, openAIModelToModelDescription, openRouterModelFamilySortFn, openRouterModelToModelDescription, perplexityAIModelDescriptions, perplexityAIModelSort, togetherAIModelsToModelDescriptions } from './models.data';
+import { azureModelToModelDescription, groqModelSortFn, groqModelToModelDescription, lmStudioModelToModelDescription, localAIModelToModelDescription, mistralModelsSort, mistralModelToModelDescription, oobaboogaModelToModelDescription, openAIModelFilter, openAIModelToModelDescription, openRouterModelFamilySortFn, openRouterModelToModelDescription, perplexityAIModelDescriptions, perplexityAIModelSort, togetherAIModelsToModelDescriptions } from './models.data';
 import { llmsChatGenerateWithFunctionsOutputSchema, llmsListModelsOutputSchema, ModelDescriptionSchema } from '../llm.server.types';
 import { wilreLocalAIModelsApplyOutputSchema, wireLocalAIModelsAvailableOutputSchema, wireLocalAIModelsListOutputSchema } from './localai.wiretypes';
 
@@ -168,7 +168,8 @@ export const llmOpenAIRouter = createTRPCRouter({
 
         case 'groq':
           models = openAIModels
-            .map(groqModelToModelDescription);
+            .map(groqModelToModelDescription)
+            .sort(groqModelSortFn);
           break;
 
         case 'lmstudio':