DFIQ card redesign and AI query UI (#3157)

* Redesign DFIQ card
* Generate LLM queries
* Add system settings API endpoint
* Add system settings to frontend store

---------

Co-authored-by: Janosch <[email protected]>

docker/dev/build/docker-entrypoint.sh

@@ -20,7 +20,7 @@ if [ "$1" = 'timesketch' ]; then
   cp /usr/local/src/timesketch/data/bigquery_matcher.yaml /etc/timesketch/
   ln -s /usr/local/src/timesketch/data/sigma_config.yaml /etc/timesketch/sigma_config.yaml
   ln -s /usr/local/src/timesketch/data/sigma /etc/timesketch/
-  ln -s /usr/local/src/timesketch/data/scenarios /etc/timesketch/
+  ln -s /usr/local/src/timesketch/data/dfiq /etc/timesketch/
   ln -s /usr/local/src/timesketch/data/context_links.yaml /etc/timesketch/context_links.yaml
   ln -s /usr/local/src/timesketch/data/plaso_formatters.yaml /etc/timesketch/plaso_formatters.yaml
 

timesketch/api/v1/resources/nl2q.py

@@ -200,24 +200,30 @@ def post(self, sketch_id):
 
         question = form.get("question")
         prompt = self.build_prompt(question, sketch_id)
+        result_schema = {
+            "name": "AI generated search query",
+            "query_string": None,
+            "error": None,
+        }
         try:
             llm = manager.LLMManager().get_provider(llm_provider)()
         except Exception as e:  # pylint: disable=broad-except
             logger.error("Error LLM Provider: {}".format(e))
-            abort(
-                HTTP_STATUS_CODE_INTERNAL_SERVER_ERROR,
-                "Error in loading the LLM Provider. Please contact your "
-                "Timesketch administrator.",
+            result_schema["error"] = (
+                "Error loading LLM Provider. Please try again later!"
             )
+            return jsonify(result_schema)
 
         try:
             prediction = llm.generate(prompt)
         except Exception as e:  # pylint: disable=broad-except
             logger.error("Error NL2Q prompt: {}".format(e))
-            abort(
-                HTTP_STATUS_CODE_INTERNAL_SERVER_ERROR,
-                "An error occurred generating the NL2Q prediction via the "
-                "defined LLM. Please contact your Timesketch administrator.",
+            result_schema["error"] = (
+                "An error occurred generating the query via the defined LLM. "
+                "Please try again later!"
             )
-        result = {"question": question, "llm_query": prediction}
-        return jsonify(result)
+            return jsonify(result_schema)
+        # The model sometimes output tripple backticks that needs to be removed.
+        result_schema["query_string"] = prediction.strip("```")
+
+        return jsonify(result_schema)

timesketch/api/v1/resources/settings.py

@@ -0,0 +1,39 @@
+# Copyright 2024 Google Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""System settings."""
+
+from flask import current_app
+from flask import jsonify
+from flask_restful import Resource
+from flask_login import login_required
+
+
+class SystemSettingsResource(Resource):
+    """Resource to get system settings."""
+
+    @login_required
+    def get(self):
+        """GET system settings.
+
+        Returns:
+            JSON object with system settings.
+        """
+        # Settings from timesketch.conf to expose to the frontend clients.
+        settings_to_return = ["LLM_PROVIDER"]
+        result = {}
+
+        for setting in settings_to_return:
+            result[setting] = current_app.config.get(setting)
+
+        return jsonify(result)

timesketch/api/v1/resources_test.py

@@ -1222,7 +1222,11 @@ def test_nl2q_prompt(self, mock_aggregator, mock_llm_manager):
         self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)
         self.assertDictEqual(
             response.json,
-            {"question": "Question for LLM?", "llm_query": "LLM generated query"},
+            {
+                "name": "AI generated search query",
+                "query_string": "LLM generated query",
+                "error": None,
+            },
         )
 
     @mock.patch("timesketch.api.v1.utils.run_aggregator")
@@ -1253,6 +1257,8 @@ def test_nl2q_no_prompt(self, mock_aggregator):
             content_type="application/json",
         )
         self.assertEqual(response.status_code, HTTP_STATUS_CODE_INTERNAL_SERVER_ERROR)
+        # data = json.loads(response.get_data(as_text=True))
+        # self.assertIsNotNone(data.get("error"))
 
     @mock.patch("timesketch.api.v1.utils.run_aggregator")
     @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
@@ -1315,7 +1321,9 @@ def test_nl2q_wrong_llm_provider(self, mock_aggregator):
             data=json.dumps(data),
             content_type="application/json",
         )
-        self.assertEqual(response.status_code, HTTP_STATUS_CODE_INTERNAL_SERVER_ERROR)
+        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)
+        data = json.loads(response.get_data(as_text=True))
+        self.assertIsNotNone(data.get("error"))
 
     @mock.patch("timesketch.api.v1.resources.OpenSearchDataStore", MockDataStore)
     def test_nl2q_no_llm_provider(self):
@@ -1379,4 +1387,19 @@ def test_nl2q_llm_error(self, mock_aggregator, mock_llm_manager):
             data=json.dumps(data),
             content_type="application/json",
         )
-        self.assertEqual(response.status_code, HTTP_STATUS_CODE_INTERNAL_SERVER_ERROR)
+        self.assertEqual(response.status_code, HTTP_STATUS_CODE_OK)
+        data = json.loads(response.get_data(as_text=True))
+        self.assertIsNotNone(data.get("error"))
+
+
+class SystemSettingsResourceTest(BaseTest):
+    """Test system settings resource."""
+
+    resource_url = "/api/v1/settings/"
+
+    def test_system_settings_resource(self):
+        """Authenticated request to get system settings."""
+        self.login()
+        response = self.client.get(self.resource_url)
+        expected_response = {"LLM_PROVIDER": "test"}
+        self.assertEqual(response.json, expected_response)

timesketch/api/v1/routes.py

@@ -76,6 +76,7 @@
 from .resources.contextlinks import ContextLinkConfigResource
 from .resources.unfurl import UnfurlResource
 from .resources.nl2q import Nl2qResource
+from .resources.settings import SystemSettingsResource
 
 from .resources.scenarios import ScenarioTemplateListResource
 from .resources.scenarios import ScenarioListResource
@@ -196,6 +197,7 @@
     (ContextLinkConfigResource, "/contextlinks/"),
     (UnfurlResource, "/unfurl/"),
     (Nl2qResource, "/sketches/<int:sketch_id>/nl2q/"),
+    (SystemSettingsResource, "/settings/"),
     # Scenario templates
     (ScenarioTemplateListResource, "/scenarios/"),
     # Scenarios

timesketch/frontend-ng/src/components/Explore/EventList.vue

@@ -23,13 +23,19 @@ limitations under the License.
     </v-dialog>
 
     <div v-if="!eventList.objects.length && !searchInProgress" class="ml-3">
-      <p class="ml-n2 mt-n4">
+      <p>
+        Your search <span v-if="currentQueryString">'{{ currentQueryString }}'</span
+        ><span v-if="filterChips.length"> in combination with the selected filter terms</span> did not match any events.
+      </p>
+      <p>
         <v-dialog v-model="saveSearchMenu" v-if="!disableSaveSearch" width="500">
           <template v-slot:activator="{ on, attrs }">
-            <v-btn small text rounded color="secondary" v-bind="attrs" v-on="on">
-              <v-icon left small title="Save current search">mdi-content-save-outline</v-icon>
-              save this search
-            </v-btn>
+            <div v-bind="attrs" v-on="on">
+              <v-btn small depressed>
+                <v-icon left small title="Save current search">mdi-content-save-outline</v-icon>
+                Save search
+              </v-btn>
+            </div>
           </template>
 
           <v-card class="pa-4">
@@ -62,9 +68,6 @@ limitations under the License.
           </v-card>
         </v-dialog>
       </p>
-      <p>
-        Your search <span v-if="currentQueryString">'{{ currentQueryString }}'</span><span v-if="filterChips.length"> in combination with the selected filter terms</span> did not match any events.
-      </p>
       <p>Suggestions:</p>
       <ul>
         <li>Try different keywords<span v-if="filterChips.length"> or filter terms</span>.</li>

timesketch/frontend-ng/src/components/Scenarios/QuestionApproach.vue

@@ -28,7 +28,7 @@ limitations under the License.
         :key="opensearchQuery.value"
         :searchchip="opensearchQuery"
         type="link"
-        class="mb-1"
+        class="mb-1 mt-2"
       ></ts-search-chip>
     </div>