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()
|
@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.
|
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)
|
success, researcher, error = get_researcher_by_id(mcp.researchers, research_id)
|
||||||
if not success:
|
if not success:
|
||||||
return error
|
yield format_sse_event("message", error)
|
||||||
|
|
||||||
logger.info(f"Generating report for research ID: {research_id}")
|
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()
|
sources = researcher.get_research_sources()
|
||||||
costs = researcher.get_costs()
|
costs = researcher.get_costs()
|
||||||
|
|
||||||
return create_success_response({
|
final_data_payload = {
|
||||||
"report": report,
|
"report": report,
|
||||||
"source_count": len(sources),
|
"source_count": len(sources),
|
||||||
"costs": costs
|
"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:
|
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()
|
@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.
|
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)
|
success, researcher, error = get_researcher_by_id(mcp.researchers, research_id)
|
||||||
if not success:
|
if not success:
|
||||||
return error
|
yield format_sse_event("message", error)
|
||||||
|
|
||||||
sources = researcher.get_research_sources()
|
sources = researcher.get_research_sources()
|
||||||
source_urls = researcher.get_source_urls()
|
source_urls = researcher.get_source_urls()
|
||||||
|
|
||||||
return create_success_response({
|
final_data_payload = {
|
||||||
"sources": format_sources_for_response(sources),
|
"sources": format_sources_for_response(sources),
|
||||||
"source_urls": source_urls
|
"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()
|
@mcp.tool()
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue