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`.
This commit is contained in:
Vishwanath Martur 2024-12-28 13:58:12 +05:30
parent eda10b1212
commit fdefd717b0
6 changed files with 24 additions and 12 deletions

View File

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

View File

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

View File

@ -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,
});
return anthropic('claude-3-5-sonnet-20240620');
}
export function getGroqModel(apiKey: string) {
const groq = createGroqInstance(apiKey);
return groq('groq-model');
}
export function getGroqLlamaModel(apiKey: string) {
const groq = createGroqInstance(apiKey);
return groq('llama3.3-70b');
}

View File

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

View File

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

View File

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