Add groq support

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

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;

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');
 }

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,

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);

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();