File size: 3,874 Bytes
caf16cf c1e2ddb caf16cf c1e2ddb caf16cf d462865 caf16cf d462865 caf16cf d462865 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import os
from typing import Literal
from langgraph.graph import StateGraph, START
from langgraph.types import Command
from dotenv import load_dotenv
import gradio as gr
from langchain_huggingface import HuggingFaceEndpoint
# Load environment variables
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")
# Define HuggingFaceEndpoint
llm = HuggingFaceEndpoint(
repo_id="mistralai/Mistral-7B-Instruct-v0.3",
huggingfacehub_api_token=HF_TOKEN.strip(),
temperature=0.7,
max_new_tokens=150,
)
# Define state
class State(dict):
disaster_type: str
severity: str
resources: list[str]
recommendation: str
# Create the graph
builder = StateGraph(State)
# Define nodes
def disaster_type_node(state: State) -> Command[Literal["severity_node"]]:
"""Set disaster type and proceed to severity selection."""
return Command(update={"disaster_type": state["disaster_type"]}, goto="severity_node")
def severity_node(state: State) -> Command[Literal["resources_node"]]:
"""Set severity level and proceed to resource selection."""
return Command(update={"severity": state["severity"]}, goto="resources_node")
def resources_node(state: State) -> Command[Literal["generate_response_node"]]:
"""Set available resources and proceed to generating response."""
return Command(update={"resources": state["resources"]}, goto="generate_response_node")
def generate_response_node(state: State):
"""Generate a disaster response plan based on inputs."""
resources = ", ".join(state["resources"])
prompt = f"""
Disaster: {state['disaster_type']}
Severity: {state['severity']}
Resources: {resources}
You are an emergency response assistant. Provide a detailed response plan for the given situation in 100 words.
"""
response = llm(prompt)
return {"recommendation": response}
# Add nodes to the graph
builder.add_edge(START, "disaster_type_node")
builder.add_node("disaster_type_node", disaster_type_node)
builder.add_node("severity_node", severity_node)
builder.add_node("resources_node", resources_node)
builder.add_node("generate_response_node", generate_response_node)
# Compile the graph
graph = builder.compile()
# Gradio Interface
def process_disaster(disaster_type, severity, resources):
"""Run the disaster response graph with user inputs."""
state = {"disaster_type": disaster_type, "severity": severity, "resources": resources}
result = graph.invoke(state)
# Generate Mermaid code for graph visualization
mermaid_code = graph.get_graph().draw_mermaid()
# Return the response plan, placeholder markdown, the link, and the Mermaid visualization
return (
result["recommendation"],
"## Mermaid Graph",
"Check out this [mermaid link](https://mermaid.live/) to display a graph with the following data",
f"```mermaid\n{mermaid_code}\n```"
)
iface = gr.Interface(
fn=process_disaster,
inputs=[
gr.Dropdown(
["Flood", "Fire", "Earthquake", "Tornado", "Other"],
label="Select Disaster Type",
),
gr.Dropdown(
["Low", "Medium", "High", "Critical"],
label="Select Severity Level",
),
gr.CheckboxGroup(
[
"Water and Food Supplies",
"Medical Aid",
"Shelter and Evacuation",
"Search and Rescue Teams",
"Other",
],
label="Select Available Resources (Multi-Select)",
),
],
outputs=[
gr.Textbox(label="Generated Disaster Response Plan"),
gr.Markdown(), # Placeholder for the "Mermaid Graph" heading
gr.Markdown(), # Placeholder for the link text
gr.Markdown(label="Mermaid Graph Visualization") # Mermaid visualization
],
title="Disaster Response Assistant",
)
iface.launch() |