Refactor SNOMED-CT_Assistant.py to improve entity mapping result display with markdown table.

SNOMED-CT_Assistant.py

@@ -129,16 +129,39 @@ def chat_input(prompt, med_text):
             print("entity: ", entity)
             results = query_chroma_db(entity, 10)
             results_dict = get_dict_from_chroma_results(results)
-            entity_mapping_prompt = generate_entity_mapping_prompt(entity, results_dict)
-            st.session_state.messages.append({"role": "user", "content": entity_mapping_prompt})
-            entity_mapping_response = client.chat.completions.create(
-                model=model_tag, messages=st.session_state.messages, temperature=0.5)
-            mapping_msg = entity_mapping_response.choices[0].message.content
-            st.session_state.messages.append({"role": "assistant", "content": mapping_msg})
-            st.chat_message("assistant").write(mapping_msg)
+            results_table = entity_mapping_result_to_table(entity, results_dict)
+            st.session_state.messages.append({"role": "assistant", "content": results_table})
+            st.chat_message("assistant").write(results_table)
+            # entity_mapping_prompt = generate_entity_mapping_prompt(entity, results_dict)
+            # st.session_state.messages.append({"role": "user", "content": entity_mapping_prompt})
+            # entity_mapping_response = client.chat.completions.create(
+            #     model=model_tag, messages=st.session_state.messages, temperature=0.5)
+            # mapping_msg = entity_mapping_response.choices[0].message.content
+            # st.session_state.messages.append({"role": "assistant", "content": mapping_msg})
+            # st.chat_message("assistant").write(mapping_msg)
         
-            
-            
+        
+# Conver entity mapping result to markdown table
+def entity_mapping_result_to_table(entity, results_dict):
+    
+    ids = results_dict['ids']
+    concept_ids = results_dict['concept_ids']
+    distances = results_dict['distances']
+    descriptions = results_dict['descriptions']
+    
+    # header
+    header = "| Identified Entity | Distance | IDs | SNOMED CT - Concept IDs | SNOMED CT - Descriptions |"
+    seperator = "| --- | --- | --- | --- | --- |"
+    
+    # table
+    rows = []
+    for id, distance, concept_id, description in zip(ids, distances, concept_ids, descriptions):
+        row = f"| {entity} | {distance:.3f} | {id} | {concept_id} | {description} |"
+        rows.append(row)
+    
+    # merge
+    markdown_table = "\n".join([header, seperator] + rows)
+    return markdown_table
 
 if "messages" not in st.session_state:
     st.session_state["messages"] = [{"role": "system", "content": system_prompt},