From cfa16dbdc766e8122ba5377ada828d2a393d39f9 Mon Sep 17 00:00:00 2001
From: "Kar l5.G" <kar@siliconpin.com>
Date: Wed, 29 Jan 2025 21:01:05 +0530
Subject: [PATCH] list-models

---
 src/routes/api/aiTest.js | 42 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 38 insertions(+), 4 deletions(-)

diff --git a/src/routes/api/aiTest.js b/src/routes/api/aiTest.js
index 32935f3..0586612 100644
--- a/src/routes/api/aiTest.js
+++ b/src/routes/api/aiTest.js
@@ -1,6 +1,6 @@
 const aiFollowupQuestion = async (req, res) => {
   const { MongoClient } = require('mongodb');
-  const fetch = require('node-fetch');
+  // const fetch = require('node-fetch');
   const { v4: uuidv4 } = require('uuid');
   const url = process.env.MONGODB_URL;
   const dbName = process.env.MONGO_DB_NAME;
@@ -10,7 +10,18 @@ const aiFollowupQuestion = async (req, res) => {
   const database = client.db(dbName);
   const conversationsCollection = database.collection('conversations');
 
-  const MAX_TOKENS = 200; // Set max_tokens as a constant
+  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"
+    ],
+    deepseek: [
+      "deepseek-chat", "deepseek-reasoner"
+    ]
+    // 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', {
@@ -47,6 +58,31 @@ const aiFollowupQuestion = async (req, res) => {
   // Similar function for Gemini can be added here...
 
   try {
+
+    const { prompt, sessionId, provider, model = "default-model", key } = req.body;
+
+    // Check if key is provided and valid
+    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({
@@ -55,8 +91,6 @@ const aiFollowupQuestion = async (req, res) => {
       });
     }
 
-    const { prompt, sessionId, provider, model = "default-model" } = req.body;
-
     if (!conversationsCollection) {
       return res.status(500).json({
         success: false,