|
import json |
|
|
|
import pytest |
|
from langchain import PromptTemplate, LLMChain |
|
from langchain.chat_models import ChatOpenAI |
|
from FOMCTexts.prompts import PROMPT_EXTRACT_DATE |
|
|
|
|
|
PROMPT_DATE = PromptTemplate.from_template(PROMPT_EXTRACT_DATE) |
|
date_extractor = LLMChain(prompt=PROMPT_DATE, llm=ChatOpenAI(temperature=0, model_name='gpt-3.5-turbo')) |
|
|
|
|
|
@pytest.mark.parametrize( |
|
"statement, expected_dict_date", |
|
[ |
|
('Summarize in two paragraphs the monetary policy outlook discussed by the participants during the meeting of June 2023', {"year": "2023", "month": "06"}), |
|
("What were the key takeaways from the conference in September 2022?", {"year": "2022", "month": "09"}), |
|
("Provide an update on the project status in March 2023.", {"year": "2023", "month": "03"}), |
|
("Who attended the meeting on May 2021?", {"year": "2021", "month": "05"}), |
|
("Please summarize the findings of the study conducted in January 2022.", {"year": "2022", "month": "01"}), |
|
("Who were the participants in the event on July 2020?", {"year": "2020", "month": "07"}), |
|
("Provide an overview of the report in August 2019.", {"year": "2019", "month": "08"}), |
|
("Who presented the findings of the research in November 2021?", {"year": "2021", "month": "11"}), |
|
("What were the discussions during the meeting in December 2024?", {"year": "2024", "month": "12"}) |
|
] |
|
) |
|
def test_date_extraction(statement, expected_dict_date): |
|
dict_date = json.loads(date_extractor.run(statement).replace('\n', '').replace(' ', '')) |
|
for k in dict_date: |
|
assert dict_date[k] == expected_dict_date[k] |
|
|