259 lines
8.9 KiB
Python
259 lines
8.9 KiB
Python
"""
|
|
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)
|