aboutsummaryrefslogtreecommitdiff
path: root/autogpts/autogpt/autogpt/agents/features/context.py
blob: 405c6e5590d30630b263942242960682c8bb0818 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
import contextlib
from pathlib import Path
from typing import Iterator, Optional

from autogpt.agents.protocols import CommandProvider, MessageProvider
from autogpt.command_decorator import command
from autogpt.core.resource.model_providers import ChatMessage
from autogpt.core.utils.json_schema import JSONSchema
from autogpt.file_storage.base import FileStorage
from autogpt.models.command import Command
from autogpt.models.context_item import ContextItem, FileContextItem, FolderContextItem
from autogpt.utils.exceptions import InvalidArgumentError


class AgentContext:
    items: list[ContextItem]

    def __init__(self, items: Optional[list[ContextItem]] = None):
        self.items = items or []

    def __bool__(self) -> bool:
        return len(self.items) > 0

    def __contains__(self, item: ContextItem) -> bool:
        return any([i.source == item.source for i in self.items])

    def add(self, item: ContextItem) -> None:
        self.items.append(item)

    def close(self, index: int) -> None:
        self.items.pop(index - 1)

    def clear(self) -> None:
        self.items.clear()

    def format_numbered(self, workspace: FileStorage) -> str:
        return "\n\n".join(
            [f"{i}. {c.fmt(workspace)}" for i, c in enumerate(self.items, 1)]
        )


class ContextComponent(MessageProvider, CommandProvider):
    """Adds ability to keep files and folders open in the context (prompt)."""

    def __init__(self, workspace: FileStorage):
        self.context = AgentContext()
        self.workspace = workspace

    def get_messages(self) -> Iterator[ChatMessage]:
        if self.context:
            yield ChatMessage.system(
                "## Context\n"
                f"{self.context.format_numbered(self.workspace)}\n\n"
                "When a context item is no longer needed and you are not done yet, "
                "you can hide the item by specifying its number in the list above "
                "to `hide_context_item`.",
            )

    def get_commands(self) -> Iterator[Command]:
        yield self.open_file
        yield self.open_folder
        if self.context:
            yield self.close_context_item

    @command(
        parameters={
            "file_path": JSONSchema(
                type=JSONSchema.Type.STRING,
                description="The path of the file to open",
                required=True,
            )
        }
    )
    async def open_file(self, file_path: str | Path) -> str:
        """Opens a file for editing or continued viewing;
        creates it if it does not exist yet.
        Note: If you only need to read or write a file once,
        use `write_to_file` instead.

        Args:
            file_path (str | Path): The path of the file to open

        Returns:
            str: A status message indicating what happened
        """
        if not isinstance(file_path, Path):
            file_path = Path(file_path)

        created = False
        if not self.workspace.exists(file_path):
            await self.workspace.write_file(file_path, "")
            created = True

        # Try to make the file path relative
        with contextlib.suppress(ValueError):
            file_path = file_path.relative_to(self.workspace.root)

        file = FileContextItem(path=file_path)
        self.context.add(file)
        return (
            f"File {file_path}{' created,' if created else ''} has been opened"
            " and added to the context ✅"
        )

    @command(
        parameters={
            "path": JSONSchema(
                type=JSONSchema.Type.STRING,
                description="The path of the folder to open",
                required=True,
            )
        }
    )
    def open_folder(self, path: str | Path) -> str:
        """Open a folder to keep track of its content

        Args:
            path (str | Path): The path of the folder to open

        Returns:
            str: A status message indicating what happened
        """
        if not isinstance(path, Path):
            path = Path(path)

        if not self.workspace.exists(path):
            raise FileNotFoundError(
                f"open_folder {path} failed: no such file or directory"
            )

        # Try to make the path relative
        with contextlib.suppress(ValueError):
            path = path.relative_to(self.workspace.root)

        folder = FolderContextItem(path=path)
        self.context.add(folder)
        return f"Folder {path} has been opened and added to the context ✅"

    @command(
        parameters={
            "number": JSONSchema(
                type=JSONSchema.Type.INTEGER,
                description="The 1-based index of the context item to hide",
                required=True,
            )
        }
    )
    def close_context_item(self, number: int) -> str:
        """Hide an open file, folder or other context item, to save tokens.

        Args:
            number (int): The 1-based index of the context item to hide

        Returns:
            str: A status message indicating what happened
        """
        if number > len(self.context.items) or number == 0:
            raise InvalidArgumentError(f"Index {number} out of range")

        self.context.close(number)
        return f"Context item {number} hidden ✅"