File size: 5,781 Bytes
3e595a5 7828fde c2d68af 7828fde 3e595a5 c2d68af 3e595a5 95243ee c2d68af 95243ee 90f1f18 95243ee f52aa71 95243ee 3e595a5 |
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 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
from modules.data_class import DataState
from modules.instrctions import MEDICAL_INTAKE_SYSINT, WELCOME_MSG
from modules.llm_in_use import get_llm
from modules.tools import patient_id, symptom, pain, medical_hist, family_hist, social_hist, review_system, pain_manage, functional, plan, confirm_data, get_data, clear_data, save_data
from datetime import date
from typing import Literal
from langgraph.graph import StateGraph, START, END
from langchain_core.messages.ai import AIMessage
llm = get_llm()
# Order-tools will be handled by the order node.
intake_tools = [patient_id, symptom, pain, medical_hist, family_hist, social_hist, review_system, pain_manage, functional, plan, confirm_data, get_data, clear_data, save_data]
# The LLM needs to know about all of the tools, so specify everything here.
llm_with_tools = llm.bind_tools(intake_tools)
def human_node(state: DataState) -> DataState:
"""Display the last model message to the user, and receive the user's input."""
last_msg = state["messages"][-1]
print("Model:", last_msg.content)
user_input = input("User: ")
# If it looks like the user is trying to quit, flag the conversation
# as over.
if user_input in {"q", "quit", "exit", "goodbye"}:
state["finished"] = True
return state | {"messages": [("user", user_input)]}
def maybe_exit_human_node(state: DataState) -> Literal["chatbot_healthassistant", "__end__"]:
"""Route to the chatbot, unless it looks like the user is exiting."""
if state.get("finished", False):
return END
else:
return "chatbot_healthassistant"
def chatbot_with_tools(state: DataState) -> DataState:
"""The chatbot with tools. A simple wrapper around the model's own chat interface."""
defaults = {"data": {"ID": {
"name": "",
"DOB": date(1900, 1, 1), # Default placeholder date
"gender": "",
"contact": "",
"emergency_contact": ""
},
"symptom": {
"main_symptom": "",
"symptom_length": ""
},
"pain": {
"pain_location": "",
"pain_side": "",
"pain_intensity": 0,
"pain_description": "",
"start_time": date(1900, 1, 1),
"radiation": False,
"triggers": "",
"symptom": ""
},
"medical_hist": {
"medical_condition": "",
"first_time": date(1900, 1, 1),
"surgery_history": "",
"medication": "",
"allergy": ""
},
"family_hist": {
"family_history": "",
},
"social_hist": {
"occupation": "",
"smoke": False,
"alcohol": False,
"drug": False,
"support_system": "",
"living_condition": "",
},
"review_system": {
"weight_change": "",
"fever": False,
"chill": False,
"night_sweats": False,
"sleep": "",
"gastrointestinal": "",
"urinary": "",
},
"pain_manage": {
"pain_medication": "",
"specialist": False,
"other_therapy": "",
"effectiveness": False,
},
"functional": {
"life_quality": "",
"limit_activity": "",
"mood": "",
},
"plan": {
"goal": "",
"expectation": "",
"alternative_treatment": "",
}
}, "finished": False}
if state["messages"]:
new_output = llm_with_tools.invoke([MEDICAL_INTAKE_SYSINT] + state["messages"])
else:
new_output = AIMessage(content=WELCOME_MSG)
# Set up some defaults if not already set, then pass through the provided state,
# overriding only the "messages" field.
return defaults | state | {"messages": [new_output]}
def maybe_route_to_tools(state: DataState) -> str:
"""Route between chat and tool nodes if a tool call is made."""
if not (msgs := state.get("messages", [])):
raise ValueError(f"No messages found when parsing state: {state}")
msg = msgs[-1]
if state.get("finished", False):
# When an order is placed, exit the app. The system instruction indicates
# that the chatbot should say thanks and goodbye at this point, so we can exit
# cleanly.
return END
elif hasattr(msg, "tool_calls") and len(msg.tool_calls) > 0:
# Route to `tools` node for any automated tool calls first.
if any(
tool["name"] for tool in msg.tool_calls
):
# return "datacreation"
# else:
return "documenting"
else:
return "patient"
|