.

timesfm_backend.py

@@ -1,6 +1,9 @@
 # timesfm_backend.py
-import time, json, logging
+import time
+import json
+import logging
 from typing import Any, Dict, List, Optional
+
 import numpy as np
 import torch
 
@@ -9,9 +12,9 @@ from config import settings
 
 logger = logging.getLogger(__name__)
 
-# --- TimesFM import (fallback-safe) ---
+# ---------------- TimesFM import (fallback-safe) ----------------
 try:
-    from timesfm import TimesFm
+    from timesfm import TimesFm  # Google TimesFM 2.5+
     _TIMESFM_AVAILABLE = True
 except Exception as e:
     logger.warning("timesfm not available (%s) — using naive fallback.", e)
@@ -19,29 +22,41 @@ except Exception as e:
     _TIMESFM_AVAILABLE = False
 
 
-# --- helpers ---
+# ---------------- helpers ----------------
 def _parse_series(series: Any) -> np.ndarray:
+    """
+    Accepts: list[float|int], list[dict{'y'|'value'}], or dict with 'values'/'y'.
+    Returns: 1D float32 numpy array.
+    """
     if series is None:
         raise ValueError("series is required")
+
     if isinstance(series, dict):
+        # allow {"values":[...]} or {"y":[...]}
         series = series.get("values") or series.get("y")
 
     vals: List[float] = []
     if isinstance(series, (list, tuple)):
         if series and isinstance(series[0], dict):
             for item in series:
-                if "y" in item: vals.append(float(item["y"]))
-                elif "value" in item: vals.append(float(item["value"]))
+                if "y" in item:
+                    vals.append(float(item["y"]))
+                elif "value" in item:
+                    vals.append(float(item["value"]))
         else:
             vals = [float(x) for x in series]
     else:
         raise ValueError("series must be a list/tuple or dict with 'values'/'y'")
+
     if not vals:
         raise ValueError("series is empty")
     return np.asarray(vals, dtype=np.float32)
 
 
 def _fallback_forecast(y: np.ndarray, horizon: int) -> np.ndarray:
+    """
+    Naive fallback: mean of last 4 (or all if <4), repeated H times.
+    """
     if horizon <= 0:
         return np.zeros((0,), dtype=np.float32)
     k = 4 if y.shape[0] >= 4 else y.shape[0]
@@ -50,15 +65,19 @@ def _fallback_forecast(y: np.ndarray, horizon: int) -> np.ndarray:
 
 
 def _extract_json_from_text(s: str) -> Optional[Dict[str, Any]]:
+    """
+    Try to parse JSON from a plain string or a fenced ```json block.
+    Returns dict or None.
+    """
     s = s.strip()
-    # whole-string JSON
+    # whole-string JSON object/array
     if (s.startswith("{") and s.endswith("}")) or (s.startswith("[") and s.endswith("]")):
         try:
             obj = json.loads(s)
             return obj if isinstance(obj, dict) else None
         except Exception:
             pass
-    # fenced ```json ... ```
+    # fenced code blocks
     if "```" in s:
         parts = s.split("```")
         for i in range(1, len(parts), 2):
@@ -75,10 +94,10 @@ def _extract_json_from_text(s: str) -> Optional[Dict[str, Any]]:
 
 def _merge_openai_message_json(payload: Dict[str, Any]) -> Dict[str, Any]:
     """
-    OpenAI chat format:
-      messages: [{role, content}, ...]
-    content can be a string or a list of parts: [{"type":"text","text":"..."}]
-    We scan from last to first user message and merge first JSON dict found.
+    OpenAI chat format compatibility:
+      payload["messages"] may hold user JSON in the last user message.
+      content can be a plain string or a list of parts [{"type":"text","text":...}].
+    If a JSON object is found, merge its keys into payload.
     """
     msgs = payload.get("messages")
     if not isinstance(msgs, list):
@@ -87,32 +106,39 @@ def _merge_openai_message_json(payload: Dict[str, Any]) -> Dict[str, Any]:
     for m in reversed(msgs):
         if not isinstance(m, dict) or m.get("role") != "user":
             continue
-        c = m.get("content")
-        # Text parts array
-        if isinstance(c, list):
-            texts = [p.get("text") for p in c if isinstance(p, dict) and p.get("type") == "text" and isinstance(p.get("text"), str)]
-            for t in reversed(texts):
-                obj = _extract_json_from_text(t)
-                if isinstance(obj, dict):
-                    return {**payload, **obj}
-        # Plain string
-        if isinstance(c, str):
-            obj = _extract_json_from_text(c)
+        content = m.get("content")
+        texts: List[str] = []
+        if isinstance(content, list):
+            texts = [
+                p.get("text")
+                for p in content
+                if isinstance(p, dict) and p.get("type") == "text" and isinstance(p.get("text"), str)
+            ]
+        elif isinstance(content, str):
+            texts = [content]
+
+        for t in reversed(texts):
+            obj = _extract_json_from_text(t)
             if isinstance(obj, dict):
                 return {**payload, **obj}
+        break  # only inspect last user
     return payload
 
 
-# --- backend ---
+# ---------------- backend ----------------
 class TimesFMBackend(ChatBackend):
     """
     Accepts OpenAI chat-completions requests.
     Pulls timeseries config from:
       - top-level keys, OR
-      - payload['data'] (CloudEvents), OR
-      - last user message JSON (OpenAI format, string or text-part).
-    Keys: series: list[float|{y|value}], horizon: int, freq: str (optional)
+      - payload['data'] (CloudEvents wrapper), OR
+      - last user message JSON (OpenAI format).
+    Keys:
+      series: list[float|int|{y|value}]
+      horizon: int (>0)
+      freq: optional str
     """
+
     def __init__(self, model_id: Optional[str] = None, device: Optional[str] = None):
         self.model_id = model_id or "google/timesfm-2.5-200m-pytorch"
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
@@ -122,7 +148,12 @@ class TimesFMBackend(ChatBackend):
         if self._model is not None or not _TIMESFM_AVAILABLE:
             return
         try:
-            self._model = TimesFm(context_len=512, horizon_len=128, input_patch_len=32)
+            # Set lengths compatible with the 2.5 checkpoints.
+            self._model = TimesFm(
+                context_len=512,
+                horizon_len=128,
+                input_patch_len=32,
+            )
             self._model.load_from_checkpoint(self.model_id)
             try:
                 self._model.to(self.device)  # type: ignore[attr-defined]
@@ -139,6 +170,7 @@ class TimesFMBackend(ChatBackend):
             payload = {**payload, **payload["data"]}
         if isinstance(payload.get("timeseries"), dict):
             payload = {**payload, **payload["timeseries"]}
+
         # merge JSON embedded in last user message (OpenAI format)
         payload = _merge_openai_message_json(payload)
 
@@ -153,8 +185,8 @@ class TimesFMBackend(ChatBackend):
         note = None
         if _TIMESFM_AVAILABLE and self._model is not None:
             try:
-                x = torch.tensor(y, dtype=torch.float32, device=self.device).unsqueeze(0)  # [1,T]
-                preds = self._model.forecast_on_batch(x, horizon)  # -> [1,H]
+                x = torch.tensor(y, dtype=torch.float32, device=self.device).unsqueeze(0)  # [1, T]
+                preds = self._model.forecast_on_batch(x, horizon)  # -> [1, H]
                 fc = preds[0].detach().cpu().numpy().astype(float).tolist()
             except Exception as e:
                 logger.exception("TimesFM forecast failed; fallback used. %s", e)
@@ -164,12 +196,23 @@ class TimesFMBackend(ChatBackend):
             fc = _fallback_forecast(y, horizon).tolist()
             note = "fallback_used_timesfm_missing"
 
-        return {"model": self.model_id, "horizon": horizon, "freq": freq, "forecast": fc, "note": note}
+        return {
+            "model": self.model_id,
+            "horizon": horizon,
+            "freq": freq,
+            "forecast": fc,
+            "note": note,
+        }
 
     async def stream(self, request: Dict[str, Any]):
+        """
+        OA-compatible streaming shim:
+        Emits exactly one chat.completion.chunk with compact JSON content.
+        """
         rid = f"chatcmpl-timesfm-{int(time.time())}"
         now = int(time.time())
         payload = dict(request) if isinstance(request, dict) else {}
+
         try:
             result = await self.forecast(payload)
         except Exception as e:
@@ -179,24 +222,36 @@ class TimesFMBackend(ChatBackend):
                 "object": "chat.completion.chunk",
                 "created": now,
                 "model": self.model_id,
-                "choices": [{"index": 0, "delta": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
+                "choices": [
+                    {"index": 0, "delta": {"role": "assistant", "content": content}, "finish_reason": "stop"}
+                ],
             }
             return
 
         content = json.dumps(
-            {"model": result["model"], "horizon": result["horizon"], "freq": result["freq"],
-             "forecast": result["forecast"], "note": result.get("note"), "backend": "timesfm"},
-            separators=(",", ":"), ensure_ascii=False
+            {
+                "model": result["model"],
+                "horizon": result["horizon"],
+                "freq": result["freq"],
+                "forecast": result["forecast"],
+                "note": result.get("note"),
+                "backend": "timesfm",
+            },
+            separators=(",", ":"),
+            ensure_ascii=False,
         )
         yield {
             "id": rid,
             "object": "chat.completion.chunk",
             "created": now,
             "model": self.model_id,
-            "choices": [{"index": 0, "delta": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
+            "choices": [
+                {"index": 0, "delta": {"role": "assistant", "content": content}, "finish_reason": "stop"}
+            ],
         }
 
 
+# ---------------- images stub ----------------
 class StubImagesBackend(ImagesBackend):
     async def generate_b64(self, request: Dict[str, Any]) -> str:
         logger.warning("Image generation not supported in TimesFM backend.")