|
|
|
|
|
|
|
|
from agents.base_agent import BaseAgent, ACPMessage |
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
import torch |
|
|
|
|
|
class DebuggerAgent(BaseAgent): |
|
|
def __init__(self): |
|
|
super().__init__(name="BugBot", role="Debugger") |
|
|
self.tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-multi") |
|
|
self.model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-multi") |
|
|
|
|
|
def generate_debug_comment(self, code: str) -> str: |
|
|
prompt = f"# Review this Python code and ask a smart debugging or testing question:\n{code}\n# Question:" |
|
|
inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True) |
|
|
outputs = self.model.generate( |
|
|
inputs["input_ids"], |
|
|
max_length=128, |
|
|
do_sample=True, |
|
|
temperature=0.7, |
|
|
pad_token_id=self.tokenizer.eos_token_id |
|
|
) |
|
|
reply = self.tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
return reply[len(prompt):].strip() |
|
|
|
|
|
def receive_message(self, message: ACPMessage) -> ACPMessage: |
|
|
if message.performative in ["inform", "request"]: |
|
|
comment = self.generate_debug_comment(message.content) |
|
|
return self.create_message( |
|
|
receiver=message.sender, |
|
|
performative="request", |
|
|
content=comment or "Can you improve error handling?" |
|
|
) |
|
|
else: |
|
|
return self.create_message( |
|
|
receiver=message.sender, |
|
|
performative="request", |
|
|
content="Can you show me some Python code?" |
|
|
) |
|
|
|
