Update agents/programmer.py

agents/programmer.py

@@ -1,4 +1,3 @@
-# app/agents/programmer.py
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from agents.base_agent import BaseAgent, ACPMessage
 
@@ -6,25 +5,32 @@ class ProgrammerAgent(BaseAgent):
     def __init__(self):
         super().__init__(name="Aymaan", role="Helpful Coder 🤖")
         self.model_id = "microsoft/phi-1_5"
+
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
         self.model = AutoModelForCausalLM.from_pretrained(self.model_id)
 
     def receive_message(self, message: ACPMessage) -> ACPMessage:
-        prompt_lower = message.content.lower().strip()
-
-        # Short greeting handling
-        greetings = {
-            "hi": "Hey there! 😊",
-            "hello": "Hello! 👋",
-            "how are you?": "I'm doing great! How about you?",
-            "hey": "Hey! Need any help?"
-        }
-        if prompt_lower in greetings:
-            return ACPMessage(sender=self.name, receiver=message.sender, performative="inform", content=greetings[prompt_lower])
+        # Greeting shortcut
+        if self.is_greeting(message.content):
+            return ACPMessage(
+                sender=self.name,
+                receiver=message.sender,
+                performative="inform",
+                content="Hey there! I'm Aymaan, your coding buddy. 😊"
+            )
 
-        # Normal flow
         prompt = f"{message.content}"
-        inputs = self.tokenizer(prompt, return_tensors="pt")
-        outputs = self.model.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=0.7)
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        outputs = self.model.generate(
+            **inputs,
+            max_new_tokens=100,
+            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())
+        return ACPMessage(
+            sender=self.name,
+            receiver=message.sender,
+            performative="inform",
+            content=reply.strip()
+        )