Refactor API methods to use AsyncGenerator for SSE.
Updated `write_report` and `get_research_sources` to stream results via server-sent events (SSE) using `AsyncGenerator`. This allows incremental data delivery and better error reporting, improving client-server communication.
This commit is contained in:
parent
1b60eb0ae6
commit
895671189e
1 changed files with 25 additions and 9 deletions
|
@ -240,7 +240,7 @@ async def quick_search(query: str) -> AsyncGenerator[str, None]:
|
|||
|
||||
|
||||
@mcp.tool()
|
||||
async def write_report(research_id: str, custom_prompt: Optional[str] = None) -> Dict[str, Any]:
|
||||
async def write_report(research_id: str, custom_prompt: Optional[str] = None) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Generate a report based on previously conducted research.
|
||||
|
||||
|
@ -253,7 +253,7 @@ async def write_report(research_id: str, custom_prompt: Optional[str] = None) ->
|
|||
"""
|
||||
success, researcher, error = get_researcher_by_id(mcp.researchers, research_id)
|
||||
if not success:
|
||||
return error
|
||||
yield format_sse_event("message", error)
|
||||
|
||||
logger.info(f"Generating report for research ID: {research_id}")
|
||||
|
||||
|
@ -265,17 +265,29 @@ async def write_report(research_id: str, custom_prompt: Optional[str] = None) ->
|
|||
sources = researcher.get_research_sources()
|
||||
costs = researcher.get_costs()
|
||||
|
||||
return create_success_response({
|
||||
final_data_payload = {
|
||||
"report": report,
|
||||
"source_count": len(sources),
|
||||
"costs": costs
|
||||
})
|
||||
}
|
||||
|
||||
# Sende das finale Ergebnis als 'tool_result' Event
|
||||
yield format_sse_event("message", final_data_payload)
|
||||
logger.info(f"Sent final research result for ID: {research_id}")
|
||||
except Exception as e:
|
||||
return handle_exception(e, "Report generation")
|
||||
logger.error(f"Error during deep_research for query: {e}", exc_info=True)
|
||||
# Sende ein Fehler-Event an den Client
|
||||
error_payload = {
|
||||
"research_id": research_id,
|
||||
"status": "Error occurred",
|
||||
"error_message": str(e),
|
||||
"error_details": "Check server logs for more information."
|
||||
}
|
||||
yield format_sse_event("message", error_payload)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def get_research_sources(research_id: str) -> Dict[str, Any]:
|
||||
async def get_research_sources(research_id: str) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Get the sources used in the research.
|
||||
|
||||
|
@ -287,15 +299,19 @@ async def get_research_sources(research_id: str) -> Dict[str, Any]:
|
|||
"""
|
||||
success, researcher, error = get_researcher_by_id(mcp.researchers, research_id)
|
||||
if not success:
|
||||
return error
|
||||
yield format_sse_event("message", error)
|
||||
|
||||
sources = researcher.get_research_sources()
|
||||
source_urls = researcher.get_source_urls()
|
||||
|
||||
return create_success_response({
|
||||
final_data_payload = {
|
||||
"sources": format_sources_for_response(sources),
|
||||
"source_urls": source_urls
|
||||
})
|
||||
}
|
||||
|
||||
# Sende das finale Ergebnis als 'tool_result' Event
|
||||
yield format_sse_event("message", final_data_payload)
|
||||
logger.info(f"Sent final research result for ID: {research_id}")
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue