Add moirai agent replication scripts (#244)

Author: cuthalionn

project/moirai-agent/README.md

@@ -0,0 +1 @@
+This folder temporarily hosts the scripts to replicate early Moirai-agent results on contextual and non-contextual forecasting tasks.

project/moirai-agent/ctx_forecast/README.md

@@ -0,0 +1,47 @@
+## GIFT-CTX Replication
+This repo contains instructions to replicate the results of Moirai Agent on the GIFT-CTX benchmark.
+
+### Set up environment
+First, set up the environment by running:
+```
+pip install -r requirement.txt
+```
+And export your OpenAI key by:
+```
+export OPENAI_API_KEY="..."
+```
+
+### Prepare the dataset
+- Download the `gift_ctx.parquet` data here: `https://huggingface.co/datasets/Salesforce/GIFT-CTX`
+- Plot the historical data before runtime:
+```
+python gen_image.py --in_file gift_ctx.parquet --out_file gift_ctx_image.parquet --img_root img
+```
+Change the parameters to your project setting:
+```
+--in_file: path to the original GIFT-CTX dataset
+--out_file: path to the new dataset with image path added
+--img_root: path to the image folder
+```
+Note that Moirai Agent can work fine with just text input. To maximize its performance, we recommend to also provide the plot of historical data, which can be rendered beforehand to reduce runtime.
+### Run Moirai Agent
+To run Moirai Agent, first start the tools with:
+```
+bash run_tools.sh
+```
+For now, we provide agent with 2 tools: a python sandbox, and a timeseries foundation model.
+
+To evaluate Moirai Agent on the GIFT-CTX benchmark with the default setting and replicate the reported results, run:
+```
+bash run.sh
+```
+Note that the results reported in the blog post were obtained with GPT-5.1 medium reasoning effort on Jan 07 2026. Please expect some minor differences with the official reported numbers because of the non-deterministic nature of LLMs. We provided our log files in `./results`.
+
+To change the input path file, config file, output dir, input mode (text only, text with image), and parallelism, run with arguments:
+```
+bash run.sh [your_parquet_path] [your_config_path] [output_dir] [input_mode] [#jobs]
+```
+
+There are several configurations you can change in Moirai Agent such as LLM and its parameters, tools, and system prompt, which you can adjust in `src/ctx_forecast/config.py`. 
+
+

project/moirai-agent/ctx_forecast/gen_image.py

@@ -0,0 +1,41 @@
+import argparse
+import glob
+import os
+
+import pandas as pd
+from src.utils.image import plot_historical_and_save
+from tqdm import tqdm
+
+
+def main(in_file, out_file, img_root):
+    df = pd.read_parquet(in_file)
+    hist_values = df["history_values"].tolist()
+    img_paths = []
+    for i, v in enumerate(hist_values):
+        v = [float(x) for x in v.split(",")]
+
+        start = df["history_start"].iloc[i]
+
+        freq = df["frequency"].iloc[i]
+        timestamps = pd.date_range(start=start, periods=len(v), freq=freq)
+        ts_list = timestamps.strftime("%Y-%m-%d %H:%M:%S").tolist()
+
+        item = {"Date": ts_list, "Value": v}
+
+        img_path = os.path.join(img_root, f"image_{i}.jpg")
+        img_paths.append(img_path)
+        plot_historical_and_save(item, img_path)
+    df["image_path"] = img_paths
+
+    df.to_parquet(out_file)
+    return
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="A simple script to plot time series")
+    parser.add_argument("--in_file", default="./gift_ctx.parquet", required=True)
+    parser.add_argument("--out_file", default="./gift_ctx_image.parquet", required=True)
+    parser.add_argument("--img_root", default="./images/", required=True)
+
+    args = parser.parse_args()
+    main(args.in_file, args.out_file, args.img_root)

project/moirai-agent/ctx_forecast/requirements.txt

@@ -0,0 +1,179 @@
+accelerate==1.10.1
+aiohappyeyeballs==2.6.1
+aiohttp==3.13.0
+aiosignal==1.4.0
+annotated-types==0.7.0
+anyio==4.11.0
+argon2-cffi==25.1.0
+argon2-cffi-bindings==25.1.0
+arrow==1.3.0
+asttokens==3.0.0
+async-lru==2.0.5
+attrs==25.4.0
+babel==2.17.0
+beautifulsoup4==4.14.2
+bleach==6.2.0
+blinker==1.9.0
+boto3==1.42.16
+botocore==1.42.16
+certifi==2025.10.5
+cffi==2.0.0
+charset-normalizer==3.4.3
+chronos-forecasting==2.2.2
+click==8.3.0
+comm==0.2.3
+contourpy==1.3.3
+cycler==0.12.1
+datasets==4.1.1
+debugpy==1.8.17
+decorator==5.2.1
+defusedxml==0.7.1
+dill==0.4.0
+distro==1.9.0
+einops==0.8.1
+executing==2.2.1
+fastjsonschema==2.21.2
+filelock==3.19.1
+Flask==3.1.2
+fonttools==4.60.1
+fqdn==1.5.1
+frozenlist==1.8.0
+fsspec==2025.9.0
+gluonts==0.16.2
+h11==0.16.0
+hf-xet==1.1.10
+httpcore==1.0.9
+httpx==0.28.1
+httpx-sse==0.4.2
+huggingface-hub==0.35.3
+idna==3.10
+inquirerpy==0.3.4
+ipdb==0.13.13
+ipykernel==6.30.1
+ipython==9.6.0
+ipython_pygments_lexers==1.1.1
+ipywidgets==8.1.7
+isoduration==20.11.0
+itsdangerous==2.2.0
+jaxtyping==0.3.3
+jedi==0.19.2
+Jinja2==3.1.6
+jiter==0.11.0
+jmespath==1.0.1
+joblib==1.5.3
+json5==0.12.1
+jsonpointer==3.0.0
+jsonschema==4.25.1
+jsonschema-specifications==2025.9.1
+jupyter-events==0.12.0
+jupyter-lsp==2.3.0
+jupyter_client==8.6.3
+jupyter_core==5.8.1
+jupyter_server==2.17.0
+jupyter_server_terminals==0.5.3
+jupyterlab==4.4.9
+jupyterlab_pygments==0.3.0
+jupyterlab_server==2.27.3
+jupyterlab_widgets==3.0.15
+jupyterthemes==0.20.0
+kiwisolver==1.4.9
+lark==1.3.0
+lesscpy==0.15.1
+lightning==2.5.5
+lightning-utilities==0.15.2
+MarkupSafe==3.0.3
+matplotlib==3.10.6
+matplotlib-inline==0.1.7
+mcp==1.16.0
+mistune==3.1.4
+mpmath==1.3.0
+multidict==6.7.0
+multiprocess==0.70.16
+nbclient==0.10.2
+nbconvert==7.16.6
+nbformat==5.10.4
+nest-asyncio==1.6.0
+networkx==3.5
+notebook==7.4.7
+notebook_shim==0.2.4
+numpy==2.1.3
+openai==2.2.0
+orjson==3.11.3
+packaging==25.0
+pandas==2.3.3
+pandocfilters==1.5.1
+parso==0.8.5
+pexpect==4.9.0
+pfzy==0.3.4
+pillow==11.3.0
+platformdirs==4.4.0
+ply==3.11
+prometheus_client==0.23.1
+prompt_toolkit==3.0.52
+propcache==0.4.0
+psutil==7.1.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pyarrow==21.0.0
+pycparser==2.23
+pydantic==2.11.10
+pydantic-settings==2.11.0
+pydantic_core==2.33.2
+Pygments==2.19.2
+pyparsing==3.2.5
+python-dateutil==2.9.0.post0
+python-dotenv==1.1.1
+python-json-logger==4.0.0
+python-multipart==0.0.20
+pytorch-lightning==2.5.5
+pytz==2025.2
+PyYAML==6.0.3
+pyzmq==27.1.0
+referencing==0.36.2
+regex==2025.9.18
+requests==2.32.5
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rfc3987-syntax==1.1.0
+rpds-py==0.27.1
+s3transfer==0.16.0
+safetensors==0.6.2
+scikit-learn==1.8.0
+scipy==1.16.2
+Send2Trash==1.8.3
+setuptools==80.9.0
+six==1.17.0
+sniffio==1.3.1
+soupsieve==2.8
+sse-starlette==3.0.2
+stack-data==0.6.3
+starlette==0.48.0
+sympy==1.14.0
+terminado==0.18.1
+threadpoolctl==3.6.0
+tinycss2==1.4.0
+tokenizers==0.22.1
+toolz==0.12.1
+torch==2.8.0
+torchmetrics==1.8.2
+tornado==6.5.2
+tqdm==4.67.1
+traitlets==5.14.3
+transformers==4.57.0
+triton==3.4.0
+types-python-dateutil==2.9.0.20251008
+typing-inspection==0.4.2
+typing_extensions==4.15.0
+tzdata==2025.2
+uri-template==1.3.0
+urllib3==2.5.0
+uvicorn==0.37.0
+wadler_lindig==0.1.7
+wcwidth==0.2.14
+webcolors==24.11.1
+webencodings==0.5.1
+websocket-client==1.9.0
+Werkzeug==3.1.3
+widgetsnbextension==4.0.14
+xxhash==3.6.0
+yarl==1.22.0

project/moirai-agent/ctx_forecast/run.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+set -e
+
+# --------------------------------------------------
+# Positional args with defaults
+# --------------------------------------------------
+TASK_FILE=${1:-"gift_ctx_images.parquet"}
+CONFIG_FILE=${2:-"src/ctx_forecast/config.py"}
+OUTPUT_DIR=${3:-"results/moirai_agent"}
+INPUT_MODE=${4:-"text+image"}
+JOBS=${5:-$(nproc 2>/dev/null || echo 8)}
+NUM_PARTS=${6:-245}
+
+# --------------------------------------------------
+# Run all parts in parallel
+# --------------------------------------------------
+seq 0 $((NUM_PARTS - 1)) | parallel -j "$JOBS" \
+  python -m src.ctx_forecast.tsf_agent \
+    --config_file "$CONFIG_FILE" \
+    --config_name CONFIG \
+    --input_mode "$INPUT_MODE" \
+    --task_file "$TASK_FILE" \
+    --part_idx {} \
+    --num_parts "$NUM_PARTS" \
+    --output_dir "$OUTPUT_DIR"
+
+echo "Finish all parts"
+echo "Calculating metrics"
+
+# --------------------------------------------------
+# Gather metrics
+# --------------------------------------------------
+python -m src.ctx_forecast.metrics_gather \
+  --results_dir "${OUTPUT_DIR}/${INPUT_MODE}/results"
+
+echo "Done! Results are available at ${OUTPUT_DIR}/${INPUT_MODE}"
+

project/moirai-agent/ctx_forecast/run_tools.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+# build and launch python-sandbox docker
+
+bash src/tools/python_sandbox/start_background.sh
+
+
+# launch time series forecasting tools
+CUDA_VISIBLE_DEVICES=7 python3 -m src.tools.tsf_services --model_name chronos_v2 &
+# CUDA_VISIBLE_DEVICES=0 .venv/bin/python3 -m src.tools.tsf_services --model_name moirai &

project/moirai-agent/ctx_forecast/src/ctx_forecast/config.py

@@ -0,0 +1,37 @@
+max_iterations = 6
+CONFIG = {
+    "llm": {
+        "model_name": "gpt-5.1",
+        "model_params_type": {
+            "temperature": 1.0,
+            "top_p": 1.0,
+            "reasoning": {"effort": "medium", "summary": "auto"},
+            "max_output_tokens": 8192,
+        },
+    },
+    "servers": {
+        "forecast": {
+            "command": "python",
+            "args": ["-m", "src.ctx_forecast.mcp_server"],
+            "enabled_tools": ["chronos_forecast_service", "python_sandbox_service"],
+        },
+    },
+    "max_iterations": max_iterations,
+    "system_prompt": "You are an intelligent assistant that can solve complex problems by thinking step-by-step and using available tools when needed. "
+    "\n## Your Task "
+    "\n- Solve a contextual time-series forecasting problem, where historical values and contextual information are provided. "
+    "\n- The future values depends on both of the observed values and the contextual information. "
+    "\n- Contextual information may specify or imply unexpected effects in the future that will affect the normal results reasoned from history values. Then, you should combine the forecasting tool's results with the contextual information to reason the future."
+    "\n- Contextual information may specify or imply occasional or abnormal factors in the history that will not persist in the future. Then, identify and eliminate the misleading factors from the history, the future should be reasonsed from modified history values and relevant knowledge. "
+    "\n- Contextual information may totally dominate the future, when the history values are less informative. Then, discover the underlying correlation or math structure defined by context and history values, use it for prediction making. "
+    "\n- You should figure out the real cause of the future first, then make predictions accordingly. "
+    "\n## How You Work "
+    "\n 1. **Think First**: Analyze the problem and determine what information or actions you need "
+    "\n 2. **Use Tools When Needed**: Call appropriate functions/tools to perform numerical modeling. "
+    "\n 3. **Reason with Results**: Process the tool outputs and use them to inform your next steps "
+    "\n 4. **Iterate**: Continue thinking and using tools until you can provide a complete answer "
+    "\n ## Important Guidelines "
+    "\n - Only when the original or modified history values show clear patterns, you may call forecasting tools. Otherwise, fatal numerical errors will happen. "
+    "\n - Only when exact mathematical structures are inferred from the context, you may write and execute codes in the python-sandbox. Always include a print function in your codes to return valid messages. "
+    f"\n - Ensure that all reasoning and tool usage is complete within a maximum of {max_iterations} steps. Each step should either advance your understanding or gather necessary information. Be systematic and thorough in your approach. A final and fully cited answer has to be output before step {max_iterations}. ",
+}