Skip to content

feat(ui): adds genai suggestions for tags #6073

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Jun 23, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
164 changes: 164 additions & 0 deletions src/dispatch/ai/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,10 @@
from dispatch.incident.models import Incident
from dispatch.plugin import service as plugin_service
from dispatch.signal import service as signal_service
from dispatch.tag.models import Tag, TagRecommendationResponse
from dispatch.tag_type.models import TagType
from dispatch.case import service as case_service
from dispatch.incident import service as incident_service

from .exceptions import GenAIException

Expand Down Expand Up @@ -390,3 +394,163 @@ def generate_incident_summary(incident: Incident, db_session: Session) -> str:
except Exception as e:
log.exception(f"Error trying to generate summary for incident {incident.name}: {e}")
return "Incident summary not generated. An error occurred."


def get_tag_recommendations(
*, db_session, project_id: int, case_id: int | None = None, incident_id: int | None = None
) -> TagRecommendationResponse:
"""Gets tag recommendations for a project."""
genai_plugin = plugin_service.get_active_instance(
db_session=db_session, project_id=project_id, plugin_type="artificial-intelligence"
)

# we check if the artificial intelligence plugin is enabled
if not genai_plugin:
message = (
"AI tag suggestions are not available. No AI plugin is configured for this project."
)
log.warning(message)
return TagRecommendationResponse(recommendations=[], error_message=message)

storage_plugin = plugin_service.get_active_instance(
db_session=db_session, plugin_type="storage", project_id=project_id
)

# get resources from the case or incident
resources = ""
if case_id:
case = case_service.get(db_session=db_session, case_id=case_id)
if not case:
raise ValueError(f"Case with id {case_id} not found")
if case.visibility == Visibility.restricted:
message = "AI tag suggestions are not available for restricted cases."
return TagRecommendationResponse(recommendations=[], error_message=message)

resources += f"Case title: {case.name}\n"
resources += f"Description: {case.description}\n"
resources += f"Resolution: {case.resolution}\n"
resources += f"Resolution Reason: {case.resolution_reason}\n"
resources += f"Case type: {case.case_type.name}\n"

if storage_plugin and case.case_document and case.case_document.resource_id:
case_doc = storage_plugin.instance.get(
file_id=case.case_document.resource_id,
mime_type="text/plain",
)
resources += f"Case document: {case_doc}\n"

elif incident_id:
incident = incident_service.get(db_session=db_session, incident_id=incident_id)
resources += f"Incident: {incident.name}\n"
resources += f"Description: {incident.description}\n"
resources += f"Resolution: {incident.resolution}\n"
resources += f"Incident type: {incident.incident_type.name}\n"

if storage_plugin and incident.incident_document and incident.incident_document.resource_id:
incident_doc = storage_plugin.instance.get(
file_id=incident.incident_document.resource_id,
mime_type="text/plain",
)
resources += f"Incident document: {incident_doc}\n"

if (
storage_plugin
and incident.incident_review_document
and incident.incident_review_document.resource_id
):
incident_review_doc = storage_plugin.instance.get(
file_id=incident.incident_review_document.resource_id,
mime_type="text/plain",
)
resources += f"Incident review document: {incident_review_doc}\n"

else:
raise ValueError("Either case_id or incident_id must be provided")
# get all tags for the project with the tag_type that has genai_suggestions set to True
tags: list[Tag] = (
db_session.query(Tag)
.filter(Tag.project_id == project_id)
.filter(Tag.tag_type.has(TagType.genai_suggestions.is_(True)))
.all()
)

# Check if there are any tags available for AI suggestions
if not tags:
message = (
"AI tag suggestions are not available. No tag types are configured "
"for AI suggestions in this project."
)
return TagRecommendationResponse(recommendations=[], error_message=message)

# add to the resources each tag name, id, tag_type_id, and description
tag_list = "Tags you can use:\n" + (
"\n".join(
[
f"tag_name: {tag.name}\n"
f"tag_id: {tag.id}\n"
f"description: {tag.description}\n"
f"tag_type_id: {tag.tag_type_id}\n"
f"tag_type_name: {tag.tag_type.name}\n"
f"tag_type_description: {tag.tag_type.description}\n"
for tag in tags
]
)
+ "\n"
)

prompt = """
You are a security professional that can help with tag recommendations.
You will be given details about a security event and a list of tags you can use.
You will need to recommend tags for the security event using the descriptions of the tags.
Please identify the top three tags of each tag_type_id that best apply to the security event.
Provide the output in JSON format organized by tag_type_id in the following format:
{"recommendations":
[
{
"tag_type_id": 1,
"tags": [
{
"id": 1,
"name": "tag_name",
"reason": "your reasoning for including this tag"
}
]
}
]
}
Do not output anything except for the JSON.
"""

prompt += f"** Tags you can use: {tag_list} \n ** Security event details: {resources}"

tokenized_prompt, num_tokens, encoding = num_tokens_from_string(
prompt, genai_plugin.instance.configuration.chat_completion_model
)

# we check if the prompt exceeds the token limit
model_token_limit = get_model_token_limit(
genai_plugin.instance.configuration.chat_completion_model
)
if num_tokens > model_token_limit:
prompt = truncate_prompt(tokenized_prompt, num_tokens, encoding, model_token_limit)

try:
result = genai_plugin.instance.chat_completion(prompt=prompt)

# Clean the JSON string by removing markdown formatting and newlines
# Remove markdown code block markers
cleaned_result = result.strip()
if cleaned_result.startswith("```json"):
cleaned_result = cleaned_result[7:] # Remove ```json
if cleaned_result.endswith("```"):
cleaned_result = cleaned_result[:-3] # Remove ```

# Replace escaped newlines with actual newlines, then clean whitespace
cleaned_result = cleaned_result.replace("\\n", "\n")
cleaned_result = " ".join(cleaned_result.split())

return TagRecommendationResponse.model_validate_json(cleaned_result)
except Exception as e:
log.exception(f"Error generating tag recommendations: {e}")
message = "AI tag suggestions encountered an error. Please try again later."
return TagRecommendationResponse(recommendations=[], error_message=message)
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
"""Add GenAI suggestions column to tag_type table

Revision ID: 7fc3888c7b9a
Revises: 8f324b0f365a
Create Date: 2025-06-04 14:49:20.592746

"""

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = "7fc3888c7b9a"
down_revision = "8f324b0f365a"
branch_labels = None
depends_on = None


def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.add_column("tag_type", sa.Column("genai_suggestions", sa.Boolean(), nullable=True))
# ### end Alembic commands ###


def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column("tag_type", "genai_suggestions")
# ### end Alembic commands ###
67 changes: 61 additions & 6 deletions src/dispatch/database/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
from dispatch.search.fulltext.composite_search import CompositeSearch
from dispatch.signal.models import Signal, SignalInstance
from dispatch.tag.models import Tag
from dispatch.tag_type.models import TagType

from dispatch.task.models import Task
from typing import Annotated

Expand Down Expand Up @@ -134,12 +134,23 @@ def format_for_sqlalchemy(self, query: SQLAlchemyQuery, default_model):
operator = self.operator
value = self.value

# Special handling for TagType.id filtering on Tag model
# Needed since TagType.id is not a column on the Tag model
# Convert TagType.id filter to tag_type_id filter on Tag model
if (
filter_spec.get("model") == "TagType"
and filter_spec.get("field") == "id"
and default_model
and getattr(default_model, "__tablename__", None) == "tag"
):
filter_spec = {"model": "Tag", "field": "tag_type_id", "op": filter_spec.get("op")}

model = get_model_from_spec(filter_spec, query, default_model)

function = operator.function
arity = operator.arity

field_name = self.filter_spec["field"]
field_name = filter_spec["field"]
field = Field(model, field_name)
sqlalchemy_field = field.get_sqlalchemy_field()

Expand Down Expand Up @@ -400,8 +411,18 @@ def apply_filters(query, filter_spec, model_cls=None, do_auto_join=True):

filter_spec = {
'or': [
{'model': 'Foo', 'field': 'id', 'op': '==', 'value': '1'},
{'model': 'Bar', 'field': 'id', 'op': '==', 'value': '2'},
{
'model': 'Foo',
'field': 'id',
'op': '==',
'value': '1'
},
{
'model': 'Bar',
'field': 'id',
'op': '==',
'value': '2'
},
]
}

Expand Down Expand Up @@ -456,7 +477,7 @@ def apply_filter_specific_joins(model: Base, filter_spec: dict, query: orm.query
(Signal, "TagType"): (Signal.tags, True),
(SignalInstance, "Entity"): (SignalInstance.entities, True),
(SignalInstance, "EntityType"): (SignalInstance.entities, True),
(Tag, "TagType"): (TagType, False),
# (Tag, "TagType"): (TagType, False), # Disabled: filtering by tag_type_id directly
(Tag, "Project"): (Project, False),
(IndividualContact, "Project"): (Project, False),
}
Expand Down Expand Up @@ -733,7 +754,41 @@ def search_filter_sort_paginate(
# e.g. websearch_to_tsquery
# https://www.postgresql.org/docs/current/textsearch-controls.html
try:
query, pagination = apply_pagination(query, page_number=page, page_size=items_per_page)
# Check if this model is likely to have duplicate results from many-to-many joins
# Models with many secondary relationships (like Tag) can cause count inflation
models_needing_distinct = ["Tag"] # Add other models here as needed

if model in models_needing_distinct and items_per_page is not None:
# Use custom pagination that handles DISTINCT properly
from collections import namedtuple

Pagination = namedtuple(
"Pagination", ["page_number", "page_size", "num_pages", "total_results"]
)

# Get total count using distinct ID to avoid duplicates
# Remove ORDER BY clause for counting since it's not needed and causes issues with DISTINCT
count_query = query.with_entities(model_cls.id).distinct().order_by(None)
total_count = count_query.count()

# Apply DISTINCT to the main query as well to avoid duplicate results
# Remove ORDER BY clause since it can conflict with DISTINCT when ordering by joined table columns
query = query.distinct().order_by(None)

# Apply pagination to the distinct query
offset = (page - 1) * items_per_page if page > 1 else 0
query = query.offset(offset).limit(items_per_page)

# Calculate number of pages
num_pages = (
(total_count + items_per_page - 1) // items_per_page if items_per_page > 0 else 1
)

pagination = Pagination(page, items_per_page, num_pages, total_count)
else:
# Use standard pagination for other models
query, pagination = apply_pagination(query, page_number=page, page_size=items_per_page)

except ProgrammingError as e:
log.debug(e)
return {
Expand Down
2 changes: 2 additions & 0 deletions src/dispatch/static/dispatch/src/case/DetailsTab.vue
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,8 @@
:model-id="id"
:project="project"
show-copy
:showGenAISuggestions="true"
modelType="case"
/>
</v-col>
<v-col cols="12">
Expand Down
3 changes: 3 additions & 0 deletions src/dispatch/static/dispatch/src/incident/DetailsTab.vue
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,10 @@
:project="project"
model="incident"
:model-id="id"
:visibility="visibility"
show-copy
:showGenAISuggestions="true"
modelType="incident"
/>
</v-col>
<v-col cols="12">
Expand Down
8 changes: 4 additions & 4 deletions src/dispatch/static/dispatch/src/search/utils.js
Original file line number Diff line number Diff line change
Expand Up @@ -133,8 +133,8 @@ export default {
}
} else {
each(value, function (value) {
// filter null values
if (!value) {
// filter null/undefined values but allow false
if (value === null || value === undefined) {
return
}
if (["commander", "participant", "assignee"].includes(key) && has(value, "email")) {
Expand Down Expand Up @@ -166,8 +166,8 @@ export default {
value: value.name,
})
} else if (has(value, "model")) {
// avoid filter null values
if (value.value) {
// avoid filter null/undefined values but allow false
if (value.value !== null && value.value !== undefined) {
subFilter.push({
model: value.model,
field: value.field,
Expand Down
Loading
Loading