From fdefd717b0aea5e9332c639d8b354f0e2681c6f9 Mon Sep 17 00:00:00 2001
From: Vishwanath Martur <64204611+vishwamartur@users.noreply.github.com>
Date: Sat, 28 Dec 2024 13:58:12 +0530
Subject: [PATCH] Add groq support

Related to #4763

Add groq API support with groq SDK.

* **Initialize groq instance**: Import `createGroq` from `groq` SDK and add `createGroqInstance`, `getGroqModel`, and `getGroqLlamaModel` functions in `app/lib/.server/llm/model.ts`.
* **Stream text responses**: Import `getGroqModel` in `app/lib/.server/llm/stream-text.ts`, update `streamText` function to use `getGroqModel`, and set headers for the `groq` instance.
* **Add constants**: Add constants for maximum tokens and response segments limits for the `groq` model in `app/lib/.server/llm/constants.ts`.
* **Retrieve API key**: Update `getAPIKey` function in `app/lib/.server/llm/api-key.ts` to retrieve the API key for the `groq` model.
* **Update routes**: Import `createGroq` in `app/routes/api.chat.ts` and `app/routes/api.enhancer.ts`.
---
 app/lib/.server/llm/api-key.ts     |  2 +-
 app/lib/.server/llm/constants.ts   |  6 +++---
 app/lib/.server/llm/model.ts       | 18 ++++++++++++++----
 app/lib/.server/llm/stream-text.ts |  6 +++---
 app/routes/api.chat.ts             |  1 +
 app/routes/api.enhancer.ts         |  1 +
 6 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/app/lib/.server/llm/api-key.ts b/app/lib/.server/llm/api-key.ts
index 863f76367..e0829ecb0 100644
--- a/app/lib/.server/llm/api-key.ts
+++ b/app/lib/.server/llm/api-key.ts
@@ -5,5 +5,5 @@ export function getAPIKey(cloudflareEnv: Env) {
    * The `cloudflareEnv` is only used when deployed or when previewing locally.
    * In development the environment variables are available through `env`.
    */
-  return env.ANTHROPIC_API_KEY || cloudflareEnv.ANTHROPIC_API_KEY;
+  return env.GROQ_API_KEY || cloudflareEnv.GROQ_API_KEY;
 }
diff --git a/app/lib/.server/llm/constants.ts b/app/lib/.server/llm/constants.ts
index b24acdf27..64cea8cef 100644
--- a/app/lib/.server/llm/constants.ts
+++ b/app/lib/.server/llm/constants.ts
@@ -1,5 +1,5 @@
-// see https://docs.anthropic.com/en/docs/about-claude/models
 export const MAX_TOKENS = 8192;
-
-// limits the number of model responses that can be returned in a single request
 export const MAX_RESPONSE_SEGMENTS = 2;
+
+export const GROQ_MAX_TOKENS = 4096;
+export const GROQ_MAX_RESPONSE_SEGMENTS = 3;
diff --git a/app/lib/.server/llm/model.ts b/app/lib/.server/llm/model.ts
index f0d695c47..7055ee9f7 100644
--- a/app/lib/.server/llm/model.ts
+++ b/app/lib/.server/llm/model.ts
@@ -1,9 +1,19 @@
-import { createAnthropic } from '@ai-sdk/anthropic';
+import { createGroq } from 'groq';
 
-export function getAnthropicModel(apiKey: string) {
-  const anthropic = createAnthropic({
+export function createGroqInstance(apiKey: string) {
+  return createGroq({
     apiKey,
   });
+}
+
+export function getGroqModel(apiKey: string) {
+  const groq = createGroqInstance(apiKey);
+
+  return groq('groq-model');
+}
+
+export function getGroqLlamaModel(apiKey: string) {
+  const groq = createGroqInstance(apiKey);
 
-  return anthropic('claude-3-5-sonnet-20240620');
+  return groq('llama3.3-70b');
 }
diff --git a/app/lib/.server/llm/stream-text.ts b/app/lib/.server/llm/stream-text.ts
index cf937fd00..dce20a492 100644
--- a/app/lib/.server/llm/stream-text.ts
+++ b/app/lib/.server/llm/stream-text.ts
@@ -1,6 +1,6 @@
 import { streamText as _streamText, convertToCoreMessages } from 'ai';
 import { getAPIKey } from '~/lib/.server/llm/api-key';
-import { getAnthropicModel } from '~/lib/.server/llm/model';
+import { getGroqModel } from '~/lib/.server/llm/model';
 import { MAX_TOKENS } from './constants';
 import { getSystemPrompt } from './prompts';
 
@@ -23,11 +23,11 @@ export type StreamingOptions = Omit<Parameters<typeof _streamText>[0], 'model'>;
 
 export function streamText(messages: Messages, env: Env, options?: StreamingOptions) {
   return _streamText({
-    model: getAnthropicModel(getAPIKey(env)),
+    model: getGroqModel(getAPIKey(env)),
     system: getSystemPrompt(),
     maxTokens: MAX_TOKENS,
     headers: {
-      'anthropic-beta': 'max-tokens-3-5-sonnet-2024-07-15',
+      'groq-beta': 'max-tokens-3-5-sonnet-2024-07-15',
     },
     messages: convertToCoreMessages(messages),
     ...options,
diff --git a/app/routes/api.chat.ts b/app/routes/api.chat.ts
index b685ac853..b9b2ea7c7 100644
--- a/app/routes/api.chat.ts
+++ b/app/routes/api.chat.ts
@@ -3,6 +3,7 @@ import { MAX_RESPONSE_SEGMENTS, MAX_TOKENS } from '~/lib/.server/llm/constants';
 import { CONTINUE_PROMPT } from '~/lib/.server/llm/prompts';
 import { streamText, type Messages, type StreamingOptions } from '~/lib/.server/llm/stream-text';
 import SwitchableStream from '~/lib/.server/llm/switchable-stream';
+import { createGroq } from 'groq';
 
 export async function action(args: ActionFunctionArgs) {
   return chatAction(args);
diff --git a/app/routes/api.enhancer.ts b/app/routes/api.enhancer.ts
index 5c8175ca3..a079186fa 100644
--- a/app/routes/api.enhancer.ts
+++ b/app/routes/api.enhancer.ts
@@ -2,6 +2,7 @@ import { type ActionFunctionArgs } from '@remix-run/cloudflare';
 import { StreamingTextResponse, parseStreamPart } from 'ai';
 import { streamText } from '~/lib/.server/llm/stream-text';
 import { stripIndents } from '~/utils/stripIndent';
+import { createGroq } from 'groq';
 
 const encoder = new TextEncoder();
 const decoder = new TextDecoder();