feat: add multiseries support (#64)

Author: AzulGarza

tests/test_agent.py

@@ -22,19 +22,19 @@ def _response_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse
 def test_forecast_returns_expected_output(query):
     df = generate_series(n_series=1, freq="D", min_length=30)
     expected_output = {
-        "tsfeatures_results": ["mean: 0.5"],
         "tsfeatures_analysis": "ok",
         "selected_model": "ZeroModel",
         "model_details": "details",
-        "cross_validation_results": ["ZeroModel: 0.1"],
         "model_comparison": "cmp",
         "is_better_than_seasonal_naive": True,
         "reason_for_selection": "reason",
-        "forecast": ["2025-01-01: 1.0"],
         "forecast_analysis": "analysis",
         "user_query_response": query,
     }
     tc = TimeCopilot(llm=build_stub_llm(expected_output))
+    tc.fcst_df = None
+    tc.eval_df = None
+    tc.features_df = None
     result = tc.forecast(df=df, h=2, freq="D", seasonality=7, query=query)
 
     assert result.output == ForecastAgentOutput(**expected_output)

tests/test_live.py

@@ -3,29 +3,37 @@
 Keeping it separate from the other tests because costs and requires a live LLM.
 """
 
+import logfire
 import pytest
+from dotenv import load_dotenv
 from utilsforecast.data import generate_series
 
 from timecopilot import TimeCopilot
 
+load_dotenv()
+logfire.configure(send_to_logfire="if-token-present")
+logfire.instrument_pydantic_ai()
+
 
 @pytest.mark.live
-def test_forecast_returns_expected_output():
+@pytest.mark.parametrize("n_series", [1, 2])
+def test_forecast_returns_expected_output(n_series):
+    h = 2
     df = generate_series(
-        n_series=1,
+        n_series=n_series,
         freq="D",
         min_length=30,
         static_as_categorical=False,
     )
-    forecasting_agent = TimeCopilot(
+    tc = TimeCopilot(
         llm="openai:gpt-4o-mini",
         retries=3,
     )
-    result = forecasting_agent.forecast(
+    result = tc.forecast(
         df=df,
-        query="Please forecast the series with a horizon of 2 and frequency D.",
+        query=f"Please forecast the series with a horizon of {h} and frequency D.",
     )
-    assert len(result.output.forecast) == 2
+    assert len(result.fcst_df) == n_series * h
     assert result.output.is_better_than_seasonal_naive
     assert result.output.forecast_analysis is not None
     assert result.output.reason_for_selection is not None

timecopilot/agent.py

@@ -87,12 +87,6 @@
 class ForecastAgentOutput(BaseModel):
     """The output of the forecasting agent."""
 
-    tsfeatures_results: list[str] = Field(
-        description=(
-            "The time series features that were considered as a list of strings of "
-            "feature names and their values separated by colons."
-        )
-    )
     tsfeatures_analysis: str = Field(
         description=(
             "Analysis of what the time series features reveal about the data "
@@ -108,12 +102,6 @@ class ForecastAgentOutput(BaseModel):
             "strengths, and typical use cases."
         )
     )
-    cross_validation_results: list[str] = Field(
-        description=(
-            "The cross-validation results as a string of model names "
-            "and their scores separated by colons."
-        )
-    )
     model_comparison: str = Field(
         description=(
             "Detailed comparison of model performances, explaining why certain "
@@ -126,12 +114,6 @@ class ForecastAgentOutput(BaseModel):
     reason_for_selection: str = Field(
         description="Explanation for why the selected model was chosen"
     )
-    forecast: list[str] = Field(
-        description=(
-            "The forecasted values for the time series as a list of strings of "
-            "periods and their values separated by colons."
-        )
-    )
     forecast_analysis: str = Field(
         description=(
             "Detailed interpretation of the forecast, including trends, patterns, "
@@ -324,7 +306,7 @@ def __init__(
 
         3. Final Model Selection and Forecasting:
            - Choose the best performing model with clear justification
-           - Generate and analyze the forecast
+           - Generate the forecast using just the selected model
            - Interpret trends and patterns in the forecast
            - Discuss reliability and potential uncertainties
            - Address any specific aspects from the user's prompt
@@ -364,10 +346,17 @@ def __init__(
         )
 
         @self.forecasting_agent.system_prompt
-        async def add_time_series(ctx: RunContext[ExperimentDataset]) -> str:
+        async def add_time_series(
+            ctx: RunContext[ExperimentDataset],
+        ) -> str:
+            df_agg = ctx.deps.df.groupby("unique_id").agg(list)
             output = (
-                f"The time series is: {ctx.deps.df['y'].tolist()}, "
-                f"the date column is: {ctx.deps.df['ds'].tolist()}"
+                "these are the time series in json format where the key is the "
+                "identifier of the time series and the values is also a json "
+                "of two elements: "
+                "the first element is the date column and the second element is the "
+                "value column."
+                f"{df_agg.to_json(orient='index')}"
             )
             return output
 
@@ -384,15 +373,27 @@ async def tsfeatures_tool(
                         f"{', '.join(TSFEATURES.keys())}"
                     )
                 callable_features.append(TSFEATURES[feature])
-            features_df = _get_feats(
-                index=ctx.deps.df["unique_id"].iloc[0],
-                ts=ctx.deps.df,
-                features=callable_features,
-                freq=ctx.deps.seasonality,
-            )
-            return ",".join(
-                [f"{col}: {features_df[col].iloc[0]}" for col in features_df.columns]
+            features_df: pd.DataFrame | None = None
+            for uid in ctx.deps.df["unique_id"].unique():
+                features_df_uid = _get_feats(
+                    index=uid,
+                    ts=ctx.deps.df,
+                    features=callable_features,
+                    freq=ctx.deps.seasonality,
+                )
+                if features_df is None:
+                    features_df = features_df_uid
+                else:
+                    features_df = pd.concat([features_df, features_df_uid])
+            features_df = features_df.rename_axis("unique_id")  # type: ignore
+            self.features_df = features_df
+            output = (
+                "these are the time series features in json format where the key is "
+                "the identifier of the time series and the values is also a json of "
+                "feature names and their values."
+                f"{features_df.to_json(orient='index')}"
             )
+            return output
 
         @self.forecasting_agent.tool
         async def cross_validation_tool(
@@ -429,6 +430,7 @@ async def cross_validation_tool(
                 ["metric"],
                 as_index=False,
             ).mean(numeric_only=True)
+            self.eval_df = eval_df
             return ", ".join(
                 [
                     f"{model.alias}: {eval_df[model.alias].iloc[0]}"
@@ -437,19 +439,25 @@ async def cross_validation_tool(
             )
 
         @self.forecasting_agent.tool
-        async def forecast_tool(ctx: RunContext[ExperimentDataset], model: str) -> str:
+        async def forecast_tool(
+            ctx: RunContext[ExperimentDataset],
+            model: str,
+        ) -> str:
             callable_model = MODELS[model]
             fcst_df = callable_model.forecast(
                 df=ctx.deps.df,
                 h=ctx.deps.h,
                 freq=ctx.deps.freq,
             )
-            output = ",".join(
-                [
-                    f"{row['ds'].strftime('%Y-%m-%d')}: {row[model]}"
-                    for _, row in fcst_df.iterrows()
-                ]
+            df_agg = fcst_df.groupby("unique_id").agg(list)
+            output = (
+                "these are the forecasted values in json format where the key is the "
+                "identifier of the time series and the values is also a json of two "
+                "elements: the first element is the date column and the second "
+                "element is the value column."
+                f"{df_agg.to_json(orient='index')}"
             )
+            self.fcst_df = fcst_df
             return output
 
         @self.forecasting_agent.output_validator
@@ -519,5 +527,7 @@ def forecast(
             user_prompt=query,
             deps=dataset,
         )
-
+        result.fcst_df = self.fcst_df
+        result.eval_df = self.eval_df
+        result.features_df = self.features_df
         return result