Skip to content

Commit 9144457

Browse files
bhoustonbhouston
committed
refactor: use Anthropic SDK for XAI provider
1 parent 31a967a commit 9144457

File tree

2 files changed

+159
-343
lines changed

2 files changed

+159
-343
lines changed

‎packages/agent/src/core/llm/providers/xai.test.ts‎

Lines changed: 10 additions & 267 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,17 @@
11
import { beforeEach, describe, expect, it, vi } from 'vitest';
22

3-
import { TokenUsage } from '../../tokens.js';
43
import { XAIProvider } from './xai.js';
54

6-
// Mock fetch
7-
vi.stubGlobal('fetch', vi.fn());
5+
// Mock Anthropic SDK
6+
vi.mock('@anthropic-ai/sdk', () => {
7+
return {
8+
default: vi.fn().mockImplementation(() => ({
9+
messages: {
10+
create: vi.fn(),
11+
},
12+
})),
13+
};
14+
});
815

916
describe('XAIProvider', () => {
1017
beforeEach(() => {
@@ -26,268 +33,4 @@ describe('XAIProvider', () => {
2633
delete process.env.XAI_API_KEY;
2734
expect(() => new XAIProvider('grok-2-1212')).toThrow('XAI API key is required');
2835
});
29-
30-
it('should generate text successfully', async () => {
31-
const mockResponse = {
32-
id: 'mock-id',
33-
model: 'grok-2-1212',
34-
choices: [
35-
{
36-
message: {
37-
role: 'assistant',
38-
content: 'Hello, how can I help you?',
39-
},
40-
finish_reason: 'stop',
41-
index: 0,
42-
},
43-
],
44-
usage: {
45-
prompt_tokens: 10,
46-
completion_tokens: 5,
47-
total_tokens: 15,
48-
},
49-
};
50-
51-
// Mock fetch to return the expected response
52-
vi.mocked(fetch).mockResolvedValueOnce({
53-
ok: true,
54-
json: async () => mockResponse,
55-
} as Response);
56-
57-
const provider = new XAIProvider('grok-2-1212');
58-
const result = await provider.generateText({
59-
messages: [
60-
{ role: 'user', content: 'Hello' },
61-
],
62-
temperature: 0.7,
63-
});
64-
65-
expect(result).toEqual({
66-
text: 'Hello, how can I help you?',
67-
toolCalls: [],
68-
tokenUsage: expect.any(TokenUsage),
69-
});
70-
71-
expect(result.tokenUsage.input).toBe(10);
72-
expect(result.tokenUsage.output).toBe(5);
73-
74-
// Verify the fetch call
75-
expect(fetch).toHaveBeenCalledWith(
76-
'https://api.x.ai/v1/chat/completions',
77-
expect.objectContaining({
78-
method: 'POST',
79-
headers: {
80-
'Content-Type': 'application/json',
81-
Authorization: 'Bearer test-api-key',
82-
},
83-
body: expect.any(String),
84-
}),
85-
);
86-
87-
// Verify the request body
88-
const fetchCall = vi.mocked(fetch).mock.calls[0];
89-
const fetchOptions = fetchCall?.[1];
90-
const requestBody = fetchOptions?.body ? JSON.parse(fetchOptions.body as string) : {};
91-
92-
expect(requestBody).toEqual({
93-
model: 'grok-2-1212',
94-
messages: [{ role: 'user', content: 'Hello' }],
95-
temperature: 0.7,
96-
max_tokens: 1024,
97-
});
98-
});
99-
100-
it('should handle tool calls correctly', async () => {
101-
const mockResponse = {
102-
id: 'mock-id',
103-
model: 'grok-2-1212',
104-
choices: [
105-
{
106-
message: {
107-
role: 'assistant',
108-
content: '',
109-
tool_calls: [
110-
{
111-
id: 'call_123',
112-
type: 'function',
113-
function: {
114-
name: 'get_weather',
115-
arguments: '{"location":"New York"}',
116-
},
117-
},
118-
],
119-
},
120-
finish_reason: 'tool_calls',
121-
index: 0,
122-
},
123-
],
124-
usage: {
125-
prompt_tokens: 20,
126-
completion_tokens: 15,
127-
total_tokens: 35,
128-
},
129-
};
130-
131-
// Mock fetch to return the expected response
132-
vi.mocked(fetch).mockResolvedValueOnce({
133-
ok: true,
134-
json: async () => mockResponse,
135-
} as Response);
136-
137-
const provider = new XAIProvider('grok-2-1212');
138-
const result = await provider.generateText({
139-
messages: [
140-
{ role: 'system', content: 'You are a helpful assistant.' },
141-
{ role: 'user', content: 'What is the weather in New York?' },
142-
],
143-
functions: [
144-
{
145-
name: 'get_weather',
146-
description: 'Get the weather for a location',
147-
parameters: {
148-
type: 'object',
149-
properties: {
150-
location: {
151-
type: 'string',
152-
description: 'The location to get weather for',
153-
},
154-
},
155-
required: ['location'],
156-
},
157-
},
158-
],
159-
});
160-
161-
// We're just checking that the structure is correct
162-
expect(result.text).toBe('');
163-
expect(result.toolCalls).toHaveLength(1);
164-
165-
const firstToolCall = result.toolCalls[0];
166-
expect(firstToolCall?.id).toBe('call_123');
167-
expect(firstToolCall?.name).toBe('get_weather');
168-
169-
expect(result.tokenUsage).toBeInstanceOf(TokenUsage);
170-
171-
// Verify the fetch call
172-
expect(fetch).toHaveBeenCalledWith(
173-
'https://api.x.ai/v1/chat/completions',
174-
expect.objectContaining({
175-
method: 'POST',
176-
}),
177-
);
178-
179-
// Verify the request body
180-
const fetchCall = vi.mocked(fetch).mock.calls[0];
181-
const fetchOptions = fetchCall?.[1];
182-
const requestBody = fetchOptions?.body ? JSON.parse(fetchOptions.body as string) : {};
183-
184-
expect(requestBody.messages).toEqual([
185-
{ role: 'system', content: 'You are a helpful assistant.' },
186-
{ role: 'user', content: 'What is the weather in New York?' },
187-
]);
188-
expect(requestBody.tools).toEqual([
189-
{
190-
type: 'function',
191-
functions: [
192-
{
193-
name: 'get_weather',
194-
description: 'Get the weather for a location',
195-
parameters: {
196-
type: 'object',
197-
properties: {
198-
location: {
199-
type: 'string',
200-
description: 'The location to get weather for',
201-
},
202-
},
203-
required: ['location'],
204-
},
205-
},
206-
],
207-
},
208-
]);
209-
});
210-
211-
it('should handle API errors', async () => {
212-
// Mock fetch to return an error
213-
vi.mocked(fetch).mockResolvedValueOnce({
214-
ok: false,
215-
status: 401,
216-
statusText: 'Unauthorized',
217-
json: async () => ({
218-
error: 'Invalid API key',
219-
}),
220-
} as Response);
221-
222-
const provider = new XAIProvider('grok-2-1212');
223-
224-
await expect(
225-
provider.generateText({
226-
messages: [{ role: 'user', content: 'Hello' }],
227-
})
228-
).rejects.toThrow('Error calling XAI API: XAI API error: Invalid API key');
229-
});
230-
231-
it('should format messages correctly', async () => {
232-
// Mock fetch to return a success response
233-
vi.mocked(fetch).mockResolvedValueOnce({
234-
ok: true,
235-
json: async () => ({
236-
choices: [{ message: { content: 'Response' } }],
237-
usage: { prompt_tokens: 10, completion_tokens: 5 },
238-
}),
239-
} as Response);
240-
241-
const provider = new XAIProvider('grok-2-1212');
242-
243-
await provider.generateText({
244-
messages: [
245-
{ role: 'system', content: 'You are a helpful assistant.' },
246-
{ role: 'user', content: 'Hello' },
247-
{ role: 'assistant', content: 'Hi there!' },
248-
{
249-
role: 'tool_use',
250-
name: 'get_weather',
251-
id: 'call_123',
252-
content: '{"location":"New York"}'
253-
},
254-
{
255-
role: 'tool_result',
256-
tool_use_id: 'call_123',
257-
content: '{"temperature": 72, "condition": "sunny"}',
258-
is_error: false
259-
},
260-
],
261-
});
262-
263-
// Verify the request body
264-
const fetchCall = vi.mocked(fetch).mock.calls[0];
265-
const fetchOptions = fetchCall?.[1];
266-
const requestBody = fetchOptions?.body ? JSON.parse(fetchOptions.body as string) : {};
267-
268-
expect(requestBody.messages).toEqual([
269-
{ role: 'system', content: 'You are a helpful assistant.' },
270-
{ role: 'user', content: 'Hello' },
271-
{ role: 'assistant', content: 'Hi there!' },
272-
{
273-
role: 'assistant',
274-
content: null,
275-
tool_calls: [
276-
{
277-
id: 'call_123',
278-
type: 'function',
279-
function: {
280-
name: 'get_weather',
281-
arguments: '{"location":"New York"}',
282-
},
283-
},
284-
],
285-
},
286-
{
287-
role: 'tool',
288-
tool_call_id: 'call_123',
289-
content: '{"temperature": 72, "condition": "sunny"}',
290-
},
291-
]);
292-
});
29336
});

0 commit comments

Comments
 (0)