Tell the API caller is the data was truncated

Author: cccs-sgaron

assemblyline_ui/api/v4/file.py

@@ -17,7 +17,8 @@
 from assemblyline_ui.api.base import api_login, make_api_response, make_subapi_blueprint, stream_file_response
 from assemblyline_ui.config import ALLOW_ZIP_DOWNLOADS, ALLOW_RAW_DOWNLOADS, FILESTORE, STORAGE, config, \
     CLASSIFICATION as Classification, ARCHIVESTORE
-from assemblyline_ui.helper.ai import AiApiException, summarize_code_snippet as ai_code, summarized_al_submission
+from assemblyline_ui.helper.ai import APIException, EmptyAIResponse, \
+    summarize_code_snippet as ai_code, summarized_al_submission
 from assemblyline_ui.helper.result import format_result
 from assemblyline_ui.helper.user import load_user_settings
 from assemblyline.datastore.collection import Index
@@ -325,7 +326,10 @@ def summarized_results(sha256, **kwargs):
     /api/v4/file/ai/123456...654321/
 
     Result example:
-    <AI summary of the AL results>
+    {
+      "content": <AI summary of the AL results>,
+      "truncated": false
+    }
     """
     if not config.ui.ai.enabled:
         return make_api_response({}, "AI Support is disabled on this system.", 400)
@@ -341,7 +345,7 @@ def summarized_results(sha256, **kwargs):
         # TODO: Caching maybe?
         ai_summary = summarized_al_submission(data)
         return make_api_response(ai_summary)
-    except AiApiException as e:
+    except (APIException, EmptyAIResponse) as e:
         return make_api_response("", str(e), 400)
 
 
@@ -365,7 +369,10 @@ def summarize_code_snippet(sha256, **kwargs):
     /api/v4/file/code_summary/123456...654321/
 
     Result example:
-    <AI summary of the code snippet>
+    {
+      "content": <AI summary of the code snippet>,
+      "truncated": false
+    }
     """
     if not config.ui.ai.enabled:
         return make_api_response({}, "AI Support is disabled on this system.", 400)
@@ -406,7 +413,7 @@ def summarize_code_snippet(sha256, **kwargs):
             # TODO: Caching maybe?
             ai_summary = ai_code(data)
             return make_api_response(ai_summary)
-        except AiApiException as e:
+        except (APIException, EmptyAIResponse) as e:
             return make_api_response("", str(e), 400)
     else:
         return make_api_response({}, "You are not allowed to view this file.", 403)

assemblyline_ui/api/v4/submission.py

@@ -1,15 +1,15 @@
 import time
-from assemblyline.datastore.collection import Index
-from assemblyline_core.dispatching.client import DispatchClient
-from assemblyline_ui.helper.ai import AiApiException, summarized_al_submission
 
 from flask import request
 from werkzeug.exceptions import BadRequest
 
 from assemblyline.datastore.exceptions import MultiKeyError, SearchException
+from assemblyline.datastore.collection import Index
 from assemblyline.odm.models.user import ROLES
+from assemblyline_core.dispatching.client import DispatchClient
 from assemblyline_ui.api.base import api_login, make_api_response, make_subapi_blueprint
 from assemblyline_ui.config import STORAGE, LOGGER, FILESTORE, config, CLASSIFICATION as Classification
+from assemblyline_ui.helper.ai import EmptyAIResponse, APIException, summarized_al_submission
 from assemblyline_ui.helper.result import cleanup_heuristic_sections, format_result
 from assemblyline_ui.helper.submission import get_or_create_summary
 
@@ -500,7 +500,11 @@ def get_ai_summary(sid, **kwargs):
     None
 
     Result example:
-    < THE AI SUMMARY IN MARKDOWN FORMAT >
+    {
+      "content": < THE AI SUMMARY IN MARKDOWN FORMAT >,
+      "truncated": false
+    }
+
 
     """
     if not config.ui.ai.enabled:
@@ -517,7 +521,7 @@ def get_ai_summary(sid, **kwargs):
         # TODO: Caching maybe?
         ai_summary = summarized_al_submission(data)
         return make_api_response(ai_summary)
-    except AiApiException as e:
+    except (APIException, EmptyAIResponse) as e:
         return make_api_response("", str(e), 400)
 
 

assemblyline_ui/helper/ai.py

@@ -6,7 +6,11 @@
 from assemblyline_ui.config import config, LOGGER
 
 
-class AiApiException(Exception):
+class APIException(Exception):
+    pass
+
+
+class EmptyAIResponse(Exception):
     pass
 
 
@@ -30,22 +34,23 @@ def _call_ai_backend(data, params: AIQueryParams, action):
     except Exception as e:
         message = f"An exception occured while trying to {action} with AI on server {config.ui.ai.chat_url}. [{e}]"
         LOGGER.warning(message)
-        raise AiApiException(message)
+        raise APIException(message)
 
     if not resp.ok:
         msg_data = resp.json()
         msg = msg_data.get('error', {}).get('message', None) or msg_data
         message = f"The AI API denied the request to {action} with the following message: {msg}"
         LOGGER.warning(message)
-        raise AiApiException(message)
+        raise APIException(message)
 
     # Get AI responses
-    responses = resp.json().get('choices', [])
+    responses = resp.json()['choices']
     if responses:
-        content = responses[0].get('message', {}).get('content', None)
-        return content or None
+        content = responses[0]['message']['content']
+        reason = responses[0]['finish_reason']
+        return {'content': content, 'truncated': reason == 'length'}
 
-    return None
+    raise EmptyAIResponse("There was no response returned by the AI")
 
 
 def detailed_al_submission(report):