You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

50 lines
2.2 KiB

import json
from services import Orchestrator
from services.crawler import CrawlerService
from fastapi import WebSocket, WebSocketDisconnect, APIRouter
from services.llm import LLMService
from services.socket import ConnectionManager
manager = ConnectionManager()
router = APIRouter()
@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
user_unique_id = await manager.connect(websocket)
try:
data = await websocket.receive_text()
data = json.loads(data)
user_state = manager.get_user_state(user_unique_id)
if not user_state:
orchestrator = Orchestrator(crawler_service=CrawlerService(url=data['url'],
crawl_type=data['crawlType']),
llm_service=LLMService(model_type='openai',
model_name="gpt-4o-mini",
crawl_type='summarize'))
status_message = dict(type="status", message="Processing your request...")
await manager.send_message(unique_id=user_unique_id,
message=json.dumps(status_message))
manager.modify_user_state(user_unique_id, "crawling")
async for update in orchestrator.stream_website_data(
user_id=user_unique_id, manager=manager,
description=data['description'], url=data['url'],
site_type=data['siteType']):
message = dict(type="message", message=update,
is_complete=False)
await manager.send_message(user_unique_id, json.dumps(message))
status_message = dict(type="status", message="Disconnecting YOU NOW")
message = dict(type="message", message="..",
is_complete=True)
await manager.send_message(user_unique_id, json.dumps(message))
await manager.send_message(user_unique_id, json.dumps(status_message))
await manager.disconnect(user_unique_id)
except WebSocketDisconnect:
await manager.disconnect(user_unique_id)