Skip to content

Commit d72435f

Browse files
CoveMBCoveMB
authored
Merge branch 'master' into remix-encoding-fix
2 parents 6933f06 + e568983 commit d72435f

File tree

9 files changed

+257
-321
lines changed

9 files changed

+257
-321
lines changed

‎packages/ui/api/ai.ts‎

Lines changed: 25 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,22 @@
1-
import { OpenAIStream } from 'ai';
21
import * as solidityFunctions from './ai-assistant/function-definitions/solidity.ts';
32
import * as cairoFunctions from './ai-assistant/function-definitions/cairo.ts';
43
import * as cairoAlphaFunctions from './ai-assistant/function-definitions/cairo-alpha.ts';
54
import * as stellarFunctions from './ai-assistant/function-definitions/stellar.ts';
65
import * as stylusFunctions from './ai-assistant/function-definitions/stylus.ts';
76
import { saveChatInRedisIfDoesNotExist } from './services/redis.ts';
8-
import { getOpenAiInstance } from './services/open-ai.ts';
9-
import { getEnvironmentVariableOr } from './utils/env.ts';
10-
import type { AiChatBodyRequest, Chat } from './ai-assistant/types/assistant.ts';
7+
import type { ChatMessages } from './services/open-ai.ts';
8+
import { createOpenAiCompletionStream } from './services/open-ai.ts';
9+
import type { AiChatBodyRequest } from './ai-assistant/types/assistant.ts';
1110
import type { SupportedLanguage } from './ai-assistant/types/languages.ts';
1211
import type {
1312
AllContractsAIFunctionDefinitions,
1413
SimpleAiFunctionDefinition,
1514
} from './ai-assistant/types/function-definition.ts';
15+
import { Cors } from './utils/cors.ts';
1616

1717
const getFunctionsContext = <TLanguage extends SupportedLanguage = SupportedLanguage>(
1818
language: TLanguage,
19-
): SimpleAiFunctionDefinition[] => {
19+
): { type: 'function'; function: SimpleAiFunctionDefinition }[] => {
2020
const functionPerLanguages: AllContractsAIFunctionDefinitions = {
2121
solidity: solidityFunctions,
2222
cairo: cairoFunctions,
@@ -25,10 +25,15 @@ const getFunctionsContext = <TLanguage extends SupportedLanguage = SupportedLang
2525
stylus: stylusFunctions,
2626
};
2727

28-
return Object.values(functionPerLanguages[language] ?? {});
28+
return (Object.values(functionPerLanguages[language] ?? {}) as SimpleAiFunctionDefinition[]).map(
29+
functionDefinition => ({
30+
type: 'function',
31+
function: functionDefinition,
32+
}),
33+
);
2934
};
3035

31-
const buildAiChatMessages = (request: AiChatBodyRequest): Chat[] => {
36+
const buildAiChatMessages = (request: AiChatBodyRequest): ChatMessages => {
3237
const validatedMessages = request.messages.filter((message: { role: string; content: string }) => {
3338
return message.content.length < 500;
3439
});
@@ -51,27 +56,23 @@ export default async (req: Request): Promise<Response> => {
5156
try {
5257
const aiChatBodyRequest: AiChatBodyRequest = await req.json();
5358

54-
const openai = getOpenAiInstance();
59+
const chatMessages = buildAiChatMessages(aiChatBodyRequest);
5560

56-
const aiChatMessages = buildAiChatMessages(aiChatBodyRequest);
57-
58-
const response = await openai.chat.completions.create({
59-
model: getEnvironmentVariableOr('OPENAI_MODEL', 'gpt-4o-mini'),
60-
messages: aiChatMessages,
61-
functions: getFunctionsContext(aiChatBodyRequest.language),
62-
temperature: 0.7,
63-
stream: true,
64-
});
65-
66-
const stream = OpenAIStream(response, {
67-
onCompletion: saveChatInRedisIfDoesNotExist(aiChatBodyRequest.chatId, aiChatMessages),
61+
const openAiStream = createOpenAiCompletionStream({
62+
streamParams: {
63+
messages: chatMessages,
64+
tools: getFunctionsContext(aiChatBodyRequest.language),
65+
},
66+
chatId: aiChatBodyRequest.chatId,
67+
onAiStreamCompletion: async ({ chatId, chatMessages }, finalStreamResult) =>
68+
finalStreamResult && (await saveChatInRedisIfDoesNotExist(chatId, chatMessages)(finalStreamResult)),
6869
});
6970

70-
return new Response(stream, {
71+
return new Response(openAiStream, {
7172
headers: new Headers({
72-
'Access-Control-Allow-Origin': '*',
73-
'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
74-
'Content-Type': 'text/html; charset=utf-8',
73+
...Cors,
74+
'Content-Type': 'application/x-ndjson; charset=utf-8',
75+
'Cache-Control': 'no-transform',
7576
}),
7677
});
7778
} catch (e) {

‎packages/ui/api/services/open-ai.ts‎

Lines changed: 90 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import OpenAI from 'openai';
2-
import { getEnvironmentVariablesOrFail } from '../utils/env.ts';
2+
import { getEnvironmentVariableOr, getEnvironmentVariablesOrFail } from '../utils/env.ts';
33

44
export const getOpenAiInstance = () => {
55
const { OPENAI_API_KEY } = getEnvironmentVariablesOrFail('OPENAI_API_KEY');
@@ -8,3 +8,92 @@ export const getOpenAiInstance = () => {
88
apiKey: OPENAI_API_KEY,
99
});
1010
};
11+
12+
export type ChatMessages = OpenAI.Chat.ChatCompletionCreateParams['messages'];
13+
14+
type ProcessOpenAiStreamParams = {
15+
openAiStream: AsyncIterable<OpenAI.Chat.Completions.ChatCompletionChunk> & { controller?: AbortController };
16+
chatMessages: ChatMessages;
17+
chatId: string;
18+
};
19+
type OnAiStreamCompletion = (
20+
processOpenAiStreamParams: ProcessOpenAiStreamParams,
21+
processOpenAiStreamResults: string,
22+
) => Promise<unknown>;
23+
24+
const processOpenAIStream = (
25+
{ openAiStream, chatMessages, chatId }: ProcessOpenAiStreamParams,
26+
onAiStreamCompletion?: OnAiStreamCompletion,
27+
) =>
28+
new ReadableStream({
29+
async pull(controller: ReadableStreamDefaultController<Uint8Array>) {
30+
const textEncoder = new TextEncoder();
31+
let finalResponse = '';
32+
const finalToolCall = {
33+
name: '',
34+
arguments: '',
35+
};
36+
37+
try {
38+
for await (const chunk of openAiStream) {
39+
const delta = chunk?.choices?.[0]?.delta;
40+
const isToolCallBuilding = Boolean(delta.tool_calls?.[0]?.index === 0);
41+
const isFunctionCallFinished = Boolean(chunk?.choices?.[0]?.finish_reason === 'tool_calls');
42+
43+
if (delta.content) {
44+
finalResponse += delta.content;
45+
controller.enqueue(textEncoder.encode(delta.content));
46+
} else if (isToolCallBuilding && !isFunctionCallFinished) {
47+
finalToolCall.name += delta.tool_calls?.[0]?.function?.name || '';
48+
finalToolCall.arguments += delta.tool_calls?.[0]?.function?.arguments || '';
49+
} else if (isFunctionCallFinished)
50+
controller.enqueue(
51+
textEncoder.encode(
52+
JSON.stringify({
53+
function_call: finalToolCall,
54+
}),
55+
),
56+
);
57+
}
58+
59+
controller.close();
60+
} catch (error) {
61+
console.error('OpenAI streaming error:', error);
62+
controller.error(error);
63+
return;
64+
} finally {
65+
if (onAiStreamCompletion)
66+
await onAiStreamCompletion(
67+
{ openAiStream, chatId, chatMessages },
68+
finalResponse || JSON.stringify(finalToolCall),
69+
);
70+
}
71+
},
72+
cancel() {
73+
openAiStream.controller?.abort?.();
74+
},
75+
});
76+
77+
export const createOpenAiCompletionStream = ({
78+
streamParams,
79+
chatId,
80+
onAiStreamCompletion,
81+
}: {
82+
streamParams: Omit<OpenAI.Chat.ChatCompletionCreateParams, 'model' | 'stream'>;
83+
chatId: string;
84+
onAiStreamCompletion?: OnAiStreamCompletion;
85+
}) => {
86+
const openai = getOpenAiInstance();
87+
88+
const openAiStream = openai.chat.completions.stream({
89+
model: getEnvironmentVariableOr('OPENAI_MODEL', 'gpt-4o-mini'),
90+
temperature: 0.7,
91+
stream: true,
92+
...streamParams,
93+
});
94+
95+
return processOpenAIStream(
96+
{ openAiStream, chatMessages: streamParams.messages, chatId: chatId },
97+
onAiStreamCompletion,
98+
);
99+
};

‎packages/ui/api/services/redis.ts‎

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
import { Redis } from '@upstash/redis';
22
import RedisMock from './dev-mocks/redis.ts';
33
import { getEnvironmentVariableOr, getEnvironmentVariablesOrFail } from '../utils/env.ts';
4-
import type { Chat } from '../ai-assistant/types/assistant.ts';
54

65
type RedisChat = {
76
id: string;
8-
messages: Chat[];
7+
messages: unknown[];
98
updatedAt: number;
109
createdAt?: number;
1110
};
@@ -21,25 +20,26 @@ export const getRedisInstance = () => {
2120
});
2221
};
2322

24-
export const saveChatInRedisIfDoesNotExist = (chatId: string, chatMessages: Chat[]) => async (completion: string) => {
25-
const redis = getRedisInstance();
26-
27-
const updatedAt = Date.now();
28-
const payload: RedisChat = {
29-
id: chatId,
30-
updatedAt,
31-
messages: [
32-
...chatMessages,
33-
{
34-
content: completion,
35-
role: 'assistant',
36-
},
37-
],
38-
};
23+
export const saveChatInRedisIfDoesNotExist =
24+
(chatId: string, chatMessages: unknown[]) => async (completion: string) => {
25+
const redis = getRedisInstance();
3926

40-
const exists = await redis.exists(`chat:${chatId}`);
27+
const updatedAt = Date.now();
28+
const payload: RedisChat = {
29+
id: chatId,
30+
updatedAt,
31+
messages: [
32+
...chatMessages,
33+
{
34+
content: completion,
35+
role: 'assistant',
36+
},
37+
],
38+
};
4139

42-
if (!exists) payload.createdAt = updatedAt;
40+
const exists = await redis.exists(`chat:${chatId}`);
4341

44-
await redis.hset(`chat:${chatId}`, payload);
45-
};
42+
if (!exists) payload.createdAt = updatedAt;
43+
44+
await redis.hset(`chat:${chatId}`, payload);
45+
};

‎packages/ui/api/utils/cors.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
export const Cors = {
2+
'Access-Control-Allow-Origin': '*',
3+
'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
4+
'Access-Control-Allow-Headers': 'Content-Type, Authorization',
5+
};

0 commit comments

Comments
 (0)