Develop a Multi-Agent AI Workflow for Biological Network Modeling, Protein Interactions, Metabolism, and Cell Signaling Simulation

class CellSignalingSimulationAgent:
   def run(self, df_signal: pd.DataFrame) -> AgentResult:
       peak_receptor = float(df_signal["receptor_active"].max())
       peak_kinase = float(df_signal["kinase_active"].max())
       peak_tf = float(df_signal["tf_active"].max())


       t_receptor = float(df_signal.loc[df_signal["receptor_active"].idxmax(), "time"])
       t_kinase = float(df_signal.loc[df_signal["kinase_active"].idxmax(), "time"])
       t_tf = float(df_signal.loc[df_signal["tf_active"].idxmax(), "time"])


       final_state = df_signal.iloc[-1].to_dict()


       summary = {
           "peak_receptor_activity": round(peak_receptor, 4),
           "peak_kinase_activity": round(peak_kinase, 4),
           "peak_tf_activity": round(peak_tf, 4),
           "time_to_peak_receptor": round(t_receptor, 4),
           "time_to_peak_kinase": round(t_kinase, 4),
           "time_to_peak_tf": round(t_tf, 4),
           "final_state": {k: round(float(v), 4) for k, v in final_state.items()},
       }
       return AgentResult(name="CellSignalingSimulationAgent", summary=summary)


class PrincipalInvestigatorAgent:
   def __init__(self, client, model=OPENAI_MODEL):
       self.client = client
       self.model = model


   def synthesize(self, results: List[AgentResult]) -> str:
       payload = {r.name: r.summary for r in results}


       prompt = f"""
You are a principal investigator in computational systems biology.


Given the outputs of four specialized AI agents:
1. gene regulatory network analysis
2. protein interaction prediction
3. metabolic pathway optimization
4. cell signaling simulation


Write a rigorous but readable report with these sections:
- Executive Summary
- Key Findings by Agent
- Cross-System Biological Interpretation
- Hypotheses Worth Testing in Wet Lab
- Model Limitations
- Next Computational Extensions


Use concise scientific language.
Do not fabricate datasets beyond what is shown.
When useful, connect regulation, signaling, metabolism, and protein interactions into a single systems biology story.


Agent outputs:
{json.dumps(payload, indent=2)}
"""


       try:
           resp = self.client.chat.completions.create(
               model=self.model,
               messages=[
                   {"role": "user", "content": prompt},
               ],
               temperature=0.4,
           )
           return resp.choices[0].message.content
       except Exception as e:
           return f"OpenAI synthesis failed: {e}"


genes, W = generate_gene_regulatory_network(n_genes=14, edge_prob=0.20)
X_expr = simulate_gene_expression(W, n_steps=80, noise=0.08)
grn_agent = GeneRegulatoryNetworkAgent()
grn_result = grn_agent.run(genes, W, X_expr)


proteins, prot_features, prot_families, prot_localization = generate_protein_features(n_proteins=40, feature_dim=10)
ppi_rows = generate_ppi_dataset(proteins, prot_features, prot_families, prot_localization)
ppi_agent = ProteinInteractionPredictionAgent()
ppi_result = ppi_agent.run(ppi_rows)


metabolites, reactions = generate_metabolic_network()
met_agent = MetabolicOptimizationAgent()
met_result, met_trace = met_agent.run(reactions, oxygen_budget=3.5, substrate_budget=4.2)


df_signal = simulate_cell_signaling(T=220, dt=0.05, ligand_level=1.2)
sig_agent = CellSignalingSimulationAgent()
sig_result = sig_agent.run(df_signal)


all_results = [grn_result, ppi_result, met_result, sig_result]


for r in all_results:
   pretty(r.name, json.dumps(r.summary, indent=2))


fig = plt.figure(figsize=(18, 14))


ax1 = plt.subplot(2, 2, 1)
im = ax1.imshow(W, cmap="coolwarm", aspect="auto")
ax1.set_title("Gene Regulatory Weight Matrix")
ax1.set_xticks(range(len(genes)))
ax1.set_yticks(range(len(genes)))
ax1.set_xticklabels(genes, rotation=90)
ax1.set_yticklabels(genes)
plt.colorbar(im, ax=ax1, fraction=0.046, pad=0.04)


ax2 = plt.subplot(2, 2, 2)
for i in range(min(6, X_expr.shape[1])):
   ax2.plot(X_expr[:, i], label=genes[i])
ax2.set_title("Sample Gene Expression Dynamics")
ax2.set_xlabel("Time step")
ax2.set_ylabel("Expression")
ax2.legend(loc="upper right", fontsize=8)


ax3 = plt.subplot(2, 2, 3)
ax3.plot(df_signal["time"], df_signal["receptor_active"], label="Receptor")
ax3.plot(df_signal["time"], df_signal["kinase_active"], label="Kinase")
ax3.plot(df_signal["time"], df_signal["tf_active"], label="Transcription Factor")
ax3.plot(df_signal["time"], df_signal["phosphatase"], label="Phosphatase")
ax3.set_title("Cell Signaling Simulation")
ax3.set_xlabel("Time")
ax3.set_ylabel("Activity")
ax3.legend()


ax4 = plt.subplot(2, 2, 4)
ax4.plot(met_trace)
ax4.set_title("Metabolic Search Objective Trace")
ax4.set_xlabel("Iteration")
ax4.set_ylabel("Objective score")


plt.tight_layout()
plt.show()


G_grn = nx.DiGraph()
for g in genes:
   G_grn.add_node(g)
for i in range(len(genes)):
   for j in range(len(genes)):
       if abs(W[i, j]) > 0.4:
           G_grn.add_edge(genes[i], genes[j], weight=W[i, j])


plt.figure(figsize=(10, 8))
pos = nx.spring_layout(G_grn, seed=42)
edge_colors = ["green" if G_grn[u][v]["weight"] > 0 else "red" for u, v in G_grn.edges()]
nx.draw_networkx(G_grn, pos, with_labels=True, node_size=900, font_size=9, arrows=True, edge_color=edge_colors)
plt.title("Gene Regulatory Network Graph (green=activation, red=repression)")
plt.axis("off")
plt.show()


top_ppi = ppi_result.summary["top_predicted_interactions"][:12]
G_ppi = nx.Graph()
for row in top_ppi:
   a, b, p = row["protein_a"], row["protein_b"], row["pred_prob"]
   G_ppi.add_edge(a, b, weight=p)


plt.figure(figsize=(10, 8))
pos = nx.spring_layout(G_ppi, seed=7)
widths = [2 + 4 * G_ppi[u][v]["weight"] for u, v in G_ppi.edges()]
nx.draw_networkx(G_ppi, pos, with_labels=True, node_size=1000, font_size=9, width=widths)
plt.title("Top Predicted Protein Interaction Subnetwork")
plt.axis("off")
plt.show()


grn_table = pd.DataFrame(grn_result.summary["most_dynamic_genes"])
ppi_table = pd.DataFrame(ppi_result.summary["top_predicted_interactions"])
met_table = pd.DataFrame(met_result.summary["dominant_reactions"])
sig_table = pd.DataFrame([sig_result.summary])


pretty("Most Dynamic Genes", grn_table.to_string(index=False))
pretty("Top Predicted PPIs", ppi_table.to_string(index=False))
pretty("Dominant Metabolic Reactions", met_table.to_string(index=False))


pi_agent = PrincipalInvestigatorAgent(client=client, model=OPENAI_MODEL)
final_report = pi_agent.synthesize(all_results)


pretty("OPENAI SYSTEMS BIOLOGY REPORT", final_report)


artifact = {
   "grn": grn_result.summary,
   "ppi": ppi_result.summary,
   "metabolic": met_result.summary,
   "signaling": sig_result.summary,
   "llm_report": final_report,
}


with open("bio_agents_tutorial_results.json", "w") as f:
   json.dump(artifact, f, indent=2)


print("nSaved results to: bio_agents_tutorial_results.json")