treequest 0.3.2__tar.gz → 0.3.3__tar.gz

{treequest-0.3.2 → treequest-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: treequest
-Version: 0.3.2
+Version: 0.3.3
 Summary: A Tree Search Library with Flexible API for LLM Inference-Time Scaling
 Author-email: Kou Misaki <kou.misaki@sakana.ai>, Takuya Akiba <takiba@sakana.ai>
 License-File: LICENSE
@@ -57,24 +57,27 @@ import treequest as tq
 # Each node is associated with a user-definable `state`.
 State = str
 
+
 # 1. Define a function to be used for node generation.
 def generate(parent_state: State | None) -> tuple[State, float]:
     """Generates new states and scores based on the parent state."""
-    if parent_state is None: # None represents the expansion from root.
+    if parent_state is None:  # None represents the expansion from root.
         new_state = "Initial state"
     else:
         new_state = f"State after {parent_state}"
 
-    score = random.random() # A score for the new state; It should be normalized to the [0, 1] range.
+    # A score for the new state; It should be normalized to the [0, 1] range.
+    score = random.random()
     return new_state, score
 
+
 # 2. Instantiate the algorithm and a search tree object.
 algo = tq.ABMCTSA()
 search_tree = algo.init_tree()
 
 # 3. Run the search with a generation budget (10 in this case).
 for _ in range(10):
-    search_tree = algo.step(search_tree, {'Action A': generate})
+    search_tree = algo.step(search_tree, {"Action A": generate})
 
 # 4. Extract the best score and state.
 best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
@@ -82,7 +85,8 @@ print(f"Best state: {best_state}, Score: {best_node_score}")
 
 # 5. Visualize the search tree.
 output_file_basename = Path("ab_mcts_a_search_tree")
-tq.render(search_tree, output_file_basename, format="html")  # Generates `ab_mcts_a_search_tree.html`
+# Generates `ab_mcts_a_search_tree.html`
+tq.render(search_tree, output_file_basename, format="html")
 ```
 
 Alternatively, you can use an ask–tell interface with batched AB-MCTS sampling steps:
@@ -92,8 +96,9 @@ import treequest as tq
 
 State = str
 
-def generate(parent_state: State | None) -> tuple[State, float]:
-    ...
+
+def generate(parent_state: State | None) -> tuple[State, float]: ...
+
 
 generate_fns = {"Action A": generate}
 actions = list(generate_fns.keys())
@@ -160,11 +165,13 @@ import dataclasses
 
 import treequest as tq
 
+
 @dataclasses.dataclass
 class State:
     llm_answer: str
     score: float
 
+
 def generate(parent_state: State | None) -> tuple[State, float]:
     """Generate a new node by calling an LLM."""
     if parent_state is None:
@@ -174,12 +181,14 @@ def generate(parent_state: State | None) -> tuple[State, float]:
 
     return state, state.score
 
+
 def initial_generation() -> State:
     """
     Call LLM API to generate an initial answer.
     """
     ...
 
+
 def refine_answer(llm_answer: str, score: float) -> State:
     """
     Call LLM API to refine an answer.
@@ -190,11 +199,11 @@ def refine_answer(llm_answer: str, score: float) -> State:
 algo = tq.ABMCTSM()
 search_tree = algo.init_tree()
 for i in range(20):
-    search_tree = algo.step(search_tree, {'Action Label': generate})
+    search_tree = algo.step(search_tree, {"Action Label": generate})
     # Logging best node during the search.
     if (i + 1) % 5 == 0:
         best_interim_state, _ = tq.top_k(search_tree, algo, k=1)[0]
-        print(f"Iteration {i+1}: Best state so far = {best_interim_state}")
+        print(f"Iteration {i + 1}: Best state so far = {best_interim_state}")
 
 best_state, _ = tq.top_k(search_tree, algo, k=1)[0]
 print(f"Best Answer: {best_state.llm_answer}, Best Score: {best_state.score}")
@@ -209,6 +218,7 @@ from functools import partial
 
 import treequest as tq
 
+
 def generate(llm_name: str, parent_state=None):
     """
     Call LLM API using litellm, vllm, etc., to generate a new node
@@ -216,9 +226,12 @@ def generate(llm_name: str, parent_state=None):
     ...
     return new_state, new_score
 
+
 llm_names = ["o4-mini", "gemini-2.5-pro"]
 # Create dict of different actions backed by different LLMs.
-generate_fns = {llm_name: partial(generate, llm_name=llm_name) for llm_name in llm_names}
+generate_fns = {
+    llm_name: partial(generate, llm_name=llm_name) for llm_name in llm_names
+}
 
 algo = tq.StandardMCTS()
 search_tree = algo.init_tree()
@@ -286,10 +299,12 @@ from pathlib import Path
 import treequest as tq
 from PIL import Image
 
+
 class State:
     text: str
     image: Image.Image
 
+
 def state_formatter_html(state: State) -> str:
     """Formats the state for HTML visualization."""
     buffer = BytesIO()
@@ -299,13 +314,15 @@ def state_formatter_html(state: State) -> str:
     # In case state.text cannot be trusted, we recommend e.g., html.escape(state.text) to avoid html injection.
     return f'<p>{state.text}</p><br/><img src="data:image/webp;base64,{img_str}" width=100% />'
 
+
 generate_fns = {
     "action_1": ...  # Define your node generation functions (actions) here
 }
 
 algo = tq.ABMCTSA()
 search_tree = algo.init_tree()
-output_dir = Path("search_tree"); output_dir.mkdir(exist_ok=True)
+output_dir = Path("search_tree")
+output_dir.mkdir(exist_ok=True)
 for i in range(10):
     search_tree = algo.step(search_tree, generate_fns)
     # Rendering "latest.html" for checking search progress on the fly
@@ -314,11 +331,16 @@ for i in range(10):
             search_tree,
             output_basename=output_dir / "latest",
             format="html",
-            state_formatter=state_formatter_html, # Use html custom formetter
+            state_formatter=state_formatter_html,  # Use html custom formetter
         )
 
 tq.render(search_tree, output_basename=output_dir / "final", format="pdf")
-tq.render(search_tree, output_basename=output_dir / "final", format="html", state_formatter=state_formatter_html)
+tq.render(
+    search_tree,
+    output_basename=output_dir / "final",
+    format="html",
+    state_formatter=state_formatter_html,
+)
 ```
 
 > **IMPORTANT:** When using HTML format, ensure that the HTML file is securely handled, especially if the state formatter includes raw HTML content. Avoid opening untrusted HTML files in your browser. For example, XSS (cross site scripting) attacks can occur if the state includes malicious HTML/JavaScript code.

{treequest-0.3.2 → treequest-0.3.3}/README.md

@@ -21,24 +21,27 @@ import treequest as tq
 # Each node is associated with a user-definable `state`.
 State = str
 
+
 # 1. Define a function to be used for node generation.
 def generate(parent_state: State | None) -> tuple[State, float]:
     """Generates new states and scores based on the parent state."""
-    if parent_state is None: # None represents the expansion from root.
+    if parent_state is None:  # None represents the expansion from root.
         new_state = "Initial state"
     else:
         new_state = f"State after {parent_state}"
 
-    score = random.random() # A score for the new state; It should be normalized to the [0, 1] range.
+    # A score for the new state; It should be normalized to the [0, 1] range.
+    score = random.random()
     return new_state, score
 
+
 # 2. Instantiate the algorithm and a search tree object.
 algo = tq.ABMCTSA()
 search_tree = algo.init_tree()
 
 # 3. Run the search with a generation budget (10 in this case).
 for _ in range(10):
-    search_tree = algo.step(search_tree, {'Action A': generate})
+    search_tree = algo.step(search_tree, {"Action A": generate})
 
 # 4. Extract the best score and state.
 best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
@@ -46,7 +49,8 @@ print(f"Best state: {best_state}, Score: {best_node_score}")
 
 # 5. Visualize the search tree.
 output_file_basename = Path("ab_mcts_a_search_tree")
-tq.render(search_tree, output_file_basename, format="html")  # Generates `ab_mcts_a_search_tree.html`
+# Generates `ab_mcts_a_search_tree.html`
+tq.render(search_tree, output_file_basename, format="html")
 ```
 
 Alternatively, you can use an ask–tell interface with batched AB-MCTS sampling steps:
@@ -56,8 +60,9 @@ import treequest as tq
 
 State = str
 
-def generate(parent_state: State | None) -> tuple[State, float]:
-    ...
+
+def generate(parent_state: State | None) -> tuple[State, float]: ...
+
 
 generate_fns = {"Action A": generate}
 actions = list(generate_fns.keys())
@@ -124,11 +129,13 @@ import dataclasses
 
 import treequest as tq
 
+
 @dataclasses.dataclass
 class State:
     llm_answer: str
     score: float
 
+
 def generate(parent_state: State | None) -> tuple[State, float]:
     """Generate a new node by calling an LLM."""
     if parent_state is None:
@@ -138,12 +145,14 @@ def generate(parent_state: State | None) -> tuple[State, float]:
 
     return state, state.score
 
+
 def initial_generation() -> State:
     """
     Call LLM API to generate an initial answer.
     """
     ...
 
+
 def refine_answer(llm_answer: str, score: float) -> State:
     """
     Call LLM API to refine an answer.
@@ -154,11 +163,11 @@ def refine_answer(llm_answer: str, score: float) -> State:
 algo = tq.ABMCTSM()
 search_tree = algo.init_tree()
 for i in range(20):
-    search_tree = algo.step(search_tree, {'Action Label': generate})
+    search_tree = algo.step(search_tree, {"Action Label": generate})
     # Logging best node during the search.
     if (i + 1) % 5 == 0:
         best_interim_state, _ = tq.top_k(search_tree, algo, k=1)[0]
-        print(f"Iteration {i+1}: Best state so far = {best_interim_state}")
+        print(f"Iteration {i + 1}: Best state so far = {best_interim_state}")
 
 best_state, _ = tq.top_k(search_tree, algo, k=1)[0]
 print(f"Best Answer: {best_state.llm_answer}, Best Score: {best_state.score}")
@@ -173,6 +182,7 @@ from functools import partial
 
 import treequest as tq
 
+
 def generate(llm_name: str, parent_state=None):
     """
     Call LLM API using litellm, vllm, etc., to generate a new node
@@ -180,9 +190,12 @@ def generate(llm_name: str, parent_state=None):
     ...
     return new_state, new_score
 
+
 llm_names = ["o4-mini", "gemini-2.5-pro"]
 # Create dict of different actions backed by different LLMs.
-generate_fns = {llm_name: partial(generate, llm_name=llm_name) for llm_name in llm_names}
+generate_fns = {
+    llm_name: partial(generate, llm_name=llm_name) for llm_name in llm_names
+}
 
 algo = tq.StandardMCTS()
 search_tree = algo.init_tree()
@@ -250,10 +263,12 @@ from pathlib import Path
 import treequest as tq
 from PIL import Image
 
+
 class State:
     text: str
     image: Image.Image
 
+
 def state_formatter_html(state: State) -> str:
     """Formats the state for HTML visualization."""
     buffer = BytesIO()
@@ -263,13 +278,15 @@ def state_formatter_html(state: State) -> str:
     # In case state.text cannot be trusted, we recommend e.g., html.escape(state.text) to avoid html injection.
     return f'<p>{state.text}</p><br/><img src="data:image/webp;base64,{img_str}" width=100% />'
 
+
 generate_fns = {
     "action_1": ...  # Define your node generation functions (actions) here
 }
 
 algo = tq.ABMCTSA()
 search_tree = algo.init_tree()
-output_dir = Path("search_tree"); output_dir.mkdir(exist_ok=True)
+output_dir = Path("search_tree")
+output_dir.mkdir(exist_ok=True)
 for i in range(10):
     search_tree = algo.step(search_tree, generate_fns)
     # Rendering "latest.html" for checking search progress on the fly
@@ -278,11 +295,16 @@ for i in range(10):
             search_tree,
             output_basename=output_dir / "latest",
             format="html",
-            state_formatter=state_formatter_html, # Use html custom formetter
+            state_formatter=state_formatter_html,  # Use html custom formetter
         )
 
 tq.render(search_tree, output_basename=output_dir / "final", format="pdf")
-tq.render(search_tree, output_basename=output_dir / "final", format="html", state_formatter=state_formatter_html)
+tq.render(
+    search_tree,
+    output_basename=output_dir / "final",
+    format="html",
+    state_formatter=state_formatter_html,
+)
 ```
 
 > **IMPORTANT:** When using HTML format, ensure that the HTML file is securely handled, especially if the state formatter includes raw HTML content. Avoid opening untrusted HTML files in your browser. For example, XSS (cross site scripting) attacks can occur if the state includes malicious HTML/JavaScript code.

{treequest-0.3.2 → treequest-0.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "treequest"
-version = "0.3.2"
+version = "0.3.3"
 description = "A Tree Search Library with Flexible API for LLM Inference-Time Scaling"
 readme = "README.md"
 authors = [
@@ -44,16 +44,22 @@ build-backend = "hatchling.build"
 [dependency-groups]
 dev = [
     "line-profiler>=4.2.0",
-    "mypy>=1.15.0",
+    # Pin ruff and mypy so that their new releases do not change CI results.
+    # Keep ruff in sync with .pre-commit-config.yaml.
+    "mypy==2.4.0",
     "pre-commit>=4.2.0",
     "pytest>=8.3.4",
     "pytest-xdist>=3.6.1",
-    "ruff>=0.9.9",
+    "ruff==0.17.1",
     "tqdm>=4.64.0",
     "types-tqdm>=4.64.0",
     "types-PyYAML>=6.0.0",
 ]
 
+[tool.ruff.lint]
+# The rules Ruff enabled by default before 0.16, which expanded its default rule set.
+select = ["E4", "E7", "E9", "F"]
+
 [tool.hatch.build.targets.wheel]
 packages = ["src/treequest"]
 include = [

treequest-0.3.3/src/treequest/algos/ab_mcts_m/_ab_mcts_m_imports.py

@@ -0,0 +1,58 @@
+import importlib
+from pathlib import Path
+
+from treequest.imports import try_import
+
+_ARVIZ_IMPORT_ATTEMPTS = 3
+
+
+def _import_arviz() -> None:
+    # TODO: Remove this workaround after arviz fixes the race in its daily warning
+    # https://github.com/arviz-devs/arviz/issues/2530
+    # arviz 0.23.3 and 0.23.4 update the warning's stamp file through a fixed
+    # temporary file, so importing arviz from several processes at once can raise
+    # FileNotFoundError. The process that won has already updated the stamp, so
+    # importing again normally succeeds.
+    for attempt in range(1, _ARVIZ_IMPORT_ATTEMPTS + 1):
+        try:
+            importlib.import_module("arviz")
+            return
+        except FileNotFoundError as e:
+            # The process that lost the race fails to rename the temporary file.
+            lost_race = Path(e.filename or "").name == "daily_warning.tmp"
+            if not lost_race or attempt == _ARVIZ_IMPORT_ATTEMPTS:
+                raise
+            # Leave sys.modules as it is: Python has already dropped the failed
+            # package, the leftover arviz.preview is safe to reuse, and another
+            # thread may be importing arviz by now.
+
+
+with try_import() as _import:
+    import jax
+    from packaging.version import Version
+
+    # TODO: Remove this hotfix after numpyro fixes incompatibility with jax>=0.7
+    # https://github.com/pyro-ppl/numpyro/issues/2051
+    if Version(jax.__version__) >= Version("0.7.0"):
+        import jax.experimental.pjit as _pjit
+        from jax.extend.core.primitives import jit_p
+
+        _pjit.pjit_p = jit_p  # type: ignore
+    import numpy as np
+    import numpyro  # type: ignore
+    import pandas as pd  # type: ignore
+
+    _import_arviz()  # Import arviz before pymc does; see _import_arviz.
+    import pymc as pm  # type: ignore
+    from pymc.sampling.jax import sample_numpyro_nuts  # type: ignore
+    from xarray import DataArray
+
+__all__ = [
+    "jax",
+    "np",
+    "numpyro",
+    "pd",
+    "pm",
+    "sample_numpyro_nuts",
+    "DataArray",
+]

treequest-0.3.2/src/treequest/algos/ab_mcts_m/_ab_mcts_m_imports.py

@@ -1,29 +0,0 @@
-from treequest.imports import try_import
-
-with try_import() as _import:
-    import jax
-    from packaging.version import Version
-
-    # TODO: Remove this hotfix after numpyro fixes incompatibility with jax>=0.7
-    # https://github.com/pyro-ppl/numpyro/issues/2051
-    if Version(jax.__version__) >= Version("0.7.0"):
-        import jax.experimental.pjit as _pjit
-        from jax.extend.core.primitives import jit_p
-
-        _pjit.pjit_p = jit_p  # type: ignore
-    import numpy as np
-    import numpyro  # type: ignore
-    import pandas as pd  # type: ignore
-    import pymc as pm  # type: ignore
-    from pymc.sampling.jax import sample_numpyro_nuts  # type: ignore
-    from xarray import DataArray
-
-__all__ = [
-    "jax",
-    "np",
-    "numpyro",
-    "pd",
-    "pm",
-    "sample_numpyro_nuts",
-    "DataArray",
-]