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:
ThomasTaroni 2025-06-01 13:00:41 +02:00
parent 1b60eb0ae6
commit 895671189e

View file

@ -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()