Initial implementation of correctness check (microsoft#8)

Author: ultmaster

README.md

@@ -127,7 +127,7 @@ To uninstall, you can run the following commands:
 jupyter server extension disable coml
 
 # Uninstall the Python package
-pip uninstall coml
+pip uninstall mlcopilot
 ```
 
 In development mode, you will also need to remove the symlink created by `jupyter labextension develop` command.

coml/__init__.py

@@ -17,6 +17,8 @@ def load_ipython_extension(ipython):
 
     ipython.register_magics(CoMLMagics)
 
+    print(f"CoML {__version__} loaded.")
+
 
 def _jupyter_labextension_paths():
     return [{"src": "labextension", "dest": "coml"}]

coml/core.py

@@ -10,9 +10,11 @@
 from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage
 
 from .prompt_utils import (
+    CHECK_INSTRUCTION,
     EXPLAIN_INSTRUCTION,
     FIX_INSTRUCTION,
     GENERATE_INSTRUCTION,
+    SANITY_CHECK_INSTRUCTION,
     SUGGEST_INSTRUCTION,
     FixContext,
     GenerateContext,
@@ -21,9 +23,11 @@
     InteractionIncomplete,
     cached_fix_fewshots,
     cached_generate_fewshots,
+    render_check_context,
     render_fix_context,
     render_generate_context,
     render_ipython_cells,
+    render_sanity_check_context,
 )
 
 
@@ -183,3 +187,45 @@ def explain(self, code: str) -> str:
         response = self.llm(messages)
         debug_messages(response)
         return response.content
+
+    def static_check(
+        self, code: str, context: GenerateContext | FixContext
+    ) -> tuple[bool, str]:
+        # Check the quality of code by looking at it (i.e., rubberduck)
+        messages = [
+            SystemMessage(content=CHECK_INSTRUCTION),
+            HumanMessage(content=render_check_context(code, context)),
+        ]
+        debug_messages(*messages)
+        response = self.llm(messages)
+        debug_messages(response)
+        reason, last_line = response.content.rstrip().rsplit("\n", 1)
+        if "INCORRECT" in last_line.upper():
+            return False, reason
+        if "CORRECT" in last_line.upper():
+            return True, reason
+        raise ValueError("Unable to parse the response.")
+
+    def output_sanity_check(
+        self,
+        code: str,
+        context: GenerateContext | FixContext,
+        error: str | None,
+        output: str | None,
+    ) -> tuple[bool, str]:
+        # Run a sanity check of the output of the code
+        messages = [
+            SystemMessage(content=SANITY_CHECK_INSTRUCTION),
+            HumanMessage(
+                content=render_sanity_check_context(code, context, error, output)
+            ),
+        ]
+        debug_messages(*messages)
+        response = self.llm(messages)
+        debug_messages(response)
+        reason, last_line = response.content.rstrip().rsplit("\n", 1)
+        if "INCORRECT" in last_line.upper():
+            return False, reason
+        if "CORRECT" in last_line.upper():
+            return True, reason
+        raise ValueError("Unable to parse the response.")

coml/ipython_utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import base64
 import json
 import re
@@ -80,6 +82,27 @@ def run_code_in_next_cell(python_code: str, metadata: Any = None) -> None:
         )
 
 
+def update_running_cell_metadata(metadata: Any) -> None:
+    if is_jupyter_lab_environ():
+        input(
+            json.dumps(
+                {"command": "update_running_cell_metadata", "metadata": metadata}
+            )
+        )
+    else:
+        encoded_metadata = base64.b64encode(json.dumps(metadata).encode()).decode()
+        display(
+            Javascript(
+                """
+            const cell = comlGetCurrentCell();
+            cell.metadata.coml = Object.assign(cell.metadata.coml || {}, JSON.parse(atob(\""""
+                + encoded_metadata
+                + """\")));
+        """
+            )
+        )
+
+
 def get_ipython_history(ipython: InteractiveShell) -> list[str]:
     codes = []
     for code in ipython.user_ns["In"]:
@@ -98,6 +121,11 @@ def get_ipython_history(ipython: InteractiveShell) -> list[str]:
     return codes
 
 
+def get_running_cell() -> dict[str, Any] | None:
+    """See `get_last_cell` for the output format."""
+    return json.loads(input(json.dumps({"command": "running_cell"})))
+
+
 def get_last_cell() -> dict[str, Any] | None:
     """The implementation is in nbclassic_init.js. This is a *hacked* RPC channel.
 

coml/js/nbclassic_init.js

@@ -1,33 +1,52 @@
-function getDebugInformation() {
+window.comlGetRunningCellIndex = function () {
     const runningCells = $(".running");
     if (runningCells.length === 0) {
         console.warn("No running cell");
         return null;
     }
     const cellIndex = Jupyter.notebook.get_cell_elements().index(runningCells[0]);
-    if (cellIndex <= 0) {
-        console.warn("No previous cell");
+    if (cellIndex < 0) {
+        console.error("Running cell not found in cell list.");
         return null;
     }
-    const cell = IPython.notebook.get_cell(cellIndex - 1);
-    const cellDump = cell.toJSON();
-    return cellDump;
+    return cellIndex;
 }
 
-IPython.CodeCell.prototype.native_handle_input_request = IPython.CodeCell.prototype.native_handle_input_request || IPython.CodeCell.prototype._handle_input_request;
-IPython.CodeCell.prototype._handle_input_request = function (msg) {
-    try {
-        // only apply the hack if the command is valid JSON
-        const command = JSON.parse(msg.content.prompt);
-        const kernel = IPython.notebook.kernel;
-        if (command["command"] === "last_cell") {
-            kernel.send_input_reply(JSON.stringify(getDebugInformation()));
-        } else {
-            console.log("Not a command", msg);
+window.comlGetCurrentCell = function () {
+    const cell = comlGetRunningCellIndex();
+    if (cell === null) {
+        return null;
+    }
+    return IPython.notebook.get_cell(cell);
+}
+
+window.comlGetLastCell = function () {
+    const cellIndex = comlGetRunningCellIndex();
+    if (cellIndex === null) {
+        return null;
+    }
+    return IPython.notebook.get_cell(comlGetRunningCellIndex() - 1);
+}
+
+if (window.IPython && IPython.CodeCell) {
+    window.IPythonAvailable = true;
+    IPython.CodeCell.prototype.native_handle_input_request = IPython.CodeCell.prototype.native_handle_input_request || IPython.CodeCell.prototype._handle_input_request;
+    IPython.CodeCell.prototype._handle_input_request = function (msg) {
+        try {
+            // only apply the hack if the command is valid JSON
+            const command = JSON.parse(msg.content.prompt);
+            const kernel = IPython.notebook.kernel;
+            if (command["command"] === "last_cell") {
+                kernel.send_input_reply(JSON.stringify(comlGetLastCell().toJSON()));
+            } else if (command["command"] === "running_cell") {
+                kernel.send_input_reply(JSON.stringify(comlGetCurrentCell().toJSON()));
+            } else {
+                console.log("Not a command", msg);
+                this.native_handle_input_request(msg);
+            }
+        } catch(err) {
+            console.log("Not a command", msg, err);
             this.native_handle_input_request(msg);
         }
-    } catch(err) {
-        console.log("Not a command", msg, err);
-        this.native_handle_input_request(msg);
     }
-}
+}

coml/linter.py

@@ -0,0 +1,49 @@
+import json
+import tempfile
+from io import StringIO
+from typing import Literal, Tuple
+
+from pylint.lint import Run as PylintRun
+from pylint.reporters import JSONReporter
+
+LinterResult = Literal["error", "warning", "info", "ok"]
+
+
+def lint(previous_code: str, new_code: str) -> Tuple[LinterResult, str]:
+    # https://stackoverflow.com/q/75507725/6837658
+    pylint_options = [
+        "--disable=C0103",  # Invalid name
+        "--disable=C0114",  # Missing module docstring
+        "--disable=C0304",  # Final new line missing
+    ]
+    previous_lines = previous_code.count("\n") + 1
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(previous_code + "\n" + new_code)
+        f.flush()
+        f.seek(0)
+
+        reporter_buffer = StringIO()
+        results = PylintRun(
+            [f.name] + pylint_options,
+            reporter=JSONReporter(reporter_buffer),
+            do_exit=False,
+        )
+        # Score is here.
+        # score = results.linter.stats.global_note
+        file_results = json.loads(reporter_buffer.getvalue())
+        file_results = [e for e in file_results if e["line"] > previous_lines]
+
+        details = []
+        for error in file_results:
+            line = f"{error['line'] - previous_lines}:{error['column']}: {error['message-id']}: {error['message']}"
+            details.append(line)
+        details_joined = "\n".join(details)
+
+        if any(e["type"] in ("fatal", "error") for e in file_results):
+            return "error", details_joined
+        elif any(e["type"] == "warning" for e in file_results):
+            return "warning", details_joined
+        elif file_results:
+            return "info", details_joined
+        else:
+            return "ok", "No issues found."