updates

app.py

@@ -4,117 +4,46 @@ import os
 import time
 import gradio as gr
 from openai import OpenAI
-import xml.etree.ElementTree as ET
-import re
-import pandas as pd
 import prompts
 import traceback
 from io import StringIO
+import pandas as pd
+from typing import Dict, Any
+
+from typing import List, Optional
+from pydantic import BaseModel, Field
+from structures import ClinicalInfo
+
 
 client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
 
 model_name = "gpt-4o-2024-08-06"
-
+# import pdb; pdb.set_trace()
 try:
     demo = client.beta.assistants.create(
         name="Information Extractor",
-        instructions="Extract information from this note.",
+        instructions="Extract information from this note and return it as a JSON object.",
         model=model_name,
         tools=[{"type": "file_search"}],
     )
+
 except Exception as e:
     print(f"Error creating assistant: {str(e)}")
     raise
 
-def parse_xml_response(xml_string: str) -> pd.DataFrame:
-    """
-    Parse the XML response from the model and extract all fields into a dictionary,
-    then convert it to a pandas DataFrame with a nested index.
-    """
-    try:
-        # Extract only the XML content between the outermost tags
-        xml_content = re.findall(r'<[^>]+>.*?</[^>]+>', xml_string, re.DOTALL)
-        if not xml_content:
-            print("No valid XML content found.")
-            return pd.DataFrame()
-
-        # Wrap the content in a root element to ensure there's only one root
-        xml_string = f"<root>{''.join(xml_content)}</root>"
-
-        # Parse the XML
-        root = ET.fromstring(xml_string)
-
-        result = {}
-
-        for element in root:
-            tag = element.tag
-            if tag in ['patient_name', 'date_of_birth', 'sex', 'weight', 'date_of_death']:
-                result[tag] = {
-                    'reasoning': element.find('reasoning').text.strip() if element.find('reasoning') is not None else None,
-                    **{child.tag: child.text.strip() if child.text else None 
-                       for child in element if child.tag != 'reasoning'}
-                }
-            elif tag in ['traditional_chemo', 'other_cancer_treatments', 'other_conmeds']:
-                if tag not in result:
-                    result[tag] = []
-                reasoning = element.find('reasoning')
-                for item in element:
-                    if item.tag in ['drug', 'treatment', 'medication']:
-                        date_element = element.find('date')
-                        result[tag].append({
-                            'reasoning': reasoning.text.strip() if reasoning is not None else None,
-                            'name': item.text.strip() if item.text else None,
-                            'date': date_element.text.strip() if date_element is not None and date_element.text else None
-                        })
-            elif tag in ['surgery', 'surgery_outcome', 'metastasis_at_time_of_diagnosis']:
-                result[tag] = {
-                    'reasoning': element.find('reasoning').text.strip() if element.find('reasoning') is not None else None,
-                    **{child.tag: child.text.strip() if child.text else None 
-                       for child in element if child.tag != 'reasoning'}
-                }
-            elif tag == 'compounding_pharmacy':
-                result[tag] = {
-                    'reasoning': element.find('reasoning').text.strip() if element.find('reasoning') is not None else None,
-                    'pharmacy': element.find('pharmacy').text.strip() if element.find('pharmacy') is not None else None
-                }
-            elif tag == 'adverse_effects':
-                if tag not in result:
-                    result[tag] = []
-                effect = {
-                    'reasoning': element.find('reasoning').text.strip() if element.find('reasoning') is not None else None
-                }
-                for child in element:
-                    if child.tag != 'reasoning':
-                        effect[child.tag] = child.text.strip() if child.text else None
-                if effect:
-                    result[tag].append(effect)
-
-        # Convert to nested DataFrame
-        df_data = {}
-        for key, value in result.items():
-            if isinstance(value, dict):
-                for sub_key, sub_value in value.items():
-                    df_data[(key, '1', sub_key)] = [sub_value]
-            elif isinstance(value, list):
-                for i, item in enumerate(value):
-                    for sub_key, sub_value in item.items():
-                        df_data[(key, f"{i+1}", sub_key)] = [sub_value]
-            else:
-                df_data[(key, '1', '')] = [value]
+def parse_response(prompt):
+    chat_completion = client.beta.chat.completions.parse(
+        messages=[
+            {
+                "role": "user",
+                "content": prompt,
+            }
+        ],
+        model=model_name,
+        response_format=ClinicalInfo,
+    )
+    return chat_completion.choices[0].message.parsed.model_dump()
 
-        # Create multi-index DataFrame
-        df = pd.DataFrame(df_data)
-        df.columns = pd.MultiIndex.from_tuples(df.columns)
-        
-        return df
-    except ET.ParseError as e:
-        print(f"XML parsing error: {str(e)}")
-        print(f"Problematic XML content: {xml_string[:500]}...")  # Print first 500 chars of XML
-        return pd.DataFrame()
-    except Exception as e:
-        print(f"Error in parse_xml_response: {str(e)}")
-        print(f"Traceback: {traceback.format_exc()}")
-        return pd.DataFrame()
 
 def get_response(file_id, assistant_id, max_retries=3):
     for attempt in range(max_retries):
@@ -126,16 +55,24 @@ def get_response(file_id, assistant_id, max_retries=3):
                         "content": prompts.info_prompt,
                         "attachments": [
                             {"file_id": file_id, "tools": [{"type": "file_search"}]}
-                        ],
-                    }
+                        ],}
                 ]
             )
-            run = client.beta.threads.runs.create_and_poll(
-                thread_id=thread.id, assistant_id=assistant_id
+            # import pdb; pdb.set_trace()
+            run = client.beta.threads.runs.create(
+                thread_id=thread.id,
+                assistant_id=assistant_id,
+                instructions="Please provide your response as a valid JSON object.",
             )
+            run = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
+            while run.status != "completed":
+                time.sleep(1)
+                run = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
+            
             messages = list(
                 client.beta.threads.messages.list(thread_id=thread.id, run_id=run.id)
             )
+            
             assert len(messages) == 1, f"Expected 1 message, got {len(messages)}"
             message_content = messages[0].content[0].text
             annotations = message_content.annotations
@@ -150,6 +87,36 @@ def get_response(file_id, assistant_id, max_retries=3):
                 time.sleep(5)
             else:
                 raise Exception("Max retries reached. Unable to get response from the model.")
+            
+def clinical_info_to_dataframe(clinical_info: Dict[str, Any]) -> pd.DataFrame:
+    """
+    Convert ClinicalInfo dictionary to a DataFrame.
+    """
+    data = []
+    for field, value in clinical_info.items():
+        if isinstance(value, dict):
+            for sub_field, sub_value in value.items():
+                data.append({
+                    'Category': field,
+                    'Field': sub_field,
+                    'Value': str(sub_value)
+                })
+        elif isinstance(value, list):
+            for i, item in enumerate(value):
+                for sub_field, sub_value in item.items():
+                    data.append({
+                        'Category': f"{field}_{i+1}",
+                        'Field': sub_field,
+                        'Value': str(sub_value)
+                    })
+        elif value is None:
+            data.append({
+                'Category': field,
+                'Field': 'value',
+                'Value': 'None'
+            })
+    return pd.DataFrame(data)
+
 
 def process(file_content):
     try:
@@ -162,18 +129,18 @@ def process(file_content):
         message_file = client.files.create(file=open(file_name, "rb"), purpose="assistants")
 
         response = get_response(message_file.id, demo.id)  # This now includes retry logic
-        df = parse_xml_response(response)
+        response_prompt = f"Please parse the following response into the correct format: {response}"
+        clinical_info = parse_response(response_prompt)
+        df = clinical_info_to_dataframe(clinical_info)
         
         if df.empty:
             return "<p>No valid information could be extracted from the provided file.</p>"
 
-        # Transpose the DataFrame
-        df_transposed = df.T.reset_index()
-        df_transposed.columns = ['Category', 'Index', 'Field', 'Value']
-        df_transposed = df_transposed.sort_values(['Category', 'Index', 'Field'])
+        # Sort the DataFrame
+        df = df.sort_values(['Category', 'Field'])
 
         # Convert to HTML with some basic styling
-        html = df_transposed.to_html(index=False, classes='table table-striped table-bordered', escape=False)
+        html = df.to_html(index=False, classes='table table-striped table-bordered', escape=False)
         
         # Add some custom CSS for better readability
         html = f"""
@@ -227,8 +194,7 @@ def gradio_interface():
 def run_in_terminal():
     print("Clinical Note Information Extractor")
     print("This tool extracts key information from clinical notes in PDF format.")
-    print("Enter the path to your PDF file:")
-    file_path = input().strip()
+    file_path = "../clinicalnotes_raw/0b7wtxiunxwploe6tnnluh0l84qg.pdf"
 
     if not os.path.exists(file_path):
         print(f"Error: File not found at {file_path}")

structures.py

@@ -0,0 +1,73 @@
+from typing import Dict, Optional, List
+from typing_extensions import Literal
+from pydantic import BaseModel, Field as FieldInfo
+
+class Name(BaseModel):
+    reasoning: str
+    first_name: str
+    last_name: str
+
+class DateInfo(BaseModel):
+    reasoning: str
+    date: str
+
+class SexInfo(BaseModel):
+    reasoning: str
+    sex: str
+
+class ChemoInfo(BaseModel):
+    reasoning: str
+    drug: str
+    date: Optional[str] = None
+
+class TreatmentInfo(BaseModel):
+    reasoning: str
+    treatment: str
+    date: Optional[str] = None
+
+class MedicationInfo(BaseModel):
+    reasoning: str
+    medication: str
+    date: Optional[str] = None
+
+class SurgeryInfo(BaseModel):
+    reasoning: str
+    resection: str
+
+class SurgeryOutcomeInfo(BaseModel):
+    reasoning: str
+    outcome: str
+
+class MetastasisInfo(BaseModel):
+    reasoning: str
+    metastasis: str
+
+class PharmacyInfo(BaseModel):
+    reasoning: str
+    pharmacy: str
+
+class AdverseEffectInfo(BaseModel):
+    reasoning: str
+    medication: str
+    dosage: Optional[str] = None
+    date: Optional[str] = None
+    description: str
+
+class WeightInfo(BaseModel):
+    reasoning: str
+    weight: str
+
+class ClinicalInfo(BaseModel):
+    patient_name: Optional[Name] = None
+    date_of_birth: Optional[DateInfo] = None
+    sex: Optional[SexInfo] = None
+    traditional_chemo: Optional[List[ChemoInfo]] = None
+    other_cancer_treatments: Optional[List[TreatmentInfo]] = None
+    other_conmeds: Optional[List[MedicationInfo]] = None
+    surgery: Optional[SurgeryInfo] = None
+    surgery_outcome: Optional[SurgeryOutcomeInfo] = None
+    metastasis_at_time_of_diagnosis: Optional[MetastasisInfo] = None
+    compounding_pharmacy: Optional[PharmacyInfo] = None
+    adverse_effects: Optional[List[AdverseEffectInfo]] = None
+    date_of_death: Optional[DateInfo] = None
+    weight: Optional[WeightInfo] = None