strtus/Vercel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add additional auth check to openai streaming requests

Browse Source
This commit is contained in:
Sam Becker 2024-03-21 22:47:21 -05:00
parent 8a03ea8217
commit f6edaf07f1
1 changed files with 37 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

71
src/services/openai.ts
View File

@ -5,6 +5,7 @@ import { createStreamableValue, render } from 'ai/rsc';
import { kv } from '@vercel/kv';
import { Ratelimit } from '@upstash/ratelimit';
import { AI_TEXT_GENERATION_ENABLED, HAS_VERCEL_KV } from '@/site/config';
import { safelyRunAdminServerAction } from '@/auth';
const RATE_LIMIT_IDENTIFIER = 'openai-image-query';
const RATE_LIMIT_MAX_QUERIES_PER_HOUR = 100;
@ -25,43 +26,45 @@ export const streamOpenAiImageQuery = async (
imageBase64: string,
query: string,
) => {
if (ratelimit) {
const { success } = await ratelimit.limit(RATE_LIMIT_IDENTIFIER);
if (!success) {
console.error('OpenAI rate limit exceeded');
throw new Error('OpenAI rate limit exceeded');
return safelyRunAdminServerAction(async () => {
if (ratelimit) {
const { success } = await ratelimit.limit(RATE_LIMIT_IDENTIFIER);
if (!success) {
console.error('OpenAI rate limit exceeded');
throw new Error('OpenAI rate limit exceeded');
}
}
}
const stream = createStreamableValue('');
const stream = createStreamableValue('');
if (provider) {
render({
provider,
model: 'gpt-4-vision-preview',
messages: [{
'role': 'user',
'content': [
{
'type': 'text',
'text': query,
}, {
'type': 'image_url',
'image_url': {
'url': imageBase64,
if (provider) {
render({
provider,
model: 'gpt-4-vision-preview',
messages: [{
'role': 'user',
'content': [
{
'type': 'text',
'text': query,
}, {
'type': 'image_url',
'image_url': {
'url': imageBase64,
},
},
},
],
}],
text: ({ content, done }): any => {
if (done) {
stream.done(content);
} else {
stream.update(content);
}
},
});
}
],
}],
text: ({ content, done }): any => {
if (done) {
stream.done(content);
} else {
stream.update(content);
}
},
});
}
return stream.value;
return stream.value;
});
};