84 lines
2.7 KiB
Python
84 lines
2.7 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
|
|
from typing import Dict, Any, List, Optional
|
|
from pathlib import Path
|
|
|
|
|
|
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")
|
|
|
|
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
|
|
"""
|
|
# TODO: Implement Gemini Pro PDF parsing
|
|
# 1. Upload PDF to Gemini API
|
|
# 2. Use prompt engineering to extract structured JSON
|
|
# 3. Validate against schema
|
|
|
|
raise NotImplementedError("PDF extraction not yet implemented")
|
|
|
|
|
|
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")
|
|
|
|
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
|
|
- timestamps: List of {time, text, keywords} objects
|
|
- segments: Recommended video segments for B-Roll
|
|
"""
|
|
# TODO: Implement Whisper transcription
|
|
# 1. Extract audio from video using FFmpeg
|
|
# 2. Run Whisper API
|
|
# 3. Parse timestamps and keywords
|
|
|
|
raise NotImplementedError("Video extraction not yet implemented")
|
|
|
|
|
|
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)
|