Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add grok model support #15

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 39 additions & 2 deletions scripts/tweets2character.js
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,40 @@ const runChatCompletion = async (messages, useGrammar = false, model) => {
const content = data.content[0].text;
const parsed = parseJsonFromMarkdown(content) || JSON.parse(content);
return parsed;
} else if(model === 'grok'){
const modelName = 'grok-beta';
const response = await fetch('https://api.x.ai/v1/chat/completions', {
method: 'POST',
headers: {
'Content-Type': 'application/json',
'Authorization': `Bearer ${process.env.GROK_API_KEY}`,
},
body: JSON.stringify({
model: modelName,
temperature: 0,
stream: false,
messages: [
{
role: "user",
content: messages[0].content
}
],
}),
});

if (response.status === 429) {
await new Promise(resolve => setTimeout(resolve, 30000));
return runChatCompletion(messages, useGrammar, model);
}

if (!response.ok) {
throw new Error(`HTTP error! status: ${response.status}`);
}

const data = await response.json();
const content = data.choices[0].message.content.trim();
const parsed = parseJsonFromMarkdown(content) || JSON.parse(content);
return parsed;
}
};

Expand Down Expand Up @@ -407,6 +441,7 @@ process.on('unhandledRejection', (reason, promise) => {
program
.option('--openai <api_key>', 'OpenAI API key')
.option('--claude <api_key>', 'Claude API key')
.option('--grok <api_key>', 'Grok API key')
.parse(process.argv);

const limitConcurrency = async (tasks, concurrencyLimit) => {
Expand Down Expand Up @@ -472,6 +507,8 @@ const validateApiKey = (apiKey, model) => {
return apiKey.trim().startsWith('sk-');
} else if (model === 'claude') {
return apiKey.trim().length > 0;
} else if (model === 'grok') {
return apiKey.trim().startsWith('xai-');
}
return false;
};
Expand All @@ -494,7 +531,7 @@ const resumeOrStartNewSession = async (projectCache, archivePath) => {
}

if (!projectCache.unfinishedSession) {
projectCache.model = await promptUser('Select model (openai/claude): ');
projectCache.model = await promptUser('Select model (openai/claude/grok): ');
projectCache.basicUserInfo = await promptUser('Enter additional user info that might help the summarizer (real name, nicknames and handles, age, past employment vs current, etc): ');
projectCache.unfinishedSession = {
currentChunk: 0,
Expand Down Expand Up @@ -531,7 +568,7 @@ const main = async () => {
let projectCache = loadProjectCache(archivePath) || {};

projectCache = await resumeOrStartNewSession(projectCache, archivePath);

const apiKey = await getApiKey(projectCache.model);
if (!apiKey) {
throw new Error(`Failed to get a valid API key for ${projectCache.model}`);
Expand Down