fix: bugs

This commit is contained in:
Marco Vinciguerra 2024-10-12 09:57:59 +02:00
parent 257f393761
commit 026a70bd3a
3 changed files with 40 additions and 6 deletions

View File

@ -60,6 +60,7 @@ keywords = [
"web scraping tool", "web scraping tool",
"webscraping", "webscraping",
"graph", "graph",
"llm"
] ]
classifiers = [ classifiers = [
"Intended Audience :: Developers", "Intended Audience :: Developers",

View File

@ -60,7 +60,7 @@ class GenerateAnswerCSVNode(BaseNode):
self.additional_info = node_config.get("additional_info") self.additional_info = node_config.get("additional_info")
def execute(self, state): async def execute(self, state):
""" """
Generates an answer by constructing a prompt from the user's input and the scraped Generates an answer by constructing a prompt from the user's input and the scraped
content, querying the language model, and parsing its response. content, querying the language model, and parsing its response.
@ -157,7 +157,7 @@ class GenerateAnswerCSVNode(BaseNode):
) )
merge_chain = merge_prompt | self.llm_model | output_parser merge_chain = merge_prompt | self.llm_model | output_parser
answer = merge_chain.ainvoke({"context": batch_results, "question": user_prompt}) answer = await merge_chain.ainvoke({"context": batch_results, "question": user_prompt})
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})
return state return state

View File

@ -1,3 +1,6 @@
"""
GenerateAnswerNode Module
"""
from typing import List, Optional from typing import List, Optional
from langchain.prompts import PromptTemplate from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser from langchain_core.output_parsers import JsonOutputParser
@ -15,6 +18,26 @@ from ..prompts import (
) )
class GenerateAnswerNode(BaseNode): class GenerateAnswerNode(BaseNode):
"""
Initializes the GenerateAnswerNode class.
Args:
input (str): The input data type for the node.
output (List[str]): The output data type(s) for the node.
node_config (Optional[dict]): Configuration dictionary for the node,
which includes the LLM model, verbosity, schema, and other settings.
Defaults to None.
node_name (str): The name of the node. Defaults to "GenerateAnswer".
Attributes:
llm_model: The language model specified in the node configuration.
verbose (bool): Whether verbose mode is enabled.
force (bool): Whether to force certain behaviors, overriding defaults.
script_creator (bool): Whether the node is in script creation mode.
is_md_scraper (bool): Whether the node is scraping markdown data.
additional_info (Optional[str]): Any additional information to be
included in the prompt templates.
"""
def __init__( def __init__(
self, self,
input: str, input: str,
@ -34,7 +57,17 @@ class GenerateAnswerNode(BaseNode):
self.is_md_scraper = node_config.get("is_md_scraper", False) self.is_md_scraper = node_config.get("is_md_scraper", False)
self.additional_info = node_config.get("additional_info") self.additional_info = node_config.get("additional_info")
def execute(self, state: dict) -> dict: async def execute(self, state: dict) -> dict:
"""
Executes the GenerateAnswerNode.
Args:
state (dict): The current state of the graph. The input keys will be used
to fetch the correct data from the state.
Returns:
dict: The updated state with the output key containing the generated answer.
"""
self.logger.info(f"--- Executing {self.node_name} Node ---") self.logger.info(f"--- Executing {self.node_name} Node ---")
input_keys = self.get_input_keys(state) input_keys = self.get_input_keys(state)
@ -90,7 +123,7 @@ class GenerateAnswerNode(BaseNode):
chain = prompt | self.llm_model chain = prompt | self.llm_model
if output_parser: if output_parser:
chain = chain | output_parser chain = chain | output_parser
answer = chain.ainvoke({"question": user_prompt}) answer = await chain.ainvoke({"question": user_prompt})
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})
return state return state
@ -110,7 +143,7 @@ class GenerateAnswerNode(BaseNode):
chains_dict[chain_name] = chains_dict[chain_name] | output_parser chains_dict[chain_name] = chains_dict[chain_name] | output_parser
async_runner = RunnableParallel(**chains_dict) async_runner = RunnableParallel(**chains_dict)
batch_results = async_runner.invoke({"question": user_prompt}) batch_results = await async_runner.ainvoke({"question": user_prompt})
merge_prompt = PromptTemplate( merge_prompt = PromptTemplate(
template=template_merge_prompt, template=template_merge_prompt,
@ -121,7 +154,7 @@ class GenerateAnswerNode(BaseNode):
merge_chain = merge_prompt | self.llm_model merge_chain = merge_prompt | self.llm_model
if output_parser: if output_parser:
merge_chain = merge_chain | output_parser merge_chain = merge_chain | output_parser
answer = merge_chain.ainvoke({"context": batch_results, "question": user_prompt}) answer = await merge_chain.ainvoke({"context": batch_results, "question": user_prompt})
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})
return state return state