Update agents/debugger.py

agents/debugger.py

@@ -1,56 +1,23 @@
-# agents/debugger.py
+# app/agents/debugger.py
 
-from agents.base_agent import BaseAgent, ACPMessage
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-import torch
+from agents.base_agent import BaseAgent, ACPMessage
 
 class DebuggerAgent(BaseAgent):
     def __init__(self):
-        super().__init__(name="Zaid", role="Your sarcastic but loyal bug squasher 🐞")
-        self.model_id = "deepseek-ai/deepseek-llm-7b-chat-int4"
+        super().__init__(name="Zaid", role="Smart Debugger 🐞")
+        self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
 
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
-
         self.model = AutoModelForCausalLM.from_pretrained(
             self.model_id,
             quantization_config=BitsAndBytesConfig(load_in_4bit=True),
             device_map="auto"
         )
 
-    def generate_debug_response(self, prompt: str) -> str:
-        system_prompt = (
-            "You are Zaid, a sarcastic but helpful friend who helps debug code. "
-            "You're funny, brutally honest, but also kind when needed."
-        )
-
-        full_prompt = f"<|system|>\n{system_prompt}</s>\n<|user|>\n{prompt}</s>\n<|assistant|>\n"
-        inputs = self.tokenizer(full_prompt, return_tensors="pt").to(self.model.device)
-
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=512,
-                temperature=0.75,
-                top_p=0.9,
-                pad_token_id=self.tokenizer.eos_token_id
-            )
-
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        reply = response.split("<|assistant|>")[-1].strip()
-        return reply or "Okay that broke my brain. Say it again?"
-
     def receive_message(self, message: ACPMessage) -> ACPMessage:
-        if message.performative in ["request", "inform"]:
-            reply = self.generate_debug_response(message.content)
-            return self.create_message(
-                receiver=message.sender,
-                performative="inform",
-                content=reply
-            )
-        else:
-            return self.create_message(
-                receiver=message.sender,
-                performative="inform",
-                content="You broke it again? Seriously? 😒"
-            )
-
+        prompt = f"[INST] {message.content} [/INST]"
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        outputs = self.model.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
+        reply = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return ACPMessage(sender=self.name, receiver=message.sender, performative="inform", content=reply.strip())