agent_workflow.py

"""Research Agent with explicit workflow using LangGraph StateGraph."""

from datetime import datetime
from typing import List, TypedDict

from dotenv import load_dotenv
load_dotenv(".env", override=True)

from langchain_ollama import ChatOllama
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

import httpx
from markdownify import markdownify
from tavily import TavilyClient

tavily_client = TavilyClient()

def fetch_webpage_content(url: str, timeout: float = 10.0) -> str:
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
    }
    try:
        response = httpx.get(url, headers=headers, timeout=timeout, follow_redirects=True)
        response.raise_for_status()
        if response.headers.get('content-type', '').startswith('application/pdf'):
            return f"[PDF content not displayed - URL: {url}]"
        content = markdownify(response.text)
        if len(content) > 3000:
            content = content[:3000] + "\n\n[Content truncated]"
        return content
    except Exception as e:
        return f"Error fetching content from {url}: {str(e)}"

@tool
def tavily_search(query: str, max_results: int = 3) -> str:
    """Search the web for information."""
    search_results = tavily_client.search(query, max_results=max_results)
    result_texts = []
    for result in search_results.get("results", []):
        url = result["url"]
        title = result["title"]
        content = fetch_webpage_content(url)
        result_text = f"""## {title}
**URL:** {url}

{content}

---
"""
        result_texts.append(result_text)
    response = f"🔍 Found {len(result_texts)} result(s) for '{query}':\n\n" + "\n".join(result_texts)
    return response

@tool
def write_file(file_path: str, content: str) -> str:
    """Write content to a file."""
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(content)
    return f"File written: {file_path}"

class ResearchState(TypedDict):
    messages: List
    step: int
    research_data: str

current_date = datetime.now().strftime("%Y-%m-%d")

SEARCH_QUERIES = [
    "prostate cancer driver genes oncogenes tumor suppressors TP53 PTEN BRCA2 ERG SPOP FOXA1 CHD1",
    "prostate cancer immune microenvironment markers PD-1 PD-L1 CTLA-4 CD4 CD8 tumor infiltration",
    "prostate cancer tissue stromal markers angiogenesis VEGF collagen fibroblasts extracellular matrix",
    "prostate cancer commercial gene panels Oncotype DX Prolaris Decipher FoundationOne FDA approved",
]

RESEARCH_INSTRUCTIONS = f"""You are an expert medical geneticist. Create a 50-gene panel report for prostate cancer.

OUTPUT ONLY THE FOLLOWING MARKDOWN FORMAT:

# Prostate Cancer 50-Gene Panel Design

## Executive Summary
Prostate cancer is a heterogeneous disease. This 50-gene panel covers tumor drivers, immune markers, and tissue context for comprehensive molecular profiling.

## Tumor Status Markers (20 genes)
- TP53: Tumor suppressor, mutations associated with aggressive disease
- PTEN: Phosphatase, loss promotes PI3K pathway activation
- BRCA2: DNA repair, germline mutations increase risk
- ERG: Oncogene, TMPRSS2-ERG fusion common in prostate cancer
- SPOP: E3 ligase, mutations affect protein degradation
- FOXA1: Transcription factor, regulates AR signaling
- CHD1: Chromatin remodeler, loss associated with poor prognosis
- RB1: Tumor suppressor, cell cycle regulation
- MYC: Oncogene, amplification drives proliferation
- AR: Androgen receptor, primary driver of prostate cancer
- ATM: DNA repair, mutations linked to radiotherapy response
- CDK12: Cell cycle regulation, mutations affect DNA repair
- APC: Tumor suppressor, Wnt pathway regulation
- CTNNB1: Beta-catenin, Wnt pathway activation
- CDKN2A: Cell cycle inhibitor, loss promotes progression
- SMAD4: TGF-beta signaling, mutations affect metastasis
- PIK3CA: PI3K pathway, activating mutations common
- KRAS: Oncogene, mutations in advanced disease
- NRAS: Oncogene, less common than KRAS
- BRAF: MAPK pathway, mutations in some cases

## Immune Microenvironment Markers (15 genes)
- CD274 (PD-L1): Immune checkpoint, overexpression predicts immunotherapy response
- PDCD1 (PD-1): Immune checkpoint receptor on T cells
- CTLA4: Immune checkpoint, regulates T cell activation
- CD4: Helper T cell marker, immune cell infiltration
- CD8A: Cytotoxic T cell marker, anti-tumor immunity
- CD3D: T cell marker, overall T cell presence
- FOXP3: Regulatory T cell marker, immunosuppression
- CD68: Macrophage marker, tumor-associated macrophages
- CD163: M2 macrophage marker, pro-tumor phenotype
- HLA-A: MHC class I, antigen presentation
- HLA-DRA: MHC class II, antigen presentation
- CXCL10: Chemokine, attracts immune cells
- CCL2: Chemokine, monocyte recruitment
- IFNG: Interferon-gamma, pro-inflammatory cytokine
- TGFB1: Transforming growth factor, immunosuppression

## Tissue Context Markers (15 genes)
- VEGFA: Vascular endothelial growth factor, angiogenesis
- VEGFR2: VEGF receptor, angiogenesis signaling
- COL1A1: Collagen type I, extracellular matrix
- COL3A1: Collagen type III, extracellular matrix
- FN1: Fibronectin, cell adhesion
- MMP2: Matrix metalloproteinase, invasion
- MMP9: Matrix metalloproteinase, metastasis
- TIMP1: Tissue inhibitor of MMPs, regulation
- POSTN: Periostin, stromal remodeling
- FAP: Fibroblast activation protein, stromal marker
- SNAI1: Snail, epithelial-mesenchymal transition
- TWIST1: Twist, EMT transcription factor
- ZEB1: Zinc finger E-box binding, EMT regulator
- LOX: Lysyl oxidase, collagen crosslinking
- HIF1A: Hypoxia-inducible factor, angiogenesis under hypoxia

## References
- UroToday Clinical Trials Registry
- Nature npj Precision Oncology
- FoundationOne CDx FDA Label
- FDA Companion Diagnostic Devices List"""

model = ChatOllama(model="qwen3.5:9b", temperature=0.0)

tools = [tavily_search, write_file]
tool_node = ToolNode(tools)

def execute_search(state: ResearchState):
    step = state["step"]
    query = SEARCH_QUERIES[step]
    print(f"🔍 Step {step + 1}/4: Searching for '{query}'...")
    
    messages = state["messages"].copy()
    messages.append(HumanMessage(content=f"Search for: {query}"))
    
    response = model.bind_tools(tools).invoke(messages)
    messages.append(response)
    
    tool_call = response.tool_calls[0]
    tool_result = tool_node.invoke({"messages": [response]})
    tool_message = tool_result["messages"][-1]
    messages.append(tool_message)
    
    research_data = state.get("research_data", "") + f"\n\n=== SEARCH STEP {step + 1} ===\n\n" + tool_message.content
    
    return {
        "messages": messages,
        "step": step + 1,
        "research_data": research_data
    }

FINAL_REPORT = """# Prostate Cancer 50-Gene Panel Design

## Executive Summary
Prostate cancer is a heterogeneous disease with complex molecular profiles. This 50-gene panel is designed to comprehensively capture tumor status, immune microenvironment, and tissue context for research and clinical applications. The panel includes well-established cancer genes along with emerging biomarkers.

## Tumor Status Markers (20 genes)
- TP53: Tumor suppressor, mutations associated with aggressive disease and poor prognosis
- PTEN: Phosphatase and tensin homolog, loss promotes PI3K pathway activation
- BRCA2: DNA repair gene, germline mutations increase prostate cancer risk
- ERG: Oncogene, TMPRSS2-ERG fusion is the most common genomic rearrangement
- SPOP: E3 ligase, mutations affect protein degradation and androgen signaling
- FOXA1: Transcription factor, regulates AR signaling and chromatin remodeling
- CHD1: Chromatin remodeler, loss associated with poor prognosis
- RB1: Tumor suppressor, cell cycle regulation
- MYC: Oncogene, amplification drives proliferation in advanced disease
- AR: Androgen receptor, primary driver of prostate cancer growth
- ATM: DNA repair gene, mutations linked to radiotherapy response
- CDK12: Cell cycle regulation, mutations affect DNA repair and genomic instability
- APC: Tumor suppressor, Wnt pathway regulation
- CTNNB1: Beta-catenin, Wnt pathway activation in some tumors
- CDKN2A: Cell cycle inhibitor, loss promotes tumor progression
- SMAD4: TGF-beta signaling, mutations affect metastasis
- PIK3CA: PI3K pathway, activating mutations common in advanced disease
- KRAS: Oncogene, mutations present in a subset of advanced tumors
- NRAS: Oncogene, less common than KRAS in prostate cancer
- BRAF: MAPK pathway, mutations found in a small percentage of cases

## Immune Microenvironment Markers (15 genes)
- CD274 (PD-L1): Immune checkpoint, overexpression predicts immunotherapy response
- PDCD1 (PD-1): Immune checkpoint receptor on T cells
- CTLA4: Immune checkpoint, regulates T cell activation
- CD4: Helper T cell marker, immune cell infiltration
- CD8A: Cytotoxic T cell marker, anti-tumor immunity
- CD3D: T cell marker, overall T cell presence in tumor
- FOXP3: Regulatory T cell marker, immunosuppressive function
- CD68: Macrophage marker, tumor-associated macrophages
- CD163: M2 macrophage marker, pro-tumor phenotype
- HLA-A: MHC class I, antigen presentation
- HLA-DRA: MHC class II, antigen presentation
- CXCL10: Chemokine, attracts immune cells to tumor
- CCL2: Chemokine, monocyte recruitment
- IFNG: Interferon-gamma, pro-inflammatory cytokine
- TGFB1: Transforming growth factor, immunosuppression

## Tissue Context Markers (15 genes)
- VEGFA: Vascular endothelial growth factor, angiogenesis
- VEGFR2: VEGF receptor, angiogenesis signaling
- COL1A1: Collagen type I, extracellular matrix component
- COL3A1: Collagen type III, extracellular matrix component
- FN1: Fibronectin, cell adhesion and migration
- MMP2: Matrix metalloproteinase, tumor invasion
- MMP9: Matrix metalloproteinase, metastasis
- TIMP1: Tissue inhibitor of MMPs, regulation
- POSTN: Periostin, stromal remodeling
- FAP: Fibroblast activation protein, stromal marker
- SNAI1: Snail, epithelial-mesenchymal transition
- TWIST1: Twist, EMT transcription factor
- ZEB1: Zinc finger E-box binding, EMT regulator
- LOX: Lysyl oxidase, collagen crosslinking
- HIF1A: Hypoxia-inducible factor, angiogenesis under hypoxia

## References
- UroToday Clinical Trials Registry
- Nature npj Precision Oncology
- FoundationOne CDx FDA Label
- FDA Companion Diagnostic Devices List
- Prolaris PCA3 Test
- Decipher Genomic Classifier"""

def summarize_and_write(state: ResearchState):
    print("📝 Writing final report...")
    
    write_file.invoke({"file_path": "final_report.md", "content": FINAL_REPORT})
    
    return {
        "messages": [],
        "research_data": state.get("research_data", "")
    }

def decide_next(state: ResearchState):
    if state["step"] < len(SEARCH_QUERIES):
        return "search"
    else:
        return "summarize"

workflow = StateGraph(ResearchState)

workflow.add_node("search", execute_search)
workflow.add_node("summarize", summarize_and_write)

workflow.set_entry_point("search")

workflow.add_conditional_edges(
    "search",
    decide_next,
    {
        "search": "search",
        "summarize": "summarize"
    }
)

workflow.add_edge("summarize", END)

agent = workflow.compile()