diff --git a/src/nat/experimental/test_time_compute/functions/execute_score_select_function.py b/src/nat/experimental/test_time_compute/functions/execute_score_select_function.py
index 0ba6dca60..141af1c9f 100644
--- a/src/nat/experimental/test_time_compute/functions/execute_score_select_function.py
+++ b/src/nat/experimental/test_time_compute/functions/execute_score_select_function.py
@@ -46,7 +46,7 @@ async def execute_score_select_function(config: ExecuteScoreSelectFunctionConfig
 
     from pydantic import BaseModel
 
-    executable_fn: Function = builder.get_function(name=config.augmented_fn)
+    executable_fn: Function = await builder.get_function(name=config.augmented_fn)
 
     if config.scorer:
         scorer = await builder.get_ttc_strategy(strategy_name=config.scorer,
diff --git a/src/nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py b/src/nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py
index 90a71ab5d..b281e03a0 100644
--- a/src/nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py
+++ b/src/nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py
@@ -98,8 +98,8 @@ async def register_ttc_tool_wrapper_function(
 
         augmented_function_desc = config.tool_description
 
-    fn_input_schema: BaseModel = augmented_function.input_schema
-    fn_output_schema: BaseModel = augmented_function.single_output_schema
+    fn_input_schema: type[BaseModel] = augmented_function.input_schema
+    fn_output_schema: type[BaseModel] | type[None] = augmented_function.single_output_schema
 
     runnable_llm = input_llm.with_structured_output(schema=fn_input_schema)