Add get_uris_of_similar_diseases

utils.py

@@ -1,5 +1,5 @@
 # %%
-from typing import List, Dict, Any
+from typing import List, Dict, tuple, Any
 import os
 from sqlalchemy import create_engine, text
 import requests
@@ -83,6 +83,18 @@ def get_clinical_records_by_ids(clinical_record_ids: List[str]) -> List[Dict[str
         clinical_records.append(clinical_record_info)
     return clinical_records
 
+def get_uris_of_similar_diseases(uri_list: List[str]) -> List[tuple[str, str, float]]:
+    uri_list = tuple(uri_list)
+    with engine.connect() as conn:
+        with conn.begin():
+            sql = f"""
+                    SELECT e1.uri AS uri1, e2.uri AS uri2, VECTOR_COSINE(e1.embedding, e2.embedding) AS distance
+                    FROM Test.EntityEmbeddings e1, Test.EntityEmbeddings e2
+                    WHERE e1.uri IN {uri_list} AND e2.uri IN {uri_list} AND e1.uri != e2.uri
+                """
+            result = conn.execute(text(sql))
+            data = result.fetchall()
+    return data
 
 if __name__ == "__main__":
     username = 'demo'