fix(node): fixed generate answer node pydantic schema

This commit is contained in:
Marco Perini 2024-06-12 01:40:49 +02:00
parent 15421eff70
commit ab00f23d85

View File

@ -93,35 +93,20 @@ class GenerateAnswerNode(BaseNode):
# Use tqdm to add progress bar
for i, chunk in enumerate(tqdm(doc, desc="Processing chunks", disable=not self.verbose)):
if self.node_config.get("schema", None) is None and len(doc) == 1:
if len(doc) == 1:
prompt = PromptTemplate(
template=template_no_chunks,
input_variables=["question"],
partial_variables={"context": chunk.page_content,
"format_instructions": format_instructions})
elif self.node_config.get("schema", None) is not None and len(doc) == 1:
prompt = PromptTemplate(
template=template_no_chunks_with_schema,
input_variables=["question"],
partial_variables={"context": chunk.page_content,
"format_instructions": format_instructions,
"schema": self.node_config.get("schema", None)
})
elif self.node_config.get("schema", None) is None and len(doc) > 1:
else:
prompt = PromptTemplate(
template=template_chunks,
input_variables=["question"],
partial_variables={"context": chunk.page_content,
"chunk_id": i + 1,
"format_instructions": format_instructions})
elif self.node_config.get("schema", None) is not None and len(doc) > 1:
prompt = PromptTemplate(
template=template_chunks_with_schema,
input_variables=["question"],
partial_variables={"context": chunk.page_content,
"chunk_id": i + 1,
"format_instructions": format_instructions,
"schema": self.node_config.get("schema", None)})
# Dynamically name the chains based on their index
chain_name = f"chunk{i+1}"
@ -147,4 +132,4 @@ class GenerateAnswerNode(BaseNode):
# Update the state with the generated answer
state.update({self.output[0]: answer})
return state
return state