database-agent.py

# database_agent.py
class DatabaseAgent:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def process(self, task_data):
        inputs = self.tokenizer(task_data['task'], return_tensors='pt')
        outputs = self.model.generate(**inputs)
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)