-
Notifications
You must be signed in to change notification settings - Fork 4
Expand file tree
/
Copy pathmain.py
More file actions
123 lines (98 loc) · 3.29 KB
/
Copy pathmain.py
File metadata and controls
123 lines (98 loc) · 3.29 KB
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
import typing as t
from dataclasses import dataclass, field
from pathlib import Path
from textwrap import dedent
import cyclopts
import dreadnode as dn
from dreadnode.agent import Agent
from dreadnode.agent.events import AgentEnd
from dreadnode.agent.hooks import Hook
from dreadnode.agent.tools import tool
from dreadnode.data_types import Markdown
from kernel import PythonKernel
from rich.console import Console
console = Console()
# CLI
app = cyclopts.App()
@cyclopts.Parameter(name="*", group="args")
@dataclass
class Args:
model: str
"""Model to use for inference"""
task: str
"""Task to perform"""
image: str = "jupyter/datascience-notebook:latest"
"""Docker image to use for the container"""
volumes: t.Annotated[
list[str],
cyclopts.Parameter(
name=["--volume", "-v"],
help="Additional volumes to mount in the container (e.g. /path/to/dir:/path/in/container)",
),
] = field(default_factory=list)
max_steps: int = 50
"""Maximum number of steps to take"""
@cyclopts.Parameter(name="*", group="dreadnode")
@dataclass
class DreadnodeArgs:
server: str | None = None
"""Dreadnode server URL"""
token: str | None = None
"""Dreadnode API token"""
project: str | None = "python-agent"
"""Dreadnode project name"""
console: t.Annotated[bool, cyclopts.Parameter(negative=False)] = False
"""Show span information in the console"""
@tool()
async def complete_task(success: bool, markdown_summary: str) -> None: # noqa: FBT001
"""
Mark your task as complete with a success/failure status and markdown summary.
"""
dn.log_metric("task_success", success, to="run")
dn.log_output("task_summary", Markdown(markdown_summary), to="run")
def upload_work_hook(
work_dir: Path,
) -> Hook:
async def upload_work(event: AgentEnd) -> None:
dn.log_artifact(str(work_dir))
return upload_work
@app.default
async def agent(*, args: Args, dn_args: DreadnodeArgs | None = None) -> None:
"""
General agent with access to a dockerized jupyter environment.
"""
dn_args = dn_args or DreadnodeArgs()
dn.configure(
server=dn_args.server,
token=dn_args.token,
project=dn_args.project,
console=dn_args.console,
)
instructions = dedent(f"""\
Work to complete the following task. You have access to a dockerized jupyter environment.
You can run code in the environment and use the results to help you complete the task.
Unless otherwise specified, use `~/work` to store files and data. Additional volumes are listed below.
<volumes>
{args.volumes}
</volumes>
<task>
{args.task}
</task>
""")
async with PythonKernel(
image=args.image,
volumes=args.volumes,
) as kernel:
agent = Agent(
name="python-agent",
model=args.model,
description="An agent with access to a dockerized jupyter environment.",
instructions=instructions,
tools=[kernel],
hooks=[upload_work_hook(work_dir=kernel.work_dir)],
)
async with agent.stream(args.task) as events:
async for event in events:
console.print(event)
if __name__ == "__main__":
app()