fix: parse_node

This commit is contained in:
Marco Vinciguerra 2024-07-17 22:58:21 +02:00
parent 68f58cc4dd
commit 07f1e23d23
4 changed files with 29 additions and 17 deletions

View File

@ -29,7 +29,7 @@ graph_config = {
smart_scraper_graph = SmartScraperGraph( smart_scraper_graph = SmartScraperGraph(
prompt="List me all the titles", prompt="List me all the titles",
source="https://sport.sky.it/nba?gr=www", source="https://perinim.github.io/projects",
config=graph_config config=graph_config
) )

View File

@ -88,7 +88,6 @@ class BaseNode(ABC):
param (dict): The dictionary to update node_config with. param (dict): The dictionary to update node_config with.
overwrite (bool): Flag indicating if the values of node_config should be overwritten if their value is not None. overwrite (bool): Flag indicating if the values of node_config should be overwritten if their value is not None.
""" """
for key, val in params.items(): for key, val in params.items():
if hasattr(self, key) and not overwrite: if hasattr(self, key) and not overwrite:
continue continue

View File

@ -121,7 +121,7 @@ class GenerateAnswerNode(BaseNode):
answer = chain.invoke({"question": user_prompt}) answer = chain.invoke({"question": user_prompt})
break break
prompt = PromptTemplate( prompt = PromptTemplate(
template=template_chunks_prompt, template=template_chunks_prompt,
input_variables=["question"], input_variables=["question"],
partial_variables={"context": chunk, partial_variables={"context": chunk,

View File

@ -50,35 +50,48 @@ class ParseNode(BaseNode):
Args: Args:
state (dict): The current state of the graph. The input keys will be used to fetch the state (dict): The current state of the graph. The input keys will be used to fetch the
correct data from the state. correct data from the state.
Returns: Returns:
dict: The updated state with the output key containing the parsed content chunks. dict: The updated state with the output key containing the parsed content chunks.
Raises: Raises:
KeyError: If the input keys are not found in the state. KeyError: If the input keys are not found in the state, indicating that the
necessary information for parsing the content is missing.
""" """
self.logger.info(f"--- Executing {self.node_name} Node ---") self.logger.info(f"--- Executing {self.node_name} Node ---")
# Fetch data using input keys # Interpret input keys based on the provided input expression
input_keys = self.get_input_keys(state) input_keys = self.get_input_keys(state)
input_data = [state[key] for key in input_keys]
docs_transformed = input_data[0]
# Parse HTML if enabled # Fetching data from the state based on the input keys
input_data = [state[key] for key in input_keys]
# Parse the document
docs_transformed = input_data[0]
if self.parse_html: if self.parse_html:
docs_transformed = Html2TextTransformer().transform_documents(input_data[0]) docs_transformed = Html2TextTransformer().transform_documents(input_data[0])
docs_transformed = docs_transformed[0] docs_transformed = docs_transformed[0]
# Get text content chunks = chunk(text=docs_transformed.page_content,
text_content = docs_transformed.page_content if type(docs_transformed) == Document else docs_transformed chunk_size= self.node_config.get("chunk_size", 4096),
token_counter=lambda x: len(x),
memoize=False)
else:
docs_transformed = docs_transformed[0]
# Chunk the text if type(docs_transformed) == Document:
chunk_size = self.node_config.get("chunk_size", 4096) - 250 chunks = chunk(text=docs_transformed.page_content,
chunks = chunk(text=text_content, chunk_size=chunk_size, token_counter=lambda x: len(x.split()), memoize=False) chunk_size= self.node_config.get("chunk_size", 4096),
token_counter=lambda x: len(x),
# Update state with chunks memoize=False)
else:
chunks = chunk(text=docs_transformed,
chunk_size= self.node_config.get("chunk_size", 4096),
token_counter=lambda x: len(x),
memoize=False)
state.update({self.output[0]: chunks}) state.update({self.output[0]: chunks})
return state return state