1from typing import List
 2import llm
 3import os
 4import pathlib
 5import subprocess
 6import tempfile
 7import shutil
 8
 9
10@llm.hookimpl
11def register_fragment_loaders(register):
12    register("repomix", repomix_loader)
13
14
15def repomix_loader(argument: str) -> List[llm.Fragment]:
16    """
17    Load repository contents as fragments using Repomix
18    
19    Argument is a git repository URL (https:// or ssh://)
20    Examples:
21        repomix:https://git.sr.ht/~amolith/willow
22        repomix:ssh://git.sr.ht:~amolith/willow
23    """
24    if not argument.startswith(("https://", "ssh://", "git@")):
25        raise ValueError(
26            f"Repository URL must start with https://, ssh://, or git@ - got: {argument}"
27        )
28    
29    # Check if repomix is available
30    if not shutil.which("repomix"):
31        raise ValueError(
32            "repomix command not found. Please install repomix first: "
33            "https://github.com/yamadashy/repomix"
34        )
35    
36    # Create a temporary directory for the cloned repository
37    with tempfile.TemporaryDirectory() as temp_dir:
38        repo_path = pathlib.Path(temp_dir) / "repo"
39        
40        try:
41            # Clone the repository
42            subprocess.run(
43                ["git", "clone", "--depth=1", argument, str(repo_path)],
44                check=True,
45                capture_output=True,
46                text=True,
47            )
48            
49            # Run repomix on the cloned repository
50            result = subprocess.run(
51                ["repomix", "--stdout", str(repo_path)],
52                check=True,
53                capture_output=True,
54                text=True,
55            )
56            
57            # Create a single fragment with the repomix output
58            fragments = [
59                llm.Fragment(
60                    content=result.stdout,
61                    source=f"repomix:{argument}"
62                )
63            ]
64            
65            return fragments
66            
67        except subprocess.CalledProcessError as e:
68            # Handle Git or repomix errors
69            if "git" in str(e.cmd):
70                raise ValueError(
71                    f"Failed to clone repository {argument}: {e.stderr}"
72                )
73            elif "repomix" in str(e.cmd):
74                raise ValueError(
75                    f"Failed to run repomix on {argument}: {e.stderr}"
76                )
77            else:
78                raise ValueError(
79                    f"Command failed: {e.stderr}"
80                )
81        except Exception as e:
82            # Handle other errors
83            raise ValueError(
84                f"Error processing repository {argument}: {str(e)}"
85            )