File size: 1,841 Bytes
3c3f089
 
a21e1b7
 
0e460a1
3816441
a21e1b7
 
3c3f089
 
b32fabf
 
 
 
b4f7005
b32fabf
b4f7005
b32fabf
87495f8
 
 
 
 
 
 
 
3816441
 
 
 
a21e1b7
440e084
a21e1b7
 
 
 
 
50e3782
a21e1b7
 
50e3782
0e460a1
a21e1b7
b32fabf
 
a21e1b7
b32fabf
3c3f089
a21e1b7
 
 
3c3f089
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import { OpenAIStream, StreamingTextResponse } from "ai";
import { Configuration, OpenAIApi } from "openai-edge";
import { HfInference } from '@huggingface/inference';
import { HuggingFaceStream } from 'ai';
import { experimental_buildLlama2Prompt } from 'ai/prompts';
import { LLMStream } from "../../../utils/llm";

export const runtime = 'edge';

export default async function(req: Request) {
    let { messages, aiProvider = 'openai', systemMessage, url } = await req.json();

    // Set up configurations for OpenAI
    const openaiConfig = new Configuration({
        apiKey: process.env.OPENAI_API_KEY
    } as any);
    const openai = new OpenAIApi(openaiConfig, url ? url : undefined);
    const Hf = new HfInference(process.env.HUGGINGFACE_API_KEY);

    // Prepend the system message if it's not already there
    if (messages.length === 0 || messages[0].role !== "system") {
        messages = [{
            role: "system",
            content: systemMessage
        }, ...messages];
    }
    if (url) {
        const stream = await LLMStream(url, messages);
        return new StreamingTextResponse(stream);
    } else if (aiProvider === 'openai') {
        const response = await openai.createChatCompletion({
            model: 'gpt-4o-mini',
            stream: true,
            messages
        });
        const stream = OpenAIStream(response);
        return new StreamingTextResponse(stream);
    } else {
        const response = Hf.textGenerationStream({
            // @ts-ignore
            model: aiProvider,
            inputs: experimental_buildLlama2Prompt(messages),
            parameters: {
                repetition_penalty: 1,
                return_full_text: false
            }
        });

        const stream = HuggingFaceStream(response);
        return new StreamingTextResponse(stream);
    }
}