feat: add the new handling exception

This commit is contained in:
Marco Vinciguerra 2025-02-17 12:09:32 +01:00
parent 83be82a11e
commit 5c0bc46c63
2 changed files with 34 additions and 21 deletions

View File

@ -80,6 +80,7 @@ models_tokens = {
"llama3.2": 128000, "llama3.2": 128000,
"llama3.2:1b": 128000, "llama3.2:1b": 128000,
"llama3.2:3b": 128000, "llama3.2:3b": 128000,
"llama3.3": 128000,
"llama3.3:70b": 128000, "llama3.3:70b": 128000,
"scrapegraph": 8192, "scrapegraph": 8192,
"mistral-small": 128000, "mistral-small": 128000,

View File

@ -3,6 +3,7 @@ GenerateAnswerNode Module
""" """
import time import time
import json
from typing import List, Optional from typing import List, Optional
from langchain.prompts import PromptTemplate from langchain.prompts import PromptTemplate
@ -120,7 +121,11 @@ class GenerateAnswerNode(BaseNode):
else: else:
if not isinstance(self.llm_model, ChatBedrock): if not isinstance(self.llm_model, ChatBedrock):
output_parser = JsonOutputParser() output_parser = JsonOutputParser()
format_instructions = output_parser.get_format_instructions() format_instructions = (
"You must respond with a JSON object. Your response should be formatted as a valid JSON "
"with a 'content' field containing your analysis. For example:\n"
'{"content": "your analysis here"}'
)
else: else:
output_parser = None output_parser = None
format_instructions = "" format_instructions = ""
@ -131,13 +136,25 @@ class GenerateAnswerNode(BaseNode):
and not self.script_creator and not self.script_creator
or self.is_md_scraper or self.is_md_scraper
): ):
template_no_chunks_prompt = TEMPLATE_NO_CHUNKS_MD template_no_chunks_prompt = (
template_chunks_prompt = TEMPLATE_CHUNKS_MD TEMPLATE_NO_CHUNKS_MD + "\n\nIMPORTANT: " + format_instructions
template_merge_prompt = TEMPLATE_MERGE_MD )
template_chunks_prompt = (
TEMPLATE_CHUNKS_MD + "\n\nIMPORTANT: " + format_instructions
)
template_merge_prompt = (
TEMPLATE_MERGE_MD + "\n\nIMPORTANT: " + format_instructions
)
else: else:
template_no_chunks_prompt = TEMPLATE_NO_CHUNKS template_no_chunks_prompt = (
template_chunks_prompt = TEMPLATE_CHUNKS TEMPLATE_NO_CHUNKS + "\n\nIMPORTANT: " + format_instructions
template_merge_prompt = TEMPLATE_MERGE )
template_chunks_prompt = (
TEMPLATE_CHUNKS + "\n\nIMPORTANT: " + format_instructions
)
template_merge_prompt = (
TEMPLATE_MERGE + "\n\nIMPORTANT: " + format_instructions
)
if self.additional_info is not None: if self.additional_info is not None:
template_no_chunks_prompt = self.additional_info + template_no_chunks_prompt template_no_chunks_prompt = self.additional_info + template_no_chunks_prompt
@ -161,8 +178,9 @@ class GenerateAnswerNode(BaseNode):
answer = self.invoke_with_timeout( answer = self.invoke_with_timeout(
chain, {"question": user_prompt}, self.timeout chain, {"question": user_prompt}, self.timeout
) )
except Timeout: except (Timeout, json.JSONDecodeError) as e:
state.update({self.output[0]: {"error": "Response timeout exceeded"}}) error_msg = "Response timeout exceeded" if isinstance(e, Timeout) else "Invalid JSON response format"
state.update({self.output[0]: {"error": error_msg, "raw_response": str(e)}})
return state return state
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})
@ -191,14 +209,9 @@ class GenerateAnswerNode(BaseNode):
batch_results = self.invoke_with_timeout( batch_results = self.invoke_with_timeout(
async_runner, {"question": user_prompt}, self.timeout async_runner, {"question": user_prompt}, self.timeout
) )
except Timeout: except (Timeout, json.JSONDecodeError) as e:
state.update( error_msg = "Response timeout exceeded during chunk processing" if isinstance(e, Timeout) else "Invalid JSON response format in chunk processing"
{ state.update({self.output[0]: {"error": error_msg, "raw_response": str(e)}})
self.output[0]: {
"error": "Response timeout exceeded during chunk processing"
}
}
)
return state return state
merge_prompt = PromptTemplate( merge_prompt = PromptTemplate(
@ -216,10 +229,9 @@ class GenerateAnswerNode(BaseNode):
{"context": batch_results, "question": user_prompt}, {"context": batch_results, "question": user_prompt},
self.timeout, self.timeout,
) )
except Timeout: except (Timeout, json.JSONDecodeError) as e:
state.update( error_msg = "Response timeout exceeded during merge" if isinstance(e, Timeout) else "Invalid JSON response format during merge"
{self.output[0]: {"error": "Response timeout exceeded during merge"}} state.update({self.output[0]: {"error": error_msg, "raw_response": str(e)}})
)
return state return state
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})