Skip to content

Commit

Permalink
Add groq support
Browse files Browse the repository at this point in the history
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`.
  • Loading branch information
vishwamartur committed Dec 28, 2024
1 parent eda10b1 commit fdefd71
Show file tree
Hide file tree
Showing 6 changed files with 23 additions and 11 deletions.
2 changes: 1 addition & 1 deletion app/lib/.server/llm/api-key.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
6 changes: 3 additions & 3 deletions app/lib/.server/llm/constants.ts
Original file line number Diff line number Diff line change
@@ -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;
18 changes: 14 additions & 4 deletions app/lib/.server/llm/model.ts
Original file line number Diff line number Diff line change
@@ -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');
}
6 changes: 3 additions & 3 deletions app/lib/.server/llm/stream-text.ts
Original file line number Diff line number Diff line change
@@ -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';

Expand All @@ -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,
Expand Down
1 change: 1 addition & 0 deletions app/routes/api.chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
1 change: 1 addition & 0 deletions app/routes/api.enhancer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down

0 comments on commit fdefd71

Please sign in to comment.