Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import operator | |
| from typing import TypedDict, Annotated | |
| import gradio as gr | |
| from gradio_client import Client, handle_file | |
| from neo4j import GraphDatabase | |
| from langchain_openai import ChatOpenAI | |
| from langchain_core.messages import HumanMessage | |
| from langchain_core.tools import tool | |
| from langchain_core.documents import Document | |
| from langchain_experimental.graph_transformers import LLMGraphTransformer | |
| from langgraph.graph import StateGraph, END | |
| # ========================================== | |
| # 1. CONFIGURATION | |
| # ========================================== | |
| AZURE_OPENAI_API_KEY = os.environ.get("AZURE_OPENAI_API_KEY") | |
| AZURE_OPENAI_API_BASE = os.environ.get("AZURE_OPENAI_API_BASE_URL") | |
| AZURE_DEPLOYMENT_NAME = os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME") | |
| NEO4J_URI = os.environ.get("NEO4J_URI") | |
| NEO4J_USERNAME = os.environ.get("NEO4J_USERNAME", "neo4j") | |
| NEO4J_PASSWORD = os.environ.get("NEO4J_PASSWORD") | |
| # ========================================== | |
| # 2. SETUP | |
| # ========================================== | |
| llm = ChatOpenAI( | |
| api_key=AZURE_OPENAI_API_KEY, | |
| base_url=AZURE_OPENAI_API_BASE, | |
| model=AZURE_DEPLOYMENT_NAME, | |
| temperature=0 | |
| ) | |
| try: ChatOpenAI.model_rebuild() | |
| except: pass | |
| def ingest_data(): | |
| if not NEO4J_URI or not NEO4J_PASSWORD: | |
| print("⚠️ Neo4j Credentials missing. Skipping ingestion.") | |
| return | |
| print("--- 🔄 Initializing Knowledge Graph ---") | |
| try: | |
| driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) | |
| prime_kg_data = [ | |
| ("Ischemic Stroke", "associated_with", "Hypertension"), | |
| ("Ischemic Stroke", "associated_with", "Atrial Fibrillation"), | |
| ("Ischemic Stroke", "treated_by", "Alteplase"), | |
| ("Ischemic Stroke", "treated_by", "Aspirin"), | |
| ("Ischemic Stroke", "treated_by", "Clopidogrel"), | |
| ("Alteplase", "mechanism", "Fibrinolytic Agent"), | |
| ("Alteplase", "contraindicated_in", "Active Bleeding"), | |
| ("Intracranial Hemorrhage", "associated_with", "Hypertension"), | |
| ("Intracranial Hemorrhage", "associated_with", "Cerebral Amyloid Angiopathy"), | |
| ("Intracranial Hemorrhage", "treated_by", "Mannitol"), | |
| ("Intracranial Hemorrhage", "treated_by", "Nicardipine"), | |
| ("Mannitol", "mechanism", "Osmotic Diuretic"), | |
| ("Ischemic Stroke", "presents_as", "Hemiparesis"), | |
| ("Ischemic Stroke", "presents_as", "Aphasia"), | |
| ("Intracranial Hemorrhage", "presents_as", "Severe Headache"), | |
| ("Intracranial Hemorrhage", "presents_as", "Vomiting") | |
| ] | |
| query = """ | |
| UNWIND $data as row | |
| MERGE (h:MedicalEntity {name: row[0]}) | |
| MERGE (t:MedicalEntity {name: row[2]}) | |
| MERGE (h)-[:PRIMEKG_REL {type: row[1]}]->(t) | |
| """ | |
| with driver.session() as session: | |
| session.run("MATCH (n) DETACH DELETE n") | |
| session.run(query, data=prime_kg_data) | |
| aha_guideline_text = """ | |
| For patients eligible for intravenous fibrinolytic therapy, Alteplase (0.9 mg/kg) is recommended within 3 hours of onset of ischemic stroke. | |
| Mechanical Thrombectomy is indicated for patients with large vessel occlusion (LVO) in the anterior circulation within 6 hours. | |
| For patients with acute Intracranial Hemorrhage, patients with systolic blood pressure >220 mmHg should be treated with aggressive blood pressure reduction. | |
| Aspirin is recommended for secondary prevention 24 to 48 hours after thrombolysis. | |
| """ | |
| docs = [Document(page_content=aha_guideline_text)] | |
| llm_transformer = LLMGraphTransformer(llm=llm) | |
| graph_documents = llm_transformer.convert_to_graph_documents(docs) | |
| with driver.session() as session: | |
| for doc in graph_documents: | |
| for node in doc.nodes: | |
| session.run("MERGE (n:GuidelineEntity {id: $id, type: $type})", id=node.id, type=node.type) | |
| for edge in doc.relationships: | |
| session.run(""" | |
| MATCH (s:GuidelineEntity {id: $source}) | |
| MATCH (t:GuidelineEntity {id: $target}) | |
| MERGE (s)-[:GUIDELINE_REL {type: $type}]->(t) | |
| """, source=edge.source.id, target=edge.target.id, type=edge.type) | |
| driver.close() | |
| print("✅ Knowledge Graph Built Successfully.") | |
| except Exception as e: | |
| print(f"❌ KG Initialization Failed: {e}") | |
| ingest_data() | |
| # ========================================== | |
| # 3. TOOLS | |
| # ========================================== | |
| def detect_stroke_type(image_path: str) -> str: | |
| """Analyzes a patient's CT scan using the remote Vision AI model.""" | |
| HF_SPACE_URL = "Siyad798/stroke-diagnosis-ai" | |
| print(f"--- 🧠 Calling Vision AI for {image_path} ---") | |
| if not os.path.exists(image_path): return json.dumps({"error": f"File not found locally: {image_path}"}) | |
| try: | |
| client = Client(HF_SPACE_URL) | |
| result = client.predict(files=[handle_file(image_path)], api_name="/predict") | |
| ai_output = {} | |
| if isinstance(result, (list, tuple)) and len(result) > 1: | |
| json_data = result[1] | |
| if isinstance(json_data, dict): ai_output = json_data | |
| elif isinstance(json_data, str) and os.path.exists(json_data): | |
| with open(json_data, 'r') as f: ai_output = json.load(f) | |
| elif isinstance(json_data, str): | |
| try: ai_output = json.loads(json_data) | |
| except: ai_output = {"raw_output": json_data} | |
| else: ai_output = result | |
| if not ai_output: ai_output = {"error": "Empty response", "diagnosis": "Unknown"} | |
| return json.dumps(ai_output) | |
| except Exception as e: return json.dumps({"error": f"Inference Failed: {str(e)}", "diagnosis": "Unknown"}) | |
| def get_treatment_guideline(stroke_type: str) -> str: | |
| """Queries AHA/ASA Guidelines from Neo4j Knowledge Graph.""" | |
| search_term = "Ischemic Stroke" if "Ischemic" in stroke_type else "Intracranial Hemorrhage" | |
| print(f"--- 🕸️ Querying KG for: {search_term} ---") | |
| if not NEO4J_URI or not NEO4J_PASSWORD: return json.dumps({"error": "Neo4j credentials not configured."}) | |
| try: | |
| driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) | |
| with driver.session() as session: | |
| q1 = "MATCH (s:StrokeType {name: $st})-[r:INDICATED_FOR]->(t:Treatment) RETURN t.name as Treatment, t.class as Class, t.dosage as Dosage, r.criteria as Criteria" | |
| treatments = session.run(q1, st=search_term).data() | |
| q2 = "MATCH (c:ClinicalCondition)-[:CONTRAINDICATES]->(t:Treatment)<-[:INDICATED_FOR]-(s:StrokeType {name: $st}) RETURN t.name as Drug, c.name as Contraindication" | |
| safety = session.run(q2, st=search_term).data() | |
| driver.close() | |
| return json.dumps({"Guidelines_Source": "AHA/ASA 2019", "Recommended_Treatments": treatments, "Safety_Warnings": safety}) | |
| except Exception as e: return json.dumps({"error": str(e)}) | |
| # ========================================== | |
| # 4. AGENT WORKFLOW | |
| # ========================================== | |
| class AgentState(TypedDict): | |
| messages: Annotated[list, operator.add] | |
| image_path: str | |
| stroke_diagnosis: dict | |
| guidelines: dict | |
| final_report: str | |
| def detector_node(state: AgentState): | |
| print("--- Node: Stroke Detector ---") | |
| result_json = detect_stroke_type.invoke(state["image_path"]) | |
| diagnosis = json.loads(result_json) | |
| msg = f"Imaging Analysis Complete. Diagnosis: {diagnosis.get('diagnosis')}." | |
| return {"messages": [HumanMessage(content=msg)], "stroke_diagnosis": diagnosis} | |
| def advisor_node(state: AgentState): | |
| print("--- Node: Guideline Advisor ---") | |
| diagnosis = state["stroke_diagnosis"].get("diagnosis", "Unknown") | |
| guidelines_json = get_treatment_guideline.invoke(diagnosis) | |
| guidelines = json.loads(guidelines_json) | |
| return {"messages": [HumanMessage(content=f"Guidelines Retrieved for {diagnosis}.")], "guidelines": guidelines} | |
| def synthesizer_node(state: AgentState): | |
| print("--- Node: Report Synthesizer ---") | |
| prompt = f""" | |
| You are a Stroke Neurologist Assistant. Write a clinical report based on: | |
| 1. **Imaging AI Result**: {json.dumps(state['stroke_diagnosis'], indent=2)} | |
| 2. **Clinical Guidelines**: {json.dumps(state['guidelines'], indent=2)} | |
| 3. **Patient Context**: {state['messages'][0].content} | |
| **Structure:** | |
| - Diagnosis & Severity | |
| - Imaging Findings | |
| - Recommended Interventions (Cite Evidence) | |
| - Safety Checks & Contraindications | |
| """ | |
| response = llm.invoke(prompt) | |
| return {"final_report": response.content} | |
| workflow = StateGraph(AgentState) | |
| workflow.add_node("detector", detector_node) | |
| workflow.add_node("advisor", advisor_node) | |
| workflow.add_node("synthesizer", synthesizer_node) | |
| workflow.set_entry_point("detector") | |
| workflow.add_edge("detector", "advisor") | |
| workflow.add_edge("advisor", "synthesizer") | |
| workflow.add_edge("synthesizer", END) | |
| app = workflow.compile() | |
| # ========================================== | |
| # 5. GRADIO UI | |
| # ========================================== | |
| def predict_stroke_report(uploaded_file, patient_context_input): | |
| if not uploaded_file: return "Please upload a CT scan file." | |
| image_path = uploaded_file.name | |
| inputs = {"messages": [HumanMessage(content=patient_context_input)], "image_path": image_path} | |
| final_report_content = "Processing..." | |
| try: | |
| for event in app.stream(inputs, stream_mode="values"): pass | |
| final_state = event | |
| final_report_content = final_state.get("final_report", "Error generating report.") | |
| except Exception as e: final_report_content = f"Critical Error: {str(e)}" | |
| return final_report_content | |
| # Safe initialization without arguments | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 🧠 AI Stroke Diagnosis & Guideline Agent") | |
| gr.Markdown("Integrates Vision AI (DenseNet) + Knowledge Graph (Neo4j) + LLM (Azure OpenAI)") | |
| with gr.Row(): | |
| with gr.Column(): | |
| file_input = gr.File(label="Upload CT (ZIP/DCM)", file_types=[".zip", ".dcm"]) | |
| patient_context = gr.Textbox(label="Patient Context", lines=4, value="Patient 55M, onset 2 hours ago. BP 150/90. No history of bleeding.") | |
| submit_btn = gr.Button("Generate Report", variant="primary") | |
| with gr.Column(): | |
| report_output = gr.Markdown(label="Clinical Report") | |
| submit_btn.click(predict_stroke_report, inputs=[file_input, patient_context], outputs=report_output) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |