forked from LAION-AI/Open-Assistant
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy path__main__.py
94 lines (80 loc) · 3.33 KB
/
__main__.py
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
# -*- coding: utf-8 -*-
"""Simple REPL frontend."""
import requests
import typer
app = typer.Typer()
@app.command()
def main(backend_url: str, api_key: str):
"""Simple REPL frontend."""
def _post(path: str, json: dict) -> dict:
response = requests.post(f"{backend_url}{path}", json=json, headers={"X-API-Key": api_key})
response.raise_for_status()
return response.json()
typer.echo("Requesting work...")
tasks = [_post("/api/v1/tasks/", {"type": "generic"})]
while tasks:
task = tasks.pop(0)
match (task["type"]):
case "summarize_story":
typer.echo("Summarize the following story:")
typer.echo(task["story"])
# acknowledge task
_post(f"/api/v1/tasks/{task['id']}/ack", {"type": "post_created", "post_id": "1234"})
summary = typer.prompt("Enter your summary")
# send interaction
new_task = _post(
"/api/v1/tasks/interaction",
{
"type": "text_reply_to_post",
"post_id": "1234",
"user_post_id": "5678",
"text": summary,
"user": {"id": "1234", "name": "John Doe"},
},
)
tasks.append(new_task)
case "rate_summary":
typer.echo("Rate the following summary:")
typer.echo(task["summary"])
typer.echo("Full text:")
typer.echo(task["full_text"])
typer.echo(f"Rating scale: {task['scale']['min']} - {task['scale']['max']}")
# acknowledge task
_post(f"/api/v1/tasks/{task['id']}/ack", {"type": "rating_created", "post_id": "1234"})
rating = typer.prompt("Enter your rating", type=int)
# send interaction
new_task = _post(
"/api/v1/tasks/interaction",
{
"type": "post_rating",
"post_id": "1234",
"rating": rating,
"user": {"id": "1234", "name": "John Doe"},
},
)
tasks.append(new_task)
case "initial_prompt":
typer.echo("Please provide an initial prompt to the assistant.")
if task["hint"]:
typer.echo(f"Hint: {task['hint']}")
# acknowledge task
_post(f"/api/v1/tasks/{task['id']}/ack", {"type": "post_created", "post_id": "1234"})
prompt = typer.prompt("Enter your prompt")
# send interaction
new_task = _post(
"/api/v1/tasks/interaction",
{
"type": "text_reply_to_post",
"post_id": "1234",
"user_post_id": "5678",
"text": prompt,
"user": {"id": "1234", "name": "John Doe"},
},
)
tasks.append(new_task)
case "task_done":
typer.echo("Task done!")
case _:
typer.echo(f"Unknown task type {task['type']}")
if __name__ == "__main__":
app()