Chinez-dev commited on
Commit
6c7dbf1
·
verified ·
1 Parent(s): 26747d1

Create agent.py

Browse files
Files changed (1) hide show
  1. agent.py +203 -0
agent.py ADDED
@@ -0,0 +1,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
2
+ import datetime
3
+ import requests
4
+ import pytz
5
+ import yaml
6
+ from tools.final_answer import FinalAnswerTool
7
+
8
+ from Gradio_UI import GradioUI
9
+
10
+ # Below is an example of a tool that does nothing. Amaze us with your creativity !
11
+ # @tool
12
+ # def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
13
+ # #Keep this format for the description / args / args description but feel free to modify the tool
14
+ # """A tool that does nothing yet
15
+ # Args:
16
+ # arg1: the first argument
17
+ # arg2: the second argument
18
+ # """
19
+ # return "What magic will you build ?"
20
+
21
+
22
+ def __init__(self):
23
+ self.patient_data = {}
24
+
25
+ @tool
26
+ def collect_symptoms(patient_id: str, symptoms: str) -> str:
27
+ """
28
+ Nurse tool to collect symptoms and ask follow-up questions.
29
+ Args:
30
+ patient_id: Unique identifier for the patient.
31
+ symptoms: Initial symptoms provided by the patient.
32
+ Returns:
33
+ Follow-up questions based on the symptoms.
34
+ """
35
+ try:
36
+ if not patient_id or not symptoms:
37
+ raise ValueError("Patient ID and symptoms must be provided.")
38
+
39
+ self.patient_data[patient_id] = {"symptoms": symptoms, "additional_info": {}}
40
+
41
+ questions = [
42
+ "How long have you had these symptoms?",
43
+ "Do you have any allergies?",
44
+ "Are you taking any medications?",
45
+ "Have you experienced these symptoms before?",
46
+ "Have you had any recent illnesses?",
47
+ "Have you noticed any other unusual changes?",
48
+ "What is your medical history related to these symptoms?"
49
+
50
+ ]
51
+
52
+ return f"Nurse: I have noted the symptoms ({symptoms}). Here are follow-up questions:\n" + "\n".join(questions)
53
+
54
+ except ValueError as e:
55
+ return f"Error: {str(e)}"
56
+ except Exception as e:
57
+ return f"Unexpected Error: {str(e)}"
58
+
59
+ @tool
60
+ def diagnose_patient(patient_id: str) -> str:
61
+ """
62
+ Doctor tool to diagnose the patient based on symptoms.
63
+
64
+ Args:
65
+ patient_id: Unique identifier for the patient.
66
+
67
+ Returns:
68
+ Diagnosis and recommendations.
69
+ """
70
+ try:
71
+ if patient_id not in CodeAgent.patient_data:
72
+ raise ValueError("No symptoms found. Nurse must collect symptoms first.")
73
+
74
+ symptoms = CodeAgent.patient_data[patient_id]["symptoms"]
75
+ diagnosis = CodeAgent.fetch_diagnosis(symptoms)
76
+ medication = CodeAgent.fetch_medication(symptoms)
77
+ advice = CodeAgent.fetch_treatment_advice(symptoms)
78
+
79
+ return (f"Doctor: Based on the symptoms: {symptoms},\n"
80
+ f"Diagnosis: {diagnosis}\n"
81
+ f"Medication: {medication}\n"
82
+ f"Advice: {advice}")
83
+
84
+ except ValueError as e:
85
+ return f"Error: {str(e)}"
86
+ except Exception as e:
87
+ return f"Unexpected Error: {str(e)}"
88
+
89
+
90
+ @tool
91
+ def fetch_diagnosis(symptoms: str) -> str:
92
+ """
93
+ AI tool to retrieve a diagnosis based on symptoms.
94
+ Args:
95
+ symptoms: The symptoms provided by the patient.
96
+ Returns:
97
+ Detailed Diagnosis information.
98
+ """
99
+ search_query = f" A detailed medical diagnosis for this {symptoms}"
100
+ return search_tool(search_query)
101
+ try:
102
+ if not symptoms:
103
+ raise ValueError("Symptoms must be provided.")
104
+
105
+ return f"AI Diagnosis: Potential cause of {symptoms} found."
106
+
107
+ except ValueError as e:
108
+ return f"Error: {str(e)}"
109
+ except Exception as e:
110
+ return f"Unexpected Error: {str(e)}"
111
+
112
+ @tool
113
+ def fetch_medication(symptoms: str) -> str:
114
+ """
115
+ AI tool to suggest medications based on symptoms.
116
+ Args:
117
+ symptoms: The symptoms provided by the patient.
118
+ Returns:
119
+ Suggested medication.
120
+ """
121
+ search_query = f" A detailed suggested medication for this {symptoms}"
122
+ return search_tool(search_query)
123
+ try:
124
+ if not symptoms:
125
+ raise ValueError("Symptoms must be provided.")
126
+
127
+ return f"Suggested medication for {symptoms} found."
128
+
129
+ except ValueError as e:
130
+ return f"Error: {str(e)}"
131
+ except Exception as e:
132
+ return f"Unexpected Error: {str(e)}"
133
+
134
+ @tool
135
+ def fetch_treatment_advice(symptoms: str) -> str:
136
+ """
137
+ AI tool to provide treatment recommendations.
138
+ Args:
139
+ symptoms: The symptoms provided by the patient.
140
+ Returns:
141
+ Recommended treatment and advice.
142
+ """
143
+ search_query = f" A detailed recommended treatment and advice for this {symptoms}"
144
+ return search_tool(search_query)
145
+ try:
146
+ if not symptoms:
147
+ raise ValueError("Symptoms must be provided.")
148
+
149
+ return f"Treatment advice for {symptoms} retrieved."
150
+
151
+ except ValueError as e:
152
+ return f"Error: {str(e)}"
153
+ except Exception as e:
154
+ return f"Unexpected Error: {str(e)}"
155
+
156
+ search_tool = DuckDuckGoSearchTool()
157
+
158
+ @tool
159
+ def get_current_time_in_timezone(timezone: str) -> str:
160
+ """A tool that fetches the current local time in a specified timezone.
161
+ Args:
162
+ timezone: A string representing a valid timezone (e.g., 'America/New_York').
163
+ """
164
+ try:
165
+ # Create timezone object
166
+ tz = pytz.timezone(timezone)
167
+ # Get current time in that timezone
168
+ local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
169
+ return f"The current local time in {timezone} is: {local_time}"
170
+ except Exception as e:
171
+ return f"Error fetching time for timezone '{timezone}': {str(e)}"
172
+
173
+
174
+ final_answer = FinalAnswerTool()
175
+
176
+ # If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
177
+ # model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'
178
+
179
+ model = HfApiModel(
180
+ max_tokens=2096,
181
+ temperature=0.5,
182
+ model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
183
+ custom_role_conversions=None,
184
+ )
185
+
186
+
187
+ # Import tool from Hub
188
+ image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
189
+
190
+ with open("prompts.yaml", 'r') as stream:
191
+ prompt_templates = yaml.safe_load(stream)
192
+
193
+ agent = CodeAgent(
194
+ model=model,
195
+ tools=[final_answer, collect_symptoms, diagnose_patient, get_current_time_in_timezone], ## add your tools here (don't remove final answer)
196
+ max_steps=6,
197
+ verbosity_level=1,
198
+ grammar=None,
199
+ planning_interval=None,
200
+ name="MedicalDoctor",
201
+ description="A tool that simulates the roles of a medical doctor and a nurse in gathering symptoms, diagnosing, and providing recommendations",
202
+ prompt_templates=prompt_templates
203
+ )