Making inference interactive

experimental/rule_inference/local.py

@@ -48,7 +48,13 @@ def infer_from_example(data):
     room = session.get("room")
     join_room(room)
 
-    rule_name, rule = agent.infer_rules()
+    rule_name, rule = agent.infer_rules(
+        lambda intermediate_result: socketio.emit(
+            "infer_progress",
+            {"rule": intermediate_result},
+            room=room,
+        )
+    )
     socketio.emit("infer_result", {"rule_name": rule_name, "rule": rule}, room=room)
     session["last_inference_result"] = {
         "rule_name": rule_name,

experimental/rule_inference/piranha_agent.py

@@ -59,7 +59,7 @@ def get_tree_from_code(self, code: str) -> Tree:
         tree = parser.parse(bytes(code, "utf8"))
         return tree
 
-    def infer_rules(self) -> Optional[Tuple[str, str]]:
+    def infer_rules(self, callback=None) -> Optional[Tuple[str, str]]:
         """Implements the inference process of the Piranha Agent.
         The function communicates with the AI model to generate a potential refactoring rule, and subsequently tests it.
         If the rule transforms the source code into the target code, the rule is returned.
@@ -85,7 +85,10 @@ def infer_rules(self) -> Optional[Tuple[str, str]]:
         pairs = finder.find_replacement_pairs()
         for nodes_before, nodes_after in pairs.values():
             inference_engine = Inference(nodes_before, nodes_after)
-            rules += "\n\n" + inference_engine.static_infer()
+            rules += inference_engine.static_infer() + "\n\n"
+
+        if callback:
+            callback(rules)
 
         prompt_holes = {
             "source_code": self.source_code,
@@ -169,7 +172,9 @@ def validate_rule(self, completion):
             )
 
         try:
-            toml_block = toml_blocks[0].replace("parenthesized_expression", "condition")
+            toml_block = (
+                toml_blocks[0].replace("parenthesized_expression", "condition").strip()
+            )
             logger.debug(f"Generated rule: {toml_block}")
             toml_dict = toml.loads(toml_block)
             return "rule.toml", toml_block

experimental/rule_inference/static/script.js

@@ -100,6 +100,13 @@ let tree;
     buttonText.textContent = "Infer"; // Reset button text
   });
 
+  socket.on("infer_progress", function (data) {
+    // This is where you could update your interface with the data.
+    const toml = data.rule;
+    document.getElementById("query-container").style.display = "block";
+    queryEditor.setValue(toml);
+  });
+
   // To start the inference, we can emit the 'infer_piranha' event.
   document
     .getElementById("submit-button")