33"""
44
55import time
6+ import json
67from typing import List , Optional
78
89from 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 \n IMPORTANT: " + format_instructions
141+ )
142+ template_chunks_prompt = (
143+ TEMPLATE_CHUNKS_MD + "\n \n IMPORTANT: " + format_instructions
144+ )
145+ template_merge_prompt = (
146+ TEMPLATE_MERGE_MD + "\n \n IMPORTANT: " + 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 \n IMPORTANT: " + format_instructions
151+ )
152+ template_chunks_prompt = (
153+ TEMPLATE_CHUNKS + "\n \n IMPORTANT: " + format_instructions
154+ )
155+ template_merge_prompt = (
156+ TEMPLATE_MERGE + "\n \n IMPORTANT: " + 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