| 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] |
|
|
