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'

// Configurations for OpenAI
const openaiConfig = new Configuration({
    apiKey: process.env.OPENAI_API_KEY,
});
const openai = new OpenAIApi(openaiConfig);

// Create a new HuggingFace Inference instance
const Hf = new HfInference(process.env.HUGGINGFACE_API_KEY);

export const runtime = 'edge';

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

    // 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 (aiProvider === 'openai') {
        const response = await openai.createChatCompletion({
            model: 'gpt-4',
            stream: true,
            messages
        });

        const stream = OpenAIStream(response);
        return new StreamingTextResponse(stream);
    } else if (aiProvider === 'huggingface') {
        const response = Hf.textGenerationStream({
            // @ts-ignore
            model: 'meta-llama/Llama-2-7b-chat-hf',
            inputs: experimental_buildLlama2Prompt(messages),
            parameters: {
              max_new_tokens: 500,
              repetition_penalty: 1,
              truncate: 4000,
              return_full_text: false
            }
          })

        const stream = HuggingFaceStream(response);
        return new StreamingTextResponse(stream);
    } else {
        throw new Error(`Unsupported AI provider: ${aiProvider}`);
    }
}