diff --git a/reai_toolkit/app/components/tabs/ai_decomp_tab.py b/reai_toolkit/app/components/tabs/ai_decomp_tab.py index 54db1fe..5e2e63d 100644 --- a/reai_toolkit/app/components/tabs/ai_decomp_tab.py +++ b/reai_toolkit/app/components/tabs/ai_decomp_tab.py @@ -1,3 +1,4 @@ +from pathlib import Path from typing import Any, Callable, Optional from loguru import logger @@ -12,6 +13,17 @@ ) +def _thumb_icon(up: bool) -> Optional[QtGui.QIcon]: + stem = "thumb_up" if up else "thumb_down" + for ext in (".svg", ".png"): + path = Path(__file__).resolve().parent.parent / "resources" / (stem + ext) + if path.exists(): + icon = QtGui.QIcon(str(path)) + if not icon.isNull(): + return icon + return None + + def _menu_exec(menu, pos): fn = getattr(menu, "exec_", None) or getattr(menu, "exec") return fn(pos) @@ -71,9 +83,13 @@ def __init__(self, on_closed: Optional[Callable[[], None]] = None) -> None: self.on_rename: Callable[[int, str], None] | None = None self.on_edit_comment: Callable[[int], None] | None = None self.on_remove_comment: Callable[[int], None] | None = None + self.on_rate_up: Callable[[], None] | None = None + self.on_rate_down: Callable[[], None] | None = None self._parent_window: QtWidgets.QWidget | None = None self._editor: _DecompEditor | None = None self._refresh_btn: QtWidgets.QPushButton | None = None + self._rate_up_btn: QtWidgets.QPushButton | None = None + self._rate_down_btn: QtWidgets.QPushButton | None = None self._highlighter: CppHighlighter | None = None def Create(self, title: Any) -> Any: @@ -107,6 +123,31 @@ def OnCreate(self, form) -> None: title = QtWidgets.QLabel("RevEng.AI — AI Decomp", self._parent_window) header.addWidget(title) header.addStretch(1) + + up_icon = _thumb_icon(up=True) + self._rate_up_btn = QtWidgets.QPushButton(self._parent_window) + if up_icon is not None: + self._rate_up_btn.setIcon(up_icon) + self._rate_up_btn.setIconSize(QtCore.QSize(16, 16)) + else: + self._rate_up_btn.setText("\U0001f44d") + self._rate_up_btn.setCheckable(True) + self._rate_up_btn.setToolTip("Rate this AI decompilation as good") + self._rate_up_btn.clicked.connect(self._on_rate_up_clicked) + header.addWidget(self._rate_up_btn) + + down_icon = _thumb_icon(up=False) + self._rate_down_btn = QtWidgets.QPushButton(self._parent_window) + if down_icon is not None: + self._rate_down_btn.setIcon(down_icon) + self._rate_down_btn.setIconSize(QtCore.QSize(16, 16)) + else: + self._rate_down_btn.setText("\U0001f44e") + self._rate_down_btn.setCheckable(True) + self._rate_down_btn.setToolTip("Rate this AI decompilation as poor") + self._rate_down_btn.clicked.connect(self._on_rate_down_clicked) + header.addWidget(self._rate_down_btn) + self._refresh_btn = QtWidgets.QPushButton("Refresh", self._parent_window) self._refresh_btn.clicked.connect(self._on_refresh_clicked) header.addWidget(self._refresh_btn) @@ -147,12 +188,31 @@ def OnClose(self, form) -> None: self._highlighter = None self._editor = None self._refresh_btn = None + self._rate_up_btn = None + self._rate_down_btn = None self._parent_window = None def _on_refresh_clicked(self) -> None: if self.on_refresh: self.on_refresh() + def _on_rate_up_clicked(self) -> None: + self.set_rating("up") + if self.on_rate_up: + self.on_rate_up() + + def _on_rate_down_clicked(self) -> None: + self.set_rating("down") + if self.on_rate_down: + self.on_rate_down() + + @execute_ui + def set_rating(self, rating: Optional[str]) -> None: + if self._rate_up_btn: + self._rate_up_btn.setChecked(rating == "up") + if self._rate_down_btn: + self._rate_down_btn.setChecked(rating == "down") + def _on_rename_requested(self, line: int, word: str) -> None: if self.on_rename: self.on_rename(line, word) diff --git a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py index 87b37ac..7eb7aea 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py +++ b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py @@ -3,6 +3,7 @@ import ida_funcs import ida_kernwin from libbs.decompilers.ida.compat import execute_read, execute_ui +from revengai.models.ai_decompilation_rating import AiDecompilationRating from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData from revengai.models.summary_data import SummaryData @@ -70,6 +71,8 @@ def run_dialog(self) -> None: self._decomp_view.on_rename = self.request_rename self._decomp_view.on_edit_comment = self.request_edit_comment self._decomp_view.on_remove_comment = self.request_remove_comment + self._decomp_view.on_rate_up = self.rate_up + self._decomp_view.on_rate_down = self.rate_down self._decomp_view.Create(self._decomp_view.TITLE) def start_decompilation(self, ea: int) -> None: @@ -84,6 +87,7 @@ def start_decompilation(self, ea: int) -> None: cached = self.ai_decomp_service.peek_decomp(ea) if self._decomp_view is not None: + self._decomp_view.set_rating(None) if cached is not None and cached.decompilation: self._current_decomp = cached self._rerender() @@ -203,6 +207,38 @@ def refresh_current(self) -> None: self.ai_decomp_service.invalidate_ea(ea) self.start_decompilation(ea) + def rate_up(self) -> None: + self._rate(AiDecompilationRating.POSITIVE) + + def rate_down(self) -> None: + self._rate(AiDecompilationRating.NEGATIVE) + + def _rate(self, rating: AiDecompilationRating) -> None: + ea = self._current_func_vaddr + if ea is None: + return + if self._current_decomp is None: + self.show_info_dialog(message="No AI decompilation to rate yet.") + if self._decomp_view is not None: + self._decomp_view.set_rating(None) + return + self.ai_decomp_service.rate_decomp( + ea=ea, + rating=rating, + on_result=lambda response: self._on_rating_result(ea, response), + ) + + def _on_rating_result( + self, ea: int, response: GenericApiReturn[bool] + ) -> None: + if ea != self._current_func_vaddr: + return + if not response.success: + if response.error_message: + self.show_error_dialog(message=response.error_message) + if self._decomp_view is not None: + self._decomp_view.set_rating(None) + def request_rename(self, display_line: int, word: str) -> None: ea = self._current_func_vaddr if ea is None or self._baseline is None or self._current_tokenised is None: diff --git a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py index fe48fab..c53a227 100644 --- a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py +++ b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py @@ -8,12 +8,16 @@ FunctionMapping, FunctionsAIDecompilationApi, ) +from revengai.models.ai_decompilation_rating import AiDecompilationRating from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData from revengai.models.patch_comment_body import PatchCommentBody from revengai.models.summary_data import SummaryData from revengai.models.task_status import TaskStatus from revengai.models.tokenised_data import TokenisedData +from revengai.models.upsert_ai_decomplation_rating_request import ( + UpsertAiDecomplationRatingRequest, +) from revengai.models.upsert_overrides_input_body import UpsertOverridesInputBody from revengai.models.workflow_progress import WorkflowProgress @@ -177,10 +181,60 @@ def remove_comment( (function_id, line, on_result), ) + def rate_decomp( + self, + ea: int, + rating: AiDecompilationRating, + on_result: Callable[[GenericApiReturn[bool]], None], + ) -> None: + function_id = self._get_function_id(start_ea=ea) + if function_id is None: + on_result( + GenericApiReturn[bool]( + success=False, error_message="Function is not part of the analysis." + ) + ) + return + self._spawn( + self._run_rate_decomp, + f"reai-aidecomp-rating-{function_id}", + (function_id, rating, on_result), + ) + @staticmethod def _spawn(target: Callable[..., Any], name: str, args: tuple) -> None: threading.Thread(target=target, args=args, name=name, daemon=True).start() + def _run_rate_decomp( + self, + function_id: int, + rating: AiDecompilationRating, + on_result: Callable[[GenericApiReturn[bool]], None], + ) -> None: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + FunctionsAIDecompilationApi(api_client).upsert_ai_decompilation_rating( + function_id=function_id, + upsert_ai_decomplation_rating_request=UpsertAiDecomplationRatingRequest( + rating=rating, reason=None + ), + ) + except ApiException as e: + on_result( + GenericApiReturn[bool]( + success=False, error_message=_format_api_error(e) + ) + ) + return + except Exception as e: + on_result( + GenericApiReturn[bool]( + success=False, error_message=f"Unexpected error submitting rating: {e}" + ) + ) + return + on_result(GenericApiReturn[bool](success=True, data=True)) + def _run_apply_overrides( self, function_id: int, diff --git a/tests/unit/ai_decomp/test_service.py b/tests/unit/ai_decomp/test_service.py index 2e5b8e3..c25a210 100644 --- a/tests/unit/ai_decomp/test_service.py +++ b/tests/unit/ai_decomp/test_service.py @@ -4,6 +4,7 @@ import pytest from revengai import ApiException +from revengai.models.ai_decompilation_rating import AiDecompilationRating from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData from revengai.models.inline_comment import InlineComment @@ -480,6 +481,48 @@ def test_mutation_unknown_function_id_reports_failure(service, sdk, netstore): sdk.patch_ai_decompilation_inline_comment.assert_not_called() +def test_rate_decomp_calls_upsert_rating(service, sdk): + sdk.upsert_ai_decompilation_rating.return_value = MagicMock() + + on_result = MagicMock() + service.rate_decomp( + ea=4096, rating=AiDecompilationRating.POSITIVE, on_result=on_result + ) + _wait_mock(on_result) + + _, kwargs = sdk.upsert_ai_decompilation_rating.call_args + assert kwargs["function_id"] == 42 + body = kwargs["upsert_ai_decomplation_rating_request"] + assert body.rating == AiDecompilationRating.POSITIVE + assert body.reason is None + assert on_result.call_args[0][0].success is True + + +def test_rate_decomp_api_error_surfaces(service, sdk): + sdk.upsert_ai_decompilation_rating.side_effect = ApiException(status=422) + + on_result = MagicMock() + service.rate_decomp( + ea=4096, rating=AiDecompilationRating.NEGATIVE, on_result=on_result + ) + _wait_mock(on_result) + + assert on_result.call_args[0][0].success is False + + +def test_rate_decomp_unknown_function_id_reports_failure(service, sdk, netstore): + netstore.get_function_mapping.return_value.inverse_function_map = {} + + on_result = MagicMock() + service.rate_decomp( + ea=4096, rating=AiDecompilationRating.POSITIVE, on_result=on_result + ) + + on_result.assert_called_once() + assert on_result.call_args[0][0].success is False + sdk.upsert_ai_decompilation_rating.assert_not_called() + + def test_invalidate_clears_all_caches_and_inflight(service): service._decomp_cache[42] = object() service._summary_cache[42] = object()