Spaces:
Sleeping
Sleeping
File size: 4,885 Bytes
ee0877c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
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
) |