Adding some documentation, fixing tests

TCT/TCT_neighborhood_finder.py

@@ -1,3 +1,4 @@
+# reviewed by yjzhang, 2026-08-19
 from collections import Counter
 
 import pandas as pd
@@ -18,12 +19,29 @@
 from .TCT_pathfinder import generate_score_results, build_query_graph
 
 def parse_results_for_neighborhood_finder(start_node_id:str, results:dict,
-        start_node_categories=None, end_node_categories=None,
-        get_node_info=True,
-        scoring_method='infores'):
+        start_node_categories:list|None=None, end_node_categories:list|None=None,
+        get_node_info:bool=True,
+        scoring_method:str='infores') -> dict:
     """
     Converts the results of two TRAPI queries into the same general json format as the other pathfinder APIs.
-    scoring_method is how the node scores are generated, and could be 'infores' or 'edges'.
+
+    Params
+    ------
+    start_node_id : str
+        A CURIE id
+    results : dict
+        Results of a TRAPI query (e.g. parallel_api_query)
+    start_node_categories : list | None
+        Categories for the starting node.
+    end_node_categories : list | None
+        Categories for the ending nodes of the query.
+    get_node_info : bool
+    scoring_method : str
+        scoring_method is how the node scores are generated, and could be 'infores' or 'edges'. Default: 'infores'
+
+    Returns
+    -------
+    A dict of the format `{'query_graph': ..., 'knowledge_graph': ..., 'results':, 'auxiliary_graphs':...}`
     """
     # nodes
     node_info = {}
@@ -108,11 +126,10 @@ def parse_results_for_neighborhood_finder(start_node_id:str, results:dict,
     return output
 
 
-# TODO: can we use multiple input nodes?
 def parse_results_for_neighborhood_finder_multiple_inputs(start_node_ids:list[str], results:dict,
-        start_node_categories=None, end_node_categories=None,
-        get_node_info=True,
-        scoring_method='infores'):
+        start_node_categories:list|None=None, end_node_categories:list|None=None,
+        get_node_info:bool=True,
+        scoring_method:str='infores') -> dict:
     """
     Converts the results of two TRAPI queries into the same general json format as the other pathfinder APIs.
     scoring_method is how the node scores are generated, and could be 'infores' or 'edges'.

TCT/TCT_network_annotator.py

@@ -6,8 +6,8 @@
 from . import TCT_neighborhood_finder
 
 
-def network_annotator(gene_list, 
-                      select_APIs, 
+def network_annotator(gene_list,
+                      select_APIs,
                       node2_categories, 
                       select_metaKG, 
                       API_predicates, 
@@ -28,7 +28,6 @@ def network_annotator(gene_list,
     dict
         Merged TRAPI JSON object.
     """
-    import json
     TCT_neighborhood_finder_result = {}
     for gene in gene_list:
         input_identifiers = name_resolver.lookup(gene, only_taxa='NCBITaxon:9606', biolink_type='biolink:Gene').curie

TCT/TCT_pathfinder.py

@@ -1,4 +1,5 @@
 # TCT Pathfinder...
+# reviewed by yjzhang, 2026-08-20
 import requests
 
 import pandas as pd
@@ -21,15 +22,15 @@
 
 
 
-def format_query_json_for_pathfinder_with_constraints(subject_ids, 
-                                     object_ids=None,
+def format_query_json_for_pathfinder_with_constraints(subject_ids:str,
+        object_ids:str,
         subject_categories=None,
         object_categories=None,
         predicates=None,
         constraints=None
-        ):
+        ) -> dict:
     """
-        format user's input into a query json for pathfinder pipeline with constraints on the intermediate node categories.
+    Format user's input into a query json for pathfinder pipeline with constraints on the intermediate node categories.
 
     Parameters
     ----------
@@ -112,7 +113,7 @@ def format_query_json_for_pathfinder_with_constraints(subject_ids,
 
     return q
 
-def build_query_graph(start_node_id, end_node_id, start_node_categories=None, end_node_categories=None, constraints_path=None):
+def build_query_graph(start_node_id:str, end_node_id:str, start_node_categories=None, end_node_categories=None, constraints_path=None):
     """
     start_node_categories and end_node_categories are lists of categories.
     """
@@ -151,7 +152,7 @@ def build_query_graph(start_node_id, end_node_id, start_node_categories=None, en
     return q
 
 
-def generate_score_results(results, method='infores'):
+def generate_score_results(results:dict, method='infores'):
     """
     Generates a score dict, and a list of "analyses".
     method can be 'infores' or 'edges'
@@ -312,9 +313,20 @@ def parse_results_for_pathfinder(start_node_id:str, end_node_id:str, result1:dic
 
 
 # define a function that uses the query_json as an template and change the ids and categories of the nodes
-def format_pathfinder_query(node1_id, node1_category, node2_id, node2_category):
+def format_pathfinder_query(node1_id:str, node1_category:str, node2_id:str, node2_category:str) -> dict:
     '''
-
+    Formats a query to the Pathfinder API.
+
+    Params
+    ------
+    node1_id : str
+    node1_category : str
+    node2_id : str
+    node2_category : str
+
+    Returns
+    -------
+    A dict formatted as a JSON query to the Pathfinder API.
     '''
     query_json = {
         "message": {
@@ -350,15 +362,44 @@ def format_pathfinder_query(node1_id, node1_category, node2_id, node2_category):
     return query_json
 
 
-def query_aragorn_pathfinder(node1_id, node1_category, node2_id, node2_category):
+def query_aragorn_pathfinder(node1_id:str, node1_category:str, node2_id:str, node2_category:str) -> str:
+    """
+    This queries the ARAGORN Pathfinder API.
+
+    Params
+    ------
+    node1_id : str
+    node1_category : str
+    node2_id : str
+    node2_category : str
+
+    Returns
+    -------
+    A string (which should be a JSON) representing the result of an ARAGORN pathfinder query.
+    """
     #aragorn_endpoint = 'https://shepherd.renci.org/aragorn/query' # dev endpoint
     aragorn_endpoint = 'https://shepherd.ci.transltr.io/aragorn/query' # new ci endpoint
     query_current = format_pathfinder_query(node1_id, node1_category, node2_id, node2_category)
     response = requests.post(aragorn_endpoint, json=query_current)
     return response
 
 
-def query_aragorn_pathfinder_with_constraints(node1_id, node1_category, node2_id, node2_category, constraints):
+def query_aragorn_pathfinder_with_constraints(node1_id:str, node1_category:str, node2_id:str, node2_category:str, constraints:list) -> str:
+    """
+    This queries the ARAGORN Pathfinder API with a list of constraints.
+
+    Params
+    ------
+    node1_id : str
+    node1_category : str
+    node2_id : str
+    node2_category : str
+    constraints : list
+
+    Returns
+    -------
+    A string (which should be a JSON) representing the result of an ARAGORN pathfinder query.
+    """
     #aragorn_endpoint = 'https://shepherd.renci.org/aragorn/query' # dev endpoint
     aragorn_endpoint = 'https://shepherd.ci.transltr.io/aragorn/query' # new ci endpoint
     query_current = format_query_json_for_pathfinder_with_constraints(
@@ -371,13 +412,42 @@ def query_aragorn_pathfinder_with_constraints(node1_id, node1_category, node2_id
     response = requests.post(aragorn_endpoint, json=query_current)
     return response
 
-def query_arax_pathfinder(node1_id, node1_category, node2_id, node2_category):
+def query_arax_pathfinder(node1_id:str, node1_category:str, node2_id:str, node2_category:str) -> str:
+    """
+    This queries the ARAX Pathfinder API.
+
+    Params
+    ------
+    node1_id : str
+    node1_category : str
+    node2_id : str
+    node2_category : str
+
+    Returns
+    -------
+    A string (which should be a JSON) representing the result of an ARAX pathfinder query.
+    """
     ARAX_endpoint = 'https://arax.ci.transltr.io/api/arax/v1.4/query'
     query_current = format_pathfinder_query(node1_id, node1_category, node2_id, node2_category)
     response = requests.post(ARAX_endpoint, json=query_current)
     return response
 
-def query_arax_pathfinder_with_constraints(node1_id, node1_category, node2_id, node2_category, constraints):
+def query_arax_pathfinder_with_constraints(node1_id:str, node1_category:str, node2_id:str, node2_category:str, constraints:list) -> str:
+    """
+    This queries the ARAX Pathfinder API with a list of constraints.
+
+    Params
+    ------
+    node1_id : str
+    node1_category : str
+    node2_id : str
+    node2_category : str
+    constraints : list
+
+    Returns
+    -------
+    A string (which should be a JSON) representing the result of an ARAX pathfinder query.
+    """
     ARAX_endpoint = 'https://arax.ci.transltr.io/api/arax/v1.4/query'
     query_current = format_query_json_for_pathfinder_with_constraints(
         subject_ids=node1_id,

TCT/name_resolver.py

@@ -3,6 +3,7 @@
 
 API docs: https://name-lookup.ci.transltr.io/docs
 """
+# reviewed by yjzhang, 2026-08-19
 import urllib.parse
 
 import requests
@@ -12,7 +13,7 @@
 URL = 'https://name-lookup.ci.transltr.io/'
 """This is the root URL for the API."""
 
-def status():
+def status() -> str:
     """
     Returns the status of the Name Resolver API.
     """
@@ -21,7 +22,7 @@ def status():
     return response.json()
 
 
-def lookup(query: str, return_top_response:bool=True, return_synonyms:bool=False, limit:int=10, **kwargs):
+def lookup(query: str, return_top_response:bool=True, return_synonyms:bool=False, limit:int=10, **kwargs) -> TranslatorNode | list[TranslatorNode]:
     """
     A wrapper around the `lookup` api endpoint. Given a query string, this returns a TranslatorNode object or a list of TranslatorNode objects corresponding to the given name.
 
@@ -72,7 +73,7 @@ def lookup(query: str, return_top_response:bool=True, return_synonyms:bool=False
         raise requests.RequestException('Response from server had error, code ' + str(response.status_code) + ' ' + str(response))
 
 
-def synonyms(query: str|list, **kwargs):
+def synonyms(query: str|list, **kwargs) -> dict[str, TranslatorNode]:
     """
     A wrapper around the `synonyms` api endpoint. Given a CURIE or a list of CURIEs, this returns a dict of CURIE id : TranslatorNode for all synonyms for the given query.
 
@@ -107,15 +108,15 @@ def synonyms(query: str|list, **kwargs):
         raise requests.RequestException('Response from server had error, code ' + str(response.status_code) + ' ' + str(response))
 
 
-def chunk_list(data:list, size:int):
+def chunk_list(data:list, size:int) -> list:
     #Extra method to help chunk large files and avoid 504 error.
     chunks = []
     for i in range(0, len(data), size):
         chunks.append(data[i: i+size])
     return chunks
 
 
-def batch_lookup(strings:list[str], size: int=25, return_top_response:bool=True, return_synonyms:bool=False, **kwargs) -> dict:
+def batch_lookup(strings:list[str], size: int=25, return_top_response:bool=True, return_synonyms:bool=False, **kwargs) -> dict[str, TranslatorNode|list[TranslatorNode]]:
     """
     A wrapper around the `bulk-lookup` api endpoint. Given a list of query strings, this returns a TranslatorNode object or a list of TranslatorNode objects corresponding to the given name.
 
@@ -174,7 +175,7 @@ def batch_lookup(strings:list[str], size: int=25, return_top_response:bool=True,
     return curies
 
 
-def batch_synonyms(strings:list[str], size:int=50, **kwargs) -> dict:
+def batch_synonyms(strings:list[str], size:int=50, **kwargs) -> dict[str, TranslatorNode]:
     """
     A wrapper around the `synonyms` API endpoint, using POST. Given a list of CURIEs, this returns a dict of CURIE:TranslatorNode, where each TranslatorNode contains all synonyms for the given CURIE.
 

TCT/node_annotator.py

@@ -3,6 +3,7 @@
 
 API docs: https://annotator.transltr.io/
 """
+# reviewed by yjzhang, 2026-08-19
 import urllib.parse
 
 import requests
@@ -25,7 +26,7 @@ def lookup_curie(curie: str, **kwargs):
     return lookup_curies([curie], **kwargs)[curie]
 
 
-def lookup_curies(curies: list[str], **kwargs):
+def lookup_curies(curies: list[str], **kwargs) -> dict[str, dict]:
     """
     A wrapper around the `curies` API endpoint. Given a list of CURIEs, this returns a dictionary where each
     CURIE is mapped to a list of annotations.

TCT/node_normalizer.py

@@ -3,6 +3,7 @@
 
 API docs: https://nodenorm.transltr.io/docs
 """
+# reviewed by yjzhang, 2026-08-19
 import urllib.parse
 
 import requests
@@ -20,6 +21,7 @@ def status():
     response.raise_for_status()
     return response.json()
 
+
 def get_normalized_nodes(query: str | list[str],
         return_equivalent_identifiers:bool=False,
         mode:str='get',
@@ -135,7 +137,7 @@ def get_preferred_names(id_list:list[str], batch_limit=500, **kwargs) -> dict[st
     return name_map
 
 
-def ID_convert_to_preferred_name_nodeNormalizer(id_list):
+def ID_convert_to_preferred_name_nodeNormalizer(id_list: list[str]) -> dict[str, str]:
     '''
     Convert a list of CURIEs to their preferred names using NodeNorm.
     Arg:

TCT/translator_node.py

@@ -1,3 +1,4 @@
+# reviewed by yjzhang, 2026-08-19
 # translator graph node
 from dataclasses import dataclass
 import typing

TCT/translator_query.py

@@ -44,12 +44,13 @@ def get_translator_API_predicates() -> tuple[dict, pandas.DataFrame, dict]:
     return APInames, metaKG, API_predicates
 
 
-def build_attribute_constraint(attribute_id, operator, value, name=None, is_not=False):
+def build_attribute_constraint(attribute_id:str, operator:str, value, name=None, is_not=False):
     """
     This creates an attribute constraint for a TRAPI query dict.
 
     Example
     -------
+
     """
     if name is None:
         name = ''

tests/test_server.py

@@ -11,14 +11,14 @@
 def test_mcp_server_exists():
     """Test that MCP server instance exists and has correct name."""
     assert mcp is not None
-    assert mcp.name == "translator-toolkit"
+    assert mcp.name == "TCT"
 
 
 def test_mcp_server_ready():
     """Test that MCP server is ready for orchestrating agent access."""
     # Check that the server has the FastMCP functionality needed for agents
     assert hasattr(mcp, 'run'), "MCP server should be runnable for agents"
-    assert mcp.name == "translator-toolkit", "MCP server should have correct name for agents"
+    assert mcp.name == "TCT", "MCP server should have correct name for agents"
 
 
 def test_mcp_tools_accessible():
@@ -27,4 +27,4 @@ def test_mcp_tools_accessible():
 
     # These should exist as tool objects that agents can call
     assert name_lookup is not None, "name_lookup tool should be accessible"
-    assert normalize_nodes is not None, "normalize_nodes tool should be accessible"
+    assert normalize_nodes is not None, "normalize_nodes tool should be accessible"