diff --git a/pyproject.toml b/pyproject.toml
index e8fdc61..5cca402 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "treequest"
-version = "0.3.0"
+version = "0.3.1"
 description = "A Tree Search Library with Flexible API for LLM Inference-Time Scaling"
 readme = "README.md"
 authors = [
@@ -57,16 +57,10 @@ dev = [
 [tool.hatch.build.targets.wheel]
 packages = ["src/treequest"]
 include = [
-    "src/treequest/vis/assets/colormaps.json",
-    "src/treequest/vis/assets/d3.LICENSE.txt",
-    "src/treequest/vis/assets/d3.v7.min.js",
-    "src/treequest/vis/assets/d3_tree.html.jinja2",
+    "src/treequest",
 ]
 
 [tool.hatch.build.targets.sdist]
 include = [
-    "src/treequest/vis/assets/colormaps.json",
-    "src/treequest/vis/assets/d3.LICENSE.txt",
-    "src/treequest/vis/assets/d3.v7.min.js",
-    "src/treequest/vis/assets/d3_tree.html.jinja2",
+    "src/treequest",
 ]