Merge pull request #6 from xyTom/support-different-prompt-for-differe…

…nt-models

支持模型选择和提示选择

src/renderer/App.tsx

@@ -2,6 +2,8 @@
 import './App.css';
 import React, { useState,useEffect,useRef } from 'react';
 import aiModels from './models';
+//disable the eslint warning for the import of the logo
+// eslint-disable-next-line
 import logo from './assets/logo.png';
 import DisplayTextResult from './components/displayTextResult';
 import LoadingSkeleton from './components/loadingSkeleton';
@@ -11,7 +13,7 @@ import ApiKeyInput from './components/apiKeyInput';
 import PromptSelect from './components/promptSelect';
 import { KeyRound } from "lucide-react"
 import { Button } from "./components/ui/button"
-import { promptOptions } from './lib/models';
+import { promptOptions, models } from './lib/models';
 
 declare global {
   interface Window {
@@ -31,19 +33,6 @@ function App() {
   const [disabled, setDisabled] = useState(false);
   const [progressItems, setProgressItems] = useState([]);
 
-  //Input Model
-  const [apiKey, setApiKey] = useState('');
-  //read the api key from local storage
-  useEffect(() => {
-    const key = localStorage.getItem('apiKey');
-    if (key) {
-      setApiKey(key);
-    }
-  }, []);
-  const [openDialog, setOpenDialog] = useState(false);
-  const handleOpenChange = (value: boolean) => {
-    setOpenDialog(value);
-  }
   //AI model selection
   const [model, setModel] = useState('gemini');
   useEffect(() => {
@@ -52,6 +41,40 @@ function App() {
       setModel(model);
     }
   }, []);
+  //read model list, if require api, create a dictionary to store the api key
+  let keyMap = new Map();
+  models.forEach((model) => {
+    if (model.requireApiKey) {
+      keyMap.set(model.value, '');
+    }
+  });
+
+  const [apiKey, setApiKey] = useState('');
+  //read the api key from local storage
+  useEffect(() => {
+    const keys = localStorage.getItem('apiKeyMap');
+    console.log('keys', keys);
+    if (keys) {
+      const keyMapObj = JSON.parse(keys);
+      console.log('keyMapObj', keyMapObj);
+      for (const [key, value] of Object.entries(keyMapObj)) {
+        console.log('key', key, 'value', value);  
+        keyMap.set(key, value);
+        console.log('keyMap', keyMap);
+      }
+    }
+    console.log('model', model);
+    if (keyMap.has(model)) {
+      console.log('setApiKey', keyMap.get(model));
+      setApiKey(keyMap.get(model));
+    }
+  }, [model]);
+
+  const [openDialog, setOpenDialog] = useState(false);
+  const handleOpenChange = (value: boolean) => {
+    setOpenDialog(value);
+  }
+
   //Prompt selection
   const [prompt, setPrompt] = useState('Auto');
 
@@ -65,16 +88,19 @@ function App() {
   const handleModelChange = (value: string) => {
     console.log('handleModelChange', value);
     setModel(value);
-    if (value === 'gpt4') {
+    //read the api key for new model from local storage
+    if (keyMap.has(value)) {
+      setApiKey(keyMap.get(value));
       isApiKeyEmpty();
     }
   }
   //handle api key change
   const handleApiKeyChange = (value: string) => {
     console.log('handleApiKeyChange', value);
+    keyMap.set(model, value);
     setApiKey(value);
     //save the api key to local storage
-    localStorage.setItem('apiKey', value);
+    localStorage.setItem('apiKeyMap', JSON.stringify(Object.fromEntries(keyMap)));
     //close the dialog
     setOpenDialog(false);
   }
@@ -94,9 +120,9 @@ function App() {
     } else {
     setResult(null);
     setLoading(true);
-    aiModels.create(model).then((model: aiModels) => {
-      const fullPrompt = promptOptions.find((p) => p.value === value).prompt;
-      return model.run(screenShotResult,fullPrompt);
+    aiModels.create(model).then((modelInstance: aiModels) => {
+      const fullPrompt = promptOptions[model as keyof typeof promptOptions].find((p) => p.value === prompt).prompt;
+      return modelInstance.run(screenShotResult,fullPrompt);
     }).then((res: string) => {
       console.log('model res', res);
       setResult(res);
@@ -121,9 +147,12 @@ function App() {
       setscreenShotResult(value);
       setResult(null);
       setLoading(true);
-      aiModels.create(model).then((model: aiModels) => {
-        const fullPrompt = promptOptions.find((p) => p.value === prompt).prompt;
-        return model.run(value,fullPrompt);
+      aiModels.create(model).then((modelInstance: aiModels) => {
+        const fullPrompt = promptOptions[model as keyof typeof promptOptions].find((p) => p.value === prompt).prompt;
+        if (models.find((m) => m.value === model)?.requireApiKey) {
+          return modelInstance.run(value, fullPrompt, apiKey);
+        }
+        return modelInstance.run(value, fullPrompt);
       }).then((res: string) => {
         console.log('model res', res);
         setLoading(false);
@@ -242,9 +271,11 @@ function App() {
         {screenShotResult && <img src={`data:image/png;base64,${screenShotResult}`} alt="screenshot" className="mb-2 rounded-lg object-center border border-gray-100 dark:border-gray-800 mx-auto" />}
 
       <div className="flex space-x-2 mb-2 justify-center">
-      <PromptSelect handlePromptChange={handlePromptChange} />
+      <PromptSelect handlePromptChange={handlePromptChange} model={model}/>
       {/*only show the api key button when the model is gpt4 */}
-      {model === 'gpt4' &&
+      {
+      models.find((m) => m.value === model)?.requireApiKey
+       &&
       <ApiKeyButton onClick={()=>setOpenDialog(true)} />}
       </div>
       {/* {result && <button onClick={() => {
@@ -265,7 +296,7 @@ function App() {
       </div>
       </header>
 
-     <ApiKeyInput apikey={apiKey} onKeySave={handleApiKeyChange} open={openDialog} onOpenChange={handleOpenChange} />
+     <ApiKeyInput apikey={apiKey} onKeySave={handleApiKeyChange} open={openDialog} onOpenChange={handleOpenChange} model={model} />
 
     </div>
   );

src/renderer/components/apiKeyInput.tsx

@@ -10,13 +10,28 @@ import {
 import { Label } from "../components/ui/label"
 import { Input } from "../components/ui/input"
 import { Button } from "../components/ui/button"
+import { models } from "../lib/models"
 
-export default function displayTextResult(props: { apikey: string, onKeySave: (apikey: string) => void, open: boolean,onOpenChange: (value: boolean) => void }) {
+export default function displayTextResult(props: { apikey: string, onKeySave: (apikey: string) => void, open: boolean,onOpenChange: (value: boolean) => void, model: string}) {
     const [apiKey, setApiKey] = React.useState(props.apikey);
+    const [baseURL, setBaseURL] = React.useState("");
     const [open, setOpen] = React.useState(props.open);
     const onKeyChange = (value: string) => {
         setApiKey(value);
     }
+    const onBaseURLChange = (value: string) => {
+        setBaseURL(value);
+        //save the base URL to local storage
+        localStorage.setItem(`${props.model}_baseURL`, value);
+    }
+    //read the base URL from local storage
+    React.useEffect(() => {
+        let baseURL = localStorage.getItem(`${props.model}_baseURL`);
+        console.log(`${props.model}_baseURL`,baseURL);
+        if (baseURL) {
+            setBaseURL(baseURL);
+        }
+    }, [props.model]);
     //when setOpen is called, update the parent state
     React.useEffect(() => {
         props.onOpenChange(open);
@@ -25,7 +40,11 @@ export default function displayTextResult(props: { apikey: string, onKeySave: (a
     React.useEffect(() => {
         setOpen(props.open);
     }, [props.open]);
-
+    //when the parent API key is updated, update the local state
+    React.useEffect(() => {
+        setApiKey(props.apikey);
+    }
+    , [props.apikey]);
     return (
     <Dialog open={open} onOpenChange={setOpen}>
         <DialogContent>
@@ -44,6 +63,17 @@ export default function displayTextResult(props: { apikey: string, onKeySave: (a
                     onChange={(event) => onKeyChange((event.target as HTMLInputElement).value)}
                     />
                 </div>
+                    {models.find((m) => m.value === props.model).requireBaseURL && (
+                    <div className="grid grid-cols-4 items-center gap-4">
+                    <Label htmlFor="baseURL" className="text-right">
+                        Base URL
+                    </Label>
+                    <Input id="baseURL" value={baseURL} className="col-span-3" 
+                    onChange={(event) => onBaseURLChange((event.target as HTMLInputElement).value)}
+                    />
+                    </div>
+                    )}
+
             </div>
         <DialogFooter>
         <Button type="submit"

src/renderer/components/promptSelect.tsx

@@ -2,11 +2,12 @@ import React from "react"
 import { Tabs, TabsList, TabsTrigger } from "../components/ui/tabs"
 import { promptOptions } from "../lib/models"
 
-export default function promptSelect(props: { handlePromptChange: Function }) {
+export default function promptSelect(props: { handlePromptChange: Function, model: string}) {
+    const options = promptOptions as { [key: string]: { value: string; label: string; prompt: string; }[] };
     return (
         <Tabs defaultValue="Auto" onValueChange={(value) => props.handlePromptChange(value)}>
             <TabsList>
-                {promptOptions.map((prompt,index) => (
+                {options[props.model].map((prompt,index) => (
                     <TabsTrigger key={index} value={prompt.value}>{prompt.label}</TabsTrigger>
                 ))}
             </TabsList>

src/renderer/lib/models.ts

@@ -3,17 +3,27 @@ export const models = [
       value: "gemini",
       label: "Google Gemini",
       requireApiKey: false,
+      requireBaseURL: false,
       modelScript: "gemini",
     },
     {
       value: "gpt4",
       label: "OpenAI GPT-4",
       requireApiKey: true,
+      requireBaseURL: false,
       modelScript: "gpt4",
     },
+    // {
+    //     value: "azuregpt4",
+    //     label: "Azure OpenAI GPT-4",
+    //     requireApiKey: true,
+    //     requireBaseURL: true,
+    //     modelScript: "azuregpt4",
+    // }
   ]
 
-export const promptOptions = [
+export const promptOptions = {
+    "gemini":[
     {
         value: "Auto",
         label: "Auto",
@@ -55,4 +65,48 @@ export const promptOptions = [
         label:"Color",
         prompt:"Please return the color information in the image, use #RRGGBB format to describe the color.",
     }
-    ]
+    ],
+    "gpt4":[
+    {
+        value: "Auto",
+        label: "Auto",
+        prompt: "This is an image uploaded by a user, I need your help to analyze the content in the image, if the main content in the image is text, please return the result of OCR text recognition (only return the recognized text, don't add other descriptions), if the main content in the image is a mathematical formula, please return the Latex code of the mathematical formula (only return the Latex code corresponding to the formula, don't add extra If the main content of the image is a math formula, please return the Latex code of the math formula (only return the Latex code of the formula, don't add additional description), if the image contains neither text nor math formula, please return the detailed description of the image (please describe the content of the image in detail).\n",
+    },
+    {
+        value: "Formula",
+        label: "Formula",
+        prompt: "Please return the Latex code for this math formula (return only the Latex code corresponding to the formula, do not add additional descriptions)",
+
+    },
+    {
+        value: "Text",
+        label: "Text",
+        prompt: "Please return the result of OCR text recognition (only return the recognized text, do not add other descriptions)",
+    },
+    {
+        value: "Code",
+        label: "Code",
+        prompt: "Please return the content of the code in the picture (please describe in detail what the code in the picture does)",
+    },
+    {
+        value: "Table",
+        label: "Table",
+        prompt: "Please return the contents of the table in the image, using the MarkDown format.",
+    },
+    {
+        value: "Solve",
+        label: "Solve",
+        prompt: "Please return to the answer to the question in the picture",
+    },
+    {
+        value: "Image",
+        label: "Image",
+        prompt: "Please return a detailed description of the image (please describe in detail what is in the image)",
+    },
+    {
+        value:"Color",
+        label:"Color",
+        prompt:"Please return the color information in the image, use #RRGGBB format to describe the color.",
+    }
+    ]
+}

src/renderer/models/index.ts

@@ -19,7 +19,13 @@ export default class aiModel {
         return new aiModel(module.default);
     }
 
-    async run(prompt: string, image: string) {
+    async run(prompt: string, image: string, apiKey?:string, baseURL?:string) {
+        if (apiKey) {
+            if (baseURL) {
+                return this.model(prompt, image, apiKey, baseURL);
+            }
+            return this.model(prompt, image, apiKey);
+        }
         return this.model(prompt, image);
     }
 }

src/renderer/models/model-azuregpt4.ts

@@ -0,0 +1,24 @@
+import AzureOpenAI from 'openai';
+async function GPT(image:string, prompt:string, APIKey:string, baseURL:string): Promise<string> {
+  const openai = new AzureOpenAI({ apiKey: APIKey, baseURL: baseURL });
+    const response = await openai.chat.completions.create({
+        model: "gpt-4-vision-preview",
+        messages: [
+          {
+            role: "user",
+            content: [
+              { type: "text", text: prompt },
+              {
+                type: "image_url",
+                image_url: {
+                    "url": image,
+                },
+              },
+            ],
+          },
+        ],
+      });
+      console.log(response.choices[0]);
+      return response.choices[0].toString();
+}
+export default GPT;