Make DeepResearch
asynchronously iterable.
Added asynchronous iteration methods (`__aiter__` and `__anext__`) to the `DeepResearch` class to enable streaming of research report chunks. Introduced functionality to generate and split reports into smaller chunks for improved report handling.
This commit is contained in:
parent
43423732e2
commit
9f87253082
1 changed files with 44 additions and 4 deletions
|
@ -9,10 +9,35 @@ class ReportGenerator:
|
||||||
self.query = query
|
self.query = query
|
||||||
self.report_type = report_type
|
self.report_type = report_type
|
||||||
self.researcher = GPTResearcher(query, report_type)
|
self.researcher = GPTResearcher(query, report_type)
|
||||||
|
self._chunks = None # Placeholder for report chunks
|
||||||
|
self._index = 0 # Index for iteration
|
||||||
|
|
||||||
async def generate_report(self):
|
def __aiter__(self):
|
||||||
"""
|
"""
|
||||||
Conducts research and generates the report along with additional information.
|
Make this class asynchronously iterable.
|
||||||
|
"""
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __anext__(self):
|
||||||
|
"""
|
||||||
|
Defines the asynchronous iteration logic.
|
||||||
|
"""
|
||||||
|
if self._chunks is None:
|
||||||
|
# If chunks are not generated yet, generate the report
|
||||||
|
self._chunks = await self._generate_report_chunks()
|
||||||
|
|
||||||
|
if self._index >= len(self._chunks):
|
||||||
|
# Stop iteration when all chunks are yielded
|
||||||
|
raise StopAsyncIteration
|
||||||
|
|
||||||
|
# Return the next chunk and increment the index
|
||||||
|
chunk = self._chunks[self._index]
|
||||||
|
self._index += 1
|
||||||
|
return chunk
|
||||||
|
|
||||||
|
async def _generate_report_chunks(self):
|
||||||
|
"""
|
||||||
|
Conducts research and generates the report in chunks.
|
||||||
"""
|
"""
|
||||||
# Conduct research
|
# Conduct research
|
||||||
research_result = await self.researcher.conduct_research()
|
research_result = await self.researcher.conduct_research()
|
||||||
|
@ -24,14 +49,29 @@ class ReportGenerator:
|
||||||
research_images = self.researcher.get_research_images()
|
research_images = self.researcher.get_research_images()
|
||||||
research_sources = self.researcher.get_research_sources()
|
research_sources = self.researcher.get_research_sources()
|
||||||
|
|
||||||
return {
|
# Construct the full response
|
||||||
|
full_report = {
|
||||||
"report": report,
|
"report": report,
|
||||||
"context": research_context,
|
"context": research_context,
|
||||||
"costs": research_costs,
|
"costs": research_costs,
|
||||||
"images": research_images,
|
"images": research_images,
|
||||||
"sources": research_sources
|
"sources": research_sources,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
# Split the report into smaller chunks for streaming
|
||||||
|
return self._split_into_chunks(full_report)
|
||||||
|
|
||||||
|
def _split_into_chunks(self, report):
|
||||||
|
"""
|
||||||
|
Splits a report dictionary into smaller chunks for streaming.
|
||||||
|
"""
|
||||||
|
# Convert the report dictionary into a list of key-value pairs,
|
||||||
|
# where each pair represents a chunk.
|
||||||
|
chunks = []
|
||||||
|
for key, value in report.items():
|
||||||
|
chunks.append(f"{key}: {value}")
|
||||||
|
return chunks # Return the list of chunks
|
||||||
|
|
||||||
def get_query_details(self):
|
def get_query_details(self):
|
||||||
"""
|
"""
|
||||||
Returns details of the query and report type.
|
Returns details of the query and report type.
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue