wehub-resource-sync
289 lines
9.2 KiB
Python
289 lines
9.2 KiB
Python
"""
|
|
Create a local agent trace for the aisuite runs viewer.
|
|
|
|
Run from the repository root:
|
|
|
|
python examples/cli/create_demo_trace.py --trace-file .aisuite/demo.jsonl
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import aisuite as ai
|
|
from aisuite.framework.message import ChatCompletionMessageToolCall, Function, Message
|
|
|
|
|
|
def tool_call(name: str, arguments: str, call_id: str) -> ChatCompletionMessageToolCall:
|
|
return ChatCompletionMessageToolCall(
|
|
id=call_id,
|
|
type="function",
|
|
function=Function(name=name, arguments=arguments),
|
|
)
|
|
|
|
|
|
def response(content: str | None = "done", tool_calls=None):
|
|
message = Message(role="assistant", content=content, tool_calls=tool_calls)
|
|
return SimpleNamespace(choices=[SimpleNamespace(message=message)])
|
|
|
|
|
|
class ScriptedProvider:
|
|
def __init__(self, responses):
|
|
self._responses = list(responses)
|
|
|
|
def chat_completions_create(self, *_args, **_kwargs):
|
|
if not self._responses:
|
|
raise RuntimeError("No scripted responses left for demo trace.")
|
|
return self._responses.pop(0)
|
|
|
|
|
|
LARGE_FILE_CONTENT = "demo artifact content\n" * 1300
|
|
LARGE_STDOUT_COMMAND = "python3 -c 'print(\"demo artifact stdout\\n\" * 1300)'"
|
|
|
|
|
|
def build_client() -> ai.Client:
|
|
client = ai.Client()
|
|
client.providers["openai"] = ScriptedProvider(
|
|
[
|
|
response(
|
|
None,
|
|
[
|
|
tool_call(
|
|
"list_files",
|
|
'{"path": "tests/toolkits"}',
|
|
"call_demo_list",
|
|
)
|
|
],
|
|
),
|
|
response("Found `test_files.py` and `test_shell.py` in `tests/toolkits`."),
|
|
response(
|
|
None,
|
|
[
|
|
tool_call(
|
|
"read_file_lines",
|
|
(
|
|
'{"path": "tests/toolkits/test_shell.py", '
|
|
'"start_line": 1, "max_lines": 10}'
|
|
),
|
|
"call_demo_lines",
|
|
)
|
|
],
|
|
),
|
|
response(
|
|
"Read the first 10 lines of `tests/toolkits/test_shell.py` "
|
|
"using the line-range file tool."
|
|
),
|
|
response(
|
|
None,
|
|
[
|
|
tool_call(
|
|
"run_shell",
|
|
'{"command": "python3 -m pytest tests/toolkits -q"}',
|
|
"call_demo_tests",
|
|
)
|
|
],
|
|
),
|
|
response("The focused toolkit test suite passed."),
|
|
response(
|
|
None,
|
|
[
|
|
tool_call(
|
|
"write_file",
|
|
json.dumps(
|
|
{
|
|
"path": ".aisuite/demo-work/large-artifact.txt",
|
|
"content": LARGE_FILE_CONTENT,
|
|
}
|
|
),
|
|
"call_demo_write_large",
|
|
)
|
|
],
|
|
),
|
|
response(
|
|
"Wrote a large demo file so the viewer can show artifactized arguments."
|
|
),
|
|
response(
|
|
None,
|
|
[
|
|
tool_call(
|
|
"run_shell",
|
|
json.dumps({"command": LARGE_STDOUT_COMMAND}),
|
|
"call_demo_large_stdout",
|
|
)
|
|
],
|
|
),
|
|
response(
|
|
"Captured large shell output so the viewer can show artifactized results."
|
|
),
|
|
response(
|
|
None,
|
|
[
|
|
tool_call(
|
|
"run_shell",
|
|
json.dumps({"command": "python3 -c 'print(\"blocked\")'"}),
|
|
"call_demo_denied_shell",
|
|
)
|
|
],
|
|
),
|
|
response("The demo policy denied the blocked shell command."),
|
|
response(
|
|
None,
|
|
[
|
|
tool_call(
|
|
"review_changes",
|
|
json.dumps(
|
|
{"input": "Review the demo trace changes and risks."}
|
|
),
|
|
"call_demo_review",
|
|
)
|
|
],
|
|
),
|
|
response(
|
|
"Reviewer found no material issues. Residual risk: this is scripted demo data."
|
|
),
|
|
response("Reviewer subagent completed and the demo trace is ready."),
|
|
]
|
|
)
|
|
return client
|
|
|
|
|
|
def approve_demo_tool(context: ai.ToolPolicyContext) -> ai.ToolPolicyDecision:
|
|
if context.tool_name == "run_shell":
|
|
command = str(context.arguments.get("command", ""))
|
|
if "blocked" in command:
|
|
return ai.ToolPolicyDecision(allowed=False, reason="blocked by demo policy")
|
|
return ai.ToolPolicyDecision(allowed=True, reason="approved by demo")
|
|
return ai.ToolPolicyDecision(allowed=True, reason="low risk")
|
|
|
|
|
|
def create_demo_trace(
|
|
trace_file: Path,
|
|
cwd: Path,
|
|
append: bool = False,
|
|
artifact_root: Path | None = None,
|
|
) -> None:
|
|
if trace_file.exists() and not append:
|
|
trace_file.unlink()
|
|
trace_file.parent.mkdir(parents=True, exist_ok=True)
|
|
sink = ai.tracing.LocalTraceSink(trace_file)
|
|
artifact_store = ai.FileArtifactStore(
|
|
artifact_root or trace_file.parent / "artifacts"
|
|
)
|
|
reviewer = ai.Agent(
|
|
name="reviewer",
|
|
model="openai:gpt-4o-mini",
|
|
instructions=(
|
|
"Review the demo trace work. Do not edit files. Keep the response concise."
|
|
),
|
|
tools=ai.toolkits.files(root=cwd, allow_write=False),
|
|
tags=["demo", "dev", "reviewer"],
|
|
metadata={"app": "aisuite_demo_trace", "role": "reviewer"},
|
|
)
|
|
agent = ai.Agent(
|
|
name="aisuite_demo_dev",
|
|
model="openai:gpt-4o-mini",
|
|
instructions=(
|
|
"You are a concise coding assistant. Use tools to inspect the "
|
|
"project and run focused checks."
|
|
),
|
|
tools=[
|
|
*ai.toolkits.files(root=cwd, allow_write=True),
|
|
*ai.toolkits.git(root=cwd),
|
|
*ai.toolkits.shell(cwd=cwd, allowed_commands=["python3"]),
|
|
ai.agent_tool(
|
|
reviewer,
|
|
name="review_changes",
|
|
description="Ask a read-only reviewer subagent to inspect the demo work.",
|
|
),
|
|
],
|
|
tags=["demo", "dev"],
|
|
metadata={"app": "aisuite_demo_trace"},
|
|
)
|
|
|
|
client = build_client()
|
|
result = ai.Runner.run_sync(
|
|
agent,
|
|
"List files in tests/toolkits.",
|
|
client=client,
|
|
run_name="demo_list_files",
|
|
group_id="aisuite-demo",
|
|
trace_sinks=[sink],
|
|
tool_policy=approve_demo_tool,
|
|
artifact_store=artifact_store,
|
|
)
|
|
result = ai.Runner.continue_sync(
|
|
result,
|
|
"Read the first 10 lines of tests/toolkits/test_shell.py.",
|
|
max_turns=5,
|
|
trace_sinks=[sink],
|
|
tool_policy=approve_demo_tool,
|
|
artifact_store=artifact_store,
|
|
)
|
|
result = ai.Runner.continue_sync(
|
|
result,
|
|
"Run the focused toolkit tests.",
|
|
max_turns=5,
|
|
trace_sinks=[sink],
|
|
tool_policy=approve_demo_tool,
|
|
artifact_store=artifact_store,
|
|
)
|
|
result = ai.Runner.continue_sync(
|
|
result,
|
|
"Write a large demo file for artifact preview testing.",
|
|
max_turns=5,
|
|
trace_sinks=[sink],
|
|
tool_policy=approve_demo_tool,
|
|
artifact_store=artifact_store,
|
|
)
|
|
result = ai.Runner.continue_sync(
|
|
result,
|
|
"Run a command with large stdout for artifact preview testing.",
|
|
max_turns=5,
|
|
trace_sinks=[sink],
|
|
tool_policy=approve_demo_tool,
|
|
artifact_store=artifact_store,
|
|
)
|
|
result = ai.Runner.continue_sync(
|
|
result,
|
|
"Attempt a blocked shell command so the viewer shows a denied event.",
|
|
max_turns=5,
|
|
trace_sinks=[sink],
|
|
tool_policy=approve_demo_tool,
|
|
artifact_store=artifact_store,
|
|
)
|
|
ai.Runner.continue_sync(
|
|
result,
|
|
"Ask the reviewer subagent to review the demo work.",
|
|
max_turns=5,
|
|
trace_sinks=[sink],
|
|
tool_policy=approve_demo_tool,
|
|
artifact_store=artifact_store,
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Create a sample aisuite trace.")
|
|
parser.add_argument("--trace-file", default=".aisuite/demo.jsonl")
|
|
parser.add_argument("--cwd", default=".")
|
|
parser.add_argument("--append", action="store_true")
|
|
parser.add_argument("--artifact-root", default=None)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
create_demo_trace(
|
|
trace_file=Path(args.trace_file),
|
|
cwd=Path(args.cwd).expanduser().resolve(),
|
|
append=args.append,
|
|
artifact_root=Path(args.artifact_root) if args.artifact_root else None,
|
|
)
|
|
print(f"Wrote demo trace to {args.trace_file}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|