quantlab-factor-agent 0.1.0__py3-none-any.whl

factor_agent/llm/openai_compatible.py

@@ -0,0 +1,113 @@
+"""The OpenAI-compatible backend: chat completions at any ``base_url`` (OpenAI, DeepSeek, DashScope,
+Gemini's compatible endpoint, a local vLLM or Ollama), the key from the environment variable the config names."""
+
+import openai
+
+from factor_agent.llm._common import key_from_environment
+from factor_agent.llm.base import ProviderQuotaExhausted, ProviderReplyMalformed, Reply, Request
+
+
+class OpenAICompatibleBackend:
+    """Complete a request with one chat-completions call.
+
+    Parameters
+    ----------
+    model : str
+        The model of every call.
+    base_url : str
+        The provider's endpoint, such as ``https://api.deepseek.com/v1``.
+    api_key_env : str, optional
+        The environment variable holding the key; None for a local server without one.
+    max_tokens : int, optional
+        The ``max_tokens`` of every call; the provider's default if None.
+    max_completion_tokens : int, optional
+        The ``max_completion_tokens`` of every call, in place of
+        ``max_tokens``: OpenAI's reasoning models take only this one, and
+        count the reasoning tokens in it.
+    reasoning_effort : str, optional
+        The ``reasoning_effort`` of every call (such as ``"low"``,
+        ``"medium"``, ``"high"``), for a reasoning model; the provider's
+        default if None.
+    timeout_seconds : float
+        Seconds a call may take.
+
+    Raises
+    ------
+    ValueError
+        If both ``max_tokens`` and ``max_completion_tokens`` are given.
+    """
+
+    def __init__(
+        self, model: str, base_url: str, api_key_env: str | None = None,
+        max_tokens: int | None = None, max_completion_tokens: int | None = None,
+        reasoning_effort: str | None = None, timeout_seconds: float = 600.0,
+    ):
+        if max_tokens is not None and max_completion_tokens is not None:
+            raise ValueError("openai_compatible: give max_tokens or max_completion_tokens, not both")
+        self.model = model
+        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
+        self.reasoning_effort = reasoning_effort
+        key = key_from_environment(api_key_env) if api_key_env else "none"
+        self._client = openai.OpenAI(api_key=key, base_url=base_url, timeout=timeout_seconds)
+
+    @classmethod
+    def from_config(cls, spec: dict) -> "OpenAICompatibleBackend":
+        """Build from a role's config: ``model``, ``base_url`` and optionally ``api_key_env``,
+        ``max_tokens`` or ``max_completion_tokens``, ``reasoning_effort`` and ``timeout_seconds``."""
+        names = ("api_key_env", "max_tokens", "max_completion_tokens", "reasoning_effort", "timeout_seconds")
+        options = {k: spec[k] for k in names if k in spec}
+        return cls(spec["model"], spec["base_url"], **options)
+
+    @property
+    def params(self) -> dict:
+        """The parameters every call sends with the prompt."""
+        params = {
+            "max_tokens": self.max_tokens, "max_completion_tokens": self.max_completion_tokens,
+            "reasoning_effort": self.reasoning_effort,
+        }
+        return {k: v for k, v in params.items() if v is not None}
+
+    def complete(self, request: Request) -> Reply:
+        """Send ``request``'s prompt as one user message.
+
+        Raises
+        ------
+        ProviderQuotaExhausted
+            If the provider answers HTTP 402 (Payment Required): the balance or quota is exhausted.
+        openai.APIError
+            If the call fails otherwise, after the SDK's retries (which never retry a 402).
+        ProviderReplyMalformed
+            If the provider answers something that is not a reply: no body
+            (``null``), no choice, no message or no token usage.
+        """
+        try:
+            response = self._client.chat.completions.create(
+                model=self.model, messages=[{"role": "user", "content": request.prompt}], **self.params,
+            )
+        except openai.APIStatusError as exc:
+            if exc.status_code == 402:
+                raise ProviderQuotaExhausted(f"openai_compatible: {self.model}: {exc.message}") from exc
+            raise
+        lacking = _lacking(response)
+        if lacking:
+            raise ProviderReplyMalformed(f"openai_compatible: {self.model} answered a reply without {lacking}")
+        usage = response.usage
+        return Reply(
+            text=response.choices[0].message.content or "",
+            usage={"input_tokens": usage.prompt_tokens, "output_tokens": usage.completion_tokens},
+            model_id=response.model, params=self.params, raw=response.to_json(indent=None),
+        )
+
+
+def _lacking(response) -> str | None:
+    """What ``response`` lacks to be a reply (its body, a choice, the message, the usage), or None."""
+    if response is None:
+        return "a body"
+    if not getattr(response, "choices", None):
+        return "a choice"
+    if getattr(response.choices[0], "message", None) is None:
+        return "a message"
+    if getattr(response, "usage", None) is None:
+        return "token usage"
+    return None

factor_agent/llm/replay.py

@@ -0,0 +1,59 @@
+"""The replay backend: a recorded run's replies served back, so a campaign reproduces without a provider."""
+
+from collections import defaultdict, deque
+
+from factor_agent.llm.base import Reply, Request
+from factor_agent.llm.log import LLMLog
+
+
+class ReplayExhausted(RuntimeError):
+    """The recorded run holds no reply for a request: it ran out of calls, or the run asked something else."""
+
+
+class ReplayBackend:
+    """Answer each request with the recorded reply to the same role, model and prompt.
+
+    A request asked several times gets its recorded replies in order.
+
+    Raises
+    ------
+    ValueError
+        If the log holds no call (a wrong path, or a run that made none).
+    """
+
+    def __init__(self, log: LLMLog):
+        self.log = log
+        self._replies: dict[tuple[str, str, str], deque] = defaultdict(deque)
+        calls = log.calls()
+        if not calls:
+            raise ValueError(f"replay: no recorded call in {log.path}")
+        for call in calls:
+            request = call["request"]
+            self._replies[request["role"], request["model"], request["prompt"]].append(call)
+
+    @classmethod
+    def from_config(cls, spec: dict) -> "ReplayBackend":
+        """Build from ``{"log": <the llm_log directory of the recorded run>}``."""
+        return cls(LLMLog(spec["log"]))
+
+    def complete(self, request: Request) -> Reply:
+        """Return the recorded reply to ``request``.
+
+        Raises
+        ------
+        ReplayExhausted
+            If the log holds no reply left for this role, model and prompt.
+        """
+        queue = self._replies.get((request.role, request.model, request.prompt))
+        if not queue:
+            raise ReplayExhausted(
+                f"replay of {self.log.path}: no recorded reply left for this {request.role} request "
+                f"(the recorded run ran out of calls, or this run asked something it never did); "
+                f"the prompt begins {request.prompt[:200]!r}"
+            )
+        call = queue.popleft()
+        reply = call["reply"]
+        return Reply(
+            text=reply["text"], usage=reply["usage"], model_id=reply["model_id"],
+            params=call["request"]["params"], raw=reply["raw"],
+        )

factor_agent/llm/scripted.py

@@ -0,0 +1,39 @@
+"""A scripted backend: replies read from a file, in order. For tests and offline smoke runs."""
+
+import json
+import threading
+from pathlib import Path
+
+from factor_agent.llm.base import Reply, Request
+
+
+class ScriptedBackend:
+    """Answer each request with the next reply of a fixed list.
+
+    ``requests`` keeps every request received, in order; calls from several
+    threads take the replies one at a time, in the order they arrive.
+
+    Raises
+    ------
+    RuntimeError
+        From ``complete`` when every reply has been used.
+    """
+
+    def __init__(self, replies: list[str]):
+        self._replies = list(replies)
+        self.requests: list[Request] = []
+        self._taking = threading.Lock()
+
+    @classmethod
+    def from_config(cls, spec: dict) -> "ScriptedBackend":
+        """Build from ``{"replies": <path to a JSON list of reply texts>}``."""
+        return cls(json.loads(Path(spec["replies"]).read_text()))
+
+    def complete(self, request: Request) -> Reply:
+        """Return the next scripted reply and keep ``request``."""
+        with self._taking:
+            if len(self.requests) >= len(self._replies):
+                raise RuntimeError(f"ScriptedBackend: no reply left for request {len(self.requests) + 1}")
+            text = self._replies[len(self.requests)]
+            self.requests.append(request)
+        return Reply(text=text, usage={}, model_id="scripted")

factor_agent/period_lag.py

@@ -0,0 +1,91 @@
+"""Period lags read from data: the lag columns ``PeriodLag(x, k)`` reads, computed when a factor is.
+
+A lag column (``Vocabulary.lag_columns``) is no column of the price
+dataset but a derived column (``factor_agent.derived``): its provider,
+``PeriodLags``, reads the field's column and the key column of its family
+(the period end of the report each bar shows) and computes it with
+quantlab's ``report_lags``, as ``FundamentalAlpha`` computes its report
+lags. On a bar showing report r, the value ``k`` periods ago is the last
+value shown during the k-th earlier report shown; it is missing unless that
+report's period ends ``k * days`` (within ``tolerance_days``) before the
+current one, so a skipped or missing report never passes for an earlier
+period. A key that is a date counts in days; a numeric key is used as it
+is.
+
+A lag of ``k`` periods reads ``(k + 2) * bars`` bars before the first bar it
+computes: the k earlier periods, the current one and one more for a late
+filing.
+
+Each lag column is described by a mapping, as ``Vocabulary.lag_columns``
+gives it: ``column`` (the field's dataset column), ``k``, ``key``,
+``days``, ``tolerance_days`` and ``bars``.
+"""
+
+from collections.abc import Iterable, Mapping
+
+import numpy as np
+import xarray as xr
+from quantlab.dataset.base import BaseDataset
+from quantlab.factor.predefined.fundamental_alpha import report_lags
+
+from factor_agent import derived
+
+
+def read(inputs: Iterable[str], lag_columns: Mapping[str, Mapping]) -> dict[str, dict]:
+    """The entries of ``lag_columns`` (``Vocabulary.lag_columns``) for the columns ``inputs``, by sorted name."""
+    return {name: dict(lag_columns[name]) for name in sorted(inputs) if name in lag_columns}
+
+
+def warmup_bars(lags: Mapping[str, Mapping]) -> int:
+    """The bars before the first computed bar the ``lags`` read; 0 without any."""
+    return max(((int(lag["k"]) + 2) * int(lag["bars"]) for lag in lags.values()), default=0)
+
+
+class PeriodLags(derived.Provider):
+    """The lag columns ``lags`` (``read``), filled from the daily panel; see the module docstring.
+
+    It reads the fields' columns and keys (``data_columns``), the keys from
+    the panel only, and the warm-up the lags need (``warmup_bars``).
+    """
+
+    def __init__(self, lags: Mapping[str, Mapping]):
+        self._lags = {name: dict(lag) for name, lag in lags.items()}
+        self.columns = tuple(self._lags)
+        self.data_columns = tuple(sorted({c for lag in self._lags.values() for c in (lag["column"], lag["key"])}))
+        self.panel_columns = tuple(sorted({lag["key"] for lag in self._lags.values()}))
+        self.warmup_bars = warmup_bars(self._lags)
+
+    def fill(
+        self, dataset: BaseDataset, panel: xr.Dataset, arrays: dict[str, np.ndarray],
+        symbols: np.ndarray, timestamps: np.ndarray,
+    ) -> dict[str, np.ndarray]:
+        groups: dict[tuple, dict[str, list]] = {}
+        for name, lag in self._lags.items():
+            group = groups.setdefault((lag["key"], float(lag["days"]), float(lag["tolerance_days"])), {})
+            group.setdefault(lag["column"], []).append((int(lag["k"]), name))
+        lagged = {}
+        for (key, days, tolerance), wanted in groups.items():
+            periods = _key_values(panel[key].sel(timestamp=timestamps, symbol=symbols).transpose("timestamp", "symbol"))
+            out = report_lags(
+                {column: arrays[column] for column in wanted}, periods,
+                {column: tuple(k for k, _ in ks) for column, ks in wanted.items()}, days, tolerance,
+            )
+            for column, ks in wanted.items():
+                for k, name in ks:
+                    lagged[name] = out[f"{column}__q{k}"]
+        return lagged
+
+
+def providers(lags: Mapping[str, Mapping]) -> list[derived.Provider]:
+    """The provider of the lag columns ``lags``, in a list; none without any."""
+    return [PeriodLags(lags)] if lags else []
+
+
+def _key_values(key: xr.DataArray) -> np.ndarray:
+    """The key as float: days since 1970 for dates (NaN where NaT), the numbers themselves otherwise."""
+    values = key.to_numpy()
+    if np.issubdtype(values.dtype, np.datetime64):
+        days = values.astype("datetime64[D]").astype(np.float64)
+        days[np.isnat(values)] = np.nan
+        return days
+    return values.astype(np.float64)