3
3
"""
4
4
5
5
import time
6
+ import json
6
7
from typing import List , Optional
7
8
8
9
from langchain .prompts import PromptTemplate
@@ -120,7 +121,11 @@ def execute(self, state: dict) -> dict:
120
121
else :
121
122
if not isinstance (self .llm_model , ChatBedrock ):
122
123
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
+ )
124
129
else :
125
130
output_parser = None
126
131
format_instructions = ""
@@ -131,13 +136,25 @@ def execute(self, state: dict) -> dict:
131
136
and not self .script_creator
132
137
or self .is_md_scraper
133
138
):
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
+ )
137
148
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
+ )
141
158
142
159
if self .additional_info is not None :
143
160
template_no_chunks_prompt = self .additional_info + template_no_chunks_prompt
@@ -161,8 +178,9 @@ def execute(self, state: dict) -> dict:
161
178
answer = self .invoke_with_timeout (
162
179
chain , {"question" : user_prompt }, self .timeout
163
180
)
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 )}})
166
184
return state
167
185
168
186
state .update ({self .output [0 ]: answer })
@@ -191,14 +209,9 @@ def execute(self, state: dict) -> dict:
191
209
batch_results = self .invoke_with_timeout (
192
210
async_runner , {"question" : user_prompt }, self .timeout
193
211
)
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 )}})
202
215
return state
203
216
204
217
merge_prompt = PromptTemplate (
@@ -216,10 +229,9 @@ def execute(self, state: dict) -> dict:
216
229
{"context" : batch_results , "question" : user_prompt },
217
230
self .timeout ,
218
231
)
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 )}})
223
235
return state
224
236
225
237
state .update ({self .output [0 ]: answer })
0 commit comments