✨ added functions and prompts for function calling

tracex/extraction/prototype/function_calls.py

@@ -1,5 +1,5 @@
 TOOLS = [
-        {
+    {
         "type": "function",
         "function": {
                 "name": "add_start_dates",
@@ -11,7 +11,7 @@
                             "type": "array",
                             "items": {
                                 "type": "string",
-                                "description": "an start date",
+                                "description": "a start date",
                             }
                         },
                     },
@@ -31,7 +31,67 @@
                             "type": "array",
                             "items": {
                                 "type": "string",
-                                "description": "an end date",
+                                "description": "a end date",
+                            }
+                        },
+                    },
+                    "required": ["output"]
+                },
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+                "name": "add_duration",
+                "description": "this function extract the duration",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "output": {
+                            "type": "array",
+                            "items": {
+                                "type": "string",
+                                "description": "a duration",
+                            }
+                        },
+                    },
+                    "required": ["output"]
+                },
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+                "name": "add_event_type",
+                "description": "this function extract the event type",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "output": {
+                            "type": "array",
+                            "items": {
+                                "type": "string",
+                                "description": "an event type",
+                            }
+                        },
+                    },
+                    "required": ["output"]
+                },
+        }
+    },
+    {
+        "type": "function",
+        "function": {
+                "name": "add_location",
+                "description": "this function extract the location",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "output": {
+                            "type": "array",
+                            "items": {
+                                "type": "string",
+                                "description": "a location",
                             }
                         },
                     },

tracex/extraction/prototype/input_handling.py

@@ -7,28 +7,28 @@
 from . import prompts as p
 
 
-def convert_text_to_csv(inp):
+def convert_text_to_csv(text):
     """Converts the input to CSV with intermediate steps."""
     steps = str(7)
     print("Converting Data: Summarizing the text. (1/" + steps + ")", end="\r")
-    bulletpoints = convert_text_to_bulletpoints(inp)
+    bulletpoints = convert_text_to_bulletpoints(text)
     print(
         "Converting Data: Extracting start date information. (2/" + steps + ")",
         end="\r",
     )
     u.pause_between_queries()
-    bulletpoints_start = add_start_dates(inp, bulletpoints)
+    bulletpoints_start = add_start_dates(text, bulletpoints)
     print(
         "Converting Data: Extracting end date information. (3/" + steps + ")   ",
         end="\r",
     )
     u.pause_between_queries()
-    bulletpoints_end = add_end_dates(inp, bulletpoints_start)
+    bulletpoints_end = add_end_dates(text, bulletpoints_start)
     print(
         "Converting Data: Extracting duration information. (4/" + steps + ") ", end="\r"
     )
     u.pause_between_queries()
-    bulletpoints_duration = add_durations(inp, bulletpoints_end)
+    bulletpoints_duration = add_durations(text, bulletpoints_end)
     print(
         "Converting Data: Extracting event types. (5/" + steps + ")          ", end="\r"
     )
@@ -47,11 +47,11 @@ def convert_text_to_csv(inp):
     return output_path
 
 
-def convert_text_to_bulletpoints(inp):
+def convert_text_to_bulletpoints(text):
     """Converts the input text to bulletpoints."""
     messages = [
         {"role": "system", "content": p.TXT_TO_BULLETPOINTS_CONTEXT},
-        {"role": "user", "content": p.TXT_TO_BULLETPOINTS_PROMPT + inp},
+        {"role": "user", "content": p.TXT_TO_BULLETPOINTS_PROMPT + text},
         {"role": "assistant", "content": p.TXT_TO_BULLETPOINTS_ANSWER},
     ]
     bulletpoints = u.query_gpt(messages)
@@ -90,10 +90,10 @@ def add_start_dates(text, df):
         output = u.query_gpt(messages)
 
         fc_message = [
-            {"role": "system", "content": p.START_DATE_CONTEXT},
+            {"role": "system", "content": p.FC_START_DATE_CONTEXT},
             {
                 "role": "user",
-                "content": p.START_DATE_FUNCTION_CALL
+                "content": p.FC_START_DATE_PROMPT
                 + "The text: "
                 + output
             },
@@ -135,10 +135,10 @@ def add_end_dates(text, df):
         output = u.query_gpt(messages)
 
         fc_message = [
-            {"role": "system", "content": p.END_DATE_CONTEXT},
+            {"role": "system", "content": p.FC_END_DATE_CONTEXT},
             {
                 "role": "user",
-                "content": p.END_DATE_FUNCTION_CALL
+                "content": p.FC_END_DATE_PROMPT
                 + "The text: "
                 + output
             },
@@ -157,13 +157,13 @@ def add_end_dates(text, df):
     return df
 
 
-def add_durations(inp, bulletpoints_start):
+def add_durations(text, bulletpoints_start):
     """Adds durations to the bulletpoints."""
     messages = [
         {"role": "system", "content": p.BULLETPOINTS_DURATION_CONTEXT},
         {
             "role": "user",
-            "content": p.BULLETPOINTS_DURATION_PROMPT + inp + "\n" + bulletpoints_start,
+            "content": p.BULLETPOINTS_DURATION_PROMPT + text + "\n" + bulletpoints_start,
         },
         {"role": "assistant", "content": p.BULLETPOINTS_DURATION_ANSWER},
     ]

tracex/extraction/prototype/prompts.py

@@ -142,7 +142,13 @@ def life_circumstances_prompt(sex):
     'experiencing mild symptoms' you should return '20200401T0000'.
     If the bullet point is 'testing positive for Covid19' you should return '20200405T0000'.
 """
+FC_START_DATE_CONTEXT = """
+You are an expert in extracting information. You easily detect the start dates and extract them as they are without changing any format.
+"""
 
+FC_START_DATE_PROMPT = """
+Please extract the following start date of the text without changing the given date format:
+"""
 
 # Adding of a end date to every bullet point
 BULLETPOINTS_END_DATE_CONTEXT = """
@@ -167,19 +173,11 @@ def life_circumstances_prompt(sex):
 
     The text 'In the next time I made sure to improve my mental wellbeing.' and the bulletpoint 'improving mental wellbeing' with the start date '20210610T0000', you should output '20210710T0000'.
 """
-END_DATE_CONTEXT = """
-You are an expert in extracting information. You easily detect the end dates and extract the as they are without changing any format.
-"""
-
-END_DATE_FUNCTION_CALL = """
-Please extract the following end date of the text without changing the given date format:
-"""
-
-START_DATE_CONTEXT = """
-You are an expert in extracting information. You easily detect the end dates and extract the as they are without changing any format.
+FC_END_DATE_CONTEXT = """
+You are an expert in extracting information. You easily detect the end dates and extract them as they are without changing any format.
 """
 
-START_DATE_FUNCTION_CALL = """
+FC_END_DATE_PROMPT = """
 Please extract the following end date of the text without changing the given date format:
 """
 
@@ -207,6 +205,13 @@ def life_circumstances_prompt(sex):
     '(visiting doctor's, 20200405T0000, 02:00:00), (testing positive for Covid19, 20200405T0000, 00:00:00), (getting hospitalized, 20200405T0000, 03:00:00)'.
     The text 'In the next time I made sure to improve my mental wellbeing.' and the bulletpoint '(improving mental wellbeing, 20210610T0000)' could be updated to '(improving mental wellbeing, 20210610T0000, 720:00:00)'.
 """
+FC_DURATION_CONTEXT = """
+You are an expert in extracting information. You easily detect durations and extract them as they are without changing any format.
+"""
+
+FC_DURATION_PROMPT = """
+Please extract the following duration of the text without changing the given date format:
+"""
 
 
 # Adding of a event type to every bullet point
@@ -227,6 +232,13 @@ def life_circumstances_prompt(sex):
     (getting hospitalized, 20200405T0000, 03:00:00, Hospital admission)'.
 """
 
+FC_EVENT_TYPE_CONTEXT = """
+You are an expert in extracting information. You easily detect event type and extract them as they are without changing any format.
+"""
+
+FC_EVENT_TYPE_PROMPT = """
+Please extract the following event type of the text without changing the given date format:
+"""
 
 # Adding of a location type to every bullet point
 BULLETPOINTS_LOCATION_CONTEXT = """
@@ -247,3 +259,10 @@ def life_circumstances_prompt(sex):
     '(visiting doctor's, 20200405T0000, 02:00:00, Doctors Visit, Doctors), (testing positive for Covid19, 20200405T0000, 00:00:00, Diagnosis, Doctors),
     (getting hospitalized, 20200405T0000, 03:00:00, Hospital admission, Hospital)'.
 """
+FC_LOCATION_CONTEXT = """
+You are an expert in extracting information. You easily detect locations and extract them as they are without changing any format.
+"""
+
+FC_LOCATION_PROMPT = """
+Please extract the following location of the text without changing the given date format:
+"""