📦 NEW: XAI models support (#76)

* 📦 NEW: XAI models support

* 🐛 FIX: Import

Author: msaaddev

examples/nextjs/.env.baseai.example

@@ -19,3 +19,4 @@ GROQ_API_KEY=
 MISTRAL_API_KEY=
 PERPLEXITY_API_KEY=
 TOGETHER_API_KEY=
+XAI_API_KEY=

packages/baseai/src/data/models.ts

@@ -90,6 +90,8 @@ export const PERPLEXITY: string = 'Perplexity';
 export const DEEPINFRA: string = 'deepinfra';
 export const BEDROCK: string = 'bedrock';
 export const AZURE_OPEN_AI: string = 'azure-openai';
+export const X_AI: string = 'xAI';
+export const OLLAMA: string = 'ollama';
 
 interface Model {
 	id: string;
@@ -570,6 +572,18 @@ export const modelsByProvider: ModelsByProviderInclCosts = {
 			promptCost: 0.2,
 			completionCost: 0.2
 		}
+	],
+	[X_AI]: [
+		{
+			id: 'grok-beta',
+			provider: X_AI,
+			promptCost: 5,
+			completionCost: 15,
+			toolSupport: {
+				toolChoice: true,
+				parallelToolCalls: false
+			}
+		}
 	]
 };
 

packages/baseai/src/dev/data/models.ts

@@ -11,6 +11,7 @@ export const DEEPINFRA: string = 'deepinfra';
 export const BEDROCK: string = 'bedrock';
 export const AZURE_OPEN_AI: string = 'azure-openai';
 export const OLLAMA: string = 'ollama';
+export const X_AI: string = 'xAI';
 
 interface Model {
 	id: string;
@@ -511,6 +512,14 @@ export const modelsByProvider: ModelsByProviderInclCosts = {
 			promptCost: 1,
 			completionCost: 3
 		}
+	],
+	[X_AI]: [
+		{
+			id: 'grok-beta',
+			provider: X_AI,
+			promptCost: 5,
+			completionCost: 15
+		}
 	]
 };
 

packages/baseai/src/dev/llms/call-llm.ts

@@ -7,7 +7,8 @@ import {
 	OLLAMA,
 	OPEN_AI,
 	PERPLEXITY,
-	TOGETHER_AI
+	TOGETHER_AI,
+	X_AI
 } from '@/dev/data/models';
 
 import { addContextFromMemory } from '@/utils/memory/lib';
@@ -25,6 +26,7 @@ import { callOllama } from './call-ollama';
 import { callOpenAI } from './call-openai';
 import { callPerplexity } from './call-perplexity';
 import { callTogether } from './call-together';
+import { callXAI } from './call-xai';
 
 export async function callLLM({
 	pipe,
@@ -110,6 +112,16 @@ export async function callLLM({
 			});
 		}
 
+		if (modelProvider === X_AI) {
+			dlog('XAI', '✅');
+			return await callXAI({
+				pipe,
+				messages,
+				llmApiKey,
+				stream
+			});
+		}
+
 		if (modelProvider === COHERE) {
 			dlog('COHERE', '✅');
 			return await callCohere({

packages/baseai/src/dev/llms/call-xai.ts

@@ -0,0 +1,50 @@
+import OpenAI from 'openai';
+import { dlog } from '../utils/dlog';
+import { X_AI } from '../data/models';
+import { handleLlmError } from './utils';
+import type { Message } from 'types/pipe';
+import type { ModelParams } from 'types/providers';
+import { addToolsToParams } from '../utils/add-tools-to-params';
+
+export async function callXAI({
+	pipe,
+	stream,
+	llmApiKey,
+	messages
+}: {
+	pipe: any;
+	stream: boolean;
+	llmApiKey: string;
+	messages: Message[];
+}) {
+	try {
+		const modelParams = buildModelParams(pipe, stream, messages);
+
+		// LLM.
+		const groq = new OpenAI({
+			apiKey: llmApiKey,
+			baseURL: 'https://api.x.ai/v1'
+		});
+
+		// Add tools (functions) to modelParams
+		addToolsToParams(modelParams, pipe);
+		dlog('modelParams', modelParams);
+
+		return await groq.chat.completions.create(modelParams as any);
+	} catch (error: any) {
+		handleLlmError({ error, provider: X_AI });
+	}
+}
+
+function buildModelParams(
+	pipe: any,
+	stream: boolean,
+	messages: Message[]
+): ModelParams {
+	return {
+		messages,
+		stream,
+		model: pipe.model.name,
+		...pipe.model.params
+	};
+}

packages/baseai/src/dev/utils/get-llm-api-key.ts

@@ -6,7 +6,8 @@ import {
 	GROQ,
 	OPEN_AI,
 	PERPLEXITY,
-	TOGETHER_AI
+	TOGETHER_AI,
+	X_AI
 } from '@/dev/data/models';
 
 export function getLLMApiKey(modelProvider: string): string {
@@ -27,6 +28,8 @@ export function getLLMApiKey(modelProvider: string): string {
 			return process.env.FIREWORKS_API_KEY || '';
 		case modelProvider.includes(PERPLEXITY):
 			return process.env.PERPLEXITY_API_KEY || '';
+		case modelProvider.includes(X_AI):
+			return process.env.XAI_API_KEY || '';
 		default:
 			throw new Error(`Unsupported model provider: ${modelProvider}`);
 	}

packages/baseai/src/init/index.ts

@@ -292,6 +292,7 @@ GROQ_API_KEY=
 MISTRAL_API_KEY=
 PERPLEXITY_API_KEY=
 TOGETHER_API_KEY=
+XAI_API_KEY=
 `;
 
 	try {

packages/baseai/src/utils/to-old-pipe-format.ts

@@ -1,14 +1,15 @@
-import { OLLAMA } from '@/dev/data/models';
 import type { Pipe, PipeOld } from './../../types/pipe';
 import {
 	ANTHROPIC,
 	COHERE,
 	FIREWORKS_AI,
 	GOOGLE,
 	GROQ,
+	OLLAMA,
 	OPEN_AI,
 	PERPLEXITY,
-	TOGETHER_AI
+	TOGETHER_AI,
+	X_AI
 } from './../data/models';
 
 type Provider =
@@ -102,7 +103,8 @@ function getProvider(providerString: string): Provider {
 		cohere: COHERE,
 		fireworks: FIREWORKS_AI,
 		perplexity: PERPLEXITY,
-		ollama: OLLAMA
+		ollama: OLLAMA,
+		xai: X_AI
 	};
 
 	const provider = providerMap[providerString.toLowerCase()];

packages/core/src/data/models.ts

@@ -11,6 +11,7 @@ export const DEEPINFRA: string = 'deepinfra';
 export const BEDROCK: string = 'bedrock';
 export const AZURE_OPEN_AI: string = 'azure-openai';
 export const OLLAMA: string = 'ollama';
+export const X_AI: string = 'xAI';
 
 interface Model {
 	id: string;
@@ -400,6 +401,14 @@ export const modelsByProvider: ModelsByProviderInclCosts = {
 			completionCost: 3,
 		},
 	],
+	[X_AI]: [
+		{
+			id: 'grok-beta',
+			provider: X_AI,
+			promptCost: 5,
+			completionCost: 15,
+		},
+	],
 };
 
 export const jsonModeModels = [

packages/core/src/utils/get-llm-api-key.ts

@@ -8,6 +8,7 @@ import {
 	OPEN_AI,
 	PERPLEXITY,
 	TOGETHER_AI,
+	X_AI,
 } from '../data/models';
 
 export function getLLMApiKey(modelProvider: string): string {
@@ -30,6 +31,8 @@ export function getLLMApiKey(modelProvider: string): string {
 			return process.env.PERPLEXITY_API_KEY || '';
 		case modelProvider.includes(OLLAMA):
 			return process.env.OLLAMA_API_KEY || '';
+		case modelProvider.includes(X_AI):
+			return process.env.XAI_API_KEY || '';
 
 		default:
 			throw new Error(`Unsupported model provider: ${modelProvider}`);