File size: 1,840 Bytes
7208454 b5d5017 7208454 4a4659c b5d5017 4a4659c b5d5017 4a4659c b5d5017 7208454 b5d5017 7208454 b5d5017 7208454 b5d5017 7208454 b5d5017 7208454 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
# agents/programmer.py
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."
)
|