""" Google NotebookLM integration service. Uses the unofficial `notebooklm-py` library to create notebooks, add feasibility-study data as sources, and generate/download assets. """ import asyncio import json import os import sys from datetime import datetime from flask import current_app if sys.platform == "win32": asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) # Asset type -> generate method name mapping ASSET_GENERATE_MAP = { "audio": "generate_audio", "video": "generate_video", "mind_map": "generate_mind_map", "report": "generate_report", "flashcards": "generate_flashcards", "quiz": "generate_quiz", "infographic": "generate_infographic", "slide_deck": "generate_slide_deck", } # Asset type -> download method + default format ASSET_DOWNLOAD_MAP = { "audio": {"method": "download_audio", "format": "mp3"}, "video": {"method": "download_video", "format": "mp4"}, "mind_map": {"method": "download_mind_map", "format": "json"}, "report": {"method": "download_report", "format": "md"}, "flashcards": {"method": "download_flashcards", "format": "json"}, "quiz": {"method": "download_quiz", "format": "json"}, "infographic": {"method": "download_infographic", "format": "png"}, "slide_deck": {"method": "download_slide_deck", "format": "pdf"}, } # Display info for UI ASSET_DISPLAY = { "audio": {"label_ar": "ملخص صوتي", "label_en": "Audio Overview", "icon": "🎧", "color": "emerald"}, "video": {"label_ar": "ملخص مرئي", "label_en": "Video Overview", "icon": "🎬", "color": "blue"}, "mind_map": {"label_ar": "خريطة ذهنية", "label_en": "Mind Map", "icon": "🧠", "color": "purple"}, "report": {"label_ar": "تقرير", "label_en": "Report", "icon": "📄", "color": "slate"}, "flashcards": {"label_ar": "بطاقات تعليمية", "label_en": "Flashcards", "icon": "📝", "color": "amber"}, "quiz": {"label_ar": "اختبار", "label_en": "Quiz", "icon": "❓", "color": "rose"}, "infographic": {"label_ar": "إنفوجرافيك", "label_en": "Infographic", "icon": "📊", "color": "cyan"}, "slide_deck": {"label_ar": "عرض تقديمي", "label_en": "Slide Deck", "icon": "📑", "color": "orange"}, } def _get_uploads_dir(report_id): """Get or create the uploads directory for a report's notebook assets.""" base = os.path.join(current_app.static_folder, "uploads", "notebooks", str(report_id)) os.makedirs(base, exist_ok=True) return base def _synthesis_to_text(report): """Convert report data into a rich text source for NotebookLM.""" parts = [ f"# دراسة جدوى: {report.project_name}", f"\n## وصف المشروع\n{report.project_description}", ] synthesis = report.synthesis_result if isinstance(synthesis, dict): for key, value in synthesis.items(): title = key.replace("_", " ").title() parts.append(f"\n## {title}") if isinstance(value, dict): for k2, v2 in value.items(): parts.append(f"### {k2.replace('_', ' ').title()}") if isinstance(v2, list): for item in v2: parts.append(f"- {item}") else: parts.append(str(v2)) elif isinstance(value, list): for item in value: if isinstance(item, dict): parts.append(json.dumps(item, ensure_ascii=False, indent=2)) else: parts.append(f"- {item}") else: parts.append(str(value)) if isinstance(report.market_analysis, dict): parts.append("\n## تحليل السوق التفصيلي") parts.append(json.dumps(report.market_analysis, ensure_ascii=False, indent=2)) if isinstance(report.financial_analysis, dict): parts.append("\n## التحليل المالي التفصيلي") parts.append(json.dumps(report.financial_analysis, ensure_ascii=False, indent=2)) if isinstance(report.competitive_analysis, dict): parts.append("\n## التحليل التنافسي التفصيلي") parts.append(json.dumps(report.competitive_analysis, ensure_ascii=False, indent=2)) return "\n".join(parts) def get_user_session_dir(user_id): """Get or create the session directory for a specific user.""" base = os.path.join( current_app.instance_path, "..", "data", "notebooklm_sessions", str(user_id) ) base = os.path.abspath(base) os.makedirs(base, exist_ok=True) return base def get_user_storage_path(user_id): """Get the storage_state.json path for a specific user.""" return os.path.join(get_user_session_dir(user_id), "storage_state.json") def check_session_valid(user_id): """Check if a user has a valid NotebookLM session file.""" path = get_user_storage_path(user_id) if not os.path.exists(path): return False try: with open(path, "r", encoding="utf-8") as f: data = json.load(f) return bool(data.get("cookies")) except Exception: return False def start_google_login_sync(user_id): """ Launch a Playwright browser (in a subprocess) for the user to log into Google. Saves the session cookies to the user's storage_state.json. Works on both Windows and Linux by running Playwright in a separate process. Returns the storage path on success. """ import subprocess, sys storage_path = get_user_storage_path(user_id) worker_script = os.path.join(os.path.dirname(__file__), "_nlm_login_worker.py") result = subprocess.run( [sys.executable, worker_script, storage_path], capture_output=True, text=True, timeout=360, ) if result.returncode == 0 and result.stdout.strip().startswith("OK:"): return storage_path else: err = result.stderr.strip() or result.stdout.strip() or "Unknown error" # Extract clean error message for line in err.splitlines(): if line.startswith("ERROR:"): raise Exception(line[6:]) raise Exception(err) async def _get_client(user_id=None, admin_account_id=None): """Get an authenticated NotebookLM client for a specific user or admin account. NOTE: NOTEBOOKLM_HOME is kept set for the entire lifetime of the client so that download operations can find the storage_state.json. """ from notebooklm import NotebookLMClient # Use extended timeout for all generation tasks - some assets (video, audio) can take # several minutes to initiate the generation process timeout = 300.0 # 5 minutes for RPC calls (was 30s default, too short for complex assets) if admin_account_id: from app.ai.nlm_admin_service import get_admin_account account = get_admin_account(admin_account_id) if account and account.session_path and os.path.exists(account.session_path): # Set NOTEBOOKLM_HOME to the directory containing storage_state.json session_dir = os.path.dirname(account.session_path) os.environ["NOTEBOOKLM_HOME"] = session_dir client = await NotebookLMClient.from_storage(timeout=timeout) return client else: raise Exception(f"Admin account {admin_account_id} has no valid session.") if user_id: storage_path = get_user_storage_path(user_id) if os.path.exists(storage_path): # Set NOTEBOOKLM_HOME and keep it set — the client needs it # for both creation AND subsequent download/generate calls. os.environ["NOTEBOOKLM_HOME"] = get_user_session_dir(user_id) client = await NotebookLMClient.from_storage(timeout=timeout) return client else: raise Exception("User has not connected Google account. Please connect first.") return await NotebookLMClient.from_storage(timeout=timeout) async def create_notebook_for_report(report, user_id=None, admin_account_id=None): """ Create a NotebookLM notebook and add the report data as a text source. Returns (notebook_id, source_id, notebook_url). """ async with await _get_client(user_id, admin_account_id) as client: nb = await client.notebooks.create(f"جدوى - {report.project_name}") notebook_id = nb.id source_text = _synthesis_to_text(report) source = await client.sources.add_text( notebook_id, title=f"دراسة جدوى: {report.project_name}", content=source_text, wait=True, ) source_id = source.id if hasattr(source, "id") else str(source) notebook_url = f"https://notebooklm.google.com/notebook/{notebook_id}" return notebook_id, source_id, notebook_url async def generate_asset(notebook_id, asset_type, language="ar", user_id=None, admin_account_id=None): """ Trigger generation of a specific asset type in a notebook. Returns when generation is complete. """ method_name = ASSET_GENERATE_MAP.get(asset_type) if not method_name: raise ValueError(f"Unknown asset type: {asset_type}") from notebooklm import ( InfographicOrientation, InfographicDetail, VideoStyle, QuizDifficulty, ) async with await _get_client(user_id, admin_account_id) as client: artifacts = client.artifacts gen_method = getattr(artifacts, method_name) # Build kwargs matching notebooklm v0.3.2 API signatures kwargs = {} if asset_type == "audio": kwargs["language"] = language or "ar" kwargs["instructions"] = "اجعل المحتوى جذاباً وشاملاً" elif asset_type == "video": kwargs["language"] = language or "ar" kwargs["video_style"] = VideoStyle.WHITEBOARD elif asset_type == "quiz": kwargs["difficulty"] = QuizDifficulty.MEDIUM elif asset_type == "flashcards": kwargs["difficulty"] = QuizDifficulty.MEDIUM elif asset_type == "infographic": kwargs["language"] = language or "ar" kwargs["orientation"] = InfographicOrientation.PORTRAIT kwargs["detail_level"] = InfographicDetail.DETAILED elif asset_type == "slide_deck": kwargs["language"] = language or "ar" elif asset_type == "report": kwargs["language"] = language or "ar" # Retry loop for rate-limit handling max_retries = 4 last_error = None for attempt in range(max_retries): try: try: status = await gen_method(notebook_id, **kwargs) except TypeError as te: current_app.logger.warning(f"NotebookLM {asset_type} TypeError: {te}") status = await gen_method(notebook_id) # Handle dict responses (e.g. mind_map returns dict) if isinstance(status, dict): return True # Check for rate limiting — retry after delay if hasattr(status, "is_rate_limited") and status.is_rate_limited: wait_secs = 30 * (attempt + 1) current_app.logger.warning(f"NotebookLM rate limited for {asset_type}, waiting {wait_secs}s (attempt {attempt+1}/{max_retries})") await asyncio.sleep(wait_secs) continue # Check for failure if hasattr(status, "is_failed") and status.is_failed: err_msg = getattr(status, "error", "") or "" if "rate limit" in err_msg.lower() or "quota" in err_msg.lower(): wait_secs = 30 * (attempt + 1) current_app.logger.warning(f"NotebookLM rate limited for {asset_type}: {err_msg}, waiting {wait_secs}s") await asyncio.sleep(wait_secs) continue raise Exception(f"فشل إنشاء {asset_type} في NotebookLM: {err_msg}".strip()) # Wait for async task completion if hasattr(status, "task_id") and status.task_id: # Use extended timeouts for different asset types: # - Video: 30 minutes (1800s) - video generation takes the longest # - Audio: 25 minutes (1500s) - audio overview is complex # - Slide Deck: 20 minutes (1200s) - multi-slide generation # - Report: 15 minutes (900s) - long-form content # - Infographic: 15 minutes (900s) - visual generation # - Mind Map: 10 minutes (600s) - medium complexity # - Quiz/Flashcards: 5 minutes (300s) - simpler assets timeout_map = { "video": 1800, "audio": 1500, "slide_deck": 1200, "report": 900, "infographic": 900, "mind_map": 600, "quiz": 300, "flashcards": 300, } completion_timeout = timeout_map.get(asset_type, 900) # Default 15 minutes current_app.logger.info(f"Waiting for {asset_type} generation with timeout {completion_timeout}s") final = await artifacts.wait_for_completion(notebook_id, status.task_id, timeout=completion_timeout) if hasattr(final, "is_failed") and final.is_failed: err_msg = getattr(final, "error", "") or "" raise Exception(f"فشل إنشاء {asset_type} في NotebookLM: {err_msg}".strip()) return True except Exception as e: last_error = e err_str = str(e).lower() if "rate limit" in err_str or "quota" in err_str or "429" in err_str: wait_secs = 30 * (attempt + 1) current_app.logger.warning(f"NotebookLM rate limited (exception) for {asset_type}: {e}, waiting {wait_secs}s") await asyncio.sleep(wait_secs) continue raise raise last_error or Exception(f"فشل إنشاء {asset_type} بعد {max_retries} محاولات بسبب تجاوز حد الاستخدام. حاول بعد عدة دقائق.") async def download_asset(notebook_id, asset_type, output_path, user_id=None, admin_account_id=None): """ Download a generated asset to a local file. Retries up to 5 times with delays to wait for the artifact to become available. Returns the actual file path written. """ info = ASSET_DOWNLOAD_MAP.get(asset_type) if not info: raise ValueError(f"Unknown asset type: {asset_type}") method_name = info["method"] file_format = info["format"] file_path = f"{output_path}.{file_format}" max_retries = 5 last_error = None for attempt in range(max_retries): try: async with await _get_client(user_id, admin_account_id) as client: dl_method = getattr(client.artifacts, method_name) kwargs = {} if asset_type in ("quiz", "flashcards"): kwargs["output_format"] = "json" elif asset_type == "mind_map": kwargs["output_format"] = "json" if file_format == "json" else None try: await dl_method(notebook_id, file_path, **{k: v for k, v in kwargs.items() if v}) except TypeError: await dl_method(notebook_id, file_path) return file_path, file_format except Exception as e: last_error = e if attempt < max_retries - 1: await asyncio.sleep(5 * (attempt + 1)) raise last_error def run_sync(coro): """Run an async coroutine synchronously.""" try: loop = asyncio.get_event_loop() if loop.is_running(): import concurrent.futures with concurrent.futures.ThreadPoolExecutor() as pool: future = pool.submit(asyncio.run, coro) return future.result(timeout=700) return loop.run_until_complete(coro) except RuntimeError: return asyncio.run(coro)