Spaces:
Running
Running
File size: 4,835 Bytes
372531f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import asyncio
import datetime
from typing import Dict, List
from fastapi import WebSocket
from backend.report_type import BasicReport, DetailedReport
from backend.chat import ChatAgentWithMemory
from gpt_researcher.utils.enum import ReportType, Tone
from multi_agents.main import run_research_task
from gpt_researcher.actions import stream_output # Import stream_output
from backend.server.server_utils import CustomLogsHandler
class WebSocketManager:
"""Manage websockets"""
def __init__(self):
"""Initialize the WebSocketManager class."""
self.active_connections: List[WebSocket] = []
self.sender_tasks: Dict[WebSocket, asyncio.Task] = {}
self.message_queues: Dict[WebSocket, asyncio.Queue] = {}
self.chat_agent = None
async def start_sender(self, websocket: WebSocket):
"""Start the sender task."""
queue = self.message_queues.get(websocket)
if not queue:
return
while True:
message = await queue.get()
if websocket in self.active_connections:
try:
if message == "ping":
await websocket.send_text("pong")
else:
await websocket.send_text(message)
except:
break
else:
break
async def connect(self, websocket: WebSocket):
"""Connect a websocket."""
await websocket.accept()
self.active_connections.append(websocket)
self.message_queues[websocket] = asyncio.Queue()
self.sender_tasks[websocket] = asyncio.create_task(
self.start_sender(websocket))
async def disconnect(self, websocket: WebSocket):
"""Disconnect a websocket."""
if websocket in self.active_connections:
self.active_connections.remove(websocket)
self.sender_tasks[websocket].cancel()
await self.message_queues[websocket].put(None)
del self.sender_tasks[websocket]
del self.message_queues[websocket]
async def start_streaming(self, task, report_type, report_source, source_urls, document_urls, tone, websocket, headers=None):
"""Start streaming the output."""
tone = Tone[tone]
# add customized JSON config file path here
config_path = "default"
report = await run_agent(task, report_type, report_source, source_urls, document_urls, tone, websocket, headers = headers, config_path = config_path)
#Create new Chat Agent whenever a new report is written
self.chat_agent = ChatAgentWithMemory(report, config_path, headers)
return report
async def chat(self, message, websocket):
"""Chat with the agent based message diff"""
if self.chat_agent:
await self.chat_agent.chat(message, websocket)
else:
await websocket.send_json({"type": "chat", "content": "Knowledge empty, please run the research first to obtain knowledge"})
async def run_agent(task, report_type, report_source, source_urls, document_urls, tone: Tone, websocket, headers=None, config_path=""):
"""Run the agent."""
start_time = datetime.datetime.now()
# Create logs handler for this research task
logs_handler = CustomLogsHandler(websocket, task)
# Initialize researcher based on report type
if report_type == "multi_agents":
report = await run_research_task(
query=task,
websocket=logs_handler, # Use logs_handler instead of raw websocket
stream_output=stream_output,
tone=tone,
headers=headers
)
report = report.get("report", "")
elif report_type == ReportType.DetailedReport.value:
researcher = DetailedReport(
query=task,
report_type=report_type,
report_source=report_source,
source_urls=source_urls,
document_urls=document_urls,
tone=tone,
config_path=config_path,
websocket=logs_handler, # Use logs_handler instead of raw websocket
headers=headers
)
report = await researcher.run()
else:
researcher = BasicReport(
query=task,
report_type=report_type,
report_source=report_source,
source_urls=source_urls,
document_urls=document_urls,
tone=tone,
config_path=config_path,
websocket=logs_handler, # Use logs_handler instead of raw websocket
headers=headers
)
report = await researcher.run()
return report
|