Update app.py

app.py

@@ -3,58 +3,25 @@ import streamlit as st
 import os
 import json
 import time
-from annotated_text import annotated_text
-
-
-
-dax_input = st.text_area("DAX input")
-
-
-def generate_annotated_array(text, json_data):
-    """
-    Generate an array of strings based on annotations from the JSON data.
-
-    Parameters:
-    - text (str): The input text to be annotated.
-    - json_data (dict): The JSON data containing annotations.
-
-    Returns:
-    - list: An array of strings with separate elements for each annotation.
-    """
-    entities = json_data['documents'][0]['entities']
-    output = []
-    index = 0
-    buffer_text = ""
-
-    while index < len(text):
-        # Find the entity that matches the current position
-        entity = next((e for e in entities if e['offset'] == index), None)
-        
-        if entity:
-            # If there's buffer_text, add it to the output
-            if buffer_text:
-                output.append(buffer_text)
-                buffer_text = ""
-            
-            # Add the annotated entity to the output
-            output.append("[" + entity['text'] + ":" + entity['category'] + "]")
-            index += entity['length']
-        else:
-            # If no entity is found, add the character to buffer_text
-            buffer_text += text[index]
-            index += 1
-
-    # Add any remaining buffer_text to the output
-    if buffer_text:
-        output.append(buffer_text)
-
-    return output
-
 import requests
+from html import escape
+from random import choice
+import plotly.graph_objects as go
+import networkx as nx
+import os
+import openai
+from langchain.agents import create_json_agent, AgentExecutor
+from langchain.agents.agent_toolkits import JsonToolkit
+from langchain.chains import LLMChain
+from langchain.chat_models import AzureChatOpenAI
+from langchain.requests import TextRequestsWrapper
+from langchain.tools.json.tool import JsonSpec
 
-
-import requests
-import time 
+openai.api_key = os.environ["OPENAI_API_KEY"]
+openai.api_version = os.environ["OPENAI_API_VERSION"]
+openai.api_base = os.environ["OPENAI_API_BASE"]
+os.environ["OPENAI_API_VERSION"] = openai.api_version = "2023-03-15-preview"
+st.set_page_config(layout="wide")
 
 def analyze_healthcare_text(text):
     # Endpoint, headers and subscription key
@@ -80,8 +47,6 @@ def analyze_healthcare_text(text):
 
     # Making the initial POST request
     response = requests.post(f"{base_url}?api-version=2022-10-01-preview", headers=headers, json=data)
-
-    time.sleep(10)
     
     # Get the operation-location from the response header
     operation_location = response.headers.get('operation-location')
@@ -89,103 +54,252 @@ def analyze_healthcare_text(text):
     # Extract JOB-ID from the operation-location
     job_id = operation_location.split('/')[-1].split('?')[0]
     
+    
     # Make a subsequent GET request to retrieve the results using the JOB-ID
-    result_response = requests.get(f"{base_url}/{job_id}?api-version=2022-10-01-preview", headers=headers)
+    success = False
+    while not success:
+        result_response = requests.get(f"{base_url}/{job_id}?api-version=2022-10-01-preview", headers=headers)
+        if result_response.json()['status'] == 'succeeded':
+            success = True
+        else:
+            time.sleep(1)
     
     # Return the JSON response from the GET request
     result = result_response.json()
     
     return result
 
-
-
-def convert_to_annotated_text(input_list):
-    """
-    Convert a list with annotated content into a nested list suitable for annotated_text format.
+def annotate_text_with_entities(original_text, entities_data):
+        # Color palette for different categories
+    PALETTE = [
+        "#ff4b4b",
+        "#ffa421",
+        "#ffe312",
+        "#21c354",
+        "#00d4b1",
+        "#00c0f2",
+        "#1c83e1",
+        "#803df5",
+        "#808495",
+    ]
     
-    Args:
-    - input_list (list): The list with content and annotations in format '[text:annotation]'.
+    # Opacities
+    OPACITIES = [
+        "33", "66",
+    ]
     
-    Returns:
-    - list: A nested list in the annotated_text format.
-    """
-    annotated_list = []
-    temp_group = []
-
-    for item in input_list:
-        # Check if the item is an annotation
-        if item.startswith('[') and item.endswith(']'):
-            content = item[1:-1].split(':')
-            temp_group.append((content[0], content[1]))
-        else:
-            if temp_group:  # if there are items in the temporary group
-                annotated_list.append(temp_group)
-                temp_group = []
-            annotated_list.append(item)
+    json_data = entities_data
+    
+    # Extract entities from the JSON data
+    entities = json_data['documents'][0]['entities']
 
-    # Add any remaining items in the temporary group to the final list
-    if temp_group:
-        annotated_list.append(temp_group)
+    # Get unique categories from entities
+    unique_categories = list(set(entity['category'] for entity in entities))
 
-    return annotated_list
+    # Create a mapping of categories to colors
+    category_to_color = {category: PALETTE[i % len(PALETTE)] for i, category in enumerate(unique_categories)}
+    
+    # If we run out of colors in the palette, we will assign random colors to the remaining categories
+    if len(unique_categories) > len(PALETTE):
+        additional_colors = ['#'+''.join([choice('0123456789ABCDEF') for _ in range(6)]) for _ in range(len(unique_categories) - len(PALETTE))]
+        for i, category in enumerate(unique_categories[len(PALETTE):]):
+            category_to_color[category] = additional_colors[i]
+
+    def create_entity_html(entity, entity_id):
+        # Get the color for the entity category
+        color = category_to_color[entity["category"]] + OPACITIES[entity["offset"] % len(OPACITIES)]
+        
+        entity_html = f'<span id="entity-{entity_id}"><span style="display: inline-flex; flex-direction: row; align-items: center; background: {color}; border-radius: 0.5rem; padding: 0.25rem 0.5rem; overflow: hidden; line-height: 1;">{entity_id}. {escape(entity["text"])}'
+        
+        # If there are links, create a dropdown menu with the links
+        if entity.get("links"):
+            options = "".join(f'<option value="{link["id"]}">{link["dataSource"]} Code {link["id"]}</option>' for link in entity["links"])
+            dropdown_html = f'''
+            <span style="border-left: 1px solid; opacity: 0.1; margin-left: 0.5rem; align-self: stretch;"></span>
+            <span style="margin-left: 0.5rem; display: flex; flex-direction: column; align-items: flex-start;">
+                <select style="font-size: 0.75rem; opacity: 0.5;">
+                    {options}
+                </select>
+                <label style="font-size: 0.6rem; margin-top: 0.25rem;">{entity["category"]}</label>
+            </span>
+            '''
+            entity_html += dropdown_html
+        else:
+            # If there are no links, just display the category label
+            entity_html += f'<span style="border-left: 1px solid; opacity: 0.1; margin-left: 0.5rem; align-self: stretch;"></span><span style="margin-left: 0.5rem; font-size: 0.75rem; opacity: 0.5;">{entity["category"]}</span>'
 
+        # Close the main span element
+        entity_html += '</span></span>'
+        
+        return entity_html
+    
+    # Create HTML representation for each entity
+    entity_htmls = [create_entity_html(entity, i) for i, entity in enumerate(entities)]
+    
+    # Replace entities in the original text with their HTML representations
+    # We iterate from the end to avoid changing the offsets of the yet-to-be-replaced entities
+    for entity, entity_html in sorted(zip(entities, entity_htmls), key=lambda x: x[0]['offset'], reverse=True):
+        start = entity['offset']
+        end = start + entity['length']
+        original_text = original_text[:start] + entity_html + original_text[end:]
+    
+    # Create a color key section
+    color_key_section = "<br><br><strong>Color Key:</strong><br>"
+    for category, color in category_to_color.items():
+        color_key_section += f'<span style="display: inline-block; background: {color}; width: 1em; height: 1em; margin-right: 0.5em; vertical-align: middle;"></span>{category}<br>'
+    
+    original_text += color_key_section
+    
+    return original_text, category_to_color
 
-if st.button("Analyze"):
-    text = dax_input
-    json_analysis = analyze_healthcare_text(text)
-    json_analysis = json_analysis["tasks"]["items"][0]["results"]
-    # save json analysis as a file
+def create_interactive_graph_from_json(json_data, category_to_color):
+     # Load the JSON data
+    entities = json_data['documents'][0]['entities']
+    relations = json_data['documents'][0].get('relations', [])
+    
+    # Create a new directed graph
+    graph = nx.DiGraph()
+    
+    # Add nodes to the graph
+    for i, entity in enumerate(entities):
+        graph.add_node(i, label=entity['text'], category=entity['category'])
+    
+    # Add edges to the graph
+    for relation in relations:
+        source_index = int(relation['entities'][0]['ref'].split('/')[-1])
+        target_index = int(relation['entities'][1]['ref'].split('/')[-1])
+        graph.add_edge(source_index, target_index, label=relation['relationType'])
+    
+    # Get positions of the nodes using spring layout
+    pos = nx.spring_layout(graph)
+    
+    # Get node positions
+    x_nodes = [pos[i][0] for i in graph.nodes]
+    y_nodes = [pos[i][1] for i in graph.nodes]
+    
+    # Get the colors for each node based on its category
+    node_colors = [category_to_color[graph.nodes[i]['category']] for i in graph.nodes]
     
+    # Get edge positions
+    x_edges = []
+    y_edges = []
+    for edge in graph.edges:
+        x_edges += [pos[edge[0]][0], pos[edge[1]][0], None]
+        y_edges += [pos[edge[0]][1], pos[edge[1]][1], None]
     
-    new_text = generate_annotated_array(text, json_analysis)
-    new_text = convert_to_annotated_text(new_text)
-    annotated_text(new_text)
- 
+    # Create edge traces
+    edge_trace = go.Scatter(x=x_edges, y=y_edges, line=dict(width=0.5, color='#888'), hoverinfo='none', mode='lines')
+    
+    # Create node traces with all nodes
+    node_trace_all = go.Scatter(x=x_nodes, y=y_nodes, text=[graph.nodes[i]['label'] for i in graph.nodes], 
+                                mode='markers+text', hoverinfo='text', marker=dict(color=node_colors, size=10))
+    
+    # Create node traces with nodes having at least one edge
+    nodes_with_edges = set([edge[0] for edge in graph.edges] + [edge[1] for edge in graph.edges])
+    x_nodes_with_edges = [pos[i][0] for i in nodes_with_edges]
+    y_nodes_with_edges = [pos[i][1] for i in nodes_with_edges]
+    
+    node_trace_with_edges = go.Scatter(x=x_nodes_with_edges, y=y_nodes_with_edges, 
+                                       text=[graph.nodes[i]['label'] for i in nodes_with_edges], 
+                                       mode='markers+text', hoverinfo='text', marker=dict(color=node_colors, size=10))
+    
+    # Create figure
+    fig = go.Figure(data=[edge_trace, node_trace_all, node_trace_with_edges],
+                    layout=go.Layout(title='Entities and Relationships in Patient Notes',
+                                     titlefont_size=16,
+                                     showlegend=False,
+                                     hovermode='closest',
+                                     margin=dict(b=20, l=5, r=5, t=40),
+                                     xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+                                     yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+                                     updatemenus=[dict(type="buttons",
+                                                       x=1.15,
+                                                       y=1.2,
+                                                       buttons=[dict(label="All Entities",
+                                                                     method="update",
+                                                                     args=[{"visible": [True, True, False]}]),
+                                                                 dict(label="Entities with Relationships",
+                                                                     method="update",
+                                                                     args=[{"visible": [True, False, True]}])])]))
+    
+    # Display the interactive plot
+    return fig
 
+# divide the page into 3 columns
+col1, col2, col3 = st.columns([2,5,2])
 
+if 'r' not in st.session_state:
+    st.session_state.r = 'value'
 
-st.title("Nuance DAX Copilot")
+if 'r_annotated' not in st.session_state:
+    st.session_state.r_annotated = 'value'
 
-os.environ["OPENAI_API_BASE"] = openai.api_type = "azure"
-os.environ["OPENAI_API_BASE"] = openai.api_base = "https://eastus-openai-sean.openai.azure.com/"
-os.environ["OPENAI_API_VERSION"] = openai.api_version = "2023-03-15-preview"
+if 'colour_to_category' not in st.session_state:
+    st.session_state.colour_to_category = 'value'
 
-openai.api_key = os.environ["OPENAI_API_KEY"]
-openai.api_version = os.environ["OPENAI_API_VERSION"]
-openai.api_base = os.environ["OPENAI_API_BASE"]
-os.environ["OPENAI_API_VERSION"] = openai.api_version = "2023-03-15-preview"
+with col1:
+    col1.subheader("DAX Express Input")
 
+    st.text("Enter your DAX Express output below:")
 
-if "messages" not in st.session_state:
-    st.session_state.messages = [{"role":"system","content":"You are an AI assistant that ansswers questions about patient encounters. You are not a doctor and should not diagnose or treat patients. However, you can suggest common practices and help doctors with their questions that will help them make better decisions. \ Use only the information below: \n Patient Note / Encounter Summary:  \n"}]
+    dax_input = st.text_area("", height=500)
+    analyze_btn = st.button("Analyze")
+
+with col2:
+    col2.subheader("Text Analytics for Health Output")
+    if st.session_state.r_annotated != 'value':
+        with st.expander("Entity Mappings"):
+            st.markdown(st.session_state.r_annotated, unsafe_allow_html=True)
+        with st.expander("Show Relationships"):
+            st.plotly_chart(create_interactive_graph_from_json(st.session_state.r, st.session_state.colour_to_category), use_container_width=True)
+        with st.expander("Show JSON"):
+            st.json(st.session_state.r)
+        
+    if analyze_btn:
+        
+        st.session_state.r = analyze_healthcare_text(dax_input)["tasks"]["items"][0]["results"]
+        r_annotated, category_to_color = annotate_text_with_entities(dax_input, st.session_state.r)
+        st.session_state.r_annotated = r_annotated
+        st.session_state.colour_to_category = category_to_color
+        with st.expander("Entity Mappings"):
+            st.markdown(r_annotated, unsafe_allow_html=True)
+        with st.expander("Show Relationships"):
+            st.plotly_chart(create_interactive_graph_from_json(st.session_state.r, category_to_color), use_container_width=True)
+        with st.expander("Show JSON"):
+            st.json(st.session_state.r)
 
+        
+        
+        
 
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+with col3:
+    col3.subheader("Copilot Concept")
 
+    question = st.text_input("Ask a question to Copilot:")
 
+    toggle = st.toggle("Grounded", False)
 
 
+    ask = st.button("Ask")
+    if toggle:
+        if ask:
+            json_spec = JsonSpec(dict_=st.session_state.r, max_value_length=7000)
+            json_toolkit = JsonToolkit(spec=json_spec)
 
+            json_agent_executor = create_json_agent(
+                llm=AzureChatOpenAI(temperature=0, deployment_name="gpt-4"), toolkit=json_toolkit, verbose=True
+            )
 
-if prompt := st.chat_input("Nuance DAX Copilot?"):
-    st.session_state.messages.append({"role": "user", "content": prompt + dax_input})
-    with st.chat_message("user"):
-        st.markdown(prompt)
+            st.write(json_agent_executor.run(question))
 
-    with st.chat_message("assistant"):
-        message_placeholder = st.empty()
-        full_response = ""
-        for response in openai.ChatCompletion.create(
-            messages=[
-                {"role": m["role"], "content": m["content"]}
-                for m in st.session_state.messages
-            ],
-            stream=True,
-            engine="gpt-4",
-        ):
-            full_response += response.choices[0].delta.get("content", "")
-            message_placeholder.markdown(full_response + "▌")
-        message_placeholder.markdown(full_response)
-    st.session_state.messages.append({"role": "assistant", "content": full_response})
+    elif ask:
+        response = openai.ChatCompletion.create(
+    engine="gpt-4",
+    messages = [{"role":"system","content":"You are an AI assistant that helps people find information."}, {"role": "user", "content" : question}],
+    temperature=0.7,
+    max_tokens=800,
+    top_p=0.95,
+    frequency_penalty=0,
+    presence_penalty=0,
+    stop=None)
+        st.write(response.choices[0].message.content)