Update app.py

app.py

@@ -2,6 +2,8 @@ import json
 import streamlit as st
 from google.oauth2 import service_account
 from google.cloud import language_v1
+import urllib.parse
+import urllib.request
 import pandas as pd
 
 # Function to query Google's Knowledge Graph API
@@ -12,9 +14,14 @@ def query_knowledge_graph(entity_id):
     except Exception as e:
         st.write(f"An error occurred: {e}")
 
-# Function to filter entities with "mid" in their metadata
-def filter_entities_with_mid(entities):
-    return [entity for entity in entities if 'mid' in entity.metadata]
+# Function to count entities with 'mid' that contains '/g/' or '/m/' in their metadata
+def count_entities(entities):
+    count = 0
+    for entity in entities:
+        metadata = entity.metadata
+        if 'mid' in metadata and ('/g/' in metadata['mid'] or '/m/' in metadata['mid']):
+            count += 1
+    return count
 
 # Function to serialize entity metadata
 def serialize_entity_metadata(metadata):
@@ -73,14 +80,22 @@ def sample_analyze_entities(text_content):
 
     response = client.analyze_entities(request={"document": document, "encoding_type": encoding_type})
     
-    # Filter entities that have a "mid" in their metadata
-    entities_with_mid = filter_entities_with_mid(response.entities)
-    
-    st.markdown(f"# We found {len(entities_with_mid)} entities with 'mid' in their metadata")
-    st.write("---")
-    
-    for i, entity in enumerate(entities_with_mid):
-        st.write(f"Entity {i+1} of {len(entities_with_mid)}")
+    # Count the entities with 'mid' and either '/g/' or '/m/' in their metadata
+    entity_count = count_entities(response.entities)
+
+    if entity_count == 0:
+        st.markdown(f"# We found {len(response.entities)} entities - but found no Google Entities")
+        st.write("---")
+    elif entity_count == 1:
+        st.markdown(f"# We found {len(response.entities)} entities - and found 1 Google Entity")
+        st.write("---")
+    else:
+        st.markdown(f"# We found {len(response.entities)} entities - and found {entity_count} Google Entities")
+        st.write("---")
+
+
+    for i, entity in enumerate(response.entities):
+        st.write(f"Entity {i+1} of {len(response.entities)}")
         st.write(f"Name: {entity.name}")
         st.write(f"Type: {language_v1.Entity.Type(entity.type_).name}")
         st.write(f"Salience Score: {entity.salience}")
@@ -89,7 +104,7 @@ def sample_analyze_entities(text_content):
             st.write("Metadata:")
             st.write(entity.metadata)
             
-            if 'mid' in entity.metadata:
+            if 'mid' in entity.metadata and ('/g/' in entity.metadata['mid'] or '/m/' in entity.metadata['mid']):
                 entity_id = entity.metadata['mid']
                 query_knowledge_graph(entity_id)
 
@@ -103,7 +118,7 @@ def sample_analyze_entities(text_content):
         st.write("---")
 
     # Add the export functionality
-    export_entities(entities_with_mid)
+    export_entities(response.entities)
 
 # User input for text analysis
 user_input = st.text_area("Enter text to analyze")