Merge pull request lichess-org#81 from FriedrichtenHagen/add_external_engine_analysis_post_endpoints

Add external engine analysis post endpoints

Author: kraktus

CHANGELOG.rst

@@ -19,6 +19,8 @@ To be released
   ``client.studies.get_by_user`` to get the metadata of all studies
   of a given user.
 * Added ``bots.handle_draw_offer`` and ``bots.handle_takeback_offer`` to handle draw and takeback offers
+* Added ``client.external_engine.analyse``, ``client.external_engine.acquire_request``, ``client.external_engine.answer_request`` to handle analysis with an external engine
+
 
 Thanks to all the contributors who helped to this release:
 - @hsheth2

README.rst

@@ -148,6 +148,9 @@ Most of the API is available:
     client.external_engine.create
     client.external_engine.update
     client.external_engine.delete
+    client.external_engine.analyse
+    client.external_engine.acquire_request
+    client.external_engine.answer_request
 
     client.fide.get_player
     client.fide.search_players

berserk/clients/__init__.py

@@ -93,6 +93,7 @@ def __init__(
         *,
         tablebase_url: str | None = None,
         explorer_url: str | None = None,
+        external_engine_url: str | None = None,
     ):
         session = session or requests.Session()
         super().__init__(session, base_url)
@@ -116,5 +117,7 @@ def __init__(
         self.tablebase = Tablebase(session, tablebase_url)
         self.opening_explorer = OpeningExplorer(session, explorer_url)
         self.bulk_pairings = BulkPairings(session, base_url)
-        self.external_engine = ExternalEngine(session, base_url)
+        self.external_engine = ExternalEngine(
+            session, base_url=base_url, external_engine_url=external_engine_url
+        )
         self.fide = Fide(session)

berserk/clients/external_engine.py

@@ -1,13 +1,32 @@
 from __future__ import annotations
 
-from typing import List, cast
+from typing import List, cast, Iterator
 
 from .base import BaseClient
+import requests
+from ..formats import NDJSON, TEXT
+from ..types.common import UciVariant
+from ..types.external_engine import ExternalEngineRequest, EngineAnalysisOutput
+
+EXTERNAL_ENGINE_URL = "https://engine.lichess.ovh"
 
 
 class ExternalEngine(BaseClient):
     """Client for external engine related endpoints."""
 
+    def __init__(
+        self,
+        session: requests.Session,
+        *,
+        base_url: str | None = None,
+        external_engine_url: str | None = None,
+    ):
+        """Create a subclient for the endpoints that use a different base url."""
+        super().__init__(session, base_url)
+        self._external_client = BaseClient(
+            session, external_engine_url or EXTERNAL_ENGINE_URL
+        )
+
     def get(self) -> List[ExternalEngine]:
         """Lists all external engines that have been registered for the user, and the credentials required to use them.
 
@@ -112,3 +131,85 @@ def delete(self, engine_id: str) -> None:
         """
         path = f"/api/external-engine/{engine_id}"
         self._r.request("DELETE", path)
+
+    def analyse(
+        self,
+        engine_id: str,
+        client_secret: str,
+        session_id: str,
+        threads: int,
+        hash_table_size: int,
+        pri_num_variations: int,
+        variant: UciVariant,
+        initial_fen: str,
+        moves: List[str],
+        movetime: int | None = None,
+        depth: int | None = None,
+        nodes: int | None = None,
+    ) -> Iterator[EngineAnalysisOutput]:
+        """
+        Analyse with external engine
+
+        Request analysis from an external engine. Response content is streamed as newline delimited JSON.
+        The properties are based on the UCI specification.
+        Analysis stops when the client goes away, the requested limit is reached, or the provider goes away.
+
+        :param engine_id: external engine id
+        :param client_secret: engine credentials
+        :param session_id: Arbitary string that identifies the analysis session. Providers may wish to clear the hash table between sessions.
+        :param threads: Number of threads to use for analysis.
+        :param hash_table_size: Hash table size to use for analysis, in MiB.
+        :param pri_num_variations: Requested number of principal variations. (1-5)
+        :param variant: uci variant
+        :param initial_fen: Initial position of the game.
+        :param moves: List of moves played from the initial position, in UCI notation.
+        :param movetime: Amount of time to analyse the position, in milliseconds.
+        :param depth: Analysis target depth
+        :param nodes: Number of nodes to analyse in the position
+        """
+        path = f"/api/external-engine/{engine_id}/analyse"
+        payload = {
+            "clientSecret": client_secret,
+            "work": {
+                "sessionId": session_id,
+                "threads": threads,
+                "hash": hash_table_size,
+                "multiPv": pri_num_variations,
+                "variant": variant,
+                "initialFen": initial_fen,
+                "moves": moves,
+                "movetime": movetime,
+                "depth": depth,
+                "nodes": nodes,
+            },
+        }
+
+        for response in self._external_client._r.post(
+            path=path,
+            payload=payload,
+            stream=True,
+            fmt=NDJSON,
+        ):
+            yield cast(EngineAnalysisOutput, response)
+
+    def acquire_request(self, provider_secret: str) -> ExternalEngineRequest:
+        """Wait for an analysis request to any of the external engines that have been registered with the given secret.
+        :param provider_secret: provider credentials
+        :return: the requested analysis
+        """
+        path = "/api/external-engine/work"
+        payload = {"providerSecret": provider_secret}
+        return cast(
+            ExternalEngineRequest,
+            self._external_client._r.post(path=path, payload=payload),
+        )
+
+    def answer_request(self, engine_id: str) -> str:
+        """Submit a stream of analysis as UCI output.
+        The server may close the connection at any time, indicating that the requester has gone away and analysis
+        should be stopped.
+        :param engine_id: engine ID
+        :return: the requested analysis
+        """
+        path = f"/api/external-engine/work/{engine_id}"
+        return self._external_client._r.post(path=path, fmt=TEXT)

berserk/types/__init__.py

@@ -17,8 +17,9 @@
 )
 
 from .bulk_pairings import BulkPairing, BulkPairingGame
+from .external_engine import ExternalEngine
 from .challenges import ChallengeJson
-from .common import ClockConfig, ExternalEngine, LightUser, OnlineLightUser, VariantKey
+from .common import ClockConfig, LightUser, OnlineLightUser, VariantKey
 from .fide import FidePlayer
 from .puzzles import PuzzleData, PuzzleRace
 from .opening_explorer import (

berserk/types/common.py

@@ -12,27 +12,6 @@ class ClockConfig(TypedDict):
     increment: int
 
 
-class ExternalEngine(TypedDict):
-    # Engine ID
-    id: str
-    # Engine display name
-    name: str
-    # Secret token that can be used to request analysis
-    clientSecret: str
-    # User this engine has been registered for
-    userId: str
-    # Max number of available threads
-    maxThreads: int
-    # Max available hash table size, in MiB
-    maxHash: int
-    # Estimated depth of normal search
-    defaultDepth: int
-    # List of supported chess variants
-    variants: str
-    # Arbitrary data that engine provider can use for identification or bookkeeping
-    providerData: NotRequired[str]
-
-
 Color: TypeAlias = Literal["white", "black"]
 
 GameType: TypeAlias = Literal[
@@ -47,6 +26,17 @@ class ExternalEngine(TypedDict):
     "fromPosition",
 ]
 
+UciVariant = Literal[
+    "chess",
+    "crazyhouse",
+    "antichess",
+    "atomic",
+    "horde",
+    "kingofthehill",
+    "racingkings",
+    "3check",
+]
+
 Speed = Literal[
     "ultraBullet", "bullet", "blitz", "rapid", "classical", "correspondence"
 ]

berserk/types/external_engine.py

@@ -0,0 +1,72 @@
+from typing import List
+
+from typing_extensions import TypedDict, NotRequired
+from .common import UciVariant
+
+
+class ExternalEngine(TypedDict):
+    # Engine ID
+    id: str
+    # Engine display name
+    name: str
+    # Secret token that can be used to request analysis
+    clientSecret: str
+    # User this engine has been registered for
+    userId: str
+    # Max number of available threads
+    maxThreads: int
+    # Max available hash table size, in MiB
+    maxHash: int
+    # Estimated depth of normal search
+    defaultDepth: int
+    # List of supported chess variants
+    variants: list[UciVariant]
+    # Arbitrary data that engine provider can use for identification or bookkeeping
+    providerData: NotRequired[str]
+
+
+class ExternalEngineWork(TypedDict):
+    # Arbitrary string that identifies the analysis session. Providers may clear the hash table between sessions
+    sessionId: str
+    # Number of threads to use for analysis
+    threads: int
+    # Hash table size to use for analysis, in MiB
+    hash: int
+    # Requested number of principle variations
+    multiPv: List[int]
+    # Uci variant
+    variant: str
+    # Initial position of the game
+    initialFen: str
+    # List of moves played from the initial position, in UCI notation
+    moves: List[str]
+    # Request an infinite search (rather than roughly aiming for defaultDepth)
+    infinite: NotRequired[bool]
+
+
+class ExternalEngineRequest(TypedDict):
+    id: str
+    work: ExternalEngineWork
+    engine: ExternalEngine
+
+
+class PrincipleVariationAnalysis(TypedDict):
+    # Current search depth of the pv
+    depth: int
+    # Variation in UCI notation
+    moves: List[str]
+    # Evaluation in centi-pawns, from White's point of view
+    cp: NotRequired[int]
+    # Evaluation in signed moves to mate, from White's point of view
+    mate: NotRequired[int]
+
+
+class EngineAnalysisOutput(TypedDict):
+    # Number of milliseconds the search has been going on
+    time: int
+    # Current search depth
+    depth: int
+    # Number of nodes visited so far
+    nodes: int
+    # Information about up to 5 pvs, with the primary pv at index 0
+    pvs: List[PrincipleVariationAnalysis]