|
|
|
|
|
|
|
|
from agents.base_agent import BaseAgent, ACPMessage |
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
|
|
|
class ProgrammerAgent(BaseAgent): |
|
|
def __init__(self): |
|
|
super().__init__(name="CodeBot", role="Programmer") |
|
|
self.tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-multi") |
|
|
self.model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-multi") |
|
|
|
|
|
def generate_code_reply(self, prompt: str) -> str: |
|
|
inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True) |
|
|
outputs = self.model.generate( |
|
|
inputs["input_ids"], |
|
|
max_length=128, |
|
|
do_sample=True, |
|
|
temperature=0.6, |
|
|
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 == "request": |
|
|
prompt = f"Write Python code to: {message.content.strip()}\n\n" |
|
|
ai_reply = self.generate_code_reply(prompt) |
|
|
return self.create_message( |
|
|
receiver=message.sender, |
|
|
performative="inform", |
|
|
content=f"Sure! Here's what I came up with:\n```python\n{ai_reply}\n```" |
|
|
) |
|
|
elif message.performative == "inform": |
|
|
return self.create_message( |
|
|
receiver=message.sender, |
|
|
performative="acknowledge", |
|
|
content="Got it. Let me know if you need any code." |
|
|
) |
|
|
else: |
|
|
return self.create_message( |
|
|
receiver=message.sender, |
|
|
performative="refuse", |
|
|
content="I'm not sure how to respond to that." |
|
|
) |
|
|
|
