from .model import InformationExtractedFromABillReceipt as PydanticModel

from langchain.chains import LLMChain
from langchain.chat_models import ChatOpenAI
from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
from langchain.prompts import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    SystemMessagePromptTemplate,
)

model = ChatOpenAI(
    temperature=0,
    n=1,
    model_kwargs= {
        'stop': None,
        'top_p': 1,
        'frequency_penalty': 0,
        'presence_penalty': 0,
    }
)

# Build categorizing chain
system_message_prompt = SystemMessagePromptTemplate.from_template(
    "You are an information extraction engine that outputs details from OCR processed "
    "documents such as date/time/place of departure and arrival.\n"
    "{format_instructions}"
)
human_message_prompt = HumanMessagePromptTemplate.from_template("{text}")
chat_prompt = ChatPromptTemplate.from_messages(
    [system_message_prompt, human_message_prompt]
)
output_parser = PydanticOutputParser(pydantic_object=PydanticModel)
fixing_parser = OutputFixingParser.from_llm(llm=model, parser=output_parser)
chain = LLMChain(
    llm=model, prompt=chat_prompt, output_parser=fixing_parser
)