Skip to content

Commit 5c0bc46

Browse files
committed
feat: add the new handling exception
1 parent 83be82a commit 5c0bc46

File tree

2 files changed

+34
-21
lines changed

2 files changed

+34
-21
lines changed

scrapegraphai/helpers/models_tokens.py

+1
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@
8080
"llama3.2": 128000,
8181
"llama3.2:1b": 128000,
8282
"llama3.2:3b": 128000,
83+
"llama3.3": 128000,
8384
"llama3.3:70b": 128000,
8485
"scrapegraph": 8192,
8586
"mistral-small": 128000,

scrapegraphai/nodes/generate_answer_node.py

+33-21
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
"""
44

55
import time
6+
import json
67
from typing import List, Optional
78

89
from langchain.prompts import PromptTemplate
@@ -120,7 +121,11 @@ def execute(self, state: dict) -> dict:
120121
else:
121122
if not isinstance(self.llm_model, ChatBedrock):
122123
output_parser = JsonOutputParser()
123-
format_instructions = output_parser.get_format_instructions()
124+
format_instructions = (
125+
"You must respond with a JSON object. Your response should be formatted as a valid JSON "
126+
"with a 'content' field containing your analysis. For example:\n"
127+
'{"content": "your analysis here"}'
128+
)
124129
else:
125130
output_parser = None
126131
format_instructions = ""
@@ -131,13 +136,25 @@ def execute(self, state: dict) -> dict:
131136
and not self.script_creator
132137
or self.is_md_scraper
133138
):
134-
template_no_chunks_prompt = TEMPLATE_NO_CHUNKS_MD
135-
template_chunks_prompt = TEMPLATE_CHUNKS_MD
136-
template_merge_prompt = TEMPLATE_MERGE_MD
139+
template_no_chunks_prompt = (
140+
TEMPLATE_NO_CHUNKS_MD + "\n\nIMPORTANT: " + format_instructions
141+
)
142+
template_chunks_prompt = (
143+
TEMPLATE_CHUNKS_MD + "\n\nIMPORTANT: " + format_instructions
144+
)
145+
template_merge_prompt = (
146+
TEMPLATE_MERGE_MD + "\n\nIMPORTANT: " + format_instructions
147+
)
137148
else:
138-
template_no_chunks_prompt = TEMPLATE_NO_CHUNKS
139-
template_chunks_prompt = TEMPLATE_CHUNKS
140-
template_merge_prompt = TEMPLATE_MERGE
149+
template_no_chunks_prompt = (
150+
TEMPLATE_NO_CHUNKS + "\n\nIMPORTANT: " + format_instructions
151+
)
152+
template_chunks_prompt = (
153+
TEMPLATE_CHUNKS + "\n\nIMPORTANT: " + format_instructions
154+
)
155+
template_merge_prompt = (
156+
TEMPLATE_MERGE + "\n\nIMPORTANT: " + format_instructions
157+
)
141158

142159
if self.additional_info is not None:
143160
template_no_chunks_prompt = self.additional_info + template_no_chunks_prompt
@@ -161,8 +178,9 @@ def execute(self, state: dict) -> dict:
161178
answer = self.invoke_with_timeout(
162179
chain, {"question": user_prompt}, self.timeout
163180
)
164-
except Timeout:
165-
state.update({self.output[0]: {"error": "Response timeout exceeded"}})
181+
except (Timeout, json.JSONDecodeError) as e:
182+
error_msg = "Response timeout exceeded" if isinstance(e, Timeout) else "Invalid JSON response format"
183+
state.update({self.output[0]: {"error": error_msg, "raw_response": str(e)}})
166184
return state
167185

168186
state.update({self.output[0]: answer})
@@ -191,14 +209,9 @@ def execute(self, state: dict) -> dict:
191209
batch_results = self.invoke_with_timeout(
192210
async_runner, {"question": user_prompt}, self.timeout
193211
)
194-
except Timeout:
195-
state.update(
196-
{
197-
self.output[0]: {
198-
"error": "Response timeout exceeded during chunk processing"
199-
}
200-
}
201-
)
212+
except (Timeout, json.JSONDecodeError) as e:
213+
error_msg = "Response timeout exceeded during chunk processing" if isinstance(e, Timeout) else "Invalid JSON response format in chunk processing"
214+
state.update({self.output[0]: {"error": error_msg, "raw_response": str(e)}})
202215
return state
203216

204217
merge_prompt = PromptTemplate(
@@ -216,10 +229,9 @@ def execute(self, state: dict) -> dict:
216229
{"context": batch_results, "question": user_prompt},
217230
self.timeout,
218231
)
219-
except Timeout:
220-
state.update(
221-
{self.output[0]: {"error": "Response timeout exceeded during merge"}}
222-
)
232+
except (Timeout, json.JSONDecodeError) as e:
233+
error_msg = "Response timeout exceeded during merge" if isinstance(e, Timeout) else "Invalid JSON response format during merge"
234+
state.update({self.output[0]: {"error": error_msg, "raw_response": str(e)}})
223235
return state
224236

225237
state.update({self.output[0]: answer})

0 commit comments

Comments
 (0)