Drug molecule screening agent

main_img = flyte.Image.from_uv_script(__file__, name="drug-molecule-screening", pre=True).with_apt_packages(
    "libxrender1", "libxext6", "libexpat1",
)

env = flyte.TaskEnvironment(
    name="drug-molecule-screening",
    image=main_img,
    resources=flyte.Resources(cpu=2, memory="6Gi"),
    secrets=[
        flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    ],
)

SCREENING_AGENT_INSTRUCTIONS = """\
You are a medicinal chemistry screening strategist. You orchestrate a virtual \
screening pipeline using durable Flyte tools. You NEVER invent molecular \
properties — only RDKit tools compute them.

Workflow:
1. If target_profile is not provided in the user message, derive a JSON \
target_profile from the therapeutic brief. Valid keys: mw, logp, hbd, hba, tpsa \
(each [min, max]). Ground choices in oral bioavailability / kinase / CNS rules \
as appropriate to the brief.
2. First pass (always): load_molecules → compute_properties → \
screen_candidates → generate_report. Pass tool outputs between steps exactly \
(molecule_dir from load_molecules into compute_properties and generate_report; \
properties_json from compute_properties into screen_candidates and \
generate_report; screening_json must be the complete, unmodified string \
returned by screen_candidates — never rebuild or summarize JSON yourself).
3. Read the JSON summary returned by generate_report. Reflect:
   - If all_criteria_met == 0: relax exactly ONE profile bound by ~10–20% \
and re-run screen_candidates then generate_report only, reusing the same \
molecule_dir and properties_json from the first pass.
   - If all molecules pass but diversity is a stated goal: note high similarity \
in your summary; do not re-run unless brief asks for stricter filters.
   - Maximum ONE rescreen iteration.
4. Finish with plain text: top candidate, rationale tied to computed metrics \
from the tool JSON, funnel interpretation, and suggested next steps (docking, \
ADMET lab tests).

If the user supplies an explicit target_profile JSON, use it as-is.

Do NOT ask the user for SMILES or molecule lists when molecules_json is empty — \
the default library is loaded automatically.
"""

screening_agent = Agent(
    name="drug-screening-agent",
    instructions=SCREENING_AGENT_INSTRUCTIONS,
    model=MODEL,
    tools=[
        load_molecules,
        compute_properties,
        screen_candidates,
        generate_report,
    ],
    max_turns=12,
)

@env.task(report=True)
async def pipeline(
    brief: str = "Screen the default drug library for orally bioavailable small molecules.",
    molecules_json: str = "",
    target_profile: str = "",
) -> str:
    """Agentic virtual drug molecule screening pipeline.

    A medicinal-chemistry agent interprets the screening brief, derives or
    applies a target profile, orchestrates the RDKit screening stages, and
    optionally re-screens when funnel results are too narrow.

    Args:
        brief: Natural-language therapeutic goal (e.g. oral kinase inhibitors,
            CNS-penetrant small molecules).
        molecules_json: JSON mapping molecule names to SMILES strings.
            Defaults to a curated library of ~15 well-known drugs.
        target_profile: Optional JSON with desired property ranges that
            overrides agent-derived criteria
            (e.g. {"mw": [150, 500], "logp": [-0.5, 5]}).

    Returns:
        Agent summary with screening rationale and key results.
    """
    prompt_parts = [
        f"Screening brief: {brief}",
        'Use molecules_json="" for the built-in default library unless provided below.',
        "Compose the four stage tools in order: load_molecules → compute_properties "
        "→ screen_candidates → generate_report. Pass each tool's full return value "
        "verbatim to the next step (especially screening_json). Re-run "
        "screen_candidates and generate_report at most once if the funnel is too narrow.",
    ]
    if molecules_json.strip():
        prompt_parts.append(f"molecules_json: {molecules_json}")
    if target_profile.strip():
        prompt_parts.append(f"Use this target_profile exactly: {target_profile}")

    result = await screening_agent.run.aio("\n".join(prompt_parts))
    return result.summary or result.error or ""

@env.task(report=True)
async def rescreen_demo() -> str:
    """Example run with a two-round execution graph (rescreen).

    Round 1 uses a strict CNS-like profile; round 2 always re-runs
    screen_candidates and generate_report with a widened LogP window,
    reusing cached molecule_dir and properties_json.
    """
    return await pipeline(
        brief=RESCREEN_DEMO_BRIEF,
        target_profile=RESCREEN_DEMO_TARGET_PROFILE,
    )