Add callbacks to LLM chat

dailalib/__init__.py

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

dailalib/api/litellm/litellm_api.py

@@ -31,6 +31,7 @@ def __init__(
         prompts: Optional[list] = None,
         fit_to_tokens: bool = False,
         chat_use_ctx: bool = True,
+        chat_event_callbacks: Optional[dict] = None,
         **kwargs
     ):
         super().__init__(**kwargs)
@@ -40,6 +41,7 @@ def __init__(
         self.model = model
         self.fit_to_tokens = fit_to_tokens
         self.chat_use_ctx = chat_use_ctx
+        self.chat_event_callbacks = chat_event_callbacks or {"send": None, "receive": None}
 
         # delay prompt import
         from .prompts import PROMPTS, DEFAULT_STYLE

dailalib/llm_chat/llm_chat_ui.py

@@ -68,7 +68,6 @@ def __init__(self, ai_api: "LiteLLMAIAPI", parent=None, context: Context = None)
         # create a context for this first message
         if ai_api.chat_use_ctx:
             ai_api.info("Collecting context for the current function...")
-            #import remote_pdb; remote_pdb.RemotePdb('localhost', 4444).set_trace()
             if context is None:
                 context = ai_api._dec_interface.gui_active_context()
             dec = ai_api._dec_interface.decompile(context.func_addr)
@@ -124,6 +123,12 @@ def send_message(self, add_text=True, role="user"):
         if not user_text:
             return
 
+        # do aiapi calback
+        if self.ai_api:
+            send_callback = self.ai_api.chat_event_callbacks.get("send", None)
+            if send_callback:
+                send_callback(user_text)
+
         # Display user message
         if add_text:
             self.add_message(user_text, is_user=True)
@@ -145,6 +150,12 @@ def receive_message(self, assistant_message):
         # Display assistant message
         self.add_message(assistant_message, is_user=False)
 
+        # do aiapi calback
+        if self.ai_api:
+            recv_callback = self.ai_api.chat_event_callbacks.get("receive", None)
+            if recv_callback:
+                recv_callback(assistant_message)
+
         # Append to chat history
         self.chat_history.append({"role": "user", "content": assistant_message})
 

setup.cfg

@@ -17,7 +17,7 @@ install_requires =
     litellm>=1.44.27
     tiktoken
     Jinja2
-    libbs>=2.0.2
+    libbs>=2.1.0
 
 python_requires = >= 3.10
 include_package_data = True