feat: ✨ AI Read File Tool, Configurable system prompts and loading lots of llms

src/retrieve.py

@@ -1,16 +1,93 @@
 import sys
 import dspy
-# import turso
+import logging
+from dspy.utils.callback import BaseCallback
+
+from logging.handlers import RotatingFileHandler
 
 from config_loader import load_config
-from experts.dnd_agent import DnDRAG
+from experts.retrieval_agent import DnDRAG
 
 CFG = load_config()
 RETRIEVE_MODEL = CFG["models"]["retrieval"]
 API_BASE = CFG["api"]["base_url"]
 API_VERSION = CFG["api"]["api_version"]
 
+class CallbackHandler(BaseCallback):
+    """Custom callback class for logging agent interactions."""
+
+    def __init__(self, logger):
+        """Initialize the callback with a logger instance."""
+        super().__init__()
+        self.logger = logger
+
+    def on_module_end(self, call_id, outputs, exception):
+        """Handle module end events for logging."""
+        step = "Reasoning" if self._is_reasoning_output(outputs) else "Acting"
+        self.logger.debug(f"== {step} Step ===")
+        for k, v in outputs.items():
+            self.logger.debug(f"  {k}: {v}")
+
+    def on_lm_start(self, call_id, instance, inputs):
+        """Handle language model start events for logging."""
+        self.logger.debug(f"LM is called with inputs: {inputs}")
+
+    def on_tool_start(self, call_id, instance, inputs):
+        """Handle tool start events for logging."""
+        self.logger.debug(f"Tool {instance} called with inputs: {inputs}")
+
+    def on_tool_end(self, call_id, outputs, exception):
+        """Handle tool end events for logging."""
+        self.logger.debug(f"Tool finished with outputs: {outputs}")
+
+    def on_lm_end(self, call_id, outputs, exception):
+        """Handle language model end events for logging."""
+        self.logger.debug(f"LM is finished with outputs: {outputs}")
+
+    def _is_reasoning_output(self, outputs):
+        return any(k.startswith("Thought") for k in outputs)
+
+def setup_logging():
+    """Set up logging configuration for Merlin."""
+    # Create a custom logger
+    logger = logging.getLogger(__name__)
+
+    # Set the minimum level for the logger
+    logger.setLevel(logging.DEBUG)
+
+    # Create a console handler
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.INFO)
+
+    # Create a file handler with rotation every 5MB
+    file_handler = RotatingFileHandler(
+        "dmv.log", maxBytes=5 * 1024 * 1024, backupCount=3
+    )
+    file_handler.setLevel(logging.DEBUG)
+
+    # Create a formatter
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+
+    # Set the formatter for the handler
+    console_handler.setFormatter(formatter)
+    file_handler.setFormatter(formatter)
+
+    # Add the handler to the logger
+    logger.addHandler(console_handler)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
 def main():
+    logger = setup_logging()
+    logger.debug("main application started")
+
+    # Add verbose callback
+    dspy.configure(verbose_errors=True)
+    dspy.configure(callbacks=[CallbackHandler(logger)])
     # 1. Setup the LLM
     print("🚀 Initializing Qwen-8B via LM Studio...")
     lm = dspy.LM(RETRIEVE_MODEL, api_base=API_BASE + API_VERSION)
@@ -32,7 +109,7 @@ def main():
             query = input("📝 Query: ").strip()
 
             # Exit conditions
-            if query.lower() in ["exit", "quit", "q"]:
+            if query.lower() in ["exit", "quit", "q", "bye"]:
                 print("Farewell, traveler. Good luck on your quest!")
                 break
 
@@ -47,11 +124,10 @@ def main():
             print(response.answer)
 
         except KeyboardInterrupt:
-            print("\n\nExiting... See you next session!")
+            print("\n\nRude?!.... Exiting...")
             sys.exit(0)
         except Exception as e:
             print(f"\n⚠️ An error occurred: {e}")
 
-
 if __name__ == "__main__":
-    main()
+    main()