truefoundry 0.1.2__py3-none-any.whl → 0.2.0rc2__py3-none-any.whl

truefoundry/autodeploy/tools/docker_run.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import atexit
+from typing import Any, Dict, Generator, List, Optional
+
+import docker
+import requests
+from pydantic import Field
+from rich.console import Console, ConsoleOptions, RenderResult
+from rich.padding import Padding
+from rich.text import Text
+
+from truefoundry.autodeploy.tools.base import (
+    Event,
+    Message,
+    RequestEvent,
+    ResponseEvent,
+    Tool,
+)
+
+
+class DockerRunLog(Event):
+    index: int
+    log: str
+
+    def render(self, console: Console):
+        # if self.index > 1:
+        #     print("\r", end="")
+        console.print(Padding.indent(renderable=Text.from_ansi(self.log), level=2))
+        # console.print("Press control-c to stop waiting for more logs", end="")
+
+
+class DockerRun(Tool):
+    description = """
+Run a docker image
+"""
+
+    class Request(RequestEvent):
+        image_tag: str
+        ports: Optional[Dict[str, int]] = Field(
+            None,
+            description="""
+Ports to expose.
+The keys of the dictionary are the ports to bind inside the container in 'port'.
+The values are the ports to open on the host""",
+        )
+        command: str
+
+        def render(self, console: Console):
+            console.print(
+                f"[bold magenta]TFY-Agent[/] is executing the Docker container. Image Tag: [bold green]{self.image_tag}[/], Exposed Port: [bold green]{str(self.ports) if self.ports is not None else 'Not exposed'}[/], Command: [bold green]{self.command}[/]"
+            )
+
+    class Response(ResponseEvent):
+        logs: Optional[List[str]] = Field(
+            None, description="Logs of the container. Only last 50 chars."
+        )
+        exit_code: Optional[int] = Field(
+            None,
+            description="""
+Exit code of the process if the container stops.
+This will not be passed if the container is still running.
+""",
+        )
+        client_error: Optional[str] = Field(None, description="Docker client error.")
+
+        def __rich_console__(
+            self, console: Console, options: ConsoleOptions
+        ) -> RenderResult:
+            none_text = "[italic magenta]None[/]"
+            error_text = (
+                f"[green]'{self.client_error}'[/]"
+                if self.client_error is not None
+                else none_text
+            )
+            exit_code_text = (
+                "[green]'0'[/]"
+                if self.exit_code == 0
+                else (
+                    none_text if self.exit_code is None else f"[red]{self.exit_code}[/]"
+                )
+            )
+
+            yield Text.from_markup("[bold magenta]Response[/](")
+            if self.logs is not None:
+                yield Text.from_markup('  [yellow]logs=[/]"')
+                yield Text.from_ansi("".join(self.logs))
+                yield Text.from_markup('"')
+            else:
+                yield Text.from_markup(f"  [yellow]logs[/]={none_text}")
+
+            yield Text.from_markup(f"  [yellow]exit_code[/]={exit_code_text}")
+            yield Text.from_markup(f"  [yellow]client_error[/]={error_text} \n)")
+
+    def __init__(self, docker_client: docker.DockerClient, environment: Dict):
+        self.containers = []
+        self.docker_client = docker_client
+        self.environment = environment
+        atexit.register(self._kill_running_containers)
+
+    def _kill_running_containers(self):
+        if self.containers:
+            container = self.containers.pop()
+            try:
+                container.remove(force=True)
+            except docker.errors.APIError:
+                pass
+
+    def run(self, request: DockerRun.Request) -> Generator[Event, Any, ResponseEvent]:
+        self._kill_running_containers()
+        yield Message(message="[bold cyan]Testing:[/] Running Docker container...")
+        try:
+            container = self.docker_client.containers.run(
+                request.image_tag,
+                detach=True,
+                remove=False,
+                stderr=True,
+                ports=request.ports,
+                environment=self.environment,
+                command=request.command,
+            )
+        except docker.errors.APIError as ex:
+            if ex.is_client_error():
+                return DockerRun.Response(client_error=str(ex))
+            raise
+        yield Message(message="[bold yellow]Docker logs:[/]")
+        self.containers.append(container)
+        exit_code = None
+
+        all_logs = []
+        logs = container.logs(stream=True)
+
+        try:
+            for i, log in enumerate(logs):
+                log = log.decode()
+                all_logs.append(log)
+                yield DockerRunLog(index=i, log=log)
+        except KeyboardInterrupt:
+            pass
+        else:
+            yield Message(message="\n[bold yellow]There are no more logs.[/]")
+        exit_code = None
+        try:
+            exit_code = container.wait(timeout=1).get("StatusCode")
+        except (
+            requests.exceptions.ReadTimeout,
+            requests.exceptions.ConnectionError,
+        ):
+            ...
+        return DockerRun.Response(logs=all_logs, exit_code=exit_code)

truefoundry/autodeploy/tools/file_type_counts.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import os
+from collections import Counter
+from typing import Any, Dict, Generator
+
+import gitignorefile
+from pydantic import Field
+from rich.console import Console
+from rich.table import Table
+
+from truefoundry.autodeploy.tools.base import (
+    Event,
+    Message,
+    RequestEvent,
+    ResponseEvent,
+    Tool,
+)
+
+
+class ShowFileCount(Event):
+    file_types: Dict[str, int]
+
+    def render(self, console: Console):
+        total_files = sum(self.file_types.values())
+        console.print(f"Found {total_files} files.")
+        table = Table()
+        table.add_column("File Type", style="cyan")
+        table.add_column("Count", justify="right", style="green")
+        for file_type, count in self.file_types.items():
+            table.add_row(file_type, str(count))
+
+        console.print(table)
+
+
+class FileTypeCounts(Tool):
+    description = """
+Get counts of different types of file present.
+"""
+
+    class Request(RequestEvent): ...
+
+    class Response(ResponseEvent):
+        file_types: Dict[str, int] = Field(
+            ...,
+            description='Counts of different types of files. Ex: {"py": 1} or {"c": 1}',
+        )
+
+    def __init__(self, project_root_path: str):
+        self.project_root_path = project_root_path
+
+    def run(
+        self, request: FileTypeCounts.Request
+    ) -> Generator[Event, Any, ResponseEvent]:
+        counter = Counter()
+
+        yield Message(
+            message="[bold cyan]Processing:[/] Scanning for various file types..."
+        )
+
+        def gitignore(_):
+            return False
+
+        gitignore_path = os.path.join(self.project_root_path, ".gitignore")
+        if os.path.exists(gitignore_path):
+            gitignore = gitignorefile.parse(path=gitignore_path)
+        for root, dirs, ps in os.walk(
+            self.project_root_path,
+        ):
+            root = root[len(self.project_root_path) :]
+            if ".git" in dirs:
+                dirs.remove(".git")
+            counter.update(
+                p.split(".")[-1] if len(p) and p[0] != "." else p
+                for p in ps
+                if not gitignore(os.path.join(root, p).strip(os.path.sep))
+            )
+        yield ShowFileCount(file_types=counter)
+        return FileTypeCounts.Response(file_types=counter)

truefoundry/autodeploy/tools/list_files.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import os
+from fnmatch import fnmatch
+from typing import Any, Generator, List, Optional
+
+import gitignorefile
+from pydantic import Field
+
+from truefoundry.autodeploy.tools.base import (
+    Event,
+    Message,
+    RequestEvent,
+    ResponseEvent,
+    Tool,
+)
+
+
+class ListFiles(Tool):
+    description = """
+List all files.
+If you want to find all json files recuresively under directory a/b
+the subdir should be a/b and pattern will be *.json
+
+If you want to find all json files recuresively under current directory
+the subdir should be . and pattern will be *.json
+"""
+
+    class Request(RequestEvent):
+        sub_dir: str
+        pattern: str = Field(..., description="Glob pattern. Avoid passing '*'")
+
+    class Response(ResponseEvent):
+        paths: List[str] = Field(
+            ...,
+            description="File paths under the given directory",
+        )
+        error: Optional[str] = Field(None)
+
+    def __init__(self, project_root_path: str):
+        self.project_root_path = project_root_path
+
+    def run(self, request: ListFiles.Request) -> Generator[Event, Any, ResponseEvent]:
+        yield Message(
+            message=f"[bold cyan]Searching:[/] 🔍 Looking for files matching the pattern [magenta]{request.pattern}[/]"
+        )
+
+        paths: List[str] = []
+
+        def gitignore(_):
+            return False
+
+        gitignore_path = os.path.join(self.project_root_path, ".gitignore")
+        if os.path.exists(gitignore_path):
+            gitignore = gitignorefile.parse(path=gitignore_path)
+
+        path = os.path.join(self.project_root_path, request.sub_dir.strip(os.path.sep))
+        if not os.path.exists(path):
+            return ListFiles.Response(
+                paths=None,
+                error=f"Incorrect sub_dir {request.sub_dir}. Does not exist",
+            )
+
+        for root, dirs, ps in os.walk(
+            path,
+        ):
+            root = root[len(path) :]
+            if ".git" in dirs:
+                dirs.remove(".git")
+            paths.extend(
+                os.path.join(root, p).lstrip(os.path.sep)
+                for p in ps
+                if fnmatch(p, request.pattern)
+                and not gitignore(os.path.join(root, p).strip(os.path.sep))
+            )
+        if len(paths) > 0:
+            yield Message(message=f"[bold green]Success:[/] Found {len(paths)} files.")
+        else:
+            yield Message(
+                message=f"[red]Alert:[/] No files found matching the pattern {request.pattern}."
+            )
+        return ListFiles.Response(paths=paths)

truefoundry/autodeploy/tools/read_file.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import os
+from typing import Any, Generator, List, Optional
+
+from pydantic import BaseModel, Field
+
+from truefoundry.autodeploy.tools.base import (
+    Event,
+    Message,
+    RequestEvent,
+    ResponseEvent,
+    Tool,
+)
+
+
+class Line(BaseModel):
+    line_number: int
+    content: str
+
+
+class ReadFile(Tool):
+    description = """
+Read contents of a file.
+Avoid reading *.lock type files as they tend to be large
+"""
+
+    class Request(RequestEvent):
+        path: str = Field(
+            ...,
+            pattern=r"^[a-zA-Z0-9\.]{1}.*$",
+            description="File path to open. ",
+        )
+
+    class Response(ResponseEvent):
+        data: Optional[List[Line]] = Field(
+            None,
+            description="Content of the file.",
+        )
+        error: Optional[str] = Field(
+            None,
+            description="Error while opening a file.",
+        )
+
+    def __init__(
+        self,
+        project_root_path: str,
+    ):
+        self.project_root_path = project_root_path
+
+    def run(self, request: ReadFile.Request) -> Generator[Event, Any, ResponseEvent]:
+        yield Message(
+            message=f"[bold cyan]Processing:[/] Reading file at [magenta]{request.path}[/] and extracting details..."
+        )
+        try:
+            with open(
+                os.path.join(self.project_root_path, request.path),
+                "r",
+                encoding="utf8",
+            ) as f:
+                response = ReadFile.Response(data=[])
+                for i, line in enumerate(f):
+                    response.data.append(Line(line_number=i + 1, content=line))
+                return response
+        except FileNotFoundError as ex:
+            return ReadFile.Response(error=str(ex))

truefoundry/autodeploy/tools/send_request.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import Any, Generator, Optional
+
+import requests
+from pydantic import Field
+
+from truefoundry.autodeploy.tools.base import (
+    Event,
+    Message,
+    RequestEvent,
+    ResponseEvent,
+    Tool,
+)
+
+
+class SendRequest(Tool):
+    description = """
+Send an HTTP request.
+"""
+
+    def __init__(self):
+        self.call_count = 0
+
+    class Request(RequestEvent):
+        method: str
+        url: str
+
+    class Response(ResponseEvent):
+        response_code: Optional[int] = Field(None, description="Response Code")
+        response_body: Optional[str] = None
+        error: Optional[str] = Field(None, description="Error.")
+
+    def run(self, request: SendRequest.Request) -> Generator[Event, Any, ResponseEvent]:
+        self.call_count += 1
+        yield Message(
+            message=f"[bold cyan]Testing:[/] Sending a [magenta]{request.method.upper()}[/] request to [magenta]{request.url}[/]..."
+        )
+        try:
+            response = requests.request(request.method.lower(), url=request.url)
+            yield Message(
+                message=f"[bold green]Success:[/] Received response with status code [magenta]{response.status_code}[/]"
+            )
+            return SendRequest.Response(
+                response_code=response.status_code,
+                response_body=response.text[-50:],
+            )
+        except Exception as ex:
+            yield Message(
+                message="[red]Alert:[/] Request could not be completed successfully."
+            )
+            return SendRequest.Response(
+                error=str(ex),
+            )

truefoundry/autodeploy/tools/write_file.py

@@ -0,0 +1,101 @@
+# deprecated
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+from pydantic import BaseModel, Field
+from rich.console import Console
+from rich.prompt import Confirm
+
+from truefoundry.autodeploy.tools.base import Tool
+from truefoundry.autodeploy.utils.diff import Diff
+
+INTERACTIVE_SESSION = True
+
+
+class WriteFile(Tool):
+    description = """
+Write contents to a file.
+Read first before write.
+"""
+
+    class Request(BaseModel):
+        path: str = Field(
+            ...,
+            pattern=r"^[a-zA-Z0-9\.]{1}.*$",
+            description="File path to write.",
+        )
+        content: str = Field(..., description="Content of the file.")
+        justification: str = Field(
+            ...,
+            description="Justification of why the new content is required. Ensure the justification, justifies the content.",
+        )
+
+    class Response(BaseModel):
+        cancellation_reason: Optional[str] = Field(
+            None, description="Operation cancelled by user"
+        )
+
+        error: Optional[str] = Field(
+            None,
+            description="Error while opening a file.",
+        )
+
+    def __init__(self, project_root_path: str):
+        self.project_root_path = project_root_path
+
+    def _user_interaction(
+        self, request: WriteFile.Request, console: Console
+    ) -> Optional[WriteFile.Response]:
+        console.log("You are about to write or edit a file.")
+
+        console.print(f"Displaying changes in {request.path}")
+        prev_content = ""
+
+        if os.path.exists(os.path.join(self.project_root_path, request.path)):
+            with open(
+                os.path.join(self.project_root_path, request.path),
+                "r",
+                encoding="utf8",
+            ) as f:
+                prev_content = f.read()
+
+        console.print(
+            Diff(
+                lhs=prev_content,
+                rhs=request.content,
+                width=os.get_terminal_size().columns,
+            )
+        )
+        response = Confirm.ask(
+            f"Writing file at {request.path}?",
+        )
+        if not response:
+            description = console.input(
+                "You chose to cancel. Can you provide a reason why? [green]>> "
+            )
+            return WriteFile.Response(
+                cancellation_reason=description, error="Operation cancelled by user."
+            )
+        console.log(f"Writing file to {request.path}")
+
+    def run(
+        self,
+        request: WriteFile.Request,
+        console: Console,
+    ) -> WriteFile.Response:
+        if INTERACTIVE_SESSION:
+            interaction_response = self._user_interaction(request, console)
+            if isinstance(interaction_response, WriteFile.Response):
+                return interaction_response
+        try:
+            with open(
+                os.path.join(self.project_root_path, request.path),
+                "w",
+                encoding="utf8",
+            ) as f:
+                f.write(request.content)
+                return WriteFile.Response()
+        except FileNotFoundError as ex:
+            return WriteFile.Response(error=str(ex))

truefoundry/autodeploy/utils/diff.py

@@ -0,0 +1,157 @@
+import difflib
+import pprint
+from typing import Iterator, List
+
+from rich.console import Console, ConsoleOptions, RenderResult
+from rich.text import Text
+
+
+def rewrite_line(line, line_to_rewrite, prev_marker):
+    marker_style_map = {
+        "+": {
+            " ": "green",
+            "+": "white on green",
+            "^": "white on green",
+        },
+        "-": {
+            " ": "red",
+            "-": "white on red",
+            "^": "white on red",
+        },
+    }
+    new_line = Text("")
+    current_span = []
+    # Differ lines start with a 2 letter code, so skip past that
+    prev_char = line[2]
+    for idx, char in enumerate(line[2:], start=2):
+        if prev_marker in ("+", "-"):
+            if char != prev_char:
+                style = marker_style_map.get(prev_marker, {}).get(prev_char, None)
+                if style is not None:
+                    new_line.append_text(Text("".join(current_span), style=style))
+                current_span = []
+            if idx - 2 < len(line_to_rewrite):
+                current_span.append(line_to_rewrite[idx - 2])
+        prev_char = char
+
+    # Lines starting with ? aren't guaranteed to be the same length as the lines before them
+    #  so some characters may be left over. Add any leftover characters to the output.
+    # subtract 2 for code at start
+    remaining_index = idx - 2
+    if prev_marker == "-":
+        new_line.append_text(Text(line_to_rewrite[remaining_index:], style="red"))
+    elif prev_marker == "+":
+        new_line.append_text(Text(line_to_rewrite[remaining_index:], style="green"))
+    return new_line
+
+
+def build_symbolic_unified_diff(diff: List[str]) -> RenderResult:
+    output_lines = []
+    style = "grey"
+    last_style = style
+    for line in diff:
+        if line.startswith("+ "):
+            style = "green"
+            output_line = f"+ {line[2:]}"
+        elif line.startswith("- "):
+            style = "red"
+            output_line = f"- {line[2:]}"
+        elif line.startswith("? "):
+            if last_style == "red":
+                output_line = line[:-1].replace("+", "-")
+            elif last_style == "green":
+                output_line = line[:-1].replace("-", "+")
+        else:
+            output_line = line
+            style = "gray"
+        output_lines.append(Text(output_line, style=style))
+        last_style = style if style != "gray" else last_style
+    return output_lines
+
+
+def build_unified_diff(diff: List[str]) -> RenderResult:
+    prev_marker = ""
+    output_lines: List[Text] = []
+    for line in diff:
+        if line.startswith("- "):
+            output_lines.append(Text(line[2:], style="red"))
+        elif line.startswith("+ "):
+            output_lines.append(Text(line[2:], style="green"))
+        elif line.startswith("? "):
+            line_to_rewrite = output_lines[-1].plain
+            output_lines[-1] = rewrite_line(line, line_to_rewrite, prev_marker)
+        else:
+            output_lines.append(Text(line[2:], style="#949494"))
+        prev_marker = line[0]
+    return output_lines
+
+
+def llm_unified_diff(diff: List[str]) -> RenderResult:
+    prev_marker = ""
+    output_lines: List[Text] = []
+    for line in diff:
+        if len(line) < 1:
+            continue
+        if line.startswith("---"):
+            output_lines.append(Text(line, style="dim"))
+        elif line.startswith("+++"):
+            output_lines.append(Text(line, style="dim"))
+        elif line.startswith("@@"):
+            output_lines.append(Text(line, style="bright_blue"))
+        elif line.startswith("-"):
+            output_lines.append(Text(line[1:], style="red"))
+        elif line.startswith("+"):
+            output_lines.append(Text(line[1:], style="green"))
+        elif line.startswith("?"):
+            line_to_rewrite = output_lines[-1].plain
+            output_lines[-1] = rewrite_line(line, line_to_rewrite, prev_marker)
+        else:
+            output_lines.append(Text(line[1:], style="#949494"))
+        prev_marker = line[0]
+    return output_lines
+
+
+class LLMDiff:
+    def __init__(self, llm_diff: str) -> None:
+        self.diff = llm_diff.splitlines()
+
+    def __rich_console__(
+        self, console: Console, options: ConsoleOptions
+    ) -> RenderResult:
+        return llm_unified_diff(diff=self.diff)
+
+
+class Diff:
+    """Constructs a Diff object to render diff-highlighted code."""
+
+    def __init__(
+        self,
+        lhs: object,
+        rhs: object,
+        width: int,
+        show_symbols: bool = False,
+    ) -> None:
+        self.width = width
+        self.lhs = lhs if isinstance(lhs, str) else pprint.pformat(lhs, width=width)
+        self.rhs = rhs if isinstance(rhs, str) else pprint.pformat(rhs, width=width)
+        self.show_symbols = show_symbols
+
+    @property
+    def sides_are_different(self) -> bool:
+        return self.lhs != self.rhs
+
+    def raw_unified_diff(self) -> Iterator[str]:
+        differ = difflib.Differ()
+        lines_lhs = self.lhs.splitlines()
+        lines_rhs = self.rhs.splitlines()
+        return differ.compare(lines_lhs, lines_rhs)
+
+    def __rich_console__(
+        self, console: Console, options: ConsoleOptions
+    ) -> RenderResult:
+        diff = self.raw_unified_diff()
+        if self.show_symbols:
+            result = build_symbolic_unified_diff(diff)
+        else:
+            result = build_unified_diff(diff)
+        yield from result