Fix logging callback (#77)

Author: mahaloz

dailalib/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "3.15.4"
+__version__ = "3.15.5"
 
 import os
 # stop LiteLLM from querying at all to the remote server

dailalib/api/ai_api.py

@@ -118,9 +118,9 @@ def _requires_function(*args, ai_api: "AIAPI" = None, **kwargs):
 
         return _requires_function
 
-    def on_query(self, query_name, model, prompt_style, function, decompilation, **kwargs):
+    def on_query(self, query_name, model, prompt_style, function, decompilation, response, **kwargs):
         for func in self.query_callbacks:
             t = threading.Thread(
-                target=func, args=(query_name, model, prompt_style, function, decompilation), kwargs=kwargs
+                target=func, args=(query_name, model, prompt_style, function, decompilation, response), kwargs=kwargs
             )
             t.start()

dailalib/api/litellm/prompts/prompt.py

@@ -109,13 +109,13 @@ def _query_model(ai_api=self.ai_api, function=function, dec_text=dec_text, **_kw
             total_time = end_time - start_time
 
             # callback to handlers of post-query
-            ai_api.on_query(
-                self.name, self.ai_api.model, self.ai_api.prompt_style, function, dec_text, total_time=total_time, cost=cost
-            )
+            callback_args = [self.name, self.ai_api.model, self.ai_api.prompt_style, function, dec_text, response]
+            callback_kwargs = {"total_time": total_time, "cost": cost, "success": False}
 
             default_response = {} if self._json_response else ""
             if not response:
                 ai_api.warning(f"Response received from AI was empty! AI failed to answer.")
+                ai_api.on_query(*callback_args, **callback_kwargs)
                 return default_response
 
             # changes response type to a dict
@@ -126,6 +126,7 @@ def _query_model(ai_api=self.ai_api, function=function, dec_text=dec_text, **_kw
 
                 json_matches = JSON_REGEX.findall(response)
                 if not json_matches:
+                    ai_api.on_query(*callback_args, **callback_kwargs)
                     return default_response
 
                 json_data = json_matches[-1]
@@ -148,6 +149,8 @@ def _query_model(ai_api=self.ai_api, function=function, dec_text=dec_text, **_kw
                 log_str += f" AI likely failed to answer coherently."
             ai_api.info(log_str)
 
+            callback_kwargs["success"] = True
+            ai_api.on_query(*callback_args, **callback_kwargs)
             if ai_api.has_decompiler_gui and response:
                 ai_api.info("Updating the decompiler with the AI response...")
                 self._gui_result_callback(response, function, ai_api, context=context)

dailalib/llm_chat/llm_chat_ui.py

@@ -137,7 +137,7 @@ def send_message(self, add_text=True, role="user"):
         if self.ai_api:
             send_callback = self.ai_api.chat_event_callbacks.get("send", None)
             if send_callback:
-                send_callback(user_text)
+                send_callback(user_text, model=self.model)
 
         # Display user message
         if add_text:
@@ -164,7 +164,7 @@ def receive_message(self, assistant_message):
         if self.ai_api:
             recv_callback = self.ai_api.chat_event_callbacks.get("receive", None)
             if recv_callback:
-                recv_callback(assistant_message)
+                recv_callback(assistant_message, model=self.model)
 
         # Append to chat history
         self.chat_history.append({"role": "user", "content": assistant_message})