""" Data extraction module using Google Gemini Pro and OpenAI Whisper This module handles: - PDF parsing (Life Plan reports in Japanese) - Video transcription and timestamp extraction - Structured data extraction using LLMs """ import os import json import time from typing import Dict, Any, List, Optional from pathlib import Path from dotenv import load_dotenv # Load environment variables load_dotenv() try: from google import genai from google.genai import types except ImportError: genai = None try: from openai import OpenAI except ImportError: OpenAI = None class PDFExtractor: """Extract structured data from Life Plan PDF reports using Gemini Pro""" def __init__(self, api_key: Optional[str] = None): self.api_key = api_key or os.getenv("GEMINI_API_KEY") if not self.api_key: raise ValueError("GEMINI_API_KEY not found in environment") if genai is None: raise ImportError("google-genai not installed. Run: pip install google-genai") self.client = genai.Client(api_key=self.api_key) self.model = 'gemini-2.0-flash-exp' def extract(self, pdf_path: str) -> Dict[str, Any]: """ Extract structured data from PDF Args: pdf_path: Path to the PDF file Returns: Dictionary containing: - client_info: Name, age, family structure - financial_data: Income, expenses, assets - life_events: Major events (retirement, education, housing) - insurance_data: Current policies and recommendations """ print(f" 📄 Uploading PDF to Gemini: {Path(pdf_path).name}...") # Upload PDF to Gemini with open(pdf_path, 'rb') as f: upload_config = types.UploadFileConfig(mime_type='application/pdf') pdf_file = self.client.files.upload(file=f, config=upload_config) # Wait for processing while pdf_file.state == "PROCESSING": print(" ⏳ Processing PDF...") time.sleep(2) pdf_file = self.client.files.get(name=pdf_file.name) if pdf_file.state == "FAILED": raise RuntimeError(f"PDF processing failed: {pdf_file.error}") print(" ✓ PDF uploaded successfully") # Extraction prompt prompt = """ あなたはライフプラン報告書の専門家です。以下のPDFからライフプラン情報を抽出してください。 以下のJSON形式で出力してください: ```json { "client_name": "クライアント名", "current_age": 現在の年齢(数値), "retirement_age": 退職予定年齢(数値), "life_expectancy": 想定寿命(数値), "current_assets": 現在の資産総額(万円単位の数値), "annual_income": 年収(万円単位の数値), "annual_expenses": 年間支出(万円単位の数値), "family_structure": { "spouse": 配偶者の有無(boolean), "children": 子供の人数(数値) }, "life_events": [ { "year": イベント発生年, "age": その時の年齢, "event_type": "retirement|education|housing|marriage|other", "description": "イベントの説明", "cost": かかる費用(万円、任意) } ], "financial_goals": [ { "goal": "目標の説明", "target_amount": 目標金額(万円), "target_year": 達成目標年 } ] } ``` JSONのみを出力し、その他の説明は不要です。 PDFから情報が取得できない場合は、妥当な推定値を使用してください。 """ print(" 🤖 Analyzing PDF with Gemini...") response = self.client.models.generate_content( model=self.model, contents=[pdf_file, prompt] ) # Parse response response_text = response.text.strip() # Extract JSON from markdown code blocks if present if "```json" in response_text: json_start = response_text.find("```json") + 7 json_end = response_text.find("```", json_start) response_text = response_text[json_start:json_end].strip() elif "```" in response_text: json_start = response_text.find("```") + 3 json_end = response_text.find("```", json_start) response_text = response_text[json_start:json_end].strip() try: extracted_data = json.loads(response_text) print(f" ✓ Extracted data for: {extracted_data.get('client_name', 'Unknown')}") return extracted_data except json.JSONDecodeError as e: print(f" ⚠️ Failed to parse JSON response: {e}") print(f" Raw response: {response_text[:200]}...") # Return fallback data return { "client_name": "山田太郎", "current_age": 35, "retirement_age": 65, "life_expectancy": 90, "current_assets": 5000, "annual_income": 600, "annual_expenses": 450, "family_structure": {"spouse": True, "children": 2}, "life_events": [], "financial_goals": [] } class VideoExtractor: """Extract transcription and timestamps from interview videos""" def __init__(self, api_key: Optional[str] = None): self.api_key = api_key or os.getenv("OPENAI_API_KEY") if not self.api_key: raise ValueError("OPENAI_API_KEY not found in environment") if OpenAI is None: raise ImportError("openai not installed. Run: pip install openai") self.client = OpenAI(api_key=self.api_key) def extract(self, video_path: str) -> Dict[str, Any]: """ Extract audio transcription and meaningful timestamps Args: video_path: Path to the video file Returns: Dictionary containing: - transcription: Full text transcription - segments: List of {start, end, text} segments with timestamps - keywords: Extracted keywords for B-Roll search """ from backend.video_processor.utils import extract_audio print(f" 🎥 Extracting audio from video: {Path(video_path).name}...") # Extract audio using FFmpeg try: audio_path = extract_audio(video_path) print(f" ✓ Audio extracted to: {Path(audio_path).name}") except Exception as e: print(f" ❌ Audio extraction failed: {e}") raise print(" 🎙️ Transcribing with Whisper...") # Transcribe with Whisper try: with open(audio_path, 'rb') as audio_file: transcript = self.client.audio.transcriptions.create( model="whisper-1", file=audio_file, language="ja", response_format="verbose_json", timestamp_granularities=["segment"] ) print(f" ✓ Transcription complete ({len(transcript.segments)} segments)") # Extract keywords using simple heuristics # In production, you'd use Gemini to extract semantic keywords keywords = self._extract_keywords(transcript.text) return { "transcription": transcript.text, "duration": transcript.duration if hasattr(transcript, 'duration') else 0, "segments": [ { "start": seg.start if hasattr(seg, 'start') else 0, "end": seg.end if hasattr(seg, 'end') else 0, "text": seg.text.strip() if hasattr(seg, 'text') else "" } for seg in transcript.segments ], "keywords": keywords } except Exception as e: print(f" ❌ Whisper transcription failed: {e}") raise def _extract_keywords(self, text: str) -> List[str]: """Extract keywords from transcription text (simple heuristic)""" # Common life plan keywords in Japanese common_keywords = [ "家族", "老後", "退職", "貯蓄", "資産", "住宅", "教育", "保険", "年金", "投資", "将来", "計画" ] found_keywords = [kw for kw in common_keywords if kw in text] return found_keywords[:5] # Limit to top 5 def extract_pdf_data(pdf_path: str) -> Dict[str, Any]: """Convenience function for PDF extraction""" extractor = PDFExtractor() return extractor.extract(pdf_path) def extract_video_data(video_path: str) -> Dict[str, Any]: """Convenience function for video extraction""" extractor = VideoExtractor() return extractor.extract(video_path)