# agents/coder.py
from transformers import AutoTokenizer, AutoModelForCausalLM

class CoderAgent:
    def __init__(self, model_name="deepseek-ai/deepseek-coder-6.7b-instruct"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype="auto",
            device_map="auto"
        )
    
    def generate(self, prompt):
        full_prompt = f"""
You're a brilliant, friendly AI coder. Explain clearly and write clean Python.
Include comments and use best practices.

Task: {prompt}
        """.strip()
        inputs = self.tokenizer(full_prompt, return_tensors="pt").to(self.model.device)
        outputs = self.model.generate(**inputs, max_new_tokens=1024, temperature=0.4)
        code = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return code[len(full_prompt):].strip()