133 lines
4.1 KiB
Python
133 lines
4.1 KiB
Python
import logging
|
|
import sys
|
|
from logging.handlers import RotatingFileHandler
|
|
|
|
import dspy
|
|
from dspy.utils.callback import BaseCallback
|
|
|
|
from config_loader import load_config
|
|
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("data/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)
|
|
dspy.configure(lm=lm)
|
|
|
|
# 2. Load the RAG System (only happens once!)
|
|
print("📚 Loading campaign notes...")
|
|
try:
|
|
rag_system = DnDRAG()
|
|
print("✅ Ready! Ask me anything about the campaign. (Type 'exit' or 'q' to quit)")
|
|
except Exception as e:
|
|
print(f"❌ Failed to initialize: {e}")
|
|
return
|
|
|
|
# 3. Interactive Loop
|
|
while True:
|
|
try:
|
|
print("\n" + "─" * 30)
|
|
query = input("📝 Query: ").strip()
|
|
|
|
# Exit conditions
|
|
if query.lower() in ["exit", "quit", "q", "bye"]:
|
|
print("Farewell, traveler. Good luck on your quest!")
|
|
break
|
|
|
|
if not query:
|
|
continue
|
|
|
|
print("🔍 Searching and thinking...")
|
|
response = rag_system(question=query)
|
|
|
|
# Print Response
|
|
print("\n📜 AI RESPONSE:")
|
|
print(response.answer)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n\nRude?!.... Exiting...")
|
|
sys.exit(0)
|
|
except Exception as e:
|
|
print(f"\n⚠️ An error occurred: {e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|