import os
from huggingface_hub import InferenceClient

class CoderAgent:
    def __init__(self, hf_token):
        self.hf_token = hf_token
        # 使用适合免费方案的模型
        self.client = InferenceClient(token=hf_token, model="google/gemma-2b-it")
        
    def generate_code(self, task_description):
        """生成代码"""
        system_message = "你是编程专家，负责代码生成、审查和优化。请生成高质量、可运行的代码。"
        
        messages = [
            {"role": "system", "content": system_message},
            {"role": "user", "content": f"请为以下任务生成Python代码:\n\n{task_description}\n\n请提供完整的代码实现，包括必要的导入语句和注释。"}
        ]
        
        response = ""
        # 使用适合免费方案的参数
        for message in self.client.chat_completion(
            messages,
            max_tokens=512,  # 降低max_tokens以适应免费方案限制
            stream=True,
            temperature=0.7,
            top_p=0.9,
        ):
            choices = message.choices
            if len(choices) and choices[0].delta.content:
                response += choices[0].delta.content
                
        return response

# 使用示例
if __name__ == "__main__":
    # 从环境变量获取HF_TOKEN
    hf_token = os.environ.get("HF_TOKEN", "")
    if not hf_token:
        print("错误：请设置HF_TOKEN环境变量")
        exit(1)
        
    coder = CoderAgent(hf_token)
    task = "创建一个简单的Flask Web应用，包含一个返回'Hello World'的路由"
    
    result = coder.generate_code(task)
    print(result)