Spaces:
Running
Running
File size: 7,457 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 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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 |
from typing import Dict, Optional
import json
from ..utils.llm import construct_subtopics
from ..actions import (
stream_output,
generate_report,
generate_draft_section_titles,
write_report_introduction,
write_conclusion
)
class ReportGenerator:
"""Generates reports based on research data."""
def __init__(self, researcher):
self.researcher = researcher
self.research_params = {
"query": self.researcher.query,
"agent_role_prompt": self.researcher.cfg.agent_role or self.researcher.role,
"report_type": self.researcher.report_type,
"report_source": self.researcher.report_source,
"tone": self.researcher.tone,
"websocket": self.researcher.websocket,
"cfg": self.researcher.cfg,
"headers": self.researcher.headers,
}
async def write_report(self, existing_headers: list = [], relevant_written_contents: list = [], ext_context=None) -> str:
"""
Write a report based on existing headers and relevant contents.
Args:
existing_headers (list): List of existing headers.
relevant_written_contents (list): List of relevant written contents.
ext_context (Optional): External context, if any.
Returns:
str: The generated report.
"""
# send the selected images prior to writing report
research_images = self.researcher.get_research_images()
if research_images:
await stream_output(
"images",
"selected_images",
json.dumps(research_images),
self.researcher.websocket,
True,
research_images
)
context = ext_context or self.researcher.context
if self.researcher.verbose:
await stream_output(
"logs",
"writing_report",
f"✍️ Writing report for '{self.researcher.query}'...",
self.researcher.websocket,
)
report_params = self.research_params.copy()
report_params["context"] = context
if self.researcher.report_type == "subtopic_report":
report_params.update({
"main_topic": self.researcher.parent_query,
"existing_headers": existing_headers,
"relevant_written_contents": relevant_written_contents,
"cost_callback": self.researcher.add_costs,
})
else:
report_params["cost_callback"] = self.researcher.add_costs
report = await generate_report(**report_params)
if self.researcher.verbose:
await stream_output(
"logs",
"report_written",
f"📝 Report written for '{self.researcher.query}'",
self.researcher.websocket,
)
return report
async def write_report_conclusion(self, report_content: str) -> str:
"""
Write the conclusion for the report.
Args:
report_content (str): The content of the report.
Returns:
str: The generated conclusion.
"""
if self.researcher.verbose:
await stream_output(
"logs",
"writing_conclusion",
f"✍️ Writing conclusion for '{self.researcher.query}'...",
self.researcher.websocket,
)
conclusion = await write_conclusion(
query=self.researcher.query,
context=report_content,
config=self.researcher.cfg,
agent_role_prompt=self.researcher.cfg.agent_role or self.researcher.role,
cost_callback=self.researcher.add_costs,
websocket=self.researcher.websocket,
)
if self.researcher.verbose:
await stream_output(
"logs",
"conclusion_written",
f"📝 Conclusion written for '{self.researcher.query}'",
self.researcher.websocket,
)
return conclusion
async def write_introduction(self):
"""Write the introduction section of the report."""
if self.researcher.verbose:
await stream_output(
"logs",
"writing_introduction",
f"✍️ Writing introduction for '{self.researcher.query}'...",
self.researcher.websocket,
)
introduction = await write_report_introduction(
query=self.researcher.query,
context=self.researcher.context,
agent_role_prompt=self.researcher.cfg.agent_role or self.researcher.role,
config=self.researcher.cfg,
websocket=self.researcher.websocket,
cost_callback=self.researcher.add_costs,
)
if self.researcher.verbose:
await stream_output(
"logs",
"introduction_written",
f"📝 Introduction written for '{self.researcher.query}'",
self.researcher.websocket,
)
return introduction
async def get_subtopics(self):
"""Retrieve subtopics for the research."""
if self.researcher.verbose:
await stream_output(
"logs",
"generating_subtopics",
f"🌳 Generating subtopics for '{self.researcher.query}'...",
self.researcher.websocket,
)
subtopics = await construct_subtopics(
task=self.researcher.query,
data=self.researcher.context,
config=self.researcher.cfg,
subtopics=self.researcher.subtopics,
)
if self.researcher.verbose:
await stream_output(
"logs",
"subtopics_generated",
f"📊 Subtopics generated for '{self.researcher.query}'",
self.researcher.websocket,
)
return subtopics
async def get_draft_section_titles(self, current_subtopic: str):
"""Generate draft section titles for the report."""
if self.researcher.verbose:
await stream_output(
"logs",
"generating_draft_sections",
f"📑 Generating draft section titles for '{self.researcher.query}'...",
self.researcher.websocket,
)
draft_section_titles = await generate_draft_section_titles(
query=self.researcher.query,
current_subtopic=current_subtopic,
context=self.researcher.context,
role=self.researcher.cfg.agent_role or self.researcher.role,
websocket=self.researcher.websocket,
config=self.researcher.cfg,
cost_callback=self.researcher.add_costs,
)
if self.researcher.verbose:
await stream_output(
"logs",
"draft_sections_generated",
f"🗂️ Draft section titles generated for '{self.researcher.query}'",
self.researcher.websocket,
)
return draft_section_titles
|