Shuffle around contents of execute_request, format_response, and search

app.py

@@ -70,7 +70,7 @@ class Work:
         )
 
     @staticmethod
-    def raw_fields() -> list[str]:
+    def get_raw_fields() -> list[str]:
         return [
             "title",
             "abstract_inverted_index",
@@ -124,18 +124,24 @@ def get_index(dir: Path, search_time_s: float) -> Dataset:
     return index
 
 
-def execute_request(request_str):
-    response = requests.get(request_str).json()
-    return response
+def execute_request(ids: list[str]) -> list[Work]:
+    # query with the /works endpoint with a specific list of IDs and fields
+    search_filter = f'openalex_id:{"|".join(ids)}'
+    search_select = ",".join(["id"] + Work.get_raw_fields())
+    response = requests.get(
+        "https://api.openalex.org/works",
+        {"filter": search_filter, "select": search_select}
+    )
+    response.raise_for_status()
 
+    # the response is not necessarily ordered, so order them
+    response = {d["id"]: Work.from_dict(d) for d in response.json()["results"]}
+    return [response[id_] for id_ in ids]
 
-def format_response(neighbors, response):
-    response = {doc['id']: doc for doc in response['results']}
-    
-    result_string = ''
-    for distance, openalex_id in neighbors:
-        work = Work.from_dict(response[openalex_id])
 
+def format_response(neighbors: list[Work], distances: list[float]) -> str:
+    result_string = ''
+    for work, distance in zip(neighbors, distances):
         # edge cases: no title
         abstract = work.abstract if work.abstract is not None else "No abstract"
 
@@ -196,21 +202,18 @@ model = get_model(MODEL_NAME, "cpu")
 index = get_index(DIR, SEARCH_TIME_S)
 
 
-def search(query):
+def search(query: str) -> str:
     global model, index
 
-    # TODO: pass in param string directly?
     query_embedding = model.encode(query)
     distances, faiss_ids = index.search("embeddings", query_embedding, 20)
-
     openalex_ids = index[faiss_ids]["idxs"]
-    search_filter = f'openalex_id:{"|".join(openalex_ids)}'
-    search_select = 'id,title,abstract_inverted_index,authorships,primary_location,publication_year,cited_by_count,doi'
 
-    neighbors = [(distance, openalex_id) for distance, openalex_id in zip(distances, openalex_ids)]
-    request_str = f'https://api.openalex.org/works?filter={search_filter}&select={search_select}'
+    works = execute_request(openalex_ids)
+    result_string = format_response(works, distances)
+
+    return result_string
 
-    return neighbors, request_str
 
 with gr.Blocks() as demo:
     gr.Markdown('# abstracts-index')
@@ -232,14 +235,10 @@ with gr.Blocks() as demo:
         results = gr.Markdown()
 
     md = MarkdownIt('js-default', {'linkify': True, 'typographer': True}) # don't render html or latex!
-    results.md = md
+    results.md = md  # TODO: does this not work anymore?
     
-    query.submit(search, inputs=[query], outputs=[neighbors_var, request_str_var]) \
-        .success(execute_request, inputs=[request_str_var], outputs=[response_var]) \
-        .success(format_response, inputs=[neighbors_var, response_var], outputs=[results])
-    btn.click(search, inputs=[query], outputs=[neighbors_var, request_str_var]) \
-        .success(execute_request, inputs=[request_str_var], outputs=[response_var]) \
-        .success(format_response, inputs=[neighbors_var, response_var], outputs=[results])
+    query.submit(search, inputs=[query], outputs=[results])
+    btn.click(search, inputs=[query], outputs=[results])
 
 demo.queue(2)
 demo.launch()