Spaces:
Running
Running
File size: 2,519 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 |
from .utils.file_formats import \
write_md_to_pdf, \
write_md_to_word, \
write_text_to_md
from .utils.views import print_agent_output
class PublisherAgent:
def __init__(self, output_dir: str, websocket=None, stream_output=None, headers=None):
self.websocket = websocket
self.stream_output = stream_output
self.output_dir = output_dir
self.headers = headers or {}
async def publish_research_report(self, research_state: dict, publish_formats: dict):
layout = self.generate_layout(research_state)
await self.write_report_by_formats(layout, publish_formats)
return layout
def generate_layout(self, research_state: dict):
sections = '\n\n'.join(f"{value}"
for subheader in research_state.get("research_data")
for key, value in subheader.items())
references = '\n'.join(f"{reference}" for reference in research_state.get("sources"))
headers = research_state.get("headers")
layout = f"""# {headers.get('title')}
#### {headers.get("date")}: {research_state.get('date')}
## {headers.get("introduction")}
{research_state.get('introduction')}
## {headers.get("table_of_contents")}
{research_state.get('table_of_contents')}
{sections}
## {headers.get("conclusion")}
{research_state.get('conclusion')}
## {headers.get("references")}
{references}
"""
return layout
async def write_report_by_formats(self, layout:str, publish_formats: dict):
if publish_formats.get("pdf"):
await write_md_to_pdf(layout, self.output_dir)
if publish_formats.get("docx"):
await write_md_to_word(layout, self.output_dir)
if publish_formats.get("markdown"):
await write_text_to_md(layout, self.output_dir)
async def run(self, research_state: dict):
task = research_state.get("task")
publish_formats = task.get("publish_formats")
if self.websocket and self.stream_output:
await self.stream_output("logs", "publishing", f"Publishing final research report based on retrieved data...", self.websocket)
else:
print_agent_output(output="Publishing final research report based on retrieved data...", agent="PUBLISHER")
final_research_report = await self.publish_research_report(research_state, publish_formats)
return {"report": final_research_report}
|