Spaces:
Running
Running
File size: 2,618 Bytes
372531f |
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 |
from .utils.views import print_agent_output
from .utils.llms import call_model
import json
sample_revision_notes = """
{
"draft": {
draft title: The revised draft that you are submitting for review
},
"revision_notes": Your message to the reviewer about the changes you made to the draft based on their feedback
}
"""
class ReviserAgent:
def __init__(self, websocket=None, stream_output=None, headers=None):
self.websocket = websocket
self.stream_output = stream_output
self.headers = headers or {}
async def revise_draft(self, draft_state: dict):
"""
Review a draft article
:param draft_state:
:return:
"""
review = draft_state.get("review")
task = draft_state.get("task")
draft_report = draft_state.get("draft")
prompt = [
{
"role": "system",
"content": "You are an expert writer. Your goal is to revise drafts based on reviewer notes.",
},
{
"role": "user",
"content": f"""Draft:\n{draft_report}" + "Reviewer's notes:\n{review}\n\n
You have been tasked by your reviewer with revising the following draft, which was written by a non-expert.
If you decide to follow the reviewer's notes, please write a new draft and make sure to address all of the points they raised.
Please keep all other aspects of the draft the same.
You MUST return nothing but a JSON in the following format:
{sample_revision_notes}
""",
},
]
response = await call_model(
prompt,
model=task.get("model"),
response_format="json",
)
return response
async def run(self, draft_state: dict):
print_agent_output(f"Rewriting draft based on feedback...", agent="REVISOR")
revision = await self.revise_draft(draft_state)
if draft_state.get("task").get("verbose"):
if self.websocket and self.stream_output:
await self.stream_output(
"logs",
"revision_notes",
f"Revision notes: {revision.get('revision_notes')}",
self.websocket,
)
else:
print_agent_output(
f"Revision notes: {revision.get('revision_notes')}", agent="REVISOR"
)
return {
"draft": revision.get("draft"),
"revision_notes": revision.get("revision_notes"),
}
|