treequest 0.2.0__tar.gz → 0.3.1__tar.gz

treequest-0.2.0/README.md → treequest-0.3.1/PKG-INFO

@@ -1,3 +1,39 @@
+Metadata-Version: 2.4
+Name: treequest
+Version: 0.3.1
+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
+Requires-Python: >=3.11
+Requires-Dist: joblib>=1.5.2
+Requires-Dist: pyyaml>=6.0.0
+Requires-Dist: scipy>=1.0.0
+Requires-Dist: ulid-py>=1.1.0
+Provides-Extra: abmcts-m
+Requires-Dist: jax>=0.5.2; extra == 'abmcts-m'
+Requires-Dist: numpy>=1.26.4; extra == 'abmcts-m'
+Requires-Dist: numpyro>=0.17.0; extra == 'abmcts-m'
+Requires-Dist: packaging>=23; extra == 'abmcts-m'
+Requires-Dist: pandas>=2.2.3; extra == 'abmcts-m'
+Requires-Dist: pymc>=5.21.0; extra == 'abmcts-m'
+Requires-Dist: xarray>=2025.1.2; extra == 'abmcts-m'
+Provides-Extra: all
+Requires-Dist: graphviz>=0.20.3; extra == 'all'
+Requires-Dist: jax>=0.5.2; extra == 'all'
+Requires-Dist: jinja2>=3.1.0; extra == 'all'
+Requires-Dist: numpy>=1.26.4; extra == 'all'
+Requires-Dist: numpyro>=0.17.0; extra == 'all'
+Requires-Dist: packaging>=23; extra == 'all'
+Requires-Dist: pandas>=2.2.3; extra == 'all'
+Requires-Dist: pillow>=10.0.0; extra == 'all'
+Requires-Dist: pymc>=5.21.0; extra == 'all'
+Requires-Dist: xarray>=2025.1.2; extra == 'all'
+Provides-Extra: vis
+Requires-Dist: graphviz>=0.20.3; extra == 'vis'
+Requires-Dist: jinja2>=3.1.0; extra == 'vis'
+Requires-Dist: pillow>=10.0.0; extra == 'vis'
+Description-Content-Type: text/markdown
+
 # TreeQuest
 [![Python](https://img.shields.io/badge/python-3.11-blue)](https://www.python.org)
 [![GitHub license](https://img.shields.io/badge/license-Apache%202.0-blue?style=flat-square)](https://www.apache.org/licenses/LICENSE-2.0)
@@ -14,6 +50,7 @@ A flexible answer tree search library featuring **AB-MCTS**, useful for (but not
 ## Quick Start
 ```python
 import random
+from pathlib import Path
 
 import treequest as tq
 
@@ -42,6 +79,10 @@ for _ in range(10):
 # 4. Extract the best score and state.
 best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
 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`
 ```
 
 Alternatively, you can use an ask–tell interface with batched AB-MCTS sampling steps:
@@ -88,16 +129,27 @@ We recommend `batch_size<=5` as a starting point.
 - Checkpointing and resuming searches.
 
 ## Installation
-### uv
 First, install [`uv`](https://github.com/astral-sh/uv?tab=readme-ov-file#installation). Then you can install TreeQuest with the following command:
 ```bash
-uv add "treequest[abmcts-m]"
+uv add "treequest[all]"
 ```
 
-### pip
 Alternatively, you can use pip to install TreeQuest:
 ```bash
-pip install "treequest[abmcts-m]"
+pip install "treequest[all]"
+```
+
+There are optional dependencies for ABMCTS-M and visualization features. You can install them with:
+```bash
+# with all dependencies
+uv add "treequest[all]"
+
+# without dependencies
+uv add treequest
+
+# with individual dependencies
+uv add "treequest[abmcts-m]"
+uv add "treequest[vis]"
 ```
 
 ## Usage
@@ -121,7 +173,7 @@ def generate(parent_state: State | None) -> tuple[State, float]:
         state = refine_answer(parent_state.llm_answer, parent_state.score)
 
     return state, state.score
-    
+
 def initial_generation() -> State:
     """
     Call LLM API to generate an initial answer.
@@ -218,7 +270,57 @@ for _ in range(30):
     search_tree = ab_mcts_m.step(search_tree, generate_fns)
 ```
 
-**NOTE**: To run AB-MCTS-M, you need to install extra dependencies with the `treequest[abmcts-m]` option.
+**NOTE**: To run AB-MCTS-M, you need to install extra dependencies with the `treequest[abmcts-m]` or `treequest[all]` option.
+
+## Visualization
+TreeQuest provides visualization utilities to render the search tree. You can visualize the search tree using the `tq.render` function as shown in the Quick Start section.
+You need to install optional dependencies for visualization either by `uv add treequest[vis]` or `uv add treequest[all]`.
+
+```python
+import html
+import base64
+from io import BytesIO
+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()
+    state.image.save(buffer, format="WEBP")
+    img_str = base64.b64encode(buffer.getvalue()).decode()
+    # HTML representation is allowed (ensure safety when using untrusted content)
+    # 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)
+for i in range(10):
+    search_tree = algo.step(search_tree, generate_fns)
+    # Rendering "latest.html" for checking search progress on the fly
+    if i % 2 == 0 or i == 9:
+        tq.render(
+            search_tree,
+            output_basename=output_dir / "latest",
+            format="html",
+            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)
+```
+
+> **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.
 
 ## Requirements
 
@@ -230,14 +332,24 @@ Contributions are welcome! Please see [CONTRIBUTING.md](./CONTRIBUTING.md) for d
 
 ## Citation
 ```bibtex
-@article{inoue2025wider,
-  title={Wider or Deeper? Scaling LLM Inference-Time Compute with Adaptive Branching Tree Search},
-  author={Inoue, Yuichi and Misaki, Kou and Imajuku, Yuki and Kuroki, So and Nakamura, Taishi and Akiba, Takuya},
-  journal={arXiv preprint arXiv:2503.04412},
-  year={2025}
+@inproceedings{inoue2025wider,
+  title={Wider or Deeper?  Scaling {LLM} Inference-Time Compute with Adaptive Branching Tree Search},
+  author={Yuichi Inoue and Kou Misaki and Yuki Imajuku and So Kuroki and Taishi Nakamura and Takuya Akiba},
+  booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
+  year={2025},
+  url={https://openreview.net/forum?id=jAsr5GHt3P}
 }
 ```
 
 ## License
 
 [Apache 2.0](./LICENSE)
+
+### Third-party notices
+
+TreeQuest bundles a few assets whose original authors retain copyright:
+
+- **D3.js v7** by Mike Bostock and contributors, distributed under the ISC License. The full license text is included in `src/treequest/vis/assets/d3.LICENSE.txt` alongside the `d3.v7.min.js` binary.
+- **Colormap samples** extracted from matplotlib and seaborn. The data and license references are embedded in `src/treequest/vis/assets/colormaps.json`.
+
+These notices must be preserved in any redistribution of TreeQuest, including any compiled artifacts.

treequest-0.2.0/PKG-INFO → treequest-0.3.1/README.md

@@ -1,23 +1,3 @@
-Metadata-Version: 2.4
-Name: treequest
-Version: 0.2.0
-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
-Requires-Python: >=3.11
-Requires-Dist: joblib>=1.5.2
-Requires-Dist: scipy>=1.0.0
-Requires-Dist: ulid-py>=1.1.0
-Provides-Extra: abmcts-m
-Requires-Dist: jax>=0.5.2; extra == 'abmcts-m'
-Requires-Dist: numpy>=1.26.4; extra == 'abmcts-m'
-Requires-Dist: numpyro>=0.17.0; extra == 'abmcts-m'
-Requires-Dist: packaging>=23; extra == 'abmcts-m'
-Requires-Dist: pandas>=2.2.3; extra == 'abmcts-m'
-Requires-Dist: pymc>=5.21.0; extra == 'abmcts-m'
-Requires-Dist: xarray>=2025.1.2; extra == 'abmcts-m'
-Description-Content-Type: text/markdown
-
 # TreeQuest
 [![Python](https://img.shields.io/badge/python-3.11-blue)](https://www.python.org)
 [![GitHub license](https://img.shields.io/badge/license-Apache%202.0-blue?style=flat-square)](https://www.apache.org/licenses/LICENSE-2.0)
@@ -34,6 +14,7 @@ A flexible answer tree search library featuring **AB-MCTS**, useful for (but not
 ## Quick Start
 ```python
 import random
+from pathlib import Path
 
 import treequest as tq
 
@@ -62,6 +43,10 @@ for _ in range(10):
 # 4. Extract the best score and state.
 best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
 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`
 ```
 
 Alternatively, you can use an ask–tell interface with batched AB-MCTS sampling steps:
@@ -108,16 +93,27 @@ We recommend `batch_size<=5` as a starting point.
 - Checkpointing and resuming searches.
 
 ## Installation
-### uv
 First, install [`uv`](https://github.com/astral-sh/uv?tab=readme-ov-file#installation). Then you can install TreeQuest with the following command:
 ```bash
-uv add "treequest[abmcts-m]"
+uv add "treequest[all]"
 ```
 
-### pip
 Alternatively, you can use pip to install TreeQuest:
 ```bash
-pip install "treequest[abmcts-m]"
+pip install "treequest[all]"
+```
+
+There are optional dependencies for ABMCTS-M and visualization features. You can install them with:
+```bash
+# with all dependencies
+uv add "treequest[all]"
+
+# without dependencies
+uv add treequest
+
+# with individual dependencies
+uv add "treequest[abmcts-m]"
+uv add "treequest[vis]"
 ```
 
 ## Usage
@@ -141,7 +137,7 @@ def generate(parent_state: State | None) -> tuple[State, float]:
         state = refine_answer(parent_state.llm_answer, parent_state.score)
 
     return state, state.score
-    
+
 def initial_generation() -> State:
     """
     Call LLM API to generate an initial answer.
@@ -238,7 +234,57 @@ for _ in range(30):
     search_tree = ab_mcts_m.step(search_tree, generate_fns)
 ```
 
-**NOTE**: To run AB-MCTS-M, you need to install extra dependencies with the `treequest[abmcts-m]` option.
+**NOTE**: To run AB-MCTS-M, you need to install extra dependencies with the `treequest[abmcts-m]` or `treequest[all]` option.
+
+## Visualization
+TreeQuest provides visualization utilities to render the search tree. You can visualize the search tree using the `tq.render` function as shown in the Quick Start section.
+You need to install optional dependencies for visualization either by `uv add treequest[vis]` or `uv add treequest[all]`.
+
+```python
+import html
+import base64
+from io import BytesIO
+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()
+    state.image.save(buffer, format="WEBP")
+    img_str = base64.b64encode(buffer.getvalue()).decode()
+    # HTML representation is allowed (ensure safety when using untrusted content)
+    # 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)
+for i in range(10):
+    search_tree = algo.step(search_tree, generate_fns)
+    # Rendering "latest.html" for checking search progress on the fly
+    if i % 2 == 0 or i == 9:
+        tq.render(
+            search_tree,
+            output_basename=output_dir / "latest",
+            format="html",
+            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)
+```
+
+> **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.
 
 ## Requirements
 
@@ -250,14 +296,24 @@ Contributions are welcome! Please see [CONTRIBUTING.md](./CONTRIBUTING.md) for d
 
 ## Citation
 ```bibtex
-@article{inoue2025wider,
-  title={Wider or Deeper? Scaling LLM Inference-Time Compute with Adaptive Branching Tree Search},
-  author={Inoue, Yuichi and Misaki, Kou and Imajuku, Yuki and Kuroki, So and Nakamura, Taishi and Akiba, Takuya},
-  journal={arXiv preprint arXiv:2503.04412},
-  year={2025}
+@inproceedings{inoue2025wider,
+  title={Wider or Deeper?  Scaling {LLM} Inference-Time Compute with Adaptive Branching Tree Search},
+  author={Yuichi Inoue and Kou Misaki and Yuki Imajuku and So Kuroki and Taishi Nakamura and Takuya Akiba},
+  booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
+  year={2025},
+  url={https://openreview.net/forum?id=jAsr5GHt3P}
 }
 ```
 
 ## License
 
 [Apache 2.0](./LICENSE)
+
+### Third-party notices
+
+TreeQuest bundles a few assets whose original authors retain copyright:
+
+- **D3.js v7** by Mike Bostock and contributors, distributed under the ISC License. The full license text is included in `src/treequest/vis/assets/d3.LICENSE.txt` alongside the `d3.v7.min.js` binary.
+- **Colormap samples** extracted from matplotlib and seaborn. The data and license references are embedded in `src/treequest/vis/assets/colormaps.json`.
+
+These notices must be preserved in any redistribution of TreeQuest, including any compiled artifacts.

{treequest-0.2.0 → treequest-0.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "treequest"
-version = "0.2.0"
+version = "0.3.1"
 description = "A Tree Search Library with Flexible API for LLM Inference-Time Scaling"
 readme = "README.md"
 authors = [
@@ -9,6 +9,7 @@ authors = [
 ]
 requires-python = ">=3.11"
 dependencies = [
+    "PyYAML>=6.0.0",
     "joblib>=1.5.2",
     "scipy>=1.0.0",
     "ulid-py>=1.1.0",
@@ -25,17 +26,41 @@ abmcts-m = [
     "packaging>=23",
 ]
 
+vis = [
+    "Pillow>=10.0.0",
+    "graphviz>=0.20.3",
+    "jinja2>=3.1.0",
+]
+
+all = [
+    "treequest[abmcts-m]",
+    "treequest[vis]",
+]
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
 
 [dependency-groups]
 dev = [
-    "graphviz>=0.20.3",
     "line-profiler>=4.2.0",
     "mypy>=1.15.0",
     "pre-commit>=4.2.0",
     "pytest>=8.3.4",
     "pytest-xdist>=3.6.1",
     "ruff>=0.9.9",
+    "tqdm>=4.64.0",
+    "types-tqdm>=4.64.0",
+    "types-PyYAML>=6.0.0",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/treequest"]
+include = [
+    "src/treequest",
+]
+
+[tool.hatch.build.targets.sdist]
+include = [
+    "src/treequest",
 ]

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/__init__.py

@@ -20,6 +20,7 @@ from .algos.base import Algorithm
 from .algos.standard_mcts import StandardMCTS
 from .algos.tree_of_thought_bfs import TreeOfThoughtsBFSAlgo
 from .ranker import top_k
+from .vis import render
 
 __all__ = [
     "StandardMCTS",
@@ -28,4 +29,5 @@ __all__ = [
     "ABMCTSA",
     "ABMCTSM",
     "Algorithm",
+    "render",
 ]

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/algos/ab_mcts_a/algo.py

@@ -8,8 +8,8 @@ from typing import Dict, Generic, List, Literal, Optional, Tuple, TypeVar, Union
 from treequest.algos.ab_mcts_a.prob_state import NodeProbState, PriorConfig
 from treequest.algos.base import Algorithm
 from treequest.algos.tree import Node, Tree
-from treequest.trial import Trial, TrialId, TrialStore
-from treequest.types import GenerateFnType, StateScoreType
+from treequest.trial import Trial, TrialStore
+from treequest.types import GenerateFnType, StateScoreType, TrialId
 
 # Type variable for state
 StateT = TypeVar("StateT")
@@ -369,7 +369,9 @@ class ABMCTSA(Algorithm[StateT, ABMCTSAAlgoState[StateT]]):
         parent_node = state.tree.get_node(finished_trial.node_to_expand)
         action = finished_trial.action
         # Add new node to the tree
-        new_node = state.tree.add_node(result, parent_node)
+        new_node = state.tree.add_node(
+            result, parent_node, trial_id=finished_trial.trial_id
+        )
 
         # Update Thompson state with the new node
         thompson_state = state.thompson_states.get(parent_node)

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/algos/ab_mcts_m/algo.py

@@ -14,12 +14,12 @@ from treequest.algos.ab_mcts_m.pymc_interface import (
 )
 from treequest.algos.base import Algorithm
 from treequest.algos.tree import Node, Tree
-from treequest.trial import Trial, TrialId, TrialStore
-from treequest.types import GenerateFnType, StateScoreType
+from treequest.trial import Trial, TrialStore
+from treequest.types import GenerateFnType, StateScoreType, TrialId
 
 StateT = TypeVar("StateT")
 
-_WORKER_ALGO = None
+_WORKER_ALGO: Optional["ABMCTSM"] = None
 
 
 def _worker_init_abmctsm(config: dict, per_worker_cpu_devices: int):
@@ -169,7 +169,7 @@ class ABMCTSM(Algorithm[StateT, ABMCTSMState[StateT]]):
         Args:
             state: Current algorithm state
             node: Node to select child from
-            generate_fn: Mapping of action names to generation functions
+            actions: List of action names
 
         Returns:
             Tuple of (selected node, action if new node was generated)
@@ -297,7 +297,9 @@ class ABMCTSM(Algorithm[StateT, ABMCTSMState[StateT]]):
         parent_node = state.tree.get_node(finished_trial.node_to_expand)
 
         # Add new node to the tree
-        new_node = state.tree.add_node(result, parent_node)
+        new_node = state.tree.add_node(
+            result, parent_node, trial_id=finished_trial.trial_id
+        )
 
         # Record observation
         state.all_observations[new_node.expand_idx] = Observation(

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/algos/base.py

@@ -2,8 +2,8 @@ import abc
 from collections.abc import Mapping
 from typing import Generic, List, TypeVar
 
-from treequest.trial import Trial, TrialId
-from treequest.types import GenerateFnType, StateScoreType
+from treequest.trial import Trial
+from treequest.types import GenerateFnType, StateScoreType, TrialId
 
 # Type variables for node state and algorithm state
 NodeStateT = TypeVar("NodeStateT")

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/algos/best_first_search.py

@@ -7,8 +7,8 @@ from typing import Generic, List, Tuple, TypeVar
 
 from treequest.algos.base import Algorithm
 from treequest.algos.tree import Node, Tree
-from treequest.trial import Trial, TrialId, TrialStoreWithNodeQueue
-from treequest.types import GenerateFnType, StateScoreType
+from treequest.trial import Trial, TrialStoreWithNodeQueue
+from treequest.types import GenerateFnType, StateScoreType, TrialId
 
 # Type variable for state
 StateT = TypeVar("StateT")
@@ -189,7 +189,9 @@ class BestFirstSearchAlgo(Algorithm[StateT, BFSState[StateT]]):
 
         parent_node = state.tree.get_node(finished_trial.node_to_expand)
 
-        new_node = state.tree.add_node(result, parent_node)
+        new_node = state.tree.add_node(
+            result, parent_node, trial_id=finished_trial.trial_id
+        )
 
         # Add the new node to the priority queue
         heappush(state.leaves, BFSHeapItem(node=new_node, score=new_score))

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/algos/multi_armed_bandit_ucb.py

@@ -7,8 +7,8 @@ from typing import Generic, List, TypeVar
 
 from treequest.algos.base import Algorithm
 from treequest.algos.tree import Tree
-from treequest.trial import Trial, TrialId, TrialStore
-from treequest.types import GenerateFnType, StateScoreType
+from treequest.trial import Trial, TrialStore
+from treequest.types import GenerateFnType, StateScoreType, TrialId
 
 # Type variable for state
 StateT = TypeVar("StateT")
@@ -172,7 +172,11 @@ class MultiArmedBanditUCBAlgo(Algorithm[StateT, UCBState[StateT]]):
         ):  # Trial is no longer valid, so we do not reflect the result to state
             return state
 
-        state.tree.add_node(result, state.tree.get_node(finished_trial.node_to_expand))
+        state.tree.add_node(
+            result,
+            state.tree.get_node(finished_trial.node_to_expand),
+            trial_id=finished_trial.trial_id,
+        )
 
         # Update scores for the selected action
         state.scores_by_action[finished_trial.action].append(new_score)

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/algos/standard_mcts.py

@@ -6,8 +6,8 @@ from typing import Dict, Generic, List, Optional, TypeVar
 
 from treequest.algos.base import Algorithm
 from treequest.algos.tree import Node, Tree
-from treequest.trial import Trial, TrialId, TrialStoreWithNodeQueue
-from treequest.types import GenerateFnType, StateScoreType
+from treequest.trial import Trial, TrialStoreWithNodeQueue
+from treequest.types import GenerateFnType, StateScoreType, TrialId
 
 # Type variable for state
 StateT = TypeVar("StateT")
@@ -128,7 +128,9 @@ class StandardMCTS(Algorithm[StateT, MCTSState[StateT]]):
 
         parent_node = state.tree.get_node(finished_trial.node_to_expand)
         # Add the new node to the tree
-        new_node = state.tree.add_node(result, parent_node)
+        new_node = state.tree.add_node(
+            result, parent_node, trial_id=finished_trial.trial_id
+        )
 
         # Backpropagation updates visit_counts and value_sums for the new node and all ancestors.
         self._backpropagate(state, new_node, new_score)

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/algos/tree.py

@@ -1,7 +1,7 @@
 import dataclasses
 from typing import Generic, List, Optional, TypeVar
 
-from treequest.types import StateScoreType
+from treequest.types import StateScoreType, TrialId
 
 # Type variable for state
 StateT = TypeVar("StateT")
@@ -19,6 +19,7 @@ class Node(Generic[StateT]):
 
     parent: Optional["Node[StateT]"] = None
     children: List["Node[StateT]"] = dataclasses.field(default_factory=list)
+    trial_id: Optional[TrialId] = None
 
     def __post_init__(self) -> None:
         if not self.is_root():
@@ -112,10 +113,19 @@ class Tree(Generic[StateT]):
         ]
 
     def add_node(
-        self, state_score: StateScoreType[StateT], parent: Node[StateT]
+        self,
+        state_score: StateScoreType[StateT],
+        parent: Node[StateT],
+        trial_id: Optional[TrialId] = None,
     ) -> Node[StateT]:
         state, score = state_score
-        node = Node(state=state, score=score, parent=parent, expand_idx=self.size - 1)
+        node = Node(
+            state=state,
+            score=score,
+            parent=parent,
+            expand_idx=self.size - 1,
+            trial_id=trial_id,
+        )
         parent.children.append(node)
         self.size += 1
         return node

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/algos/tree_of_thought_bfs.py

@@ -8,8 +8,8 @@ from typing import Generic, List, TypeVar
 
 from treequest.algos.base import Algorithm
 from treequest.algos.tree import Node, Tree
-from treequest.trial import Trial, TrialId, TrialStoreWithNodeQueue
-from treequest.types import GenerateFnType, StateScoreType
+from treequest.trial import Trial, TrialStoreWithNodeQueue
+from treequest.types import GenerateFnType, StateScoreType, TrialId
 
 # Type variable for state
 StateT = TypeVar("StateT")
@@ -228,7 +228,7 @@ class TreeOfThoughtsBFSAlgo(Algorithm[StateT, ToTBFSState[StateT]]):
 
         parent = state.tree.get_node(finished_trial.node_to_expand)
         # Add to the tree
-        state.tree.add_node(result, parent)
+        state.tree.add_node(result, parent, trial_id=finished_trial.trial_id)
         # Update current depth
         state.current_depth = max(state.current_depth, parent.depth + 1)
 

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/trial.py

@@ -6,9 +6,7 @@ from typing import Generic, Literal, TypeAlias, TypeVar, get_args
 import ulid
 
 from treequest.algos.tree import Node
-
-TrialId: TypeAlias = str
-NodeId: TypeAlias = int
+from treequest.types import NodeId, TrialId
 
 TrialStatus: TypeAlias = Literal["RUNNING", "INVALID", "COMPLETE"]
 TRIAL_STATUSES: tuple[TrialStatus, ...] = get_args(TrialStatus)

{treequest-0.2.0 → treequest-0.3.1}/src/treequest/types.py

@@ -1,5 +1,9 @@
 from collections.abc import Callable
-from typing import List, Optional, Tuple, TypeVar
+from typing import List, Optional, Tuple, TypeAlias, TypeVar
+
+# Type alias for NodeId and TrialId
+NodeId: TypeAlias = int
+TrialId: TypeAlias = str
 
 # Type variable for node state
 NodeStateT = TypeVar("NodeStateT")