diff --git a/pyproject.toml b/pyproject.toml index b5b22bb1..e8cf382e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,7 @@ keywords = [ "web scraping tool", "webscraping", "graph", + "llm" ] classifiers = [ "Intended Audience :: Developers", diff --git a/scrapegraphai/nodes/generate_answer_csv_node.py b/scrapegraphai/nodes/generate_answer_csv_node.py index ed58d4ba..11ab15b9 100644 --- a/scrapegraphai/nodes/generate_answer_csv_node.py +++ b/scrapegraphai/nodes/generate_answer_csv_node.py @@ -60,7 +60,7 @@ class GenerateAnswerCSVNode(BaseNode): 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 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 - 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}) return state diff --git a/scrapegraphai/nodes/generate_answer_node.py b/scrapegraphai/nodes/generate_answer_node.py index 332f9c30..384d811d 100644 --- a/scrapegraphai/nodes/generate_answer_node.py +++ b/scrapegraphai/nodes/generate_answer_node.py @@ -1,3 +1,6 @@ +""" +GenerateAnswerNode Module +""" from typing import List, Optional from langchain.prompts import PromptTemplate from langchain_core.output_parsers import JsonOutputParser @@ -15,6 +18,26 @@ from ..prompts import ( ) 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__( self, input: str, @@ -34,7 +57,17 @@ class GenerateAnswerNode(BaseNode): self.is_md_scraper = node_config.get("is_md_scraper", False) 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 ---") input_keys = self.get_input_keys(state) @@ -90,7 +123,7 @@ class GenerateAnswerNode(BaseNode): chain = prompt | self.llm_model if 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}) return state @@ -110,7 +143,7 @@ class GenerateAnswerNode(BaseNode): chains_dict[chain_name] = chains_dict[chain_name] | output_parser 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( template=template_merge_prompt, @@ -121,7 +154,7 @@ class GenerateAnswerNode(BaseNode): merge_chain = merge_prompt | self.llm_model if 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}) return state