refactor: merged find month and detect month in one function and added docstring

climateqa/engine/talk_to_data/input_processing.py

@@ -120,36 +120,6 @@ async def detect_year_with_openai(sentence: str) -> str:
     else:
         return ""
 
-async def detect_month_with_openai(sentence: str) -> dict[str, str]:
-    """
-    Detects month in a sentence using OpenAI's API via LangChain.
-    Returns the month as an integer string (e.g., "7" for July), or "" if not found.
-    """
-    llm = get_llm()
-    prompt = """
-    Extract the month (as a number from 1 to 12) mentioned in the following sentence.
-    Return the result as a Python list of integers. If no month is mentioned, return an empty list.
-
-    Sentence: "{sentence}"
-    """
-    prompt = ChatPromptTemplate.from_template(prompt)
-    structured_llm = llm.with_structured_output(ArrayOutput)
-    chain = prompt | structured_llm
-    response: ArrayOutput = await chain.ainvoke({"sentence": sentence})
-    months_list = eval(response['array'])
-    if len(months_list) > 0:
-        month_number = int(months_list[0])
-        month_name = calendar.month_name[month_number]
-        return {
-            "month_number": str(month_number),
-            "month_name": month_name
-        }
-    else:
-        return {
-            "month_number" : "",
-            "month_name" : ""
-        }
-
 
 async def detect_relevant_tables(user_question: str, plot: Plot, llm, table_names_list: list[str]) -> list[str]:
     """Identifies relevant tables for a plot based on user input.
@@ -259,11 +229,53 @@ async def find_year(user_input: str) -> str| None:
     return year
 
 async def find_month(user_input: str) -> dict[str, str|None]:
-    """Extracts month information from user input using LLM."""
-    print(f"---- Find month ---")
-    month_info = await detect_month_with_openai(user_input)
-    month_info = {key: None if value == "" else value for key, value in month_info.items()}
-    return month_info
+    """
+    Extracts month information from user input using an LLM.
+
+    This function analyzes the user's query to detect if a month is mentioned.
+    It returns both the month number (as a string, e.g. '7' for July) and the full English month name (e.g. 'July').
+    If no month is found, both values will be None.
+
+    Args:
+        user_input (str): The user's query text.
+
+    Returns:
+        dict[str, str|None]: A dictionary with keys:
+            - "month_number": the month number as a string (e.g. '7'), or None if not found
+            - "month_name": the full English month name (e.g. 'July'), or None if not found
+
+    Example:
+        >>> await find_month("Show me the temperature in Paris in July")
+        {'month_number': '7', 'month_name': 'July'}
+        >>> await find_month("Show me the temperature in Paris")
+        {'month_number': None, 'month_name': None}
+    """
+
+    llm = get_llm()
+    prompt = """
+    Extract the month (as a number from 1 to 12) mentioned in the following sentence.
+    Return the result as a Python list of integers. If no month is mentioned, return an empty list.
+
+    Sentence: "{sentence}"
+    """
+    prompt = ChatPromptTemplate.from_template(prompt)
+    structured_llm = llm.with_structured_output(ArrayOutput)
+    chain = prompt | structured_llm
+    response: ArrayOutput = await chain.ainvoke({"sentence": user_input})
+    months_list = ast.literal_eval(response['array'])
+    if len(months_list) > 0:
+        month_number = int(months_list[0])
+        month_name = calendar.month_name[month_number]
+        return {
+            "month_number": str(month_number),
+            "month_name": month_name
+        }
+    else:
+        return {
+            "month_number" : None,
+            "month_name" : None
+        }
+
 
 async def find_relevant_plots(state: State, llm, plots: list[Plot]) -> list[str]:
     print("---- Find relevant plots ----")
@@ -277,7 +289,26 @@ async def find_relevant_tables_per_plot(state: State, plot: Plot, llm, tables: l
 
 async def find_param(state: State, param_name: str, mode: Literal['DRIAS', 'IPCC'] = 'DRIAS') -> dict[str, Optional[str]] | Location | None:
     """
-    Perform the good method to retrieve the desired parameter.
+    Retrieves a specific parameter (location, year, month, etc.) from the user's input using the appropriate extraction method.
+
+    Args:
+        state (State): The current state containing at least the user's input under 'user_input'.
+        param_name (str): The name of the parameter to extract. Supported: 'location', 'year', 'month'.
+        mode (Literal['DRIAS', 'IPCC']): The data mode to use for location extraction.
+
+    Returns:
+        - For 'location': a Location object (dict with keys like 'location', 'latitude', etc.), or None if not found.
+        - For 'year': a dict {'year': year or None}.
+        - For 'month': a dict {'month_number': str or None, 'month_name': str or None}.
+        - None if the parameter is not recognized or not found.
+
+    Example:
+        >>> await find_param(state, 'location')
+        {'location': 'Paris', 'latitude': ..., ...}
+        >>> await find_param(state, 'year')
+        {'year': '2050'}
+        >>> await find_param(state, 'month')
+        {'month_number': '7', 'month_name': 'July'}
     """
     if param_name == 'location':
         location = await find_location(state['user_input'], mode)