iimtt-api/src/routes/api/aiTest.js

const aiFollowupQuestion = async (req, res) => {
  const { MongoClient } = require('mongodb');
  const fetch = require('node-fetch');
  const { v4: uuidv4 } = require('uuid');
  const url = process.env.MONGODB_URL;
  const dbName = process.env.MONGO_DB_NAME;
  const client = new MongoClient(url);

  await client.connect();
  const database = client.db(dbName);
  const conversationsCollection = database.collection('conversations');

  const MAX_TOKENS = 200;
  const PROVIDER_KEYS = JSON.parse(process.env.PROVIDER_KEY || '[]');

  const PROVIDER_MODELS = {
    openai: [ "gpt-4o", "gpt-4o-mini", "o1", "o1-mini", "gpt-3.5-turbo", "gpt-3.5", "gpt-3.5-turbo-mini", "gpt-3.5-mini", "gpt-3.5-turbo-2", "gpt-3.5-2", "gpt-3.5-turbo-2-mini", "gpt-3.5-2-mini" ],
    deepseek: [ "deepseek-chat", "deepseek-reasoner", "deepseek-qa", "deepseek-qa-mini" ],
    // Add more providers and models as needed
  };

  async function fetchOpenAICompletion(messages, model, max_tokens = MAX_TOKENS) {
    const response = await fetch('https://api.openai.com/v1/chat/completions', {
      method: 'POST',
      headers: {
        'Authorization': `Bearer ${process.env.OPENAI_KEY}`,
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({
        model: model,
        messages: messages,
        max_tokens: max_tokens,
      }),
    });
    return response.json();
  }

  async function fetchDeepSeekCompletion(messages, model) {
    const response = await fetch('https://api.deepseek.com/chat/completions', {
      method: 'POST',
      headers: {
        'Authorization': `Bearer ${process.env.DEEPSEEK_KEY}`,
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({
        model: model,
        messages: messages,
        stream: false,
      }),
    });
    return response.json();
  }

  // Similar function for Gemini can be added here...

  try {

    const { prompt, sessionId, provider, model = "default-model", key } = req.body;

    if (!key || !PROVIDER_KEYS.includes(key)) {
      return res.status(403).json({
        success: false,
        error: 'Unauthorized: Invalid or missing key.',
      });
    }
    
    if (provider && req.body['list-models']) {
      if (PROVIDER_MODELS[provider]) {
        return res.json({
          success: true,
          models: PROVIDER_MODELS[provider],
        });
      } else {
        return res.status(400).json({
          success: false,
          error: 'Invalid provider specified.',
        });
      }
    }

    // Check if the request is for listing AI providers
    if (req.body['list-ais']) {
      return res.json({
        success: true,
        providers: ['OpenAI', 'DeepSeek', 'Gemini', 'Manual'],
      });
    }

    if (!conversationsCollection) {
      return res.status(500).json({
        success: false,
        error: 'MongoDB is not connected yet. Please try again later.',
      });
    }

    let conversation;

    if (!sessionId) {
      const newSessionId = uuidv4();
      conversation = {
        sessionId: newSessionId,
        conversationHistory: [
          { role: 'system', content: 'answer within 10 words' },
          { role: 'user', content: prompt },
        ],
      };
      await conversationsCollection.insertOne(conversation);
    } else {
      conversation = await conversationsCollection.findOne({ sessionId: sessionId });
      if (!conversation) {
        return res.status(400).json({
          success: false,
          error: 'Invalid session ID.',
        });
      }
      conversation.conversationHistory.push({ role: 'user', content: prompt });
    }

    let aiResponse;

    switch (provider) {
      case 'openai':
        aiResponse = await fetchOpenAICompletion(conversation.conversationHistory, model);
        break;
      case 'deepseek':
        aiResponse = await fetchDeepSeekCompletion(conversation.conversationHistory, model);
        break;
      // case 'gemini':
      //   aiResponse = await fetchGeminiCompletion(conversation.conversationHistory, model);
      //   break;
      // Add more cases for other providers as needed
      default:
        return res.status(400).json({
          success: false,
          error: 'Invalid provider specified.',
        });
    }

    conversation.conversationHistory.push({ role: 'assistant', content: aiResponse.choices ? aiResponse.choices[0].message.content : aiResponse.response });

    await conversationsCollection.updateOne(
      { sessionId: conversation.sessionId },
      { $set: { conversationHistory: conversation.conversationHistory } }
    );

    res.json({
      success: true,
      data: aiResponse.choices ? aiResponse.choices[0].message.content : aiResponse.response,
      total_tokens: aiResponse.usage ? aiResponse.usage.total_tokens : undefined,
      sessionId: conversation.sessionId,
    });
  } catch (error) {
    console.error('Error generating response:', error.message);
    res.status(500).json({
      success: false,
      error: 'Something went wrong. Please try again later.',
    });
  }
};

module.exports = aiFollowupQuestion;