ref(similarity-embeddings): Add hash column, defer stacktrace hash column (#625)

Add hash column with default value
Accept hash and return parent_hash in similarity api
Defer un-used stacktrace_hash column

Author: jangjodi

src/migrations/versions/7278a2303b10_migration.py

@@ -0,0 +1,38 @@
+"""Migration
+
+Revision ID: 7278a2303b10
+Revises: 34eef02b2555
+Create Date: 2024-05-02 15:56:39.243317
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "7278a2303b10"
+down_revision = "34eef02b2555"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("grouping_records", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "hash",
+                sa.String(length=32),
+                server_default="00000000000000000000000000000000",
+                nullable=False,
+            )
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("grouping_records", schema=None) as batch_op:
+        batch_op.drop_column("hash")
+
+    # ### end Alembic commands ###

src/seer/db.py

@@ -23,7 +23,7 @@
 )
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.ext.asyncio import async_sessionmaker
-from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker
+from sqlalchemy.orm import DeclarativeBase, Mapped, deferred, mapped_column, sessionmaker
 
 
 class Base(DeclarativeBase):
@@ -229,7 +229,12 @@ class DbGroupingRecord(Base):
     project_id: Mapped[int] = mapped_column(BigInteger, nullable=False)
     message: Mapped[str] = mapped_column(String, nullable=False)
     stacktrace_embedding: Mapped[Vector] = mapped_column(Vector(768), nullable=False)
-    stacktrace_hash: Mapped[Optional[str]] = mapped_column(String(32), nullable=True)
+    stacktrace_hash: Mapped[Optional[str]] = deferred(
+        mapped_column(String(32).evaluates_none(), nullable=True)
+    )
+    hash: Mapped[str] = mapped_column(
+        String(32), nullable=False, default="00000000000000000000000000000000"
+    )
 
     __table_args__ = (
         Index(

src/seer/grouping/grouping.py

@@ -8,6 +8,7 @@
 import torch
 from pydantic import BaseModel, ValidationInfo, field_validator
 from sentence_transformers import SentenceTransformer
+from sqlalchemy import or_
 
 from seer.db import DbGroupingRecord, Session
 
@@ -18,8 +19,8 @@ class GroupingRequest(BaseModel):
     project_id: int
     stacktrace: str
     message: str
-    group_id: int | None = None
-    stacktrace_hash: str | None = None
+    hash: str
+    group_id: Optional[int] = None
     k: int = 1
     threshold: float = 0.01
 
@@ -32,19 +33,19 @@ def check_field_is_not_empty(cls, v, info: ValidationInfo):
 
 
 class GroupingRecord(BaseModel):
-    group_id: int | None
+    group_id: Optional[int]
     project_id: int
     message: str
     stacktrace_embedding: np.ndarray
-    stacktrace_hash: str | None
+    hash: str
 
     def to_db_model(self) -> DbGroupingRecord:
         return DbGroupingRecord(
             group_id=self.group_id,
             project_id=self.project_id,
             message=self.message,
             stacktrace_embedding=self.stacktrace_embedding,
-            stacktrace_hash=self.stacktrace_hash,
+            hash=self.hash,
         )
 
     class Config:
@@ -56,14 +57,14 @@ class Config:
 
 class GroupingResponse(BaseModel):
     parent_group_id: Optional[int]
+    parent_hash: str
     stacktrace_distance: float
     message_distance: float
     should_group: bool
 
 
 class SimilarityResponse(BaseModel):
     responses: List[GroupingResponse]
-    token: Optional[int]
 
 
 class SimilarityBenchmarkResponse(BaseModel):
@@ -149,25 +150,6 @@ def get_nearest_neighbors(self, issue: GroupingRequest) -> SimilarityResponse:
                  stacktrace similarity scores, message similarity scores, and grouping flags.
         """
         with Session() as session:
-            # If an exact match of the stacktrace hash is found, return this record
-            if hasattr(issue, "stacktrace_hash") and issue.stacktrace_hash:
-                existing_record = (
-                    session.query(DbGroupingRecord)
-                    .filter_by(stacktrace_hash=issue.stacktrace_hash)
-                    .first()
-                )
-                if existing_record:
-                    similarity_response = SimilarityResponse(responses=[], token=None)
-                    similarity_response.responses.append(
-                        GroupingResponse(
-                            parent_group_id=existing_record.group_id,
-                            stacktrace_distance=0.00,
-                            message_distance=0.00,
-                            should_group=True,
-                        )
-                    )
-                    return similarity_response
-
             embedding = self.encode_text(issue.stacktrace).astype("float32")
 
             results = (
@@ -180,22 +162,24 @@ def get_nearest_neighbors(self, issue: GroupingRequest) -> SimilarityResponse:
                 .filter(
                     DbGroupingRecord.project_id == issue.project_id,
                     DbGroupingRecord.stacktrace_embedding.cosine_distance(embedding) <= 0.15,
-                    DbGroupingRecord.group_id != issue.group_id,
-                    DbGroupingRecord.group_id != None,
+                    or_(
+                        DbGroupingRecord.group_id != issue.group_id,
+                        DbGroupingRecord.group_id == None,
+                    ),
+                    # TODO We can return a group as similar group to itself if it exists in the old table with no hash
+                    DbGroupingRecord.hash != issue.hash,
                 )
                 .order_by("distance")
                 .limit(issue.k)
                 .all()
             )
 
             # If no existing groups within the threshold, insert the request as a new GroupingRecord
-            token = None
             if not any(distance <= issue.threshold for _, distance in results):
-                token = self.insert_new_grouping_record(session, issue, embedding)
-
+                self.insert_new_grouping_record(session, issue, embedding)
             session.commit()
 
-        similarity_response = SimilarityResponse(responses=[], token=token)
+        similarity_response = SimilarityResponse(responses=[])
         for record, distance in results:
             message_similarity_score = difflib.SequenceMatcher(
                 None, issue.message, record.message
@@ -204,7 +188,8 @@ def get_nearest_neighbors(self, issue: GroupingRequest) -> SimilarityResponse:
 
             similarity_response.responses.append(
                 GroupingResponse(
-                    parent_group_id=record.group_id,
+                    parent_group_id=record.group_id if hasattr(record, "group_id") else None,
+                    parent_hash=record.hash,
                     stacktrace_distance=distance,
                     message_distance=1.0 - message_similarity_score,
                     should_group=should_group,
@@ -215,7 +200,7 @@ def get_nearest_neighbors(self, issue: GroupingRequest) -> SimilarityResponse:
 
     def insert_new_grouping_record(
         self, session, issue: GroupingRequest, embedding: np.ndarray
-    ) -> int:
+    ) -> None:
         """
         Inserts a new GroupingRecord into the database if the group_id does not already exist.
         If new grouping record was created, return the id.
@@ -224,22 +209,14 @@ def insert_new_grouping_record(
         :param issue: The issue to insert as a new GroupingRecord.
         :param embedding: The embedding of the stacktrace.
         """
-        existing_record = None
-        if issue.group_id:
-            existing_record = (
-                session.query(DbGroupingRecord).filter_by(group_id=issue.group_id).first()
-            )
+        existing_record = session.query(DbGroupingRecord).filter_by(hash=issue.hash).first()
 
         if existing_record is None:
             new_record = GroupingRecord(
-                group_id=issue.group_id,
+                group_id=issue.group_id if hasattr(issue, "group_id") else None,
                 project_id=issue.project_id,
                 message=issue.message,
                 stacktrace_embedding=embedding,
-                stacktrace_hash=issue.stacktrace_hash,
+                hash=issue.hash,
             ).to_db_model()
             session.add(new_record)
-            session.commit()
-            return new_record.id
-
-        return existing_record.id

src/seer/schemas/seer.py

@@ -86,7 +86,8 @@
 GroupingRequest = typing_extensions.TypedDict(
     "GroupingRequest",
     {
-        "group_id": int,
+        "hash": str,
+        "group_id": typing.Union[int, None],
         "project_id": int,
         "stacktrace": str,
         "message": str,
@@ -102,6 +103,7 @@
     "GroupingResponse",
     {
         "parent_group_id": typing.Union[int, None],
+        "parent_hash": str,
         "stacktrace_similarity": float,
         "message_similarity": float,
         "should_group": bool,

src/seer/schemas/seer_api.json

@@ -386,9 +386,20 @@
             "GroupingRequest": {
                 "properties": {
                     "group_id": {
-                        "type": "integer",
+                        "anyOf": [
+                            {
+                                "type": "integer"
+                            },
+                            {
+                                "type": "null"
+                            }
+                        ],
                         "title": "Group Id"
                     },
+                    "hash": {
+                        "type": "string",
+                        "title": "Hash"
+                    },
                     "project_id": {
                         "type": "integer",
                         "title": "Project Id"
@@ -413,7 +424,7 @@
                     }
                 },
                 "type": "object",
-                "required": ["group_id", "project_id", "stacktrace", "message"],
+                "required": ["hash", "project_id", "stacktrace", "message"],
                 "title": "GroupingRequest"
             },
             "GroupingResponse": {
@@ -429,6 +440,10 @@
                         ],
                         "title": "Parent Group Id"
                     },
+                    "parent_hash": {
+                        "type": "string",
+                        "title": "Parent Hash"
+                    },
                     "stacktrace_similarity": {
                         "type": "number",
                         "title": "Stacktrace Similarity"
@@ -444,7 +459,7 @@
                 },
                 "type": "object",
                 "required": [
-                    "parent_group_id",
+                    "parent_hash",
                     "stacktrace_similarity",
                     "message_similarity",
                     "should_group"