Replace requests with diffbot-kg client, and make some code style fixes

Author: brendancsmith

api/app/enhance.py

@@ -1,33 +1,33 @@
-import logging
 import os
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Union
-from urllib.parse import urlencode
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
-import requests
-from utils import graph
+from diffbot_kg import DiffbotEnhanceClient
+
+from .utils import graph
 
 CATEGORY_THRESHOLD = 0.50
 params = []
 
 DIFF_TOKEN = os.environ["DIFFBOT_API_KEY"]
 
+client = DiffbotEnhanceClient(DIFF_TOKEN)
+
 
-def get_datetime(value: Optional[Union[str, int, float]]) -> datetime:
-    if not value:
-        return value
-    return datetime.fromtimestamp(float(value) / 1000.0)
+def get_datetime(
+    value: Optional[Union[str, int, float]]
+) -> datetime | float | Literal["", 0] | None:
+    return datetime.fromtimestamp(float(value) / 1000.0) if value else value
 
 
-def process_entities(entity: str, type: str) -> Dict[str, Any]:
+async def process_entities(entity: str, type: str) -> Tuple[str, List[Dict]]:
     """
     Fetch relevant articles from Diffbot KG endpoint
     """
-    search_host = "https://kg.diffbot.com/kg/v3/enhance?"
-    params = {"type": type, "name": entity, "token": DIFF_TOKEN}
-    encoded_query = urlencode(params)
-    url = f"{search_host}{encoded_query}"
-    return entity, requests.get(url).json()
+    params = {"type": type, "name": entity}
+    response = await client.enhance(params)
+
+    return entity, response.entities
 
 
 def get_people_params(row: Dict) -> Optional[Dict]:
@@ -138,7 +138,6 @@ def get_people_params(row: Dict) -> Optional[Dict]:
 
 def get_organization_params(name: str, row: Dict) -> Dict:
     # Properties
-    type = row["type"]
     node_properties = {
         "employees": row.get("nbEmployees"),
         "revenue": row.get("revenue", {}).get("value"),
@@ -335,11 +334,11 @@ def store_enhanced_data(data: List[Dict[str, Any]]) -> Dict:
         except Exception:
             no_data.append(name)
             continue
-        type = entity["type"]
-        if type == "Organization":
+        etype = entity["type"]
+        if etype == "Organization":
             params = get_organization_params(name, entity)
             organizations.append(params)
-        elif type == "Person":
+        elif etype == "Person":
             params = get_people_params(entity)
             people.append(params)
     # Save processes status to entities without any response

api/app/importing.py

@@ -1,49 +1,48 @@
 import logging
 import os
-from typing import Any, Dict, List, Optional
+from typing import Dict, List, Optional
 
-import requests
-from utils import embeddings, text_splitter
+from diffbot_kg import DiffbotSearchClient
+
+from .utils import embeddings, text_splitter
 
 CATEGORY_THRESHOLD = 0.50
 params = []
 
 DIFF_TOKEN = os.environ["DIFFBOT_API_KEY"]
 
+client = DiffbotSearchClient(token=DIFF_TOKEN)
+
 
-def get_articles(
+async def get_articles(
     query: Optional[str], tag: Optional[str], size: int = 5, offset: int = 0
-) -> Dict[str, Any]:
+) -> List[Dict]:
     """
     Fetch relevant articles from Diffbot KG endpoint
     """
     try:
-        search_host = "https://kg.diffbot.com/kg/v3/dql?"
-        search_query = f'query=type%3AArticle+strict%3Alanguage%3A"en"+sortBy%3Adate'
+        search_query = 'query=type%3AArticle+strict%3Alanguage%3A"en"+sortBy%3Adate'
         if query:
             search_query += f'+text%3A"{query}"'
         if tag:
             search_query += f'+tags.label%3A"{tag}"'
-        url = (
-            f"{search_host}{search_query}&token={DIFF_TOKEN}&from={offset}&size={size}"
-        )
-        return requests.get(url).json()
+        response = await client.search({query: query, size: size, offset: offset})
+        return response.entities
     except Exception as ex:
         raise ex
 
 
 def get_tag_type(types: List[str]) -> str:
     try:
         return types[0].split("/")[-1]
-    except:
+    except Exception:
         return "Node"
 
 
-def process_params(data):
+def process_params(articles):
     params = []
     all_chunks = []
-    for row in data["data"]:
-        article = row["entity"]
+    for article in articles:
         split_chunks = [
             {"text": el, "index": f"{article['id']}-{i}"}
             for i, el in enumerate(text_splitter.split_text(article["text"])[:5])

api/app/main.py

@@ -1,26 +1,28 @@
+import asyncio
 import logging
 import os
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Dict
 
-from api_types import ArticleData, CountData, EntityData
-from chat import chain
-from enhance import process_entities, store_enhanced_data
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from graph_prefiltering import prefiltering_agent_executor
-from importing import get_articles, import_cypher_query, process_params
 from langserve import add_routes
-from processing import process_document, store_graph_documents
-from text2cypher import text2cypher_chain
-from utils import graph, remove_null_properties
+
+from .api_types import ArticleData, CountData, EntityData
+from .chat import chain
+from .enhance import process_entities, store_enhanced_data
+from .graph_prefiltering import prefiltering_agent_executor
+from .importing import get_articles, import_cypher_query, process_params
+from .processing import process_document, store_graph_documents
+from .text2cypher import text2cypher_chain
+from .utils import graph, remove_null_properties
 
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
 )
 
 # Multithreading for Diffbot API
-MAX_WORKERS = min(os.cpu_count() * 5, 20)
+MAX_WORKERS = min((os.cpu_count() or 1) * 5, 20)
 
 app = FastAPI()
 
@@ -35,21 +37,23 @@
 
 
 @app.post("/import_articles/")
-def import_articles_endpoint(article_data: ArticleData) -> int:
+async def import_articles_endpoint(article_data: ArticleData) -> int:
     logging.info(f"Starting to process article import with params: {article_data}")
     if not article_data.query and not article_data.tag:
         raise HTTPException(
             status_code=500, detail="Either `query` or `tag` must be provided"
         )
-    data = get_articles(article_data.query, article_data.tag, article_data.size)
-    logging.info(f"Articles fetched: {len(data['data'])} articles.")
+    articles = await get_articles(
+        article_data.query, article_data.tag, article_data.size
+    )
+    logging.info(f"Articles fetched: {len(articles)} articles.")
     try:
-        params = process_params(data)
+        params = process_params(articles)
     except Exception as e:
         # You could log the exception here if needed
-        raise HTTPException(status_code=500, detail=e)
+        raise HTTPException(status_code=500, detail=e) from e
     graph.query(import_cypher_query, params={"data": params})
-    logging.info(f"Article import query executed successfully.")
+    logging.info("Article import query executed successfully.")
     return len(params)
 
 
@@ -123,6 +127,11 @@ def fetch_unprocessed_count(count_data: CountData) -> int:
     return data[0]["output"]
 
 
+def run_async_function_in_executor(loop, async_func):
+    future = asyncio.run_coroutine_threadsafe(async_func(), loop)
+    return future.result()
+
+
 @app.post("/enhance_entities/")
 def enhance_entities(entity_data: EntityData) -> str:
     entities = graph.query(
@@ -133,11 +142,16 @@ def enhance_entities(entity_data: EntityData) -> str:
         params={"limit": entity_data.size},
     )
     enhanced_data = []
+    loop = asyncio.get_event_loop()
     with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
         # Submitting all tasks and creating a list of future objects
         for row in entities:
             futures = [
-                executor.submit(process_entities, el, row["label"])
+                executor.submit(
+                    run_async_function_in_executor(loop, process_entities),
+                    el,
+                    row["label"],
+                )
                 for el in row["entities"]
             ]
 

api/app/utils.py

@@ -52,7 +52,7 @@ def setup_indices():
     graph=graph,
     index_name=index_name,
     keyword_index_name=keyword_index_name,
-    search_type="hybrid",
+    search_type="hybrid",  # type: ignore
 )
 
 text_splitter = TokenTextSplitter(
@@ -121,7 +121,7 @@ def generate_full_text_query(input: str) -> str:
 
 def remove_null_properties(data: Dict[str, Any]):
     if not data:
-        return []
+        return {}
     # Process the 'nodes' part of the data
     for node in data["nodes"]:
         # Create a list of keys to remove (those that are None)