fix(schema): fixed json output

This commit is contained in:
Marco Perini 2024-06-05 15:07:21 +02:00
parent 4f53b09bf1
commit 5c9843f141
6 changed files with 14 additions and 35 deletions

View File

@ -185,10 +185,6 @@ idna==3.7
# via yarl
imagesize==1.4.1
# via sphinx
importlib-metadata==7.1.0
# via sphinx
importlib-resources==6.4.0
# via matplotlib
iniconfig==2.0.0
# via pytest
jinja2==3.1.4
@ -475,7 +471,6 @@ typing-extensions==4.12.0
# via pyee
# via sf-hamilton
# via sqlalchemy
# via starlette
# via streamlit
# via typer
# via typing-inspect
@ -507,6 +502,3 @@ win32-setctime==1.1.0
# via loguru
yarl==1.9.4
# via aiohttp
zipp==3.19.1
# via importlib-metadata
# via importlib-resources

View File

@ -8,7 +8,7 @@ from typing import List, Optional
# Imports from Langchain
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnableParallel
from tqdm import tqdm
@ -96,7 +96,7 @@ class GenerateAnswerCSVNode(BaseNode):
# Initialize the output parser
if self.node_config.get("schema", None) is not None:
output_parser = PydanticOutputParser(pydantic_object=self.node_config.get("schema", None))
output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
else:
output_parser = JsonOutputParser()
@ -150,9 +150,6 @@ class GenerateAnswerCSVNode(BaseNode):
single_chain = list(chains_dict.values())[0]
answer = single_chain.invoke({"question": user_prompt})
if type(answer) == PydanticOutputParser:
answer = answer.model_dump()
# Update the state with the generated answer
state.update({self.output[0]: answer})
return state

View File

@ -7,10 +7,11 @@ from typing import List, Optional
# Imports from Langchain
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnableParallel
from tqdm import tqdm
from ..utils.logging import get_logger
from ..models import Ollama
# Imports from the library
@ -81,8 +82,8 @@ class GenerateAnswerNode(BaseNode):
doc = input_data[1]
# Initialize the output parser
if self.node_config.get("schema",None) is not None:
output_parser = PydanticOutputParser(pydantic_object=self.node_config.get("schema", None))
if self.node_config.get("schema", None) is not None:
output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
else:
output_parser = JsonOutputParser()
@ -129,9 +130,6 @@ class GenerateAnswerNode(BaseNode):
single_chain = list(chains_dict.values())[0]
answer = single_chain.invoke({"question": user_prompt})
if type(answer) == PydanticOutputParser:
answer = answer.model_dump()
# Update the state with the generated answer
state.update({self.output[0]: answer})
return state

View File

@ -7,7 +7,7 @@ from typing import List, Optional
# Imports from Langchain
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnableParallel
from tqdm import tqdm
from ..models import Ollama
@ -82,7 +82,7 @@ class GenerateAnswerOmniNode(BaseNode):
# Initialize the output parser
if self.node_config.get("schema", None) is not None:
output_parser = PydanticOutputParser(pydantic_object=self.node_config.get("schema", None))
output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
else:
output_parser = JsonOutputParser()
@ -141,9 +141,6 @@ class GenerateAnswerOmniNode(BaseNode):
single_chain = list(chains_dict.values())[0]
answer = single_chain.invoke({"question": user_prompt})
if type(answer) == PydanticOutputParser:
answer = answer.model_dump()
# Update the state with the generated answer
state.update({self.output[0]: answer})
return state

View File

@ -7,7 +7,7 @@ from typing import List, Optional
# Imports from Langchain
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.runnables import RunnableParallel
from tqdm import tqdm
from ..models import Ollama
@ -96,8 +96,8 @@ class GenerateAnswerPDFNode(BaseNode):
doc = input_data[1]
# Initialize the output parser
if self.node_config.get("schema",None) is not None:
output_parser = PydanticOutputParser(pydantic_object=self.node_config.get("schema", None))
if self.node_config.get("schema", None) is not None:
output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
else:
output_parser = JsonOutputParser()

View File

@ -8,7 +8,7 @@ from tqdm import tqdm
# Imports from Langchain
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser, PydanticOutputParser
from langchain_core.output_parsers import JsonOutputParser
from tqdm import tqdm
from ..utils.logging import get_logger
@ -80,10 +80,8 @@ class MergeAnswersNode(BaseNode):
answers_str += f"CONTENT WEBSITE {i+1}: {answer}\n"
# Initialize the output parser
if self.node_config["schema"] is not None:
output_parser = PydanticOutputParser(
pydantic_object=self.node_config["schema"]
)
if self.node_config.get("schema", None) is not None:
output_parser = JsonOutputParser(pydantic_object=self.node_config["schema"])
else:
output_parser = JsonOutputParser()
@ -111,9 +109,6 @@ class MergeAnswersNode(BaseNode):
merge_chain = prompt_template | self.llm_model | output_parser
answer = merge_chain.invoke({"user_prompt": user_prompt})
if type(answer) == PydanticOutputParser:
answer = answer.model_dump()
# Update the state with the generated answer
state.update({self.output[0]: answer})
return state