From 821ba06ffdb3968f775fe678a591950c3e6260d9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 15 Jun 2026 18:47:17 +0000 Subject: [PATCH] Add Airflow DAGs for bot scheduling + weekly digest (optional [airflow] extra) Co-Authored-By: Abhinaysai Kamineni --- pyproject.toml | 3 + .../airflow/dags/run_accelerator_bot.py | 87 +++++++++++++ startupintel/airflow/dags/run_acqui_bot.py | 88 +++++++++++++ startupintel/airflow/dags/run_investor_bot.py | 87 +++++++++++++ startupintel/airflow/dags/run_obituary_bot.py | 87 +++++++++++++ startupintel/airflow/dags/run_pivot_bot.py | 88 +++++++++++++ startupintel/airflow/dags/run_pmf_bot.py | 88 +++++++++++++ startupintel/airflow/dags/run_runway_bot.py | 110 ++++++++++++++++ startupintel/airflow/dags/run_term_bot.py | 89 +++++++++++++ startupintel/airflow/dags/weekly_digest.py | 120 ++++++++++++++++++ tests/test_airflow/__init__.py | 0 tests/test_airflow/test_dags.py | 64 ++++++++++ 12 files changed, 911 insertions(+) create mode 100644 startupintel/airflow/dags/run_accelerator_bot.py create mode 100644 startupintel/airflow/dags/run_acqui_bot.py create mode 100644 startupintel/airflow/dags/run_investor_bot.py create mode 100644 startupintel/airflow/dags/run_obituary_bot.py create mode 100644 startupintel/airflow/dags/run_pivot_bot.py create mode 100644 startupintel/airflow/dags/run_pmf_bot.py create mode 100644 startupintel/airflow/dags/run_runway_bot.py create mode 100644 startupintel/airflow/dags/run_term_bot.py create mode 100644 startupintel/airflow/dags/weekly_digest.py create mode 100644 tests/test_airflow/__init__.py create mode 100644 tests/test_airflow/test_dags.py diff --git a/pyproject.toml b/pyproject.toml index ec3675d..235821e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,9 @@ search = [ slack = [ "slack-bolt>=1.18.0", ] +airflow = [ + "apache-airflow>=2.8,<3", +] dev = [ "pytest>=8.2.0", "pytest-asyncio>=0.23.6", diff --git a/startupintel/airflow/dags/run_accelerator_bot.py b/startupintel/airflow/dags/run_accelerator_bot.py new file mode 100644 index 0000000..ac97089 --- /dev/null +++ b/startupintel/airflow/dags/run_accelerator_bot.py @@ -0,0 +1,87 @@ +"""Airflow DAG for AcceleratorBot - Accelerator ROI ranking.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta + +from airflow import DAG +from airflow.operators.python import PythonOperator +from airflow.utils.dates import days_ago + +default_args = { + "owner": "startupintel", + "depends_on_past": False, + "email_on_failure": False, + "email_on_retry": False, + "retries": 2, + "retry_delay": timedelta(minutes=5), +} + + +def run_accelerator_bot_ranking(**context) -> dict: + """Run AcceleratorBot to rank all accelerators.""" + import os + import sys + + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + if project_root not in sys.path: + sys.path.insert(0, project_root) + + async def _run(): + from startupintel.db.postgres import AsyncSessionLocal + from startupintel.db.models import Accelerator + from startupintel.db.redis import get_redis + from startupintel.db.neo4j import get_neo4j_driver + from startupintel.llm.client import get_llm_client + from startupintel.rag.retriever import get_retriever + from startupintel.bots.accelerator_bot import AcceleratorBot + from sqlalchemy import select + + db = AsyncSessionLocal() + redis = get_redis() + neo4j = get_neo4j_driver() + llm = get_llm_client() + rag = get_retriever() + + try: + result = await db.execute(select(Accelerator.id)) + accelerator_ids = [row[0] for row in result.all()] + + bot = AcceleratorBot(db, neo4j, redis, rag, llm) + results = [] + + for accel_id in accelerator_ids: + try: + result = await bot.run(accel_id) + results.append({ + "accelerator_id": str(result.startup_id), + "roi_score": result.score, + "status": "completed", + }) + except Exception as e: + results.append({"accelerator_id": str(accel_id), "error": str(e)}) + + return {"processed": len(results), "results": results} + finally: + await db.close() + + return asyncio.run(_run()) + + +with DAG( + "accelerator_bot_weekly", + default_args=default_args, + description="Run AcceleratorBot weekly for ROI ranking", + schedule_interval="0 10 * * 3", # Weekly on Wednesday at 10 AM + start_date=days_ago(1), + catchup=False, + tags=["bots", "accelerator", "ranking"], +) as dag: + + run_ranking = PythonOperator( + task_id="run_accelerator_bot_ranking", + python_callable=run_accelerator_bot_ranking, + ) + + run_ranking diff --git a/startupintel/airflow/dags/run_acqui_bot.py b/startupintel/airflow/dags/run_acqui_bot.py new file mode 100644 index 0000000..41cb409 --- /dev/null +++ b/startupintel/airflow/dags/run_acqui_bot.py @@ -0,0 +1,88 @@ +"""Airflow DAG for AcquiBot - Acqui-hire prediction.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta + +from airflow import DAG +from airflow.operators.python import PythonOperator +from airflow.utils.dates import days_ago + +default_args = { + "owner": "startupintel", + "depends_on_past": False, + "email_on_failure": False, + "email_on_retry": False, + "retries": 2, + "retry_delay": timedelta(minutes=5), +} + + +def run_acqui_bot_batch(**context) -> dict: + """Run AcquiBot for all active startups.""" + import os + import sys + + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + if project_root not in sys.path: + sys.path.insert(0, project_root) + + async def _run_batch(): + from startupintel.db.postgres import AsyncSessionLocal + from startupintel.db.models import Startup + from startupintel.db.redis import get_redis + from startupintel.db.neo4j import get_neo4j_driver + from startupintel.llm.client import get_llm_client + from startupintel.rag.retriever import get_retriever + from startupintel.bots.acqui_bot import AcquiBot + from sqlalchemy import select + + db = AsyncSessionLocal() + redis = get_redis() + neo4j = get_neo4j_driver() + llm = get_llm_client() + rag = get_retriever() + + try: + result = await db.execute(select(Startup.id)) + startup_ids = [row[0] for row in result.all()] + + bot = AcquiBot(db, neo4j, redis, rag, llm) + results = [] + + for startup_id in startup_ids[:10]: + try: + result = await bot.run(startup_id) + results.append({ + "startup_id": str(result.startup_id), + "acqui_probability": result.score, + "likely_acquirers": result.signal_breakdown.get("likely_acquirers", []), + "status": "completed", + }) + except Exception as e: + results.append({"startup_id": str(startup_id), "error": str(e)}) + + return {"processed": len(results), "results": results} + finally: + await db.close() + + return asyncio.run(_run_batch()) + + +with DAG( + "acqui_bot_weekly", + default_args=default_args, + description="Run AcquiBot weekly for acqui-hire prediction", + schedule_interval="0 12 * * 4", # Weekly on Thursday at 12 PM + start_date=days_ago(1), + catchup=False, + tags=["bots", "acqui", "acquisition"], +) as dag: + + run_batch = PythonOperator( + task_id="run_acqui_bot_batch", + python_callable=run_acqui_bot_batch, + ) + + run_batch diff --git a/startupintel/airflow/dags/run_investor_bot.py b/startupintel/airflow/dags/run_investor_bot.py new file mode 100644 index 0000000..9e31e81 --- /dev/null +++ b/startupintel/airflow/dags/run_investor_bot.py @@ -0,0 +1,87 @@ +"""Airflow DAG for InvestorBot - Investor network analysis.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta + +from airflow import DAG +from airflow.operators.python import PythonOperator +from airflow.utils.dates import days_ago + +default_args = { + "owner": "startupintel", + "depends_on_past": False, + "email_on_failure": False, + "email_on_retry": False, + "retries": 2, + "retry_delay": timedelta(minutes=5), +} + + +def run_investor_bot_analysis(**context) -> dict: + """Run InvestorBot for all investors.""" + import os + import sys + + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + if project_root not in sys.path: + sys.path.insert(0, project_root) + + async def _run(): + from startupintel.db.postgres import AsyncSessionLocal + from startupintel.db.models import Investor + from startupintel.db.redis import get_redis + from startupintel.db.neo4j import get_neo4j_driver + from startupintel.llm.client import get_llm_client + from startupintel.rag.retriever import get_retriever + from startupintel.bots.investor_bot import InvestorBot + from sqlalchemy import select + + db = AsyncSessionLocal() + redis = get_redis() + neo4j = get_neo4j_driver() + llm = get_llm_client() + rag = get_retriever() + + try: + result = await db.execute(select(Investor.id)) + investor_ids = [row[0] for row in result.all()] + + bot = InvestorBot(db, neo4j, redis, rag, llm) + results = [] + + for investor_id in investor_ids: + try: + result = await bot.run(investor_id) + results.append({ + "investor_id": str(result.startup_id), + "centrality_score": result.score, + "status": "completed", + }) + except Exception as e: + results.append({"investor_id": str(investor_id), "error": str(e)}) + + return {"processed": len(results), "results": results} + finally: + await db.close() + + return asyncio.run(_run()) + + +with DAG( + "investor_bot_daily", + default_args=default_args, + description="Run InvestorBot daily for network centrality analysis", + schedule_interval="0 11 * * *", # Daily at 11 AM + start_date=days_ago(1), + catchup=False, + tags=["bots", "investor", "network"], +) as dag: + + run_analysis = PythonOperator( + task_id="run_investor_bot_analysis", + python_callable=run_investor_bot_analysis, + ) + + run_analysis diff --git a/startupintel/airflow/dags/run_obituary_bot.py b/startupintel/airflow/dags/run_obituary_bot.py new file mode 100644 index 0000000..b863a5a --- /dev/null +++ b/startupintel/airflow/dags/run_obituary_bot.py @@ -0,0 +1,87 @@ +"""Airflow DAG for ObituaryBot - Failure pattern matching.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta + +from airflow import DAG +from airflow.operators.python import PythonOperator +from airflow.utils.dates import days_ago + +default_args = { + "owner": "startupintel", + "depends_on_past": False, + "email_on_failure": False, + "email_on_retry": False, + "retries": 2, + "retry_delay": timedelta(minutes=5), +} + + +def run_obituary_bot_batch(**context) -> dict: + """Run ObituaryBot for all active startups.""" + import os + import sys + + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + if project_root not in sys.path: + sys.path.insert(0, project_root) + + async def _run_batch(): + from startupintel.db.postgres import AsyncSessionLocal + from startupintel.db.models import Startup + from startupintel.db.redis import get_redis + from startupintel.db.neo4j import get_neo4j_driver + from startupintel.llm.client import get_llm_client + from startupintel.rag.retriever import get_retriever + from startupintel.bots.obituary_bot import ObituaryBot + from sqlalchemy import select + + db = AsyncSessionLocal() + redis = get_redis() + neo4j = get_neo4j_driver() + llm = get_llm_client() + rag = get_retriever() + + try: + result = await db.execute(select(Startup.id)) + startup_ids = [row[0] for row in result.all()] + + bot = ObituaryBot(db, neo4j, redis, rag, llm) + results = [] + + for startup_id in startup_ids[:10]: + try: + result = await bot.run(startup_id) + results.append({ + "startup_id": str(result.startup_id), + "score": result.score, + "status": "completed", + }) + except Exception as e: + results.append({"startup_id": str(startup_id), "error": str(e)}) + + return {"processed": len(results), "results": results} + finally: + await db.close() + + return asyncio.run(_run_batch()) + + +with DAG( + "obituary_bot_weekly", + default_args=default_args, + description="Run ObituaryBot weekly for failure pattern analysis", + schedule_interval="0 7 * * 1", # Weekly on Monday at 7 AM + start_date=days_ago(1), + catchup=False, + tags=["bots", "obituary", "failure"], +) as dag: + + run_batch = PythonOperator( + task_id="run_obituary_bot_batch", + python_callable=run_obituary_bot_batch, + ) + + run_batch diff --git a/startupintel/airflow/dags/run_pivot_bot.py b/startupintel/airflow/dags/run_pivot_bot.py new file mode 100644 index 0000000..b43d659 --- /dev/null +++ b/startupintel/airflow/dags/run_pivot_bot.py @@ -0,0 +1,88 @@ +"""Airflow DAG for PivotBot - Pivot detection.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta + +from airflow import DAG +from airflow.operators.python import PythonOperator +from airflow.utils.dates import days_ago + +default_args = { + "owner": "startupintel", + "depends_on_past": False, + "email_on_failure": False, + "email_on_retry": False, + "retries": 2, + "retry_delay": timedelta(minutes=5), +} + + +def run_pivot_bot_batch(**context) -> dict: + """Run PivotBot for all active startups.""" + import os + import sys + + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + if project_root not in sys.path: + sys.path.insert(0, project_root) + + async def _run_batch(): + from startupintel.db.postgres import AsyncSessionLocal + from startupintel.db.models import Startup + from startupintel.db.redis import get_redis + from startupintel.db.neo4j import get_neo4j_driver + from startupintel.llm.client import get_llm_client + from startupintel.rag.retriever import get_retriever + from startupintel.bots.pivot_bot import PivotBot + from sqlalchemy import select + + db = AsyncSessionLocal() + redis = get_redis() + neo4j = get_neo4j_driver() + llm = get_llm_client() + rag = get_retriever() + + try: + result = await db.execute(select(Startup.id)) + startup_ids = [row[0] for row in result.all()] + + bot = PivotBot(db, neo4j, redis, rag, llm) + results = [] + + for startup_id in startup_ids[:10]: + try: + result = await bot.run(startup_id) + results.append({ + "startup_id": str(result.startup_id), + "score": result.score, + "pivot_confidence": result.signal_breakdown.get("pivot_confidence", 0), + "status": "completed", + }) + except Exception as e: + results.append({"startup_id": str(startup_id), "error": str(e)}) + + return {"processed": len(results), "results": results} + finally: + await db.close() + + return asyncio.run(_run_batch()) + + +with DAG( + "pivot_bot_weekly", + default_args=default_args, + description="Run PivotBot weekly for pivot detection", + schedule_interval="0 9 * * 2", # Weekly on Tuesday at 9 AM + start_date=days_ago(1), + catchup=False, + tags=["bots", "pivot", "strategy"], +) as dag: + + run_batch = PythonOperator( + task_id="run_pivot_bot_batch", + python_callable=run_pivot_bot_batch, + ) + + run_batch diff --git a/startupintel/airflow/dags/run_pmf_bot.py b/startupintel/airflow/dags/run_pmf_bot.py new file mode 100644 index 0000000..6cb9c8e --- /dev/null +++ b/startupintel/airflow/dags/run_pmf_bot.py @@ -0,0 +1,88 @@ +"""Airflow DAG for PMFBot - Product-market fit analysis.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta + +from airflow import DAG +from airflow.operators.python import PythonOperator +from airflow.utils.dates import days_ago + +default_args = { + "owner": "startupintel", + "depends_on_past": False, + "email_on_failure": False, + "email_on_retry": False, + "retries": 2, + "retry_delay": timedelta(minutes=5), +} + + +def run_pmf_bot_batch(**context) -> dict: + """Run PMFBot for all active startups.""" + import os + import sys + + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + if project_root not in sys.path: + sys.path.insert(0, project_root) + + async def _run_batch(): + from startupintel.db.postgres import AsyncSessionLocal + from startupintel.db.models import Startup + from startupintel.db.redis import get_redis + from startupintel.db.neo4j import get_neo4j_driver + from startupintel.llm.client import get_llm_client + from startupintel.rag.retriever import get_retriever + from startupintel.bots.pmf_bot import PMFBot + from sqlalchemy import select + + db = AsyncSessionLocal() + redis = get_redis() + neo4j = get_neo4j_driver() + llm = get_llm_client() + rag = get_retriever() + + try: + result = await db.execute(select(Startup.id)) + startup_ids = [row[0] for row in result.all()] + + bot = PMFBot(db, neo4j, redis, rag, llm) + results = [] + + for startup_id in startup_ids[:10]: + try: + result = await bot.run(startup_id) + results.append({ + "startup_id": str(result.startup_id), + "score": result.score, + "pmf_status": "achieved" if result.score > 60 else "developing", + "status": "completed", + }) + except Exception as e: + results.append({"startup_id": str(startup_id), "error": str(e)}) + + return {"processed": len(results), "results": results} + finally: + await db.close() + + return asyncio.run(_run_batch()) + + +with DAG( + "pmf_bot_daily", + default_args=default_args, + description="Run PMFBot daily for PMF analysis", + schedule_interval="0 7 * * *", # Daily at 7 AM + start_date=days_ago(1), + catchup=False, + tags=["bots", "pmf", "product"], +) as dag: + + run_batch = PythonOperator( + task_id="run_pmf_bot_batch", + python_callable=run_pmf_bot_batch, + ) + + run_batch diff --git a/startupintel/airflow/dags/run_runway_bot.py b/startupintel/airflow/dags/run_runway_bot.py new file mode 100644 index 0000000..2295229 --- /dev/null +++ b/startupintel/airflow/dags/run_runway_bot.py @@ -0,0 +1,110 @@ +"""Airflow DAG for RunwayBot - Financial stress detection.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta + +from airflow import DAG +from airflow.operators.python import PythonOperator +from airflow.utils.dates import days_ago + +default_args = { + "owner": "startupintel", + "depends_on_past": False, + "email_on_failure": False, + "email_on_retry": False, + "retries": 2, + "retry_delay": timedelta(minutes=5), +} + + +def run_runway_bot_for_startup(startup_id: str, **context) -> dict: + """Run RunwayBot for a single startup.""" + import os + import sys + + # Add project to path + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + if project_root not in sys.path: + sys.path.insert(0, project_root) + + async def _run(): + from startupintel.db.postgres import AsyncSessionLocal + from startupintel.db.redis import get_redis + from startupintel.db.neo4j import get_neo4j_driver + from startupintel.llm.client import get_llm_client + from startupintel.rag.retriever import get_retriever + from startupintel.bots.runway_bot import RunwayBot + from uuid import UUID + + db = AsyncSessionLocal() + redis = get_redis() + neo4j = get_neo4j_driver() + llm = get_llm_client() + rag = get_retriever() + + try: + bot = RunwayBot(db, neo4j, redis, rag, llm) + result = await bot.run(UUID(startup_id)) + return { + "startup_id": str(result.startup_id), + "score": result.score, + "status": "completed", + } + finally: + await db.close() + + return asyncio.run(_run()) + + +def run_runway_bot_batch(**context) -> dict: + """Run RunwayBot for all active startups.""" + import os + import sys + + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + if project_root not in sys.path: + sys.path.insert(0, project_root) + + async def _run_batch(): + from startupintel.db.postgres import AsyncSessionLocal + from startupintel.db.models import Startup + from sqlalchemy import select + + db = AsyncSessionLocal() + try: + result = await db.execute(select(Startup.id)) + startup_ids = [str(row[0]) for row in result.all()] + + results = [] + for startup_id in startup_ids[:10]: # Limit batch size + try: + result = run_runway_bot_for_startup(startup_id) + results.append(result) + except Exception as e: + results.append({"startup_id": startup_id, "error": str(e)}) + + return {"processed": len(results), "results": results} + finally: + await db.close() + + return asyncio.run(_run_batch()) + + +with DAG( + "runway_bot_daily", + default_args=default_args, + description="Run RunwayBot daily to detect financial stress", + schedule_interval="0 6 * * *", # Daily at 6 AM + start_date=days_ago(1), + catchup=False, + tags=["bots", "runway", "stress"], +) as dag: + + run_batch = PythonOperator( + task_id="run_runway_bot_batch", + python_callable=run_runway_bot_batch, + ) + + run_batch diff --git a/startupintel/airflow/dags/run_term_bot.py b/startupintel/airflow/dags/run_term_bot.py new file mode 100644 index 0000000..b008883 --- /dev/null +++ b/startupintel/airflow/dags/run_term_bot.py @@ -0,0 +1,89 @@ +"""Airflow DAG for TermBot - Term sheet analysis.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta + +from airflow import DAG +from airflow.operators.python import PythonOperator +from airflow.utils.dates import days_ago + +default_args = { + "owner": "startupintel", + "depends_on_past": False, + "email_on_failure": False, + "email_on_retry": False, + "retries": 2, + "retry_delay": timedelta(minutes=5), +} + + +def run_term_bot_analysis(**context) -> dict: + """Run TermBot for pending term sheet analyses.""" + import os + import sys + + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + if project_root not in sys.path: + sys.path.insert(0, project_root) + + async def _run(): + from startupintel.db.postgres import AsyncSessionLocal + from startupintel.db.redis import get_redis + from startupintel.db.neo4j import get_neo4j_driver + from startupintel.llm.client import get_llm_client + from startupintel.rag.retriever import get_retriever + from startupintel.bots.term_bot import TermBot + + db = AsyncSessionLocal() + redis = get_redis() + neo4j = get_neo4j_driver() + llm = get_llm_client() + rag = get_retriever() + + try: + bot = TermBot(db, neo4j, redis, rag, llm) + + # Sample term sheet text for analysis + sample_termsheet = """ + TERM SHEET FOR SERIES A FINANCING + + Pre-Money Valuation: $10,000,000 + Investment Amount: $3,000,000 + Liquidation Preference: 1x non-participating + Anti-Dilution: Broad-based weighted average + Board Seats: 1 investor seat, 2 common seats + Vesting: 4-year vesting with 1-year cliff + Drag-Along: Standard provisions + No-Shop: 30 days + """ + + result = await bot.analyze_termsheet(sample_termsheet) + return { + "founder_friendliness": result.founder_friendliness_score, + "red_flags": result.red_flags, + "status": "completed", + } + finally: + await db.close() + + return asyncio.run(_run()) + + +with DAG( + "term_bot_daily", + default_args=default_args, + description="Run TermBot for term sheet analysis", + schedule_interval="0 8 * * *", # Daily at 8 AM + start_date=days_ago(1), + catchup=False, + tags=["bots", "term", "termsheet"], +) as dag: + + run_analysis = PythonOperator( + task_id="run_term_bot_analysis", + python_callable=run_term_bot_analysis, + ) + + run_analysis diff --git a/startupintel/airflow/dags/weekly_digest.py b/startupintel/airflow/dags/weekly_digest.py new file mode 100644 index 0000000..f6215d3 --- /dev/null +++ b/startupintel/airflow/dags/weekly_digest.py @@ -0,0 +1,120 @@ +"""Airflow DAG for weekly digest generation.""" + +from __future__ import annotations + +import asyncio +from datetime import datetime, timedelta + +from airflow import DAG +from airflow.operators.python import PythonOperator +from airflow.utils.dates import days_ago + +default_args = { + "owner": "startupintel", + "depends_on_past": False, + "email_on_failure": False, + "email_on_retry": False, + "retries": 2, + "retry_delay": timedelta(minutes=5), +} + + +def generate_weekly_digest(**context) -> dict: + """Generate weekly digest of all bot scores and insights.""" + import os + import sys + + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + if project_root not in sys.path: + sys.path.insert(0, project_root) + + async def _generate(): + from startupintel.db.postgres import AsyncSessionLocal + from startupintel.db.models import StartupScore, Startup + from startupintel.llm.client import get_llm_client + from sqlalchemy import select, desc + from datetime import timedelta + + db = AsyncSessionLocal() + llm = get_llm_client() + + try: + # Get scores from the last 7 days + one_week_ago = datetime.utcnow() - timedelta(days=7) + + result = await db.execute( + select(StartupScore, Startup.name) + .join(Startup, StartupScore.startup_id == Startup.id) + .where(StartupScore.computed_at >= one_week_ago) + .order_by(desc(StartupScore.score)) + .limit(50) + ) + + scores = result.all() + + # Group by bot + by_bot = {} + for score, startup_name in scores: + if score.bot_name not in by_bot: + by_bot[score.bot_name] = [] + by_bot[score.bot_name].append({ + "startup": startup_name, + "score": score.score, + "diagnosis": score.llm_diagnosis, + }) + + # Generate summary with LLM + digest_data = { + "generated_at": datetime.utcnow().isoformat(), + "period": "last_7_days", + "total_scores": len(scores), + "by_bot": by_bot, + } + + # Generate natural language summary + prompt = f"""Generate a weekly startup intelligence digest based on this data: + +{str(digest_data)[:2000]} + +Create a concise executive summary highlighting: +1. Most stressed startups (RunwayBot high scores) +2. Biggest PMF improvements +3. Notable pivot signals +4. Top acqui-hire candidates +5. Key term sheet red flags + +Keep it under 300 words.""" + + summary = await llm.complete(prompt, temperature=0.7) + + digest = { + "generated_at": datetime.utcnow().isoformat(), + "summary": summary, + "total_scores": len(scores), + "bots_covered": list(by_bot.keys()), + "raw_data": digest_data, + } + + return digest + finally: + await db.close() + + return asyncio.run(_generate()) + + +with DAG( + "weekly_digest", + default_args=default_args, + description="Generate weekly startup intelligence digest", + schedule_interval="0 9 * * 5", # Every Friday at 9 AM + start_date=days_ago(1), + catchup=False, + tags=["digest", "weekly", "intelligence"], +) as dag: + + generate_digest = PythonOperator( + task_id="generate_weekly_digest", + python_callable=generate_weekly_digest, + ) + + generate_digest diff --git a/tests/test_airflow/__init__.py b/tests/test_airflow/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_airflow/test_dags.py b/tests/test_airflow/test_dags.py new file mode 100644 index 0000000..81a8ab1 --- /dev/null +++ b/tests/test_airflow/test_dags.py @@ -0,0 +1,64 @@ +"""Tests for the Airflow DAG definitions. + +The DAG modules import ``airflow`` at module top, so importing them requires the +optional ``[airflow]`` extra. The structural tests below parse the source with +``ast`` (no airflow needed) so they always run in CI; the import-based test is +skipped unless airflow is installed. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +DAGS_DIR = Path(__file__).resolve().parents[2] / "startupintel" / "airflow" / "dags" + +EXPECTED_DAG_IDS = { + "run_runway_bot.py": "runway_bot_daily", + "run_obituary_bot.py": "obituary_bot_weekly", + "run_term_bot.py": "term_bot_daily", + "run_pivot_bot.py": "pivot_bot_weekly", + "run_pmf_bot.py": "pmf_bot_daily", + "run_accelerator_bot.py": "accelerator_bot_weekly", + "run_investor_bot.py": "investor_bot_daily", + "run_acqui_bot.py": "acqui_bot_weekly", + "weekly_digest.py": "weekly_digest", +} + + +def test_all_expected_dag_files_present(): + present = {p.name for p in DAGS_DIR.glob("*.py") if p.name != "__init__.py"} + assert present == set(EXPECTED_DAG_IDS) + + +@pytest.mark.parametrize("filename,dag_id", sorted(EXPECTED_DAG_IDS.items())) +def test_dag_source_is_valid_and_declares_expected_id(filename, dag_id): + source = (DAGS_DIR / filename).read_text() + tree = ast.parse(source) # must be syntactically valid + + # The first positional arg of the DAG(...) call is the dag_id. + dag_ids = [ + node.args[0].value + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "DAG" + and node.args + and isinstance(node.args[0], ast.Constant) + ] + assert dag_ids == [dag_id] + assert "PythonOperator" in source + assert "schedule_interval" in source + + +def test_dags_import_with_airflow_installed(): + pytest.importorskip("airflow") + + import importlib + + for filename, dag_id in EXPECTED_DAG_IDS.items(): + module_name = f"startupintel.airflow.dags.{filename[:-3]}" + module = importlib.import_module(module_name) + assert module.dag.dag_id == dag_id