Skip to content

Commit

Permalink
Merge pull request #778 from ScrapeGraphAI/refactoring-generate-answer
Browse files Browse the repository at this point in the history
feat: update generate answer
  • Loading branch information
f-aguzzi authored Oct 31, 2024
2 parents 7440f37 + 7172b32 commit 76e2820
Showing 1 changed file with 15 additions and 2 deletions.
17 changes: 15 additions & 2 deletions scrapegraphai/nodes/generate_answer_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
GenerateAnswerNode Module
"""
from typing import List, Optional
from json.decoder import JSONDecodeError
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnableParallel
Expand Down Expand Up @@ -121,9 +122,21 @@ def execute(self, state: dict) -> dict:
partial_variables={"context": doc, "format_instructions": format_instructions}
)
chain = prompt | self.llm_model
raw_response = str((prompt | self.llm_model).invoke({"question": user_prompt}))

if output_parser:
chain = chain | output_parser
answer = chain.invoke({"question": user_prompt})
try:
answer = output_parser.parse(raw_response)
except JSONDecodeError:
lines = raw_response.split('\n')
if lines[0].strip().startswith('```'):
lines = lines[1:]
if lines[-1].strip().endswith('```'):
lines = lines[:-1]
cleaned_response = '\n'.join(lines)
answer = output_parser.parse(cleaned_response)
else:
answer = raw_response

state.update({self.output[0]: answer})
return state
Expand Down

0 comments on commit 76e2820

Please sign in to comment.