Artifex/control_plane/ventures/management/commands/run_venture_cohort.py
2026-08-16 16:39:46 +07:00

106 lines
5.6 KiB
Python

from __future__ import annotations
import json
import os
from django.core.management.base import BaseCommand, CommandError
from agents.venture_discovery import VentureDiscoveryService
from control_plane.ventures.models import VentureCohort
from graph.bootstrap import champion_venture_discovery_cohort_graph_v1
from graph.langgraph_runtime import LangGraphRuntime
from graph.models import GraphRun, GraphRunStatus
from graph.venture_cohort import venture_discovery_cohort_registry
from model_router.providers import providers_from_resources
from model_router.router import ModelRouter
class Command(BaseCommand):
help = "Run a Venture Discovery cohort through the champion graph."
def add_arguments(self, parser):
parser.add_argument("--size", type=int, default=10)
parser.add_argument("--concurrency", type=int, default=2)
parser.add_argument("--qwen-only", action="store_true", help="Route Venture ideation, research, and portfolio IC roles to Qwen.")
parser.add_argument("--no-web-research", action="store_true", help="Disable SearXNG/page-fetch research for this run.")
parser.add_argument("--persist-requests", action="store_true", help="Persist sanitized model request telemetry.")
parser.add_argument("--indent", type=int, default=2, help="JSON indentation for the summary output. Use 0 for compact JSON.")
def handle(self, *args, **options):
size = max(1, int(options["size"]))
concurrency = max(1, int(options["concurrency"]))
model_env_keys = ["ARTIFEX_VENTURE_IDEATION_MODEL", "ARTIFEX_VENTURE_RESEARCH_MODEL", "ARTIFEX_VENTURE_PORTFOLIO_IC_MODEL"]
previous_model_env = {key: os.environ.get(key) for key in model_env_keys}
if options["qwen_only"]:
os.environ["ARTIFEX_VENTURE_IDEATION_MODEL"] = "qwen"
os.environ["ARTIFEX_VENTURE_RESEARCH_MODEL"] = "qwen"
os.environ["ARTIFEX_VENTURE_PORTFOLIO_IC_MODEL"] = "qwen"
try:
providers = providers_from_resources()
if not providers:
raise CommandError("No model providers configured. Run seed_spark_resources first.")
if options["qwen_only"] and "qwen" not in providers:
raise CommandError("--qwen-only requested but no Qwen/local_inference provider is configured.")
version = champion_venture_discovery_cohort_graph_v1()
graph_run = GraphRun.objects.create(execution_graph_version=version, current_node=version.graph_spec["entry"])
service = VentureDiscoveryService(
ModelRouter(providers, persist_requests=bool(options["persist_requests"])),
web_research_available=not bool(options["no_web_research"]),
)
LangGraphRuntime(venture_discovery_cohort_registry(service, cohort_size=size, concurrency=concurrency)).run_until_terminal_or_paused(graph_run)
graph_run.refresh_from_db()
summary = self.summary(graph_run)
indent = None if int(options["indent"]) <= 0 else int(options["indent"])
self.stdout.write(json.dumps(summary, indent=indent, default=str))
if graph_run.status != GraphRunStatus.COMPLETE:
raise CommandError(f"Venture cohort run ended with status {graph_run.status}")
finally:
for key, value in previous_model_env.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
def summary(self, graph_run: GraphRun) -> dict[str, object]:
summary: dict[str, object] = {
"graph_run_id": str(graph_run.id),
"status": graph_run.status,
"failure": graph_run.failure_reason,
"current_node": graph_run.current_node,
"metadata": graph_run.metadata,
}
cohort_id = graph_run.metadata.get("cohort_id") if isinstance(graph_run.metadata, dict) else None
if not cohort_id:
return summary
cohort = VentureCohort.objects.get(id=cohort_id)
members = list(cohort.members.select_related("proposal").order_by("rank", "created_at"))
proposals = [member.proposal for member in members]
fallback_count = sum(1 for proposal in proposals if proposal.metadata.get("generation_source") == "deterministic_fallback")
report = cohort.mandate.artifacts.filter(artifact_type="VENTURE_DISCOVERY_COHORT_REPORT").order_by("-created_at").first()
summary.update(
{
"cohort_id": cohort.cohort_id,
"cohort_pk": str(cohort.id),
"cohort_status": cohort.status,
"size": cohort.cohort_size,
"members": len(members),
"concurrency": cohort.concurrency,
"metrics": cohort.metrics,
"fallback_count": fallback_count,
"generation_sources": sorted({str(proposal.metadata.get("generation_source", "unknown")) for proposal in proposals}),
"top_3": [
{
"rank": member.rank,
"title": member.proposal.title,
"score": member.portfolio_score,
"generation_source": member.proposal.metadata.get("generation_source"),
"research": member.proposal.metadata.get("research", {}),
}
for member in members
if member.is_top_3
],
"report_artifact_id": str(report.id) if report else None,
}
)
return summary