fix: async invocation

This commit is contained in:
Marco Vinciguerra 2024-10-13 11:30:39 +02:00
parent 0e4ff09a10
commit c2179abc60
15 changed files with 23 additions and 23 deletions

View File

@ -119,7 +119,7 @@ class GraphBuilder:
Returns: Returns:
dict: A JSON representation of the graph configuration. dict: A JSON representation of the graph configuration.
""" """
return self.chain.ainvoke(self.prompt) return self.chain.invoke(self.prompt)
@staticmethod @staticmethod
def convert_json_to_graphviz(json_data, format: str = 'pdf'): def convert_json_to_graphviz(json_data, format: str = 'pdf'):

View File

@ -126,7 +126,7 @@ class GenerateAnswerCSVNode(BaseNode):
) )
chain = prompt | self.llm_model | output_parser chain = prompt | self.llm_model | output_parser
answer = chain.ainvoke({"question": user_prompt}) answer = chain.invoke({"question": user_prompt})
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})
return state return state

View File

@ -143,7 +143,7 @@ class GenerateAnswerNodeKLevel(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 = merge_chain.invoke({"context": batch_results, "question": user_prompt})
state["answer"] = answer state["answer"] = answer

View File

@ -117,7 +117,7 @@ class GenerateAnswerOmniNode(BaseNode):
) )
chain = prompt | self.llm_model | output_parser chain = prompt | self.llm_model | output_parser
answer = chain.ainvoke({"question": user_prompt}) answer = chain.invoke({"question": user_prompt})
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})
return state return state
@ -149,7 +149,7 @@ class GenerateAnswerOmniNode(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 = merge_chain.invoke({"context": batch_results, "question": user_prompt})
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})
return state return state

View File

@ -325,7 +325,7 @@ class GenerateCodeNode(BaseNode):
output_parser = StrOutputParser() output_parser = StrOutputParser()
chain = prompt | self.llm_model | output_parser chain = prompt | self.llm_model | output_parser
generated_code = chain.ainvoke({}) generated_code = chain.invoke({})
return generated_code return generated_code
def semantic_comparison(self, generated_result: Any, reference_result: Any) -> Dict[str, Any]: def semantic_comparison(self, generated_result: Any, reference_result: Any) -> Dict[str, Any]:
@ -368,7 +368,7 @@ class GenerateCodeNode(BaseNode):
) )
chain = prompt | self.llm_model | output_parser chain = prompt | self.llm_model | output_parser
return chain.ainvoke({ return chain.invoke({
"generated_result": json.dumps(generated_result, indent=2), "generated_result": json.dumps(generated_result, indent=2),
"reference_result": json.dumps(reference_result_dict, indent=2) "reference_result": json.dumps(reference_result_dict, indent=2)
}) })

View File

@ -130,7 +130,7 @@ class GenerateScraperNode(BaseNode):
) )
map_chain = prompt | self.llm_model | StrOutputParser() map_chain = prompt | self.llm_model | StrOutputParser()
answer = map_chain.ainvoke({"question": user_prompt}) answer = map_chain.invoke({"question": user_prompt})
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})
return state return state

View File

@ -93,7 +93,7 @@ class HtmlAnalyzerNode(BaseNode):
output_parser = StrOutputParser() output_parser = StrOutputParser()
chain = prompt | self.llm_model | output_parser chain = prompt | self.llm_model | output_parser
html_analysis = chain.ainvoke({}) html_analysis = chain.invoke({})
state.update({self.output[0]: html_analysis, self.output[1]: reduced_html}) state.update({self.output[0]: html_analysis, self.output[1]: reduced_html})
return state return state

View File

@ -95,7 +95,7 @@ class MergeAnswersNode(BaseNode):
) )
merge_chain = prompt_template | self.llm_model | output_parser merge_chain = prompt_template | self.llm_model | output_parser
answer = merge_chain.ainvoke({"user_prompt": user_prompt}) answer = merge_chain.invoke({"user_prompt": user_prompt})
answer["sources"] = state.get("urls", []) answer["sources"] = state.get("urls", [])
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})

View File

@ -74,7 +74,7 @@ class MergeGeneratedScriptsNode(BaseNode):
) )
merge_chain = prompt_template | self.llm_model | StrOutputParser() merge_chain = prompt_template | self.llm_model | StrOutputParser()
answer = merge_chain.ainvoke({"user_prompt": user_prompt}) answer = merge_chain.invoke({"user_prompt": user_prompt})
state.update({self.output[0]: answer}) state.update({self.output[0]: answer})
return state return state

View File

@ -96,7 +96,7 @@ class PromptRefinerNode(BaseNode):
output_parser = StrOutputParser() output_parser = StrOutputParser()
chain = prompt | self.llm_model | output_parser chain = prompt | self.llm_model | output_parser
refined_prompt = chain.ainvoke({}) refined_prompt = chain.invoke({})
state.update({self.output[0]: refined_prompt}) state.update({self.output[0]: refined_prompt})
return state return state

View File

@ -91,7 +91,7 @@ class ReasoningNode(BaseNode):
output_parser = StrOutputParser() output_parser = StrOutputParser()
chain = prompt | self.llm_model | output_parser chain = prompt | self.llm_model | output_parser
refined_prompt = chain.ainvoke({}) refined_prompt = chain.invoke({})
state.update({self.output[0]: refined_prompt}) state.update({self.output[0]: refined_prompt})
return state return state

View File

@ -108,7 +108,7 @@ class RobotsNode(BaseNode):
) )
chain = prompt | self.llm_model | output_parser chain = prompt | self.llm_model | output_parser
is_scrapable = chain.ainvoke({"path": source})[0] is_scrapable = chain.invoke({"path": source})[0]
if "no" in is_scrapable: if "no" in is_scrapable:
self.logger.warning( self.logger.warning(

View File

@ -142,7 +142,7 @@ class SearchLinkNode(BaseNode):
input_variables=["content", "user_prompt"], input_variables=["content", "user_prompt"],
) )
merge_chain = merge_prompt | self.llm_model | output_parser merge_chain = merge_prompt | self.llm_model | output_parser
answer = merge_chain.ainvoke( answer = merge_chain.invoke(
{"content": chunk.page_content} {"content": chunk.page_content}
) )
relevant_links += answer relevant_links += answer

View File

@ -33,7 +33,7 @@ def syntax_focused_analysis(state: dict, llm_model) -> str:
prompt = PromptTemplate(template=TEMPLATE_SYNTAX_ANALYSIS, prompt = PromptTemplate(template=TEMPLATE_SYNTAX_ANALYSIS,
input_variables=["generated_code", "errors"]) input_variables=["generated_code", "errors"])
chain = prompt | llm_model | StrOutputParser() chain = prompt | llm_model | StrOutputParser()
return chain.ainvoke({ return chain.invoke({
"generated_code": state["generated_code"], "generated_code": state["generated_code"],
"errors": state["errors"]["syntax"] "errors": state["errors"]["syntax"]
}) })
@ -53,7 +53,7 @@ def execution_focused_analysis(state: dict, llm_model) -> str:
input_variables=["generated_code", "errors", input_variables=["generated_code", "errors",
"html_code", "html_analysis"]) "html_code", "html_analysis"])
chain = prompt | llm_model | StrOutputParser() chain = prompt | llm_model | StrOutputParser()
return chain.ainvoke({ return chain.invoke({
"generated_code": state["generated_code"], "generated_code": state["generated_code"],
"errors": state["errors"]["execution"], "errors": state["errors"]["execution"],
"html_code": state["html_code"], "html_code": state["html_code"],
@ -76,7 +76,7 @@ def validation_focused_analysis(state: dict, llm_model) -> str:
input_variables=["generated_code", "errors", input_variables=["generated_code", "errors",
"json_schema", "execution_result"]) "json_schema", "execution_result"])
chain = prompt | llm_model | StrOutputParser() chain = prompt | llm_model | StrOutputParser()
return chain.ainvoke({ return chain.invoke({
"generated_code": state["generated_code"], "generated_code": state["generated_code"],
"errors": state["errors"]["validation"], "errors": state["errors"]["validation"],
"json_schema": state["json_schema"], "json_schema": state["json_schema"],
@ -100,7 +100,7 @@ def semantic_focused_analysis(state: dict, comparison_result: Dict[str, Any], ll
input_variables=["generated_code", input_variables=["generated_code",
"differences", "explanation"]) "differences", "explanation"])
chain = prompt | llm_model | StrOutputParser() chain = prompt | llm_model | StrOutputParser()
return chain.ainvoke({ return chain.invoke({
"generated_code": state["generated_code"], "generated_code": state["generated_code"],
"differences": json.dumps(comparison_result["differences"], indent=2), "differences": json.dumps(comparison_result["differences"], indent=2),
"explanation": comparison_result["explanation"] "explanation": comparison_result["explanation"]

View File

@ -32,7 +32,7 @@ def syntax_focused_code_generation(state: dict, analysis: str, llm_model) -> str
prompt = PromptTemplate(template=TEMPLATE_SYNTAX_CODE_GENERATION, prompt = PromptTemplate(template=TEMPLATE_SYNTAX_CODE_GENERATION,
input_variables=["analysis", "generated_code"]) input_variables=["analysis", "generated_code"])
chain = prompt | llm_model | StrOutputParser() chain = prompt | llm_model | StrOutputParser()
return chain.ainvoke({ return chain.invoke({
"analysis": analysis, "analysis": analysis,
"generated_code": state["generated_code"] "generated_code": state["generated_code"]
}) })
@ -52,7 +52,7 @@ def execution_focused_code_generation(state: dict, analysis: str, llm_model) ->
prompt = PromptTemplate(template=TEMPLATE_EXECUTION_CODE_GENERATION, prompt = PromptTemplate(template=TEMPLATE_EXECUTION_CODE_GENERATION,
input_variables=["analysis", "generated_code"]) input_variables=["analysis", "generated_code"])
chain = prompt | llm_model | StrOutputParser() chain = prompt | llm_model | StrOutputParser()
return chain.ainvoke({ return chain.invoke({
"analysis": analysis, "analysis": analysis,
"generated_code": state["generated_code"] "generated_code": state["generated_code"]
}) })
@ -72,7 +72,7 @@ def validation_focused_code_generation(state: dict, analysis: str, llm_model) ->
prompt = PromptTemplate(template=TEMPLATE_VALIDATION_CODE_GENERATION, prompt = PromptTemplate(template=TEMPLATE_VALIDATION_CODE_GENERATION,
input_variables=["analysis", "generated_code", "json_schema"]) input_variables=["analysis", "generated_code", "json_schema"])
chain = prompt | llm_model | StrOutputParser() chain = prompt | llm_model | StrOutputParser()
return chain.ainvoke({ return chain.invoke({
"analysis": analysis, "analysis": analysis,
"generated_code": state["generated_code"], "generated_code": state["generated_code"],
"json_schema": state["json_schema"] "json_schema": state["json_schema"]
@ -93,7 +93,7 @@ def semantic_focused_code_generation(state: dict, analysis: str, llm_model) -> s
prompt = PromptTemplate(template=TEMPLATE_SEMANTIC_CODE_GENERATION, prompt = PromptTemplate(template=TEMPLATE_SEMANTIC_CODE_GENERATION,
input_variables=["analysis", "generated_code", "generated_result", "reference_result"]) input_variables=["analysis", "generated_code", "generated_result", "reference_result"])
chain = prompt | llm_model | StrOutputParser() chain = prompt | llm_model | StrOutputParser()
return chain.ainvoke({ return chain.invoke({
"analysis": analysis, "analysis": analysis,
"generated_code": state["generated_code"], "generated_code": state["generated_code"],
"generated_result": json.dumps(state["execution_result"], indent=2), "generated_result": json.dumps(state["execution_result"], indent=2),