fix: address review findings and ty type checking

- Move re import to module level in evaluation_agent.py
- Add missing "Entering assess node" log for consistency
- Centralize max_total_iterations derivation in state.py and workflow.py
  (was duplicated 3x in main.py, now defaults to max_validation_attempts + 1)
- Update create_initial_state defaults (was stale at 5/10)
- Fix ty warnings: remove unused type: ignore comments
- Fix ty errors: add type: ignore for LangGraph/Starlette typing limitations
- Fix return type on get_default_path (-> str | None)
- Update test_state to match new default

Author: neuromechanist

src/agents/evaluation_agent.py

@@ -4,6 +4,7 @@
 the original natural language event description.
 """
 
+import re
 from pathlib import Path
 
 from langchain_core.language_models import BaseChatModel
@@ -186,8 +187,6 @@ def _parse_decision(self, feedback: str) -> bool:
         Returns:
             True if annotation should be accepted, False if needs refinement
         """
-        import re
-
         feedback_lower = feedback.lower()
 
         # Check for explicit DECISION line

src/agents/state.py

@@ -81,8 +81,8 @@ class HedAnnotationState(TypedDict):
 def create_initial_state(
     input_description: str,
     schema_version: str = "8.4.0",
-    max_validation_attempts: int = 5,
-    max_total_iterations: int = 10,
+    max_validation_attempts: int = 3,
+    max_total_iterations: int | None = None,
     run_assessment: bool = False,
     extracted_keywords: list[str] | None = None,
     semantic_hints: list[dict] | None = None,
@@ -93,8 +93,8 @@ def create_initial_state(
     Args:
         input_description: Natural language event description to annotate
         schema_version: HED schema version to use (default: "8.4.0")
-        max_validation_attempts: Maximum validation retry attempts (default: 5)
-        max_total_iterations: Maximum total iterations to prevent infinite loops (default: 10)
+        max_validation_attempts: Maximum validation retry attempts (default: 3)
+        max_total_iterations: Maximum total iterations (default: max_validation_attempts + 1)
         run_assessment: Whether to run final assessment (default: False)
         extracted_keywords: Pre-extracted keywords from description (optional)
         semantic_hints: Pre-computed semantic search hints (optional)
@@ -103,6 +103,9 @@ def create_initial_state(
     Returns:
         Initial HedAnnotationState
     """
+    if max_total_iterations is None:
+        max_total_iterations = max_validation_attempts + 1
+
     return HedAnnotationState(
         messages=[],
         input_description=input_description,

src/agents/workflow.py

@@ -105,7 +105,7 @@ def _build_graph(self) -> StateGraph:
             Compiled StateGraph
         """
         # Create graph
-        workflow = StateGraph(HedAnnotationState)
+        workflow = StateGraph(HedAnnotationState)  # type: ignore[arg-type]  # LangGraph typing limitation
 
         # Add nodes
         if self.enable_semantic_search:
@@ -292,6 +292,7 @@ async def _assess_node(self, state: HedAnnotationState) -> dict:
         Returns:
             State update
         """
+        print("[WORKFLOW] Entering assess node")
         t0 = time.monotonic()
         result = await self.assessment_agent.assess(state)
         elapsed = time.monotonic() - t0
@@ -388,7 +389,7 @@ async def run(
         input_description: str,
         schema_version: str = "8.4.0",
         max_validation_attempts: int = 3,
-        max_total_iterations: int = 4,
+        max_total_iterations: int | None = None,
         run_assessment: bool = False,
         no_extend: bool = False,
         config: dict | None = None,
@@ -399,7 +400,7 @@ async def run(
             input_description: Natural language event description
             schema_version: HED schema version to use
             max_validation_attempts: Maximum validation retry attempts
-            max_total_iterations: Maximum total iterations to prevent infinite loops
+            max_total_iterations: Maximum total iterations (default: max_validation_attempts + 1)
             run_assessment: Whether to run final assessment (default: False)
             no_extend: If True, prohibit tag extensions (use only existing vocabulary)
             config: Optional LangGraph config (e.g., recursion_limit)
@@ -409,6 +410,9 @@ async def run(
         """
         from src.agents.state import create_initial_state
 
+        if max_total_iterations is None:
+            max_total_iterations = max_validation_attempts + 1
+
         # Create initial state
         initial_state = create_initial_state(
             input_description,
@@ -422,4 +426,4 @@ async def run(
         # Run workflow
         final_state = await self.graph.ainvoke(initial_state, config=config)  # type: ignore[attr-defined]
 
-        return final_state  # type: ignore[no-any-return]
+        return final_state

src/api/main.py

@@ -283,7 +283,7 @@ async def lifespan(app: FastAPI):
     print("Initializing HEDit annotation workflow...")
 
     # Auto-detect environment (Docker vs local)
-    def get_default_path(docker_path: str, local_path: str) -> str:
+    def get_default_path(docker_path: str, local_path: str) -> str | None:
         """Get default path based on environment.
 
         Args:
@@ -473,7 +473,7 @@ def get_default_path(docker_path: str, local_path: str) -> str:
 
 # Add CORS middleware
 app.add_middleware(
-    CORSMiddleware,
+    CORSMiddleware,  # type: ignore[arg-type]  # Starlette typing limitation
     allow_origins=allowed_origins,
     allow_credentials=True,
     allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
@@ -642,15 +642,11 @@ async def annotate(
     try:
         config = {"recursion_limit": 50}
 
-        # Derive total iteration cap from validation attempts (+1 for evaluation refinement)
-        max_total_iterations = request.max_validation_attempts + 1
-
         start_time = time.time()
         final_state = await active_workflow.run(
             input_description=request.description,
             schema_version=request.schema_version,
             max_validation_attempts=request.max_validation_attempts,
-            max_total_iterations=max_total_iterations,
             run_assessment=request.run_assessment,
             config=config,
         )
@@ -842,13 +838,11 @@ async def annotate_from_image(
 
         # Step 2: Pass description through HED annotation workflow
         config = {"recursion_limit": 50}
-        img_max_total_iters = request.max_validation_attempts + 1
 
         final_state = await active_workflow.run(
             input_description=image_description,
             schema_version=request.schema_version,
             max_validation_attempts=request.max_validation_attempts,
-            max_total_iterations=img_max_total_iters,
             run_assessment=request.run_assessment,
             config=config,
         )
@@ -996,14 +990,12 @@ async def annotate_stream(
             raise HTTPException(status_code=503, detail="Workflow not initialized")
         active_workflow = workflow
 
-    # Create initial state with iteration cap derived from validation attempts
-    max_total_iterations = request.max_validation_attempts + 1
+    # Create initial state (max_total_iterations derived from max_validation_attempts + 1)
     initial_state = create_initial_state(
         request.description,
         request.schema_version,
         request.max_validation_attempts,
-        max_total_iterations,
-        request.run_assessment,
+        run_assessment=request.run_assessment,
     )
 
     # Node name to user-friendly stage mapping
@@ -1037,7 +1029,7 @@ def send_event(event_type: str, data: dict) -> str:
 
             # Use LangGraph's astream_events for real-time streaming
             config = {"recursion_limit": 50}
-            async for event in active_workflow.graph.astream_events(
+            async for event in active_workflow.graph.astream_events(  # type: ignore[union-attr]
                 initial_state, config=config, version="v2"
             ):
                 event_type = event.get("event")
@@ -1292,13 +1284,11 @@ def send_event(event_type: str, data: dict) -> str:
             )
 
             # Step 2: Create initial state for annotation workflow
-            img_max_total_iterations = request.max_validation_attempts + 1
             initial_state = create_initial_state(
                 image_description,
                 request.schema_version,
                 request.max_validation_attempts,
-                img_max_total_iterations,
-                request.run_assessment,
+                run_assessment=request.run_assessment,
             )
 
             # Track state and progress
@@ -1308,7 +1298,7 @@ def send_event(event_type: str, data: dict) -> str:
 
             # Use LangGraph's astream_events for real-time streaming
             config = {"recursion_limit": 50}
-            async for event in active_workflow.graph.astream_events(
+            async for event in active_workflow.graph.astream_events(  # type: ignore[union-attr]
                 initial_state, config=config, version="v2"
             ):
                 event_type = event.get("event")

src/utils/openrouter_llm.py

@@ -101,7 +101,7 @@ class CachingLLMWrapper(BaseChatModel):
 
     model_config = {"arbitrary_types_allowed": True}
 
-    def __init__(self, llm: BaseChatModel, **kwargs) -> None:  # type: ignore[no-untyped-def]
+    def __init__(self, llm: BaseChatModel, **kwargs) -> None:
         super().__init__(llm=llm, **kwargs)  # type: ignore[call-arg]
 
     @property

tests/test_state.py

@@ -15,7 +15,7 @@ def test_create_initial_state():
     assert state["is_valid"] is False
     assert state["is_faithful"] is False
     assert state["is_complete"] is False
-    assert state["max_validation_attempts"] == 5
+    assert state["max_validation_attempts"] == 3
     assert state["schema_version"] == "8.4.0"
     assert state["no_extend"] is False
     assert state["tag_suggestions"] == {}