Display model name used (OWASP#705)

* Update Vertex AI prompt formatting and structure for improved RAG

* fix(chatbot): update model reference from PALM2 to Gemini 2.0 Flash for improved accuracy

* feat(prompt-client): add model name retrieval methods and update chatbot to display model name

* Remove unnecessary blank line in vertex_prompt_client.py

Signed-off-by: Spyros <northdpole@users.noreply.github.com>

* Fix indentation in OpenAI prompt client

Signed-off-by: Spyros <northdpole@users.noreply.github.com>

* Fix formatting and whitespace in vertex_prompt_client.py

Signed-off-by: Spyros <northdpole@users.noreply.github.com>

* Format return statement for better readability

Signed-off-by: Spyros <northdpole@users.noreply.github.com>

---------

Signed-off-by: Spyros <northdpole@users.noreply.github.com>
Co-authored-by: Spyros <northdpole@users.noreply.github.com>

Author: 2 people

application/frontend/src/pages/chatbot/chatbot.tsx

@@ -34,6 +34,21 @@ export const Chatbot = () => {
   const [error, setError] = useState<string>('');
   const [chat, setChat] = useState<ChatState>(DEFAULT_CHAT_STATE);
   const [user, setUser] = useState('');
+  const [modelName, setModelName] = useState<string>('');
+
+  function getModelDisplayName(modelName: string): string {
+    if (!modelName) {
+      return 'a Large Language Model';
+    }
+    // Format model names for display
+    if (modelName.startsWith('gemini')) {
+      return `Google ${modelName.replace('gemini-', 'Gemini ').replace(/-/g, ' ')}`;
+    } else if (modelName.startsWith('gpt')) {
+      return `OpenAI ${modelName.toUpperCase()}`;
+    }
+    return modelName;
+  }
+
   const hasMessages = chatMessages.length > 0;
   function login() {
     fetch(`${apiUrl}/user`, { method: 'GET' })
@@ -104,6 +119,9 @@ export const Chatbot = () => {
       .then((data) => {
         setLoading(false);
         setError('');
+        if (data.model_name) {
+          setModelName(data.model_name);
+        }
         setChatMessages((prev) => [
           ...prev,
           {
@@ -211,7 +229,7 @@ export const Chatbot = () => {
 
             <div className="chatbot-disclaimer">
               <i>
-                Answers are generated by a Google PALM2 Large Language Model, which uses the internet as
+                Answers are generated by {getModelDisplayName(modelName)} Large Language Model, which uses the internet as
                 training data, plus collected key cybersecurity standards from{' '}
                 <a href="https://opencre.org">OpenCRE</a> as the preferred source. This leads to more reliable
                 answers and adds references, but note: it is still generative AI which is never guaranteed

application/prompt_client/openai_prompt_client.py

@@ -10,6 +10,11 @@ class OpenAIPromptClient:
     def __init__(self, openai_key) -> None:
         self.api_key = openai_key
         openai.api_key = self.api_key
+        self.model_name = "gpt-3.5-turbo"
+
+    def get_model_name(self) -> str:
+        """Return the model name being used."""
+        return self.model_name
 
     def get_text_embeddings(self, text: str, model: str = "text-embedding-ada-002"):
         if len(text) > 8000:

application/prompt_client/prompt_client.py

@@ -498,4 +498,10 @@ def generate_text(self, prompt: str) -> Dict[str, str]:
         logger.debug(f"retrieved completion for {prompt}")
         table = [closest_object]
         result = f"Answer: {answer}"
-        return {"response": result, "table": table, "accurate": accurate}
+        model_name = self.ai_client.get_model_name() if self.ai_client else "unknown"
+        return {
+            "response": result,
+            "table": table,
+            "accurate": accurate,
+            "model_name": model_name,
+        }

application/prompt_client/vertex_prompt_client.py

@@ -54,6 +54,11 @@ class VertexPromptClient:
 
     def __init__(self) -> None:
         self.client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
+        self.model_name = "gemini-2.0-flash"
+
+    def get_model_name(self) -> str:
+        """Return the model name being used."""
+        return self.model_name
 
     def get_text_embeddings(self, text: str) -> List[float]:
         """Text embedding with a Large Language Model."""