| |
|
|
| 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." |
| ) |
|
|
