From bb6d6c21c7353bcf2a9e65fe8992465d5267bb8b Mon Sep 17 00:00:00 2001
From: Ling LI <ling@imperialai.ai>
Date: Mon, 29 Jun 2026 07:10:22 +0000
Subject: [PATCH] Fix SignalsAPI.stake_get: query stakeValue instead of removed
 totalStake

stake_get read `data['totalStake']` from public_user_profile, but that field
was dropped from the profile query during the v2SignalsProfile/v3UserProfile
refactors (commits c4451ec/6a9819f added it; later refactors removed it).
`totalStake` no longer exists on the V3UserProfile type at all, so stake_get
raised KeyError for every call.

Query `stakeValue` (type Nmr) directly within the signals tournament and parse
it to Decimal, returning None when the model has no stake.

Verified against the live API: stake_get("uuazed") now returns the real stake
value instead of raising.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 numerapi/signalsapi.py   | 21 ++++++++++++++++-----
 tests/test_signalsapi.py | 24 ++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/numerapi/signalsapi.py b/numerapi/signalsapi.py
index c698422..c491926 100644
--- a/numerapi/signalsapi.py
+++ b/numerapi/signalsapi.py
@@ -284,18 +284,29 @@ def ticker_universe(self) -> List[str]:
         path = self.download_dataset("signals/v1.0/live.parquet")
         return pd.read_parquet(path).numerai_ticker.tolist()
 
-    def stake_get(self, username) -> decimal.Decimal:
-        """get current stake for a given users
+    def stake_get(self, username) -> decimal.Decimal | None:
+        """get current stake for a given user
 
         Args:
             username (str)
 
         Returns:
-            decimal.Decimal: current stake
+            decimal.Decimal or None: current stake, or None if the model has
+                no stake
 
         Example:
             >>> SignalsAPI().stake_get("uuazed")
             Decimal('14.63')
         """
-        data = self.public_user_profile(username)
-        return data['totalStake']
+        query = """
+          query($model_name: String!
+                $tournament: Int) {
+            v3UserProfile(model_name: $model_name
+                          tournament: $tournament) {
+              stakeValue
+            }
+          }
+        """
+        arguments = {'model_name': username, 'tournament': self.tournament_id}
+        data = self.raw_query(query, arguments)['data']['v3UserProfile']
+        return utils.parse_float_string(data['stakeValue'])
diff --git a/tests/test_signalsapi.py b/tests/test_signalsapi.py
index 0ccaf0f..dd1db70 100644
--- a/tests/test_signalsapi.py
+++ b/tests/test_signalsapi.py
@@ -1,3 +1,6 @@
+import decimal
+from unittest.mock import patch
+
 import pytest
 import responses
 
@@ -13,6 +16,27 @@ def api_fixture():
     return api
 
 
+@patch("numerapi.signalsapi.SignalsAPI.raw_query")
+def test_stake_get(mocked, api):
+    mocked.return_value = {"data": {"v3UserProfile": {"stakeValue": "14.63"}}}
+
+    stake = api.stake_get("uuazed")
+
+    assert stake == decimal.Decimal("14.63")
+    args, _ = mocked.call_args
+    # current stake lives in `stakeValue`; `totalStake` no longer exists
+    assert "stakeValue" in args[0]
+    assert "totalStake" not in args[0]
+    assert args[1]["tournament"] == api.tournament_id == 11
+
+
+@patch("numerapi.signalsapi.SignalsAPI.raw_query")
+def test_stake_get_no_stake(mocked, api):
+    # a model with no stake returns null -> None, not a KeyError
+    mocked.return_value = {"data": {"v3UserProfile": {"stakeValue": None}}}
+    assert api.stake_get("uuazed") is None
+
+
 @pytest.mark.live_api
 def test_get_leaderboard(api):
     lb = api.get_leaderboard(1)