ThomasTaroni
6993a52d47
Enhanced docstrings for better clarity and consistency across methods. Added error handling for report generation and ensured default values for research details to avoid null errors. Improved comments and logic flow for readability and maintainability.
109 lines
3.7 KiB
Python
109 lines
3.7 KiB
Python
from gpt_researcher import GPTResearcher
|
|
|
|
|
|
class ReportGenerator:
|
|
"""
|
|
A class to handle the generation of research-based reports.
|
|
|
|
This class integrates with GPTResearcher to conduct research, retrieve results,
|
|
and format them into consumable chunks for asynchronous streaming.
|
|
"""
|
|
|
|
def __init__(self, query: str, report_type: str):
|
|
"""
|
|
Initializes the ReportGenerator instance with a query and report type.
|
|
|
|
:param query: The main topic or question for research.
|
|
:param report_type: The type of report to generate.
|
|
"""
|
|
self.query = query
|
|
self.report_type = report_type
|
|
self.researcher = GPTResearcher(query, report_type)
|
|
self._chunks = None # Placeholder for report chunks
|
|
self._index = 0 # Iterator index for streaming
|
|
|
|
def __aiter__(self):
|
|
"""
|
|
Makes the ReportGenerator instance asynchronously iterable.
|
|
|
|
:return: Self instance for iteration.
|
|
"""
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
"""
|
|
Defines the logic for asynchronous iteration over report chunks.
|
|
|
|
:return: The next chunk of the report.
|
|
:raises StopAsyncIteration: Raised when all chunks are yielded.
|
|
"""
|
|
if self._chunks is None:
|
|
# Generate report chunks on first access
|
|
self._chunks = await self._generate_report_chunks()
|
|
|
|
if self._index >= len(self._chunks):
|
|
# End iteration once all chunks are exhausted
|
|
raise StopAsyncIteration
|
|
|
|
chunk = self._chunks[self._index]
|
|
self._index += 1
|
|
return chunk
|
|
|
|
async def _generate_report_chunks(self):
|
|
"""
|
|
Conducts research using the researcher and generates the report in chunks.
|
|
|
|
:return: A list of report chunks.
|
|
"""
|
|
try:
|
|
# Asynchronous research and report generation
|
|
research_result = await self.researcher.conduct_research()
|
|
report = await self.researcher.write_report()
|
|
|
|
# Retrieve additional research details
|
|
research_context = self.researcher.get_research_context() or {}
|
|
research_costs = self.researcher.get_costs() or {}
|
|
research_images = self.researcher.get_research_images() or []
|
|
research_sources = self.researcher.get_research_sources() or []
|
|
|
|
# Construct the complete research response
|
|
full_report = {
|
|
"report": report,
|
|
"context": research_context,
|
|
"costs": research_costs,
|
|
"images": research_images,
|
|
"sources": research_sources,
|
|
}
|
|
|
|
# Generate chunks for streaming
|
|
return self._split_into_chunks(full_report)
|
|
|
|
except Exception as e:
|
|
# Handle potential errors during research and report generation
|
|
raise RuntimeError(f"Error generating report chunks: {e}")
|
|
|
|
def _split_into_chunks(self, report):
|
|
"""
|
|
Splits the report dictionary into smaller chunks for easier streaming.
|
|
|
|
:param report: A dictionary containing the full report data.
|
|
:return: A list of formatted text chunks.
|
|
"""
|
|
if not report:
|
|
raise ValueError("Cannot split an empty or None report into chunks.")
|
|
|
|
chunks = []
|
|
for key, value in report.items():
|
|
chunks.append(f"{key}: {value}")
|
|
return chunks
|
|
|
|
def get_query_details(self):
|
|
"""
|
|
Retrieves the details of the query and report type.
|
|
|
|
:return: A dictionary containing the query and report type.
|
|
"""
|
|
return {
|
|
"query": self.query,
|
|
"report_type": self.report_type
|
|
}
|