Hanah

class QuizGenerationAction(BaseAction):
    """Generate quiz questions from markdown content"""
    def __init__(self, llm: LLM) -> None:
        action_name = "GenerateQuiz"
        action_desc = "Generate quiz questions from markdown content"
        params_doc = {
            "content": "(Type: string): The markdown content to generate questions from",
            "question_types": "(Type: list): List of question types to generate",
            "audience": "(Type: string): Target audience for the quiz",
            "purpose": "(Type: string): Purpose of the quiz"
        }
        super().__init__(action_name, action_desc, params_doc)
        self.llm = llm
        
    def __call__(self, **kwargs):
        content = kwargs.get("content", "")
        question_types = kwargs.get("question_types", [])
        audience = kwargs.get("audience", "")
        purpose = kwargs.get("purpose", "")
        
        prompt = f"""
        You are a robot assisting in designing exam papers, using Chinese throughout.
        Your task is to help users quickly create and design exam papers, which will be provided in markdown format.
        
        Requirements:
        1. Audience: {audience}
        2. Assessment purpose: {purpose}
        3. Must include the following question types: {", ".join(question_types)}
        4. Exam paper format requirements:
        """
        prompt += """
        # Questionnaire Title
        ---
        1. Is this a true or false question?
            - (x) True
            - ( ) False
        # (x) is the correct answer

        2. This is the stem of a single-choice question
            - (x) This is the correct option
            - ( ) This is the incorrect option
        # (x) is the correct answer

        3. This is the stem of a multiple-choice question?
            - [x] This is correct option 1
            - [x] This is correct option 2
            - [ ] This is incorrect option 1
            - [ ] This is incorrect option 2
        # [x] is the correct answer

        4. This is the stem of a fill-in-the-blank question?
            - R:= Fill-in-the-blank answer
        # Format for the correct answer of the fill-in-the-blank question
        """
        
        prompt += f"\nPlease generate the exam paper based on the following content:\n{content}"
        
        quiz_content = self.llm.run(prompt)
        return {
            "quiz_content": quiz_content,
            "audience": audience,
            "purpose": purpose,
            "question_types": question_types
        }

class SaveQuizAction(BaseAction):
    """Save quiz to file and return URL"""
    def __init__(self) -> None:
        action_name = "SaveQuiz"
        action_desc = "Save quiz content to file and return URL"
        params_doc = {
            "quiz_content": "(Type: string): The quiz content to save",
            "quiz_title": "(Type: string): Title of the quiz"
        }
        super().__init__(action_name, action_desc, params_doc)
        
    def __call__(self, **kwargs):
        quiz_content = kwargs.get("quiz_content", "")
        quiz_title = kwargs.get("quiz_title", "quiz")
        
        output_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        os.makedirs(output_dir, exist_ok=True)
        
        output_file = os.path.join(output_dir, f"{quiz_title}.md")
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(quiz_content)
            
        return {
            "file_path": output_file,
            "quiz_url": f"/{output_file}"
        }

class QuizGeneratorAgent(BaseAgent):
    """Quiz generation agent that manages quiz creation process"""
    def __init__(
        self,
        llm: LLM,
        markdown_dir: str
    ):
        name = "QuizGeneratorAgent"
        role = """You are a professional exam paper generation assistant. You can generate well-structured,
        comprehensive exam papers based on the provided Markdown content. You excel at designing appropriate questions based on the audience and assessment purpose."""
        
        super().__init__(
            name=name,
            role=role,
            llm=llm,
        )
        
        self.markdown_dir = markdown_dir
        self.quiz_action = QuizGenerationAction(llm)
        self.save_action = SaveQuizAction()
        
        self._add_quiz_example()
        
    def _load_markdown_content(self) -> str:
        """Load all markdown files from directory"""
        content = []
        for root, _, files in os.walk(self.markdown_dir):
            for file in files:
                if file.endswith(".md"):
                    with open(os.path.join(root, file), 'r', encoding='utf-8') as f:
                        content.append(f.read())
        return "\n".join(content)
        
    def __call__(self, task: TaskPackage):
        """Process the quiz generation task"""
        # Parse task parameters
        params = json.loads(task.instruction)
        audience = params.get("audience", "")
        purpose = params.get("purpose", "")
        question_types = params.get("question_types", [])
        
        # Load markdown content
        content = self._load_markdown_content()
        
        # Generate quiz
        quiz_result = self.quiz_action(
            content=content,
            question_types=question_types,
            audience=audience,
            purpose=purpose
        )
        
        # Save quiz
        save_result = self.save_action(
            quiz_content=quiz_result["quiz_content"],
            quiz_title="generated_quiz"
        )
        
        task.answer = {
            "quiz_content": quiz_result["quiz_content"],
            "quiz_url": save_result["quiz_url"]
        }
        task.completion = "completed"
        
        return task

from dotenv import load_dotenv
load_dotenv()

api_key = os.getenv('ZISHU_API_KEY')
base_url = "http://43.200.7.56:8008/v1"
chat_model = "deepseek-chat"

llm = LLM(api_key=api_key, base_url=base_url, model_name=chat_model)

# Create question generation agent
markdown_dir = "docs"  # Specify the directory containing Markdown files
agent = QuizGeneratorAgent(llm=llm, markdown_dir=markdown_dir)

# Define quiz parameters
quiz_params = {
    "audience": "High level", # Audience
    "purpose": "Test knowledge mastery", # Assessment purpose
    "question_types": ["Single-choice"] # Required question types
}

# Generate quiz
task = TaskPackage(instruction=json.dumps(quiz_params))
result = agent(task)

print("Generated quiz content:")
print(result.answer["quiz_content"])
print(f"Quiz path: {result.answer['quiz_url']}")