Skip to content

Commit ac9ceaf

Browse files
committed
fix: prevent infinite loop in extract_toc_content
The while loop exit condition used len(chat_history), but chat_history was rebuilt every iteration with exactly 2 elements, making the check len(chat_history) > 5 never true. Replace with explicit attempt counter and max_attempts limit.
1 parent 823c11e commit ac9ceaf

1 file changed

Lines changed: 9 additions & 6 deletions

File tree

pageindex/page_index.py

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -180,19 +180,22 @@ def extract_toc_content(content, model=None):
180180
response = response + new_response
181181
if_complete = check_if_toc_transformation_is_complete(content, response, model)
182182

183+
attempt = 0
184+
max_attempts = 5
185+
183186
while not (if_complete == "yes" and finish_reason == "finished"):
187+
attempt += 1
188+
if attempt > max_attempts:
189+
raise Exception('Failed to complete table of contents after maximum retries')
190+
184191
chat_history = [
185-
{"role": "user", "content": prompt},
186-
{"role": "assistant", "content": response},
192+
{"role": "user", "content": prompt},
193+
{"role": "assistant", "content": response},
187194
]
188195
prompt = f"""please continue the generation of table of contents , directly output the remaining part of the structure"""
189196
new_response, finish_reason = ChatGPT_API_with_finish_reason(model=model, prompt=prompt, chat_history=chat_history)
190197
response = response + new_response
191198
if_complete = check_if_toc_transformation_is_complete(content, response, model)
192-
193-
# Optional: Add a maximum retry limit to prevent infinite loops
194-
if len(chat_history) > 5: # Arbitrary limit of 10 attempts
195-
raise Exception('Failed to complete table of contents after maximum retries')
196199

197200
return response
198201

0 commit comments

Comments
 (0)