Pydantic agent and some important infra

.env.example

@@ -0,0 +1,14 @@
+# LOGFIRE_ENVIRONMENT=hostname-username
+
+LOGFIRE_READ_TOKEN=
+LOGFIRE_WRITE_TOKEN=
+# LOGFIRE_CONSOLE_VERBOSE=True
+# LOGFIRE_CONSOLE_MIN_LOG_LEVEL=debug
+
+LANGFUSE_SKIP=1
+# LANGFUSE_SECRET_KEY=
+# LANGFUSE_PUBLIC_KEY=
+LANGFUSE_HOST=https://us.cloud.langfuse.com
+
+# GOOGLE_API_KEY=
+# OPENAI_API_KEY=

config/baseline.yaml

@@ -0,0 +1,9 @@
+# _include: labs/hosts.yaml
+
+agent:
+  AGENT_MODEL: "o3-2025-04-16"
+  MAX_STEPS: 300
+
+# container:
+#   IMAGE_NAME: "agent-image"
+#   CONTAINER_NAME: "agent-container"

config/smoke.yaml

@@ -0,0 +1,4 @@
+_include: baseline.yaml
+
+agent:
+  MAX_STEPS: 10

pyproject.toml

@@ -5,13 +5,18 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
+    "ipdb>=0.13.13",
+    "langfuse>=3.3.4",
     "matplotlib>=3.10.0",
     "numpy>=2.2.2",
+    "pydantic-ai>=1.0.2",
+    "python-ulid[pydantic]>=3.1.0",
 ]
 [dependency-groups]
 dev = [
     "ty>=0.0.1a18",
     "ruff>=0.12.5",
+    "argcomplete>=3.6.2",
 ]
 
 [build-system]
@@ -22,7 +27,10 @@ build-backend = "hatchling.build"
 sources = ["src"]
 
 [project.scripts]
+fork-worktree = "scripts.fork_worktree:main"
+clean-worktrees = "scripts.clean_worktrees:main"
 hello = "template.hello:main"
+agent = "template.agent.main:main"
 
 [tool.pyright]
 # disable pyright in favor of ruff

src/scripts/_clean_worktrees_parser.py

@@ -0,0 +1,15 @@
+import argparse
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the argument parser for the clean-worktrees command."""
+    parser = argparse.ArgumentParser(
+        description="Remove all git worktrees created by fork-worktree matching the 'worktree-*' pattern."
+    )
+    parser.add_argument(
+        "--force",
+        "-f",
+        action="store_true",
+        help="Skip confirmation prompt and remove all matching worktrees.",
+    )
+    return parser

src/scripts/_fork_worktree_parser.py

@@ -0,0 +1,41 @@
+import argparse
+
+DEFAULT_PREFIX = "experiment"
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the argument parser for the fork-worktree command."""
+    parser = argparse.ArgumentParser(
+        description="Create one or more worktrees branching from the current branch.",
+    )
+    parser.add_argument(
+        "-r",
+        "--repeat",
+        type=int,
+        default=1,
+        help="Number of worktrees to create (default: 1).",
+    )
+    parser.add_argument(
+        "--prefix",
+        default=DEFAULT_PREFIX,
+        help=(
+            "Prefix for new branch names. Default: 'experiment'. "
+            "Branch name format will be: <prefix>-<user>-<timestamp>-<commithash>-<customname>"
+        ),
+    )
+    parser.add_argument(
+        "--path",
+        help="Base destination directory. Defaults to directory *above* repo root.",
+    )
+    parser.add_argument(
+        "--no-open",
+        action="store_true",
+        help="Don't open each worktree with 'cursor .' after creation (default: open)",
+    )
+    parser.add_argument(
+        "--names",
+        nargs="+",
+        default=[],
+        help="Custom names for the branches. If fewer names than --repeat, names will be reused or defaults applied.",
+    )
+    return parser

src/scripts/clean_worktrees.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+
+"""clean_worktrees: Utility script to remove all git worktrees created by fork-worktree.
+
+WARNING: This script removes worktrees and their associated directories if they match the
+'worktree-*' naming convention. Use with caution.
+It will ask for confirmation before deleting anything.
+"""
+
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import argcomplete
+
+from ._clean_worktrees_parser import build_parser
+
+try:
+    REPO_ROOT_RAW = subprocess.check_output(
+        ["git", "rev-parse", "--show-toplevel"],
+        text=True,
+        stderr=subprocess.PIPE,
+    ).strip()
+    if not REPO_ROOT_RAW:
+        raise ValueError("git rev-parse returned empty string")
+    REPO_ROOT = Path(REPO_ROOT_RAW).resolve()
+except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e:
+    print("Error: Failed to determine git repository root.", file=sys.stderr)
+    print(f"Reason: {e}", file=sys.stderr)
+    print("Please run this script from within a git repository.", file=sys.stderr)
+    sys.exit(1)
+
+
+def sh(args: list, **kwargs) -> subprocess.CompletedProcess:
+    """Run a shell command, echoing it. Optionally abort on error."""
+    args_: list[str] = [str(a) for a in args]
+    print("$", " ".join(args_))
+    return subprocess.run(args_, capture_output=True, text=True, **kwargs)
+
+
+def get_worktree_paths() -> list[Path]:
+    """Get a list of paths for managed worktrees, excluding the main one."""
+    result = sh(["git", "worktree", "list", "--porcelain"], check=False)
+    if result.returncode != 0:
+        print(
+            f"Warning: 'git worktree list' failed. Stderr: {result.stderr.strip()}",
+            file=sys.stderr,
+        )
+        return []
+
+    worktree_paths = []
+    current_path_str = None
+    is_main_worktree = False
+
+    for line in result.stdout.strip().split("\n"):
+        if line.startswith("worktree "):
+            current_path_str = line.split(" ", 1)[1]
+            is_main_worktree = False
+            try:
+                if Path(current_path_str).resolve() == REPO_ROOT:
+                    is_main_worktree = True
+            except Exception:
+                pass
+
+        elif line.strip() == "":
+            if current_path_str and not is_main_worktree:
+                current_path = Path(current_path_str)
+                if current_path.name.startswith("worktree-"):
+                    worktree_paths.append(current_path)
+
+            current_path_str = None
+            is_main_worktree = False
+
+    if current_path_str and not is_main_worktree:
+        current_path = Path(current_path_str)
+        if current_path.name.startswith("worktree-"):
+            worktree_paths.append(current_path)
+
+    return worktree_paths
+
+
+def main() -> None:
+    """Entry point for the script."""
+    parser = build_parser()
+
+    if "argcomplete" in sys.modules:
+        argcomplete.autocomplete(parser)
+
+    args = parser.parse_args()
+
+    print("Looking for git worktrees...")
+
+    try:
+        worktree_paths = get_worktree_paths()
+    except Exception as e:
+        print(f"Error getting worktree list: {e}", file=sys.stderr)
+        print("Are you in a git repository?", file=sys.stderr)
+        sys.exit(1)
+
+    if not worktree_paths:
+        print("No managed 'worktree-*' worktrees found to remove.")  # Clarified message
+        sys.exit(0)
+
+    print("Found the following managed worktrees to remove:")
+    for path in worktree_paths:
+        print(f" - {path}")
+
+    if not args.force:
+        try:
+            confirm = input(
+                "Are you sure you want to remove all these worktrees and their directories? [y/N]: "
+            )
+        except EOFError:
+            confirm = "n"
+
+        if confirm.lower() != "y":
+            print("Aborting.")
+            sys.exit(0)
+    else:
+        print("\nForced removal requested. Proceeding without confirmation.")
+
+    print("\nRemoving worktrees...")
+    removed_count = 0
+    failed_count = 0
+
+    for path in worktree_paths:
+        print(f"\nProcessing worktree at: {path}")
+
+        if not path.exists():
+            print(
+                f"Warning: Directory {path} does not exist. Removing worktree entry only."
+            )
+
+        result = sh(["git", "worktree", "remove", "--force", path], check=False)
+        if result.returncode == 0:
+            print(f"Successfully removed git worktree entry for {path}")
+
+            if path.exists():
+                print(f"Removing directory {path}...")
+                try:
+                    if path.resolve() == REPO_ROOT.resolve():
+                        print(
+                            f"Safety check: Preventing removal of main repository directory ({path}).",
+                            file=sys.stderr,
+                        )
+                        failed_count += 1
+                        continue
+
+                    shutil.rmtree(path)
+                    print(f"Successfully removed directory {path}")
+                    removed_count += 1
+                except OSError as e:
+                    print(f"Error removing directory {path}: {e}", file=sys.stderr)
+                    print("You might need to remove it manually.", file=sys.stderr)
+                    failed_count += 1
+            else:
+                removed_count += 1
+
+        else:
+            print(f"Failed to remove git worktree entry for {path}. Stderr:")
+            print(result.stderr.strip(), file=sys.stderr)
+            failed_count += 1
+
+    print("\n---------------------")
+    print(f"Cleanup complete. {removed_count} worktrees removed.")
+    if failed_count > 0:
+        print(
+            f"{failed_count} worktrees failed to remove completely. Check logs above."
+        )
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()