Skip to content

Commit 075e917

Browse files
TimothyZhang7TimothyZhang7
committed
fix: retry logic broken by merge conflict
1 parent e61bdfc commit 075e917

File tree

1 file changed

+108
-83
lines changed

1 file changed

+108
-83
lines changed

‎core/framework/graph/node.py‎

Lines changed: 108 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -807,97 +807,122 @@ def executor(tool_use: ToolUse) -> ToolResult:
807807
f" ⚠ Response still truncated after {compaction_attempt} compaction attempts"
808808
)
809809

810-
total_input_tokens += response.input_tokens
811-
total_output_tokens += response.output_tokens
812-
813-
# Log the response
814-
response_preview = (
815-
response.content[:200] if len(response.content) > 200 else response.content
816-
)
817-
if len(response.content) > 200:
818-
response_preview += "..."
819-
logger.info(f" ← Response: {response_preview}")
820-
821-
# If no output_model, break immediately (no validation needed)
822-
if ctx.node_spec.output_model is None:
823-
break
810+
# Phase 2: Validation retry loop for Pydantic models
811+
max_validation_retries = ctx.node_spec.max_validation_retries if ctx.node_spec.output_model else 0
812+
validation_attempt = 0
813+
total_input_tokens = 0
814+
total_output_tokens = 0
815+
current_messages = messages.copy()
816+
817+
while True:
818+
total_input_tokens += response.input_tokens
819+
total_output_tokens += response.output_tokens
820+
821+
# Log the response
822+
response_preview = (
823+
response.content[:200] if len(response.content) > 200 else response.content
824+
)
825+
if len(response.content) > 200:
826+
response_preview += "..."
827+
logger.info(f" ← Response: {response_preview}")
824828

825-
# Try to parse and validate the response
826-
try:
827-
import json
828-
parsed = self._extract_json(response.content, ctx.node_spec.output_keys)
829+
# If no output_model, break immediately (no validation needed)
830+
if ctx.node_spec.output_model is None:
831+
break
829832

830-
if isinstance(parsed, dict):
831-
from framework.graph.validator import OutputValidator
832-
validator = OutputValidator()
833-
validation_result, validated_model = validator.validate_with_pydantic(
834-
parsed, ctx.node_spec.output_model
835-
)
833+
# Try to parse and validate the response
834+
try:
835+
import json
836+
parsed = self._extract_json(response.content, ctx.node_spec.output_keys)
836837

837-
if validation_result.success:
838-
# Validation passed, break out of retry loop
839-
model_name = ctx.node_spec.output_model.__name__
840-
logger.info(f" ✓ Pydantic validation passed for {model_name}")
841-
break
842-
else:
843-
# Validation failed
844-
validation_attempt += 1
838+
if isinstance(parsed, dict):
839+
from framework.graph.validator import OutputValidator
840+
validator = OutputValidator()
841+
validation_result, validated_model = validator.validate_with_pydantic(
842+
parsed, ctx.node_spec.output_model
843+
)
845844

846-
if validation_attempt <= max_validation_retries:
847-
# Add validation feedback to messages and retry
848-
feedback = validator.format_validation_feedback(
849-
validation_result, ctx.node_spec.output_model
850-
)
851-
logger.warning(
852-
f" ⚠ Pydantic validation failed "
853-
f"(attempt {validation_attempt}/{max_validation_retries}): "
854-
f"{validation_result.error}"
845+
if validation_result.success:
846+
# Validation passed, break out of retry loop
847+
model_name = ctx.node_spec.output_model.__name__
848+
logger.info(f" ✓ Pydantic validation passed for {model_name}")
849+
break
850+
else:
851+
# Validation failed
852+
validation_attempt += 1
853+
854+
if validation_attempt <= max_validation_retries:
855+
# Add validation feedback to messages and retry
856+
feedback = validator.format_validation_feedback(
857+
validation_result, ctx.node_spec.output_model
858+
)
859+
logger.warning(
860+
f" ⚠ Pydantic validation failed "
861+
f"(attempt {validation_attempt}/{max_validation_retries}): "
862+
f"{validation_result.error}"
863+
)
864+
logger.info(" 🔄 Retrying with validation feedback...")
865+
866+
# Add the assistant's failed response and feedback
867+
current_messages.append({
868+
"role": "assistant",
869+
"content": response.content
870+
})
871+
current_messages.append({
872+
"role": "user",
873+
"content": feedback
874+
})
875+
876+
# Re-call LLM with feedback
877+
if ctx.available_tools and self.tool_executor:
878+
response = ctx.llm.complete_with_tools(
879+
messages=current_messages,
880+
system=system,
881+
tools=ctx.available_tools,
882+
tool_executor=executor,
883+
max_tokens=ctx.max_tokens,
855884
)
856-
logger.info(" 🔄 Retrying with validation feedback...")
857-
858-
# Add the assistant's failed response and feedback
859-
current_messages.append({
860-
"role": "assistant",
861-
"content": response.content
862-
})
863-
current_messages.append({
864-
"role": "user",
865-
"content": feedback
866-
})
867-
continue # Retry the LLM call
868885
else:
869-
# Max retries exceeded
870-
latency_ms = int((time.time() - start) * 1000)
871-
err = validation_result.error
872-
logger.error(
873-
f" ✗ Pydantic validation failed after "
874-
f"{max_validation_retries} retries: {err}"
875-
)
876-
ctx.runtime.record_outcome(
877-
decision_id=decision_id,
878-
success=False,
879-
error=f"Validation failed: {validation_result.error}",
880-
tokens_used=total_input_tokens + total_output_tokens,
881-
latency_ms=latency_ms,
886+
response = ctx.llm.complete(
887+
messages=current_messages,
888+
system=system,
889+
json_mode=use_json_mode,
890+
max_tokens=ctx.max_tokens,
882891
)
883-
error_msg = (
884-
f"Pydantic validation failed after "
885-
f"{max_validation_retries} retries: {err}"
886-
)
887-
return NodeResult(
888-
success=False,
889-
error=error_msg,
890-
output=parsed,
891-
tokens_used=total_input_tokens + total_output_tokens,
892-
latency_ms=latency_ms,
893-
validation_errors=validation_result.errors,
894-
)
895-
else:
896-
# Not a dict, can't validate - break and let downstream handle
897-
break
898-
except Exception:
899-
# JSON extraction failed - break and let downstream handle
892+
continue # Retry validation
893+
else:
894+
# Max retries exceeded
895+
latency_ms = int((time.time() - start) * 1000)
896+
err = validation_result.error
897+
logger.error(
898+
f" ✗ Pydantic validation failed after "
899+
f"{max_validation_retries} retries: {err}"
900+
)
901+
ctx.runtime.record_outcome(
902+
decision_id=decision_id,
903+
success=False,
904+
error=f"Validation failed: {validation_result.error}",
905+
tokens_used=total_input_tokens + total_output_tokens,
906+
latency_ms=latency_ms,
907+
)
908+
error_msg = (
909+
f"Pydantic validation failed after "
910+
f"{max_validation_retries} retries: {err}"
911+
)
912+
return NodeResult(
913+
success=False,
914+
error=error_msg,
915+
output=parsed,
916+
tokens_used=total_input_tokens + total_output_tokens,
917+
latency_ms=latency_ms,
918+
validation_errors=validation_result.errors,
919+
)
920+
else:
921+
# Not a dict, can't validate - break and let downstream handle
900922
break
923+
except Exception:
924+
# JSON extraction failed - break and let downstream handle
925+
break
901926

902927
latency_ms = int((time.time() - start) * 1000)
903928

0 commit comments

Comments
 (0)