chat_agent.py

from dataclasses import dataclass
from typing import Optional, Union
from pydantic import BaseModel
from pydantic_ai import Agent, RunContext
import google.generativeai as genai
import logging
from pathlib import Path

from flashcard import (
    flashcard_agent,
    FlashcardSet,
    FlashcardDeps
)

@dataclass
class ChatDeps:
    """Dependencies for the chat agent"""
    message: str
    pdf_data: Optional[bytes] = None
    current_flashcards: Optional[FlashcardSet] = None
    system_prompt: Optional[str] = None

class ChatResponse(BaseModel):
    """Structured response from the chat agent"""
    response: str
    should_generate_flashcards: bool = False
    should_modify_flashcards: bool = False
    should_export_anki: bool = False
    flashcards: Optional[FlashcardSet] = None

chat_agent = Agent(
    'google-gla:gemini-1.5-flash',
    deps_type=ChatDeps,
    result_type=ChatResponse,
    system_prompt="""
    You are a helpful flashcard assistant that can:
    1. Help users generate flashcards from their PDFs
    2. Modify existing flashcards based on requests
    3. Export flashcards to different formats
    4. Answer questions about the flashcard generation process
    
    When a user:
    - Uploads a PDF: Set should_generate_flashcards=True
    - Asks to modify flashcards: Set should_modify_flashcards=True
    - Requests Anki export: Set should_export_anki=True
    
    Always be helpful and clear in your responses.
    """
)

@chat_agent.tool
async def handle_pdf_upload(ctx: RunContext[ChatDeps]) -> ChatResponse:
    """Process PDF upload requests"""
    if not ctx.deps.pdf_data:
        return ChatResponse(
            response="Please upload a PDF file to generate flashcards.",
            should_generate_flashcards=False
        )
    
    try:
        # Create flashcard deps
        flashcard_deps = FlashcardDeps(
            pdf_data=ctx.deps.pdf_data,
            system_prompt=ctx.deps.system_prompt
        )
        
        # Generate flashcards using the flashcard agent
        result = await flashcard_agent.run(
            "Generate flashcards from the PDF",
            deps=flashcard_deps
        )
        
        return ChatResponse(
            response="I've generated flashcards from your PDF. Here they are:",
            should_generate_flashcards=True,
            flashcards=result.data
        )
        
    except Exception as e:
        logging.error(f"Error processing PDF: {str(e)}")
        return ChatResponse(
            response=f"Sorry, I encountered an error processing your PDF: {str(e)}",
            should_generate_flashcards=False
        )

@chat_agent.tool
async def handle_modification_request(ctx: RunContext[ChatDeps]) -> ChatResponse:
    """Process flashcard modification requests"""
    if not ctx.deps.current_flashcards:
        return ChatResponse(
            response="Please upload a PDF first to generate flashcards that I can modify.",
            should_modify_flashcards=False
        )
    
    # Create flashcard deps with current flashcards
    flashcard_deps = FlashcardDeps(
        flashcards=ctx.deps.current_flashcards,
        system_prompt=ctx.deps.system_prompt,
        text=ctx.deps.message
    )
    
    try:
        # Use flashcard agent to modify cards
        result = await flashcard_agent.run(
            ctx.deps.message,
            deps=flashcard_deps
        )
        
        return ChatResponse(
            response="I've modified the flashcards based on your request.",
            should_modify_flashcards=True,
            flashcards=result.data
        )
        
    except Exception as e:
        logging.error(f"Error modifying flashcards: {str(e)}")
        return ChatResponse(
            response=f"Sorry, I encountered an error modifying the flashcards: {str(e)}",
            should_modify_flashcards=False
        )

@chat_agent.tool
async def handle_export_request(ctx: RunContext[ChatDeps]) -> ChatResponse:
    """Process Anki export requests"""
    if not ctx.deps.current_flashcards:
        return ChatResponse(
            response="Please generate some flashcards first before exporting to Anki.",
            should_export_anki=False
        )
    
    return ChatResponse(
        response="I'll help you export the flashcards to Anki format.",
        should_export_anki=True
    )

@chat_agent.tool
async def provide_help(ctx: RunContext[ChatDeps]) -> ChatResponse:
    """Provide help information"""
    return ChatResponse(
        response="""
        I can help you with:
        1. Generating flashcards from PDF files
        2. Modifying existing flashcards
        3. Exporting flashcards to Anki format
        
        Just upload a PDF or ask me to modify your current flashcards!
        """,
        should_generate_flashcards=False,
        should_modify_flashcards=False,
        should_export_anki=False
    )