Repository: mozilla/bugbug

Issue:
Add support for vector databases
As a start we want to support Qdrant.


Discussion:

Resolving patch:
diff --git a/bugbug/tools/code_review.py b/bugbug/tools/code_review.py
--- a/bugbug/tools/code_review.py
+++ b/bugbug/tools/code_review.py
@@ -3,17 +3,40 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this file,
 # You can obtain one at http://mozilla.org/MPL/2.0/.
 
+import re
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
+from collections import defaultdict
+from dataclasses import asdict, dataclass
+from logging import INFO, basicConfig, getLogger
+from typing import Iterable
 
+import tenacity
 from langchain.chains import ConversationChain, LLMChain
 from langchain.memory import ConversationBufferMemory
 from langchain.prompts import PromptTemplate
+from langchain_openai import OpenAIEmbeddings
+from tqdm import tqdm
 from unidiff import Hunk, PatchedFile, PatchSet
+from unidiff.errors import UnidiffParseError
 
-from bugbug import phabricator
+from bugbug import db, phabricator
 from bugbug.generative_model_tool import GenerativeModelTool
 from bugbug.utils import get_secret
+from bugbug.vectordb import VectorDB, VectorPoint
+
+basicConfig(level=INFO)
+logger = getLogger(__name__)
+
+
+@dataclass
+class InlineComment:
+    filename: str
+    start_line: int
+    end_line: int
+    content: str
+    on_removed_code: bool | None
+    id: int | None = None
+    date_created: int | None = None
 
 
 class ModelResultError(Exception):
@@ -109,6 +132,8 @@ def __init__(self, raw_diff) -> None:
 
 
 class ReviewData(ABC):
+    NIT_PATTERN = re.compile(r"[^a-zA-Z0-9]nit[\s:,]", re.IGNORECASE)
+
     @abstractmethod
     def get_review_request_by_id(self, review_id: int) -> ReviewRequest:
         raise NotImplementedError
@@ -117,9 +142,150 @@ def get_review_request_by_id(self, review_id: int) -> ReviewRequest:
     def get_patch_by_id(self, patch_id: int) -> Patch:
         raise NotImplementedError
 
+    @abstractmethod
+    def get_all_inline_comments(
+        self, comment_filter
+    ) -> Iterable[tuple[int, list[InlineComment]]]:
+        raise NotImplementedError
+
+    def load_patch_by_id(self, diff_id) -> Patch:
+        """Load a patch from local cache if it exists.
+
+        If the patch is not in the local cache it will be requested from the
+        provider and cache it locally.
+
+        Args:
+            diff_id: The ID of the patch.
+
+        Returns:
+            The patch.
+        """
+        try:
+            with open(f"patches/{diff_id}.patch", "r") as f:
+                patch = Patch(f.read())
+        except FileNotFoundError:
+            with open(f"patches/{diff_id}.patch", "w") as f:
+                patch = self.get_patch_by_id(diff_id)
+                f.write(patch.raw_diff)
+
+        return patch
+
+    def get_matching_hunk(
+        self, patched_file: PatchedFile, comment: InlineComment
+    ) -> Hunk:
+        def source_end(hunk: Hunk) -> int:
+            return hunk.source_start + hunk.source_length
+
+        def target_end(hunk: Hunk) -> int:
+            return hunk.target_start + hunk.target_length
+
+        if comment.on_removed_code is None:
+            matching_hunks = [
+                hunk
+                for hunk in patched_file
+                if hunk.target_start <= comment.start_line < target_end(hunk)
+                or hunk.source_start <= comment.start_line < source_end(hunk)
+            ]
+
+            # If there is more than one matching hunk, choose the one where the
+            # line number of the comment corresponds to an added or deleted line. We
+            # prioritize added lines over deleted lines because comments are more
+            # likely to be on added lines than deleted lines.
+            if len(matching_hunks) > 1:
+                for hunk in matching_hunks:
+                    for line in hunk:
+                        if line.is_added and line.target_line_no == comment.start_line:
+                            return hunk
+
+                    for line in hunk:
+                        if (
+                            line.is_removed
+                            and line.source_line_no == comment.start_line
+                        ):
+                            return hunk
+
+            if len(matching_hunks) != 0:
+                logger.warning(
+                    "Multiple matching hunks found for comment %s in file %s",
+                    comment.id,
+                    comment.filename,
+                )
+                return matching_hunks[0]
+
+        elif comment.on_removed_code:
+            for hunk in patched_file:
+                if hunk.source_start <= comment.start_line < source_end(hunk):
+                    return hunk
+
+        else:
+            for hunk in patched_file:
+                if hunk.target_start <= comment.start_line < target_end(hunk):
+                    return hunk
+
+    def retrieve_comments_with_hunks(self):
+        def comment_filter(comment: InlineComment):
+            comment_content = comment.content
+
+            # Ignore very short and very long comments
+            if not 50 < len(comment_content) < 500:
+                return False
+
+            # Ignore comments with URLs
+            if "https://" in comment_content or "http://" in comment_content:
+                return False
+
+            #  Ignore nit comments
+            if self.NIT_PATTERN.search(comment_content):
+                return False
+
+            # Ignore comments with code blocks
+            if "```" in comment_content:
+                return False
+
+            comment_lower = comment_content.lower()
+            if any(
+                phrase in comment_lower
+                for phrase in [
+                    "wdyt?",
+                    "what do you think?",
+                    "you explain",
+                    "understand",
+                ]
+            ):
+                return False
+
+            return True
+
+        for diff_id, comments in self.get_all_inline_comments(comment_filter):
+            try:
+                patch_set = PatchSet.from_string(
+                    self.load_patch_by_id(diff_id).raw_diff
+                )
+            except UnidiffParseError:
+                # TODO: use log instead of print
+                print(f"Failed to parse {diff_id}")
+                continue
+
+            file_map = {
+                patched_file.path: patched_file
+                for patched_file in patch_set
+                if patched_file.is_modified_file
+            }
+            for comment in comments:
+                patched_file = file_map.get(comment.filename)
+                if not patched_file:
+                    continue
+
+                hunk = self.get_matching_hunk(patched_file, comment)
+                if not hunk:
+                    continue
+
+                yield comment, hunk
+
 
 class PhabricatorReviewData(ReviewData):
     def __init__(self):
+        super().__init__()
         phabricator.set_api_key(
             get_secret("PHABRICATOR_URL"), get_secret("PHABRICATOR_TOKEN")
         )
@@ -129,26 +295,85 @@ def get_review_request_by_id(self, revision_id: int) -> ReviewRequest:
         assert len(revisions) == 1
         return ReviewRequest(revisions[0]["fields"]["diffID"])
 
+    @tenacity.retry(
+        stop=tenacity.stop_after_attempt(7),
+        wait=tenacity.wait_exponential(multiplier=1, min=16, max=64),
+        reraise=True,
+    )
     def get_patch_by_id(self, patch_id: int) -> Patch:
         assert phabricator.PHABRICATOR_API is not None
         raw_diff = phabricator.PHABRICATOR_API.load_raw_diff(int(patch_id))
         return Patch(raw_diff)
 
+    def get_all_inline_comments(
+        self, comment_filter
+    ) -> Iterable[tuple[int, list[InlineComment]]]:
+        db.download(phabricator.REVISIONS_DB)
+
+        revision_count = sum(1 for _ in phabricator.get_revisions())
+        for revision in tqdm(phabricator.get_revisions(), total=revision_count):
+            diff_comments: dict[int, list[InlineComment]] = defaultdict(list)
+
+            for transaction in revision["transactions"]:
+                if transaction["type"] != "inline":
+                    continue
+
+                # Ignore replies
+                if transaction["fields"]["replyToCommentPHID"] is not None:
+                    continue
+
+                # Ignore bot comments
+                if transaction["authorPHID"] == "PHID-USER-cje4weq32o3xyuegalpj":
+                    continue
+
+                if len(transaction["comments"]) != 1:
+                    # Follow up: https://github.com/mozilla/bugbug/issues/4218
+                    logger.warning(
+                        "Unexpected number of comments in transaction %s",
+                        transaction["id"],
+                    )
+
+                transaction_comment = transaction["comments"][0]
+                comment_id = transaction_comment["id"]
+                date_created = transaction_comment["dateCreated"]
+                comment_content = transaction_comment["content"]["raw"]
+
+                diff_id = transaction["fields"]["diff"]["id"]
+                filename = transaction["fields"]["path"]
+                start_line = transaction["fields"]["line"]
+                end_line = (
+                    transaction["fields"]["line"] + transaction["fields"]["length"] - 1
+                )
+                # Unfortunately, we do not have this information for a limitation
+                # in Phabricator's API.
+                on_removed_code = None
+
+                # TODO: we could create an extended dataclass for this
+                # instead of adding optional fields.
+                comment = InlineComment(
+                    filename,
+                    start_line,
+                    end_line,
+                    comment_content,
+                    on_removed_code,
+                    comment_id,
+                    date_created,
+                )
+
+                if not comment_filter(comment):
+                    continue
+
+                diff_comments[diff_id].append(comment)
+
+            for diff_id, comments in diff_comments.items():
+                yield diff_id, comments
+
 
 review_data_classes = {
     "phabricator": PhabricatorReviewData,
 }
 
 
-@dataclass
-class InlineComment:
-    filename: str
-    start_line: int
-    end_line: int
-    comment: str
-    on_added_code: bool
-
-
 def find_comment_scope(file: PatchedFile, line_number: int):
     for hunk in file:
         if hunk.target_start <= line_number <= hunk.target_start + hunk.target_length:
@@ -279,3 +504,33 @@ def run(self, patch: Patch) -> list[InlineComment] | None:
         )["text"]
 
         return raw_output
+
+
+class ReviewCommentsDB:
+    NAV_PATTERN = re.compile(r"\{nav, [^}]+\}")
+    WHITESPACE_PATTERN = re.compile(r"[\n\s]+")
+
+    def __init__(self, vector_db: VectorDB) -> None:
+        self.vector_db = vector_db
+        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
+
+    def clean_comment(self, comment):
+        # TODO: use the nav info instead of removing it
+        comment = self.NAV_PATTERN.sub("", comment)
+        comment = self.WHITESPACE_PATTERN.sub(" ", comment)
+        comment = comment.strip()
+
+        return comment
+
+    def add_comments_by_hunk(self, items: Iterable[tuple[Hunk, InlineComment]]):
+        self.vector_db.insert(
+            VectorPoint(
+                id=comment.id,
+                vector=self.embeddings.embed_query(str(hunk)),
+                payload=asdict(comment),
+            )
+            for comment, hunk in items
+        )
+
+    def find_similar_hunk_comments(self, hunk: Hunk):
+        return self.vector_db.search(self.embeddings.embed_query(str(hunk)))
diff --git a/bugbug/vectordb.py b/bugbug/vectordb.py
new file mode 100644
--- /dev/null
+++ b/bugbug/vectordb.py
@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this file,
+# You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Iterable
+
+from qdrant_client import QdrantClient
+from qdrant_client.http.exceptions import UnexpectedResponse
+from qdrant_client.models import Distance, PointStruct, VectorParams
+
+from bugbug.utils import get_secret
+
+
+@dataclass
+class VectorPoint:
+    id: int
+    vector: list[float]
+    payload: dict
+
+
+class VectorDB(ABC):
+    """Abstract class for a vector database.
+
+    You can implement this class to support different vector databases.
+    """
+
+    @abstractmethod
+    def setup(self):
+        ...
+
+    @abstractmethod
+    def insert(self, points: Iterable[VectorPoint]):
+        ...
+
+    @abstractmethod
+    def search(self, query: list[float]):
+        ...
+
+
+class QdrantVectorDB(VectorDB):
+    def __init__(self, collection_name: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.collection_name = collection_name
+        self.client = QdrantClient(
+            location=get_secret("QDRANT_LOCATION"), api_key=get_secret("QDRANT_API_KEY")
+        )
+
+    def setup(self):
+        config = VectorParams(
+            size=3072,
+            distance=Distance.COSINE,
+        )
+        try:
+            self.client.create_collection(self.collection_name, vectors_config=config)
+        except UnexpectedResponse as e:
+            # We only allow 409 (collection already exists) to pass through
+            if e.status_code != 409:
+                raise Exception("Failed to create collection") from e
+
+    def insert(self, points: Iterable[VectorPoint]):
+        self.client.upload_points(
+            self.collection_name,
+            (
+                PointStruct(
+                    id=point.id,
+                    vector=point.vector,
+                    payload=point.payload,
+                )
+                for point in points
+            ),
+        )
+
+    def search(self, query: list[float]):
+        return self.client.search(self.collection_name, query)
diff --git a/scripts/review_comments_retriever.py b/scripts/review_comments_retriever.py
new file mode 100644
--- /dev/null
+++ b/scripts/review_comments_retriever.py
@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this file,
+# You can obtain one at http://mozilla.org/MPL/2.0/.
+
+
+from bugbug.tools.code_review import PhabricatorReviewData, ReviewCommentsDB
+from bugbug.vectordb import QdrantVectorDB
+
+
+def main():
+    review_data = PhabricatorReviewData()
+    vector_db = QdrantVectorDB("diff_comments")
+    vector_db.setup()
+    comments_db = ReviewCommentsDB(vector_db)
+    # TODO: support resuming from where last run left off. We should run it from
+    # scratch only once. Following runs should add only new comments.
+    comments_db.add_comments_by_hunk(review_data.retrieve_comments_with_hunks())
+
+
+if __name__ == "__main__":
+    main()


Test patch: