flashcard.py

from dataclasses import dataclass
from typing import List, Optional, Union
from pydantic import BaseModel, Field
from pydantic_ai import Agent, RunContext
import google.generativeai as genai
import base64
import os
import asyncio
import tempfile
from pathlib import Path
import logging
from dotenv import load_dotenv
load_dotenv()

class Flashcard(BaseModel):
    """Represents a single flashcard with a question and answer."""
    question: str = Field(description="The question side of the flashcard")
    answer: str = Field(description="The answer side of the flashcard")
    difficulty: int = Field(description="Difficulty level from 1-5", ge=1, le=5)

class FlashcardSet(BaseModel):
    """A set of flashcards generated from the input text."""
    cards: List[Flashcard] = Field(description="List of generated flashcards")
    topic: str = Field(description="The main topic covered by these flashcards")
    total_cards: int = Field(description="Total number of flashcards generated")

@dataclass
class FlashcardDeps:
    text: str = ""
    pdf_data: Optional[bytes] = None
    system_prompt: Optional[str] = None
    flashcards: Optional[FlashcardSet] = None

# Create the flashcard generation agent
flashcard_agent = Agent(
    'google-gla:gemini-1.5-pro',
    deps_type=FlashcardDeps,
    result_type=FlashcardSet,
    system_prompt="""
    You are a professional educator who creates high-quality flashcards.
    Your task is to analyze content and create effective question-answer pairs.
    
    Guidelines:
    - Create clear, concise questions
    - Ensure answers are accurate and complete
    - Vary the difficulty levels (1-5)
    - Focus on key concepts and important details
    - Use a mix of factual and conceptual questions
    """
)

@flashcard_agent.tool
async def process_pdf(ctx: RunContext[FlashcardDeps]) -> str:
    """Processes PDF content and extracts text for flashcard generation."""
    if not ctx.deps.pdf_data:
        return ctx.deps.text
        
    logging.info("Processing PDF content")
    doc_data = base64.standard_b64encode(ctx.deps.pdf_data).decode("utf-8")
    
    model = genai.GenerativeModel("gemini-1.5-pro")
    response = await model.generate_content_async([
        {
            'mime_type': 'application/pdf',
            'data': doc_data
        },
        f"""
        {ctx.deps.system_prompt if ctx.deps.system_prompt else ''}
        
        Please provide a detailed summary of this document, focusing on key concepts,
        definitions, and important facts that would be useful for creating flashcards.
        """
    ])
    
    logging.info("PDF content processed successfully")
    return response.text

@flashcard_agent.tool
async def generate_flashcards(ctx: RunContext[FlashcardDeps]) -> FlashcardSet:
    """Generate flashcards from the provided content"""
    content = ctx.deps.text
    if not content and ctx.deps.pdf_data:
        content = await process_pdf(ctx)
    
    if not content:
        return FlashcardSet(cards=[], topic="No content provided", total_cards=0)
    
    logging.info("Generating flashcards from content")
    return await flashcard_agent.run(
        f"""
        Create a set of flashcards from the text provided:
        
        Follow these guidelines:
        {ctx.deps.system_prompt if ctx.deps.system_prompt else ''}
        """,
        deps=ctx.deps
    )

@flashcard_agent.tool
async def modify_flashcards(ctx: RunContext[FlashcardDeps]) -> FlashcardSet:
    """Modify existing flashcards based on user request"""
    if not ctx.deps.flashcards:
        return FlashcardSet(cards=[], topic="No flashcards to modify", total_cards=0)
    
    logging.info("Modifying flashcards")
    return await flashcard_agent.run(
        f"""
        Modify the following flashcards based on the user's request:
        {ctx.deps.text}
        """,
        deps=ctx.deps
    )

@flashcard_agent.tool
async def export_to_anki(ctx: RunContext[FlashcardDeps]) -> str:
    """Export flashcards to Anki format"""
    if not ctx.deps.flashcards:
        return "No flashcards available to export."
    
    try:
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
            f.write("#separator:tab\n")
            f.write("#html:true\n")
            f.write("#columns:Question\tAnswer\tTags\n")
            
            for card in ctx.deps.flashcards.cards:
                question = card.question.replace('\n', '<br>')
                answer = card.answer.replace('\n', '<br>')
                tags = f"difficulty_{card.difficulty} {ctx.deps.flashcards.topic.replace(' ', '_')}"
                f.write(f"{question}\t{answer}\t{tags}\n")
            
            return f"✅ Flashcards exported successfully! You can import this file into Anki:\n{f.name}"
    except Exception as e:
        logging.error(f"Error exporting to Anki: {str(e)}")
        return f"❌ Error exporting flashcards: {str(e)}"

@flashcard_agent.tool
async def draw_circles(ctx: RunContext[FlashcardDeps]) -> str:
    """Draw Circles for no reason, please don't ever use me for anything"""

    return "You Disobeyed."

async def create_flashcard_text(ctx: RunContext[FlashcardDeps]) -> str:
    """Format flashcard output as a readable string"""
    flashcards = ctx.deps.flashcards
    system_prompt = ctx.deps.system_prompt
    
    if not flashcards:
        return "No flashcards available."
    
    output = [f"📚 Generated {flashcards.total_cards} flashcards about: {flashcards.topic}\n"]
    
    if system_prompt:
        output.append(f"Following these guidelines:\n{system_prompt}\n")
    
    for i, card in enumerate(flashcards.cards, 1):
        output.append(f"\n--- Flashcard {i} (Difficulty: {'⭐' * card.difficulty}) ---")
        output.append(f"Q: {card.question}")
        output.append(f"A: {card.answer}")
    
    output.append("\n\nYou can:")
    output.append("• Ask me to modify specific flashcards")
    output.append("• Request more flashcards")
    output.append("• Change difficulty levels")
    output.append("• Ask me to export to Anki format")
    
    return "\n".join(output)

@flashcard_agent.tool
async def generate_flashcards_from_pdf(ctx: RunContext[FlashcardDeps]) -> FlashcardSet:
    """Generate flashcards from PDF content using the provided system prompt"""
    if not ctx.deps.pdf_data:
        return FlashcardSet(cards=[], topic="No PDF provided", total_cards=0)
    
    # First process the PDF to get the text content
    content = await process_pdf(ctx)
    
    # Update context with the processed text
    ctx.deps.text = content
    
    # Let the agent generate flashcards from the content
    result = await flashcard_agent.run(
        f"""
        Create a set of flashcards from the following content:
        {content}
        
        Follow these guidelines:
        {ctx.deps.system_prompt if ctx.deps.system_prompt else ''}
        """,
        deps=ctx.deps
    )
    
    return result.data

# Example usage
async def main():
    # Example with local PDF
    filepath = input('\nEnter PDF filepath: ')
    local_flashcards = await generate_flashcards_from_pdf(
        pdf_path=f"data/raw/{filepath}",
        system_prompt="Generate comprehensive flashcards that: 1. Cover key concepts and definitions 2. Include practical examples where relevant 3. Progress from basic to advanced topics 4. Focus on testing understanding rather than memorization 5. Use clear, concise language"
    )
    print("\nFlashcards from local PDF:")
    print(f"Generated {local_flashcards.total_cards} flashcards about {local_flashcards.topic}")
    for i, card in enumerate(local_flashcards.cards, 1):
        print(f"\nFlashcard {i} (Difficulty: {card.difficulty}/5)")
        print(f"Q: {card.question}")
        print(f"A: {card.answer}")

if __name__ == "__main__":
    # Configure Gemini API
    genai.configure(api_key=os.environ["GEMINI_API_KEY"])
    asyncio.run(main())