langroid 0.33.4__py3-none-any.whl → 0.33.7__py3-none-any.whl

langroid/parsing/repo_loader.py

@@ -0,0 +1,790 @@
+import itertools
+import json
+import logging
+import os
+import subprocess
+import tempfile
+import time
+from collections import deque
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+from urllib.parse import urlparse
+
+from dotenv import load_dotenv
+from github import Github
+from github.ContentFile import ContentFile
+from github.Label import Label
+from github.Repository import Repository
+
+from langroid.mytypes import DocMetaData, Document
+from langroid.parsing.document_parser import DocumentParser, DocumentType
+from langroid.parsing.parser import Parser, ParsingConfig
+from langroid.pydantic_v1 import BaseModel, BaseSettings, Field
+
+logger = logging.getLogger(__name__)
+
+
+def _get_decoded_content(content_file: ContentFile) -> str:
+    if content_file.encoding == "base64":
+        return content_file.decoded_content.decode("utf-8") or ""
+    elif content_file.encoding == "none":
+        return content_file.content or ""
+    else:
+        raise ValueError(f"Unsupported encoding: {content_file.encoding}")
+
+
+def _has_files(directory: str) -> bool:
+    """
+    Recursively checks if there is at least one file in a directory.
+    """
+    for dirpath, dirnames, filenames in os.walk(directory):
+        if filenames:
+            return True
+    return False
+
+
+# Pydantic model for GitHub issue data
+class IssueData(BaseModel):
+    state: str = Field(..., description="State of issue e.g. open or closed")
+    year: int = Field(..., description="Year issue was created")
+    month: int = Field(..., description="Month issue was created")
+    day: int = Field(..., description="Day issue was created")
+    assignee: Optional[str] = Field(..., description="Assignee of issue")
+    size: Optional[str] = Field(..., description="Size of issue, e.g. XS, S, M, L, XXL")
+    text: str = Field(..., description="Text of issue, i.e. description body")
+
+
+def get_issue_size(labels: List[Label]) -> str | None:
+    sizes = ["XS", "S", "M", "L", "XL", "XXL"]
+    return next((label.name for label in labels if label.name in sizes), None)
+
+
+class RepoLoaderConfig(BaseSettings):
+    """
+    Configuration for RepoLoader.
+    """
+
+    non_code_types: List[str] = [
+        "md",
+        "txt",
+        "text",
+    ]
+
+    file_types: List[str] = [
+        "py",
+        "md",
+        "yml",
+        "yaml",
+        "txt",
+        "text",
+        "sh",
+        "ini",
+        "toml",
+        "cfg",
+        "json",
+        "rst",
+        "Makefile",
+        "Dockerfile",
+    ]
+
+    exclude_dirs: List[str] = [
+        ".gitignore",
+        ".gitmodules",
+        ".gitattributes",
+        ".git",
+        ".idea",
+        ".vscode",
+        ".circleci",
+    ]
+
+
+class RepoLoader:
+    """
+    Class for recursively getting all file content in a repo.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        config: RepoLoaderConfig = RepoLoaderConfig(),
+    ):
+        """
+        Args:
+            url: full github url of repo, or just "owner/repo"
+            config: configuration for RepoLoader
+        """
+        self.url = url
+        self.config = config
+        self.clone_path: Optional[str] = None
+        self.log_file = ".logs/repo_loader/download_log.json"
+        os.makedirs(os.path.dirname(self.log_file), exist_ok=True)
+        if not os.path.exists(self.log_file):
+            with open(self.log_file, "w") as f:
+                json.dump({"junk": "ignore"}, f)
+        with open(self.log_file, "r") as f:
+            log = json.load(f)
+        if self.url in log and os.path.exists(log[self.url]):
+            logger.info(f"Repo Already downloaded in {log[self.url]}")
+            self.clone_path = log[self.url]
+
+        if "github.com" in self.url:
+            repo_name = self.url.split("github.com/")[1]
+        else:
+            repo_name = self.url
+        load_dotenv()
+        # authenticated calls to github api have higher rate limit
+        token = os.getenv("GITHUB_ACCESS_TOKEN")
+        g = Github(token)
+        self.repo = self._get_repo_with_retry(g, repo_name)
+
+    @staticmethod
+    def _get_repo_with_retry(
+        g: Github, repo_name: str, max_retries: int = 5
+    ) -> Repository:
+        """
+        Get a repo from the GitHub API, retrying if the request fails,
+        with exponential backoff.
+
+        Args:
+            g: GitHub object
+            repo_name: name of repo
+            max_retries: maximum number of retries
+        Returns:
+            Repo: GitHub repo object
+
+        """
+        base_delay = 2  # base delay in seconds
+        max_delay = 60  # maximum delay in seconds
+
+        for attempt in range(max_retries):
+            try:
+                return g.get_repo(repo_name)
+            except Exception as e:
+                delay = min(max_delay, base_delay * 2**attempt)
+                logger.info(
+                    f"Attempt {attempt+1} failed with error: {str(e)}. "
+                    f"Retrying in {delay} seconds..."
+                )
+                time.sleep(delay)
+        raise Exception(f"Failed to get repo {repo_name} after {max_retries} attempts.")
+
+    def _get_dir_name(self) -> str:
+        return urlparse(self.url).path.replace("/", "_")
+
+    def get_issues(self, k: int | None = 100) -> List[IssueData]:
+        """Get up to k issues from the GitHub repo."""
+        if k is None:
+            issues = self.repo.get_issues(state="all")
+        else:
+            issues = self.repo.get_issues(state="all")[:k]
+        issue_data_list = []
+        for issue in issues:
+            issue_data = IssueData(
+                state=issue.state,
+                year=issue.created_at.year,
+                month=issue.created_at.month,
+                day=issue.created_at.day,
+                assignee=issue.assignee.login if issue.assignee else None,
+                size=get_issue_size(issue.labels),
+                text=issue.body or "No issue description body.",
+            )
+            issue_data_list.append(issue_data)
+
+        return issue_data_list
+
+    @staticmethod
+    def _file_type(name: str) -> str:
+        """
+        Get the file type of a file name.
+        Args:
+            name: name of file, can be "a", "a.b", or ".b"
+        Returns:
+            str: file type; "a" => "a", "a.b" => "b", ".b" => "b"
+                some examples:
+                "Makefile" => "Makefile",
+                "script.py" => "py",
+                ".gitignore" => "gitignore"
+        """
+        # "a" -> ("a", ""), "a.b" -> ("a", ".b"), ".b" -> (".b", "")
+        file_parts = os.path.splitext(name)
+        if file_parts[1] == "":
+            file_type = file_parts[0]  # ("a", "") => "a"
+        else:
+            file_type = file_parts[1][1:]  # (*,".b") => "b"
+        return file_type
+
+    def _is_code(self, file_type: str) -> bool:
+        """
+        Check if a file type is code.
+
+        Args:
+            file_type: file type, e.g. "py", "md", "txt"
+        Returns:
+            bool: whether file type is code
+        """
+        return file_type not in self.config.non_code_types
+
+    def _is_allowed(self, content: ContentFile) -> bool:
+        """
+        Check if a file or directory content is allowed to be included.
+
+        Args:
+            content (ContentFile): The file or directory Content object.
+
+        Returns:
+            bool: Whether the file or directory is allowed to be included.
+        """
+        if content.type == "dir":
+            return content.name not in self.config.exclude_dirs
+        elif content.type == "file":
+            return self._file_type(content.name) in self.config.file_types
+        else:
+            return False
+
+    def default_clone_path(self) -> str:
+        return tempfile.mkdtemp(suffix=self._get_dir_name())
+
+    def clone(self, path: Optional[str] = None) -> Optional[str]:
+        """
+        Clone a GitHub repository to a local directory specified by `path`,
+        if it has not already been cloned.
+
+        Args:
+            path (str): The local directory where the repository should be cloned.
+                If not specified, a temporary directory will be created.
+
+        Returns:
+            str: The path to the local directory where the repository was cloned.
+        """
+        with open(self.log_file, "r") as f:
+            log: Dict[str, str] = json.load(f)
+
+        if (
+            self.url in log
+            and os.path.exists(log[self.url])
+            and _has_files(log[self.url])
+        ):
+            logger.warning(f"Repo Already downloaded in {log[self.url]}")
+            self.clone_path = log[self.url]
+            return self.clone_path
+
+        self.clone_path = path
+        if path is None:
+            path = self.default_clone_path()
+            self.clone_path = path
+
+        try:
+            subprocess.run(["git", "clone", self.url, path], check=True)
+            log[self.url] = path
+            with open(self.log_file, "w") as f:
+                json.dump(log, f)
+            return self.clone_path
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Git clone failed: {e}")
+        except Exception as e:
+            logger.error(f"An error occurred while trying to clone the repository:{e}")
+
+        return self.clone_path
+
+    def load_tree_from_github(
+        self, depth: int, lines: int = 0
+    ) -> Dict[str, Union[str, List[Dict[str, Any]]]]:
+        """
+        Get a nested dictionary of GitHub repository file and directory names
+        up to a certain depth, with file contents.
+
+        Args:
+            depth (int): The depth level.
+            lines (int): The number of lines of file contents to include.
+
+        Returns:
+            Dict[str, Union[str, List[Dict]]]:
+            A dictionary containing file and directory names, with file contents.
+        """
+        root_contents = self.repo.get_contents("")
+        if not isinstance(root_contents, list):
+            root_contents = [root_contents]
+        repo_structure = {
+            "type": "dir",
+            "name": "",
+            "dirs": [],
+            "files": [],
+            "path": "",
+        }
+
+        # A queue of tuples (current_node, current_depth, parent_structure)
+        queue = deque([(root_contents, 0, repo_structure)])
+
+        while queue:
+            current_node, current_depth, parent_structure = queue.popleft()
+
+            for content in current_node:
+                if not self._is_allowed(content):
+                    continue
+                if content.type == "dir" and current_depth < depth:
+                    # Create a new sub-dictionary for this directory
+                    new_dir = {
+                        "type": "dir",
+                        "name": content.name,
+                        "dirs": [],
+                        "files": [],
+                        "path": content.path,
+                    }
+                    parent_structure["dirs"].append(new_dir)
+                    contents = self.repo.get_contents(content.path)
+                    if not isinstance(contents, list):
+                        contents = [contents]
+                    queue.append(
+                        (
+                            contents,
+                            current_depth + 1,
+                            new_dir,
+                        )
+                    )
+                elif content.type == "file":
+                    file_content = "\n".join(
+                        _get_decoded_content(content).splitlines()[:lines]
+                    )
+                    file_dict = {
+                        "type": "file",
+                        "name": content.name,
+                        "content": file_content,
+                        "path": content.path,
+                    }
+                    parent_structure["files"].append(file_dict)
+
+        return repo_structure
+
+    def load(
+        self,
+        path: Optional[str] = None,
+        depth: int = 3,
+        lines: int = 0,
+    ) -> Tuple[Dict[str, Union[str, List[Dict[str, Any]]]], List[Document]]:
+        """
+        From a local folder `path` (if None, the repo clone path), get:
+          a nested dictionary (tree) of dicts, files and contents
+          a list of Document objects for each file.
+
+        Args:
+            path (str): The local folder path; if none, use self.clone_path()
+            depth (int): The depth level.
+            lines (int): The number of lines of file contents to include.
+
+        Returns:
+            Tuple of (dict, List_of_Documents):
+                A dictionary containing file and directory names, with file
+                contents, and a list of Document objects for each file.
+        """
+        if path is None:
+            if self.clone_path is None or not _has_files(self.clone_path):
+                self.clone()
+            path = self.clone_path
+        if path is None:
+            raise ValueError("Unable to clone repo")
+        return self.load_from_folder(
+            path=path,
+            depth=depth,
+            lines=lines,
+            file_types=self.config.file_types,
+            exclude_dirs=self.config.exclude_dirs,
+            url=self.url,
+        )
+
+    @staticmethod
+    def load_from_folder(
+        path: str,
+        depth: int = 3,
+        lines: int = 0,
+        file_types: Optional[List[str]] = None,
+        exclude_dirs: Optional[List[str]] = None,
+        url: str = "",
+    ) -> Tuple[Dict[str, Union[str, List[Dict[str, Any]]]], List[Document]]:
+        """
+        From a local folder `path` (required), get:
+          a nested dictionary (tree) of dicts, files and contents, restricting to
+            desired file_types and excluding undesired directories.
+          a list of Document objects for each file.
+
+        Args:
+            path (str): The local folder path, required.
+            depth (int): The depth level. Optional, default 3.
+            lines (int): The number of lines of file contents to include.
+                    Optional, default 0 (no lines => empty string).
+            file_types (List[str]): The file types to include.
+                    Optional, default None (all).
+            exclude_dirs (List[str]): The directories to exclude.
+                    Optional, default None (no exclusions).
+            url (str): Optional url, to be stored in docs as metadata. Default "".
+
+        Returns:
+            Tuple of (dict, List_of_Documents):
+                A dictionary containing file and directory names, with file contents.
+                A list of Document objects for each file.
+        """
+
+        folder_structure = {
+            "type": "dir",
+            "name": "",
+            "dirs": [],
+            "files": [],
+            "path": "",
+        }
+        # A queue of tuples (current_path, current_depth, parent_structure)
+        queue = deque([(path, 0, folder_structure)])
+        docs = []
+        exclude_dirs = exclude_dirs or []
+        while queue:
+            current_path, current_depth, parent_structure = queue.popleft()
+
+            for item in os.listdir(current_path):
+                item_path = os.path.join(current_path, item)
+                relative_path = os.path.relpath(item_path, path)
+                if (os.path.isdir(item_path) and item in exclude_dirs) or (
+                    os.path.isfile(item_path)
+                    and file_types is not None
+                    and RepoLoader._file_type(item) not in file_types
+                ):
+                    continue
+
+                if os.path.isdir(item_path) and current_depth < depth:
+                    # Create a new sub-dictionary for this directory
+                    new_dir = {
+                        "type": "dir",
+                        "name": item,
+                        "dirs": [],
+                        "files": [],
+                        "path": relative_path,
+                    }
+                    parent_structure["dirs"].append(new_dir)
+                    queue.append((item_path, current_depth + 1, new_dir))
+                elif os.path.isfile(item_path):
+                    # Add the file to the current dictionary
+                    with open(item_path, "r") as f:
+                        file_lines = list(itertools.islice(f, lines))
+                    file_content = "\n".join(line.strip() for line in file_lines)
+                    if file_content == "":
+                        continue
+
+                    file_dict = {
+                        "type": "file",
+                        "name": item,
+                        "content": file_content,
+                        "path": relative_path,
+                    }
+                    parent_structure["files"].append(file_dict)
+                    docs.append(
+                        Document(
+                            content=file_content,
+                            metadata=DocMetaData(
+                                repo=url,
+                                source=relative_path,
+                                url=url,
+                                filename=item,
+                                extension=RepoLoader._file_type(item),
+                                language=RepoLoader._file_type(item),
+                            ),
+                        )
+                    )
+        return folder_structure, docs
+
+    @staticmethod
+    def get_documents(
+        path: str | bytes,
+        parser: Parser = Parser(ParsingConfig()),
+        file_types: Optional[List[str]] = None,
+        exclude_dirs: Optional[List[str]] = None,
+        depth: int = -1,
+        lines: Optional[int] = None,
+        doc_type: str | DocumentType | None = None,
+    ) -> List[Document]:
+        """
+        Recursively get all files under a path as Document objects.
+
+        Args:
+            path (str|bytes): The path to the directory or file, or bytes content.
+                The bytes option is meant to support the case where the content
+                has already been read from a file in an upstream process
+                (e.g. from an API or a database), and we want to avoid having to
+                write it to a temporary file just to read it again.
+                (which can be very slow for large files,
+                especially in a docker container)
+            parser (Parser): Parser to use to parse files.
+            file_types (List[str], optional): List of file extensions OR
+                filenames OR file_path_names to  include.
+                Defaults to None, which includes all files.
+            exclude_dirs (List[str], optional): List of directories to exclude.
+                Defaults to None, which includes all directories.
+            depth (int, optional): Max depth of recursion. Defaults to -1,
+                which includes all depths.
+            lines (int, optional): Number of lines to read from each file.
+                Defaults to None, which reads all lines.
+            doc_type (str|DocumentType, optional): The type of document to parse.
+
+        Returns:
+            List[Document]: List of Document objects representing files.
+
+        """
+        docs = []
+        file_paths = []
+        if isinstance(path, bytes):
+            file_paths.append(path)
+        else:
+            path_obj = Path(path).resolve()
+
+            if path_obj.is_file():
+                file_paths.append(str(path_obj))
+            else:
+                path_depth = len(path_obj.parts)
+                for root, dirs, files in os.walk(path):
+                    # Exclude directories if needed
+                    if exclude_dirs:
+                        dirs[:] = [d for d in dirs if d not in exclude_dirs]
+
+                    current_depth = len(Path(root).resolve().parts) - path_depth
+                    if depth == -1 or current_depth <= depth:
+                        for file in files:
+                            file_path = str(Path(root) / file)
+                            if (
+                                file_types is None
+                                or RepoLoader._file_type(file_path) in file_types
+                                or os.path.basename(file_path) in file_types
+                                or file_path in file_types
+                            ):
+                                file_paths.append(file_path)
+
+        for file_path in file_paths:
+            docs.extend(
+                DocumentParser.chunks_from_path_or_bytes(
+                    file_path,
+                    parser,
+                    doc_type=doc_type,
+                    lines=lines,
+                )
+            )
+        return docs
+
+    def load_docs_from_github(
+        self,
+        k: Optional[int] = None,
+        depth: Optional[int] = None,
+        lines: Optional[int] = None,
+    ) -> List[Document]:
+        """
+        Directly from GitHub, recursively get all files in a repo that have one of the
+        extensions, possibly up to a max number of files, max depth, and max number
+        of lines per file (if any of these are specified).
+
+        Args:
+            k (int): max number of files to load, or None for all files
+            depth (int): max depth to recurse, or None for infinite depth
+            lines (int): max number of lines to get, from a file, or None for all lines
+
+        Returns:
+            list of Document objects, each has fields `content` and `metadata`,
+            and `metadata` has fields `url`, `filename`, `extension`, `language`
+        """
+        contents = self.repo.get_contents("")
+        if not isinstance(contents, list):
+            contents = [contents]
+        stack = list(zip(contents, [0] * len(contents)))  # stack of (content, depth)
+        # recursively get all files in repo that have one of the extensions
+        docs = []
+        i = 0
+
+        while stack:
+            if k is not None and i == k:
+                break
+            file_content, d = stack.pop()
+            if not self._is_allowed(file_content):
+                continue
+            if file_content.type == "dir":
+                if depth is None or d <= depth:
+                    items = self.repo.get_contents(file_content.path)
+                    if not isinstance(items, list):
+                        items = [items]
+                    stack.extend(list(zip(items, [d + 1] * len(items))))
+            else:
+                if depth is None or d <= depth:
+                    # need to decode the file content, which is in bytes
+                    contents = self.repo.get_contents(file_content.path)
+                    if isinstance(contents, list):
+                        contents = contents[0]
+                    text = _get_decoded_content(contents)
+                    if lines is not None:
+                        text = "\n".join(text.split("\n")[:lines])
+                    i += 1
+
+                    # Note `source` is important, it may be used to cite
+                    # evidence for an answer.
+                    # See  URLLoader
+                    # TODO we should use Pydantic to enforce/standardize this
+
+                    docs.append(
+                        Document(
+                            content=text,
+                            metadata=DocMetaData(
+                                repo=self.url,
+                                source=file_content.html_url,
+                                url=file_content.html_url,
+                                filename=file_content.name,
+                                extension=self._file_type(file_content.name),
+                                language=self._file_type(file_content.name),
+                            ),
+                        )
+                    )
+        return docs
+
+    @staticmethod
+    def select(
+        structure: Dict[str, Union[str, List[Dict[str, Any]]]],
+        includes: List[str],
+        excludes: List[str] = [],
+    ) -> Dict[str, Union[str, List[Dict[str, Any]]]]:
+        """
+        Filter a structure dictionary for certain directories and files.
+
+        Args:
+            structure (Dict[str, Union[str, List[Dict]]]): The structure dictionary.
+            includes (List[str]): A list of desired directories and files.
+                For files, either full file names or "file type" can be specified.
+                E.g.  "toml" will include all files with the ".toml" extension,
+                or "Makefile" will include all files named "Makefile".
+            excludes (List[str]): A list of directories and files to exclude.
+                Similar to `includes`, full file/dir names or "file type" can be
+                specified. Optional, defaults to empty list.
+
+
+        Returns:
+            Dict[str, Union[str, List[Dict]]]: The filtered structure dictionary.
+        """
+        filtered_structure = {
+            "type": structure["type"],
+            "name": structure["name"],
+            "dirs": [],
+            "files": [],
+            "path": structure["path"],
+        }
+
+        for dir in structure["dirs"]:
+            if (
+                dir["name"] in includes
+                or RepoLoader._file_type(dir["name"]) in includes
+            ) and (
+                dir["name"] not in excludes
+                and RepoLoader._file_type(dir["name"]) not in excludes
+            ):
+                # If the directory is in the select list, include the whole subtree
+                filtered_structure["dirs"].append(dir)
+            else:
+                # Otherwise, filter the directory's contents
+                filtered_dir = RepoLoader.select(dir, includes)
+                if (
+                    filtered_dir["dirs"] or filtered_dir["files"]
+                ):  # only add if not empty
+                    filtered_structure["dirs"].append(filtered_dir)
+
+        for file in structure["files"]:
+            if (
+                file["name"] in includes
+                or RepoLoader._file_type(file["name"]) in includes
+            ) and (
+                file["name"] not in excludes
+                and RepoLoader._file_type(file["name"]) not in excludes
+            ):
+                filtered_structure["files"].append(file)
+
+        return filtered_structure
+
+    @staticmethod
+    def ls(structure: Dict[str, Union[str, List[Dict]]], depth: int = 0) -> List[str]:
+        """
+        Get a list of names of files or directories up to a certain depth from a
+        structure dictionary.
+
+        Args:
+            structure (Dict[str, Union[str, List[Dict]]]): The structure dictionary.
+            depth (int, optional): The depth level. Defaults to 0.
+
+        Returns:
+            List[str]: A list of names of files or directories.
+        """
+        names = []
+
+        # A queue of tuples (current_structure, current_depth)
+        queue = deque([(structure, 0)])
+
+        while queue:
+            current_structure, current_depth = queue.popleft()
+
+            if current_depth <= depth:
+                names.append(current_structure["name"])
+
+                for dir in current_structure["dirs"]:
+                    queue.append((dir, current_depth + 1))
+
+                for file in current_structure["files"]:
+                    # add file names only if depth is less than the limit
+                    if current_depth < depth:
+                        names.append(file["name"])
+        names = [n for n in names if n not in ["", None]]
+        return names
+
+    @staticmethod
+    def list_files(
+        dir: str,
+        depth: int = 1,
+        include_types: List[str] = [],
+        exclude_types: List[str] = [],
+    ) -> List[str]:
+        """
+        Recursively list all files in a directory, up to a certain depth.
+
+        Args:
+            dir (str): The directory path, relative to root.
+            depth (int, optional): The depth level. Defaults to 1.
+            include_types (List[str], optional): A list of file types to include.
+                Defaults to empty list.
+            exclude_types (List[str], optional): A list of file types to exclude.
+                Defaults to empty list.
+        Returns:
+            List[str]: A list of file names.
+        """
+        depth = depth if depth >= 0 else 200
+        output = []
+
+        for root, dirs, files in os.walk(dir):
+            if root.count(os.sep) - dir.count(os.sep) < depth:
+                level = root.count(os.sep) - dir.count(os.sep)
+                sub_indent = " " * 4 * (level + 1)
+                for d in dirs:
+                    output.append("{}{}/".format(sub_indent, d))
+                for f in files:
+                    if include_types and RepoLoader._file_type(f) not in include_types:
+                        continue
+                    if exclude_types and RepoLoader._file_type(f) in exclude_types:
+                        continue
+                    output.append("{}{}".format(sub_indent, f))
+        return output
+
+    @staticmethod
+    def show_file_contents(tree: Dict[str, Union[str, List[Dict[str, Any]]]]) -> str:
+        """
+        Print the contents of all files from a structure dictionary.
+
+        Args:
+            tree (Dict[str, Union[str, List[Dict]]]): The structure dictionary.
+        """
+        contents = ""
+        for dir in tree["dirs"]:
+            contents += RepoLoader.show_file_contents(dir)
+        for file in tree["files"]:
+            path = file["path"]
+            contents += f"""
+            {path}:
+            --------------------
+            {file["content"]}
+            
+            """
+
+        return contents