MCPunk

  • mcpunk
import logging import math import os import time from concurrent.futures import ProcessPoolExecutor, as_completed from pathlib import Path from threading import Lock, Timer from typing import Literal import more_itertools from git import Repo from pydantic import ( BaseModel, ) from watchdog.events import ( DirCreatedEvent, DirDeletedEvent, DirModifiedEvent, FileCreatedEvent, FileDeletedEvent, FileModifiedEvent, FileSystemEventHandler, ) from watchdog.observers import Observer from mcpunk.file_chunk import Chunk, ChunkCategory from mcpunk.file_chunkers import ( BaseChunker, MarkdownChunker, PythonChunker, VueChunker, WholeFileChunker, ) ALL_CHUNKERS: list[type[BaseChunker]] = [ PythonChunker, MarkdownChunker, VueChunker, # Want the WholeFileChunker to be last as it's more of a "fallback" chunker WholeFileChunker, ] logger = logging.getLogger(__name__) class _ProjectFileHandler(FileSystemEventHandler): def __init__( self, project: "Project", file_watch_refresh_freq_seconds: float = 0.1, ) -> None: self.project = project self._project_lock = Lock() self._paths_pending_refresh: set[Path] = set() self._paths_pending_refresh_lock = Lock() self._refresh_freq_sec = file_watch_refresh_freq_seconds self._timer: Timer | None = None self._schedule_refresh() def _schedule_refresh(self) -> None: if self._timer is not None: self._timer.cancel() self._timer = Timer(self._refresh_freq_sec, self._refresh_paths) self._timer.daemon = True self._timer.start() def _refresh_paths(self) -> None: with self._project_lock: try: with self._paths_pending_refresh_lock: paths_pending_refresh = self._paths_pending_refresh.copy() self._paths_pending_refresh.clear() if paths_pending_refresh: logger.info(f"Refreshing {len(paths_pending_refresh)} paths") _paths_fmt = "\n\t".join(str(x) for x in paths_pending_refresh) logger.debug(f"Refreshing\n\t{_paths_fmt}") paths_to_delete = {x for x in paths_pending_refresh if not x.exists()} for p in paths_to_delete: if p.absolute() in self.project.file_map: del self.project.file_map[p.absolute()] dir_paths = {x for x in paths_pending_refresh if x.exists() and x.is_dir()} paths_to_really_refresh = paths_pending_refresh - dir_paths - paths_to_delete self.project.load_files(list(paths_to_really_refresh)) except Exception: logger.exception("Error refreshing paths") finally: self._schedule_refresh() def _path_event( self, path: Path, *, action: Literal["modified", "created", "deleted"], ) -> None: path = path.absolute() logger.debug(f"watchdog says {action}: {path}") if self._should_process(path): logger.debug(f"New path pending refresh: {path}") with self._paths_pending_refresh_lock: self._paths_pending_refresh.add(path) else: logger.debug(f"Ignoring {action} for {path}") def _should_process(self, path: Path) -> bool: if self.project.git_repo is None: return True assert path.is_absolute() # We don't want to exclude non-existent files, as they may have been deleted. if path.exists() and path.is_dir(): return False # Special case this, as we def don't want it and it seems that # `git check-ignore` doesn't consider it as ignored. if path.is_relative_to(self.project.root / ".git"): return False try: rel_path = str(path.relative_to(self.project.root)) t1 = time.monotonic() check_ignore_res: str = self.project.git_repo.git.execute( # type: ignore[call-overload] ["git", "check-ignore", str(rel_path)], with_exceptions=False, ) logger.debug(f"git check-ignore took {(time.monotonic() - t1) * 1000:.4f}ms") return check_ignore_res == "" except Exception: logger.exception(f"Error checking git ignore for {path}") return False def on_modified(self, event: FileModifiedEvent | DirModifiedEvent) -> None: self._path_event( Path(self._to_str(event.src_path)), action="modified", ) def on_created(self, event: FileCreatedEvent | DirCreatedEvent) -> None: self._path_event( Path(self._to_str(event.src_path)), action="created", ) def on_deleted(self, event: FileDeletedEvent | DirDeletedEvent) -> None: self._path_event( Path(self._to_str(event.src_path)), action="deleted", ) @staticmethod def _to_str(s: str | bytes) -> str: if isinstance(s, bytes): return s.decode("utf-8") return s class File(BaseModel): chunks: list[Chunk] abs_path: Path contents: str ext: str # File extension @classmethod def from_file_contents( cls, source_code: str, file_path: Path, max_chunk_size: int = 10_000, ) -> "File": """Extract all callables, calls and imports from the given source code file.""" chunks: list[Chunk] = [] # Try all eligible chunkers in order until one of them doesn't crash. for chunker in ALL_CHUNKERS: if chunker.can_chunk(source_code, file_path): try: chunks = chunker(source_code, file_path).chunk_file() chunks = list( more_itertools.flatten(x.split(max_size=max_chunk_size) for x in chunks), ) break except Exception: logger.exception(f"Error chunking file {file_path} with {chunker}") return File( chunks=chunks, abs_path=file_path.absolute(), contents=source_code, ext=file_path.suffix, ) def chunks_of_type(self, chunk_type: ChunkCategory) -> list[Chunk]: return [c for c in self.chunks if c.category == chunk_type] class Project: def __init__( self, *, root: Path, files_per_parallel_worker: int = 100, file_watch_refresh_freq_seconds: float = 0.1, max_chunk_size: int = 10_000, ) -> None: self.root = root.expanduser().absolute() self.files_per_parallel_worker = files_per_parallel_worker self.max_chunk_size = max_chunk_size self.file_map: dict[Path, File] = {} git_repo: Repo | None if (root / ".git").exists(): git_repo = Repo(root / ".git") else: git_repo = None self.git_repo = git_repo self._init_from_root_dir(root) # Note potential that if a file is modified here it won't be picked up. self.observer = Observer() self.observer.schedule( event_handler=_ProjectFileHandler( self, file_watch_refresh_freq_seconds=file_watch_refresh_freq_seconds, ), path=str(self.root), recursive=True, ) self.observer.start() @property def files(self) -> list[File]: return list(self.file_map.values()) def load_files(self, files: list[Path]) -> None: # How many workers to use? _cpu_count = os.cpu_count() or 1 n_workers = math.floor(len(files) / self.files_per_parallel_worker) n_workers = min(n_workers, _cpu_count // 2) # Avoid maxing out the system n_workers = max(n_workers, 1) files_analysed: list[File] if n_workers == 1: files_analysed_maybe_none = [ _analyze_file(file_path, max_chunk_size=self.max_chunk_size) for file_path in files ] files_analysed = [x for x in files_analysed_maybe_none if x is not None] else: logger.info(f"Using {n_workers} workers to process {len(files)} files") files_analysed = [] with ProcessPoolExecutor(max_workers=n_workers) as executor: future_to_file = { executor.submit( _analyze_file, file_path, max_chunk_size=self.max_chunk_size, ): file_path for file_path in files } for future in as_completed(future_to_file): file_path = future_to_file[future] try: result = future.result() if result is not None: files_analysed.append(result) except Exception: logger.exception(f"File {file_path} generated an exception") for file in files_analysed: self.file_map[file.abs_path] = file def _init_from_root_dir(self, root: Path) -> None: if not root.exists(): raise ValueError(f"Root directory {root} does not exist") files: list[Path] = [] if self.git_repo is not None: rel_paths = self.git_repo.git.ls_files().splitlines() files.extend(root / rel_path for rel_path in rel_paths) else: # Exclude specific top-level directories # TODO: make this configurable ignore_dirs = {".venv", "build", ".git", "__pycache__"} for path in root.iterdir(): if path.is_dir() and path.name not in ignore_dirs: files.extend(path.glob("**/*")) # Don't forget files in the root directory itself files.extend(root.glob("*")) files = [file for file in files if file.is_file()] self.load_files(files) def _analyze_file(file_path: Path, max_chunk_size: int = 10_000) -> File | None: try: if not file_path.exists(): logger.warning(f"File {file_path} does not exist") return None if not file_path.is_file(): logger.warning(f"File {file_path} is not a file") return None return File.from_file_contents( file_path.read_text(), file_path, max_chunk_size=max_chunk_size, ) except Exception: logger.exception(f"Error processing file {file_path}") return None