506 lines
20 KiB
Python
506 lines
20 KiB
Python
import json
|
||
import logging
|
||
import os
|
||
from typing import Any, Awaitable, Callable, Optional, Sequence
|
||
import uuid
|
||
import aiofiles
|
||
import yaml
|
||
import cv2
|
||
import base64
|
||
import asyncio
|
||
import numpy as np
|
||
import time
|
||
import subprocess
|
||
import ffmpeg
|
||
import io
|
||
from PIL import Image
|
||
from collections import deque
|
||
|
||
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.responses import FileResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
from fastapi.responses import HTMLResponse
|
||
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack
|
||
from aiortc.contrib.media import MediaRelay
|
||
from aiortc.contrib.media import MediaPlayer
|
||
|
||
from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
|
||
from autogen_agentchat.base import TaskResult
|
||
from autogen_agentchat.messages import TextMessage, UserInputRequestedEvent
|
||
from autogen_agentchat.teams import RoundRobinGroupChat
|
||
from autogen_core import CancellationToken
|
||
from autogen_core.models import ChatCompletionClient
|
||
from autogen_ext.models.openai import OpenAIChatCompletionClient
|
||
from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination
|
||
from autogen_agentchat.teams import SelectorGroupChat, RoundRobinGroupChat
|
||
from autogen_agentchat.messages import AgentEvent, ChatMessage, TextMessage, ToolCallExecutionEvent
|
||
from constant import MODEL, OPENAI_API_KEY, OPENAI_BASE_URL
|
||
from scientist_team import create_scientist_team
|
||
from engineer_team import create_engineer_team
|
||
from robot_platform import create_robot_team
|
||
from analyst_team import create_analyst_team
|
||
from utils import load_agent_configs
|
||
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
relay = MediaRelay()
|
||
|
||
model_client = OpenAIChatCompletionClient(
|
||
model=MODEL,
|
||
base_url=OPENAI_BASE_URL,
|
||
api_key=OPENAI_API_KEY,
|
||
model_info={
|
||
"vision": True,
|
||
"function_calling": True,
|
||
"json_output": True,
|
||
"family": "unknown",
|
||
},
|
||
)
|
||
|
||
|
||
async def get_team(
|
||
user_input_func: Callable[[str, Optional[CancellationToken]], Awaitable[str]],
|
||
session_dir: str
|
||
) -> RoundRobinGroupChat | SelectorGroupChat:
|
||
|
||
# Create the team.
|
||
scientist_team = create_scientist_team(model_client=model_client)
|
||
engineer_team = create_engineer_team()
|
||
robot_platform = create_robot_team()
|
||
analyst_team = create_analyst_team()
|
||
user = UserProxyAgent(
|
||
name="user",
|
||
input_func=user_input_func, # Use the user input function.
|
||
)
|
||
cur_path = os.path.dirname(os.path.abspath(__file__))
|
||
planning_agent = AssistantAgent(
|
||
"ORCHESTRATOR",
|
||
description="An agent for planning tasks, this agent should be the first to engage when given a new task.",
|
||
model_client=model_client,
|
||
system_message="""
|
||
You are a planning agent.
|
||
Your job is to break down complex Materials science research tasks into smaller, manageable subtasks.
|
||
Assign these subtasks to the appropriate sub-teams; not all sub-teams are required to participate in every task.
|
||
Your sub-teams are:
|
||
1. User: A human agent to whom you transfer information whenever you need to confirm your execution steps to a human.
|
||
2. Scientist: A professional team of material scientists who are mainly responsible for consulting on material synthesis, structure, application and properties.
|
||
- The scientist team has the following members:
|
||
2.1 Synthesis Scientist: who is good at giving perfect and correct synthesis solutions.
|
||
2.2 Structure Scientist: focusing on agents of structural topics in materials science.
|
||
2.3 Property Scientist: focuses on physical and chemistry property topics in materials science.
|
||
2.4 Application Scientist: Focus on practical applications of materials, such as devices, chips, etc.
|
||
3. Engineer: A team of professional engineers who are responsible for writing code, visualizing experimental schemes, converting experimental schemes to JSON, and more.
|
||
- The engineer team has the following members:
|
||
3.1 Structural engineer: A professional structural engineer who focus on converting natural language synthesis schemes to JSON or XML formated scheme, and then upload this JSON to S3 Storage.
|
||
3.2 Software engineer: A professional software engineers will coding with Python.
|
||
3.3 Code reviewer: A professional code reviewer will review the code written by software engineers and execute it.
|
||
3.4 Scheme Plotter: An agent responsible for converting a expriment scheme into a Mermaid flowchart.
|
||
4. Executor: A robotic platform is responsible for performing automated synthesis experiments, automated characterization experiments, and collecting experimental datas.
|
||
- The Executor team has the following members:
|
||
4.1 MobileRobot_Agent: This agent controls the mobile robot by calling the funciton sendScheme2MobileRobot to place the experimental container into the robot workstation. This agent called before RobotWorkstation_Agent.
|
||
4.2 RobotWorkstation_Agent: This agent is called by the mobile robot agent, do not plan it alone.
|
||
4.3 DataCollector_Agent: This agent collects experimental data and experimental logs from the characterization device in the robot platform and stores them.
|
||
5. Analyst: A team of data analysts who are responsible for analyzing and visualizing experimental data and logs.
|
||
- The Data Analysis team has the following members:
|
||
5.1 Expriment_Analyst: The agent of data analysts who are responsible for analyzing experimental data and logs.
|
||
5.2 Expriment_Optimizer: The agent optimizes the experimental scheme by means of component regulation and so on to make the experimental result close to the desired goal of the user.
|
||
5.3 Data_Visulizer: The agent of data visulizers who are responsible for visualizing experimental data and logs.
|
||
|
||
You only plan and delegate tasks - you do not execute them yourself.
|
||
|
||
回答时你需要初始化/更新如下任务分配表和Mermaid流程图，并按顺序执行，使用如下格式并利用：
|
||
| Team_name | Member_name | sub-task |
|
||
| ----------- | ------------- | ------------------------------------ |
|
||
| <team_name> | <member_name> | <status: brief sub-task description> |
|
||
|
||
```mermaid
|
||
graph TD
|
||
User[User]
|
||
subgraph <team_name>
|
||
A1[<member_name>]
|
||
end
|
||
style xxx # 推荐多样的风格
|
||
...
|
||
User --> A1
|
||
...
|
||
```
|
||
|
||
每次回答时，你需要清晰明确的指出已经完成的子任务下一步子任务，使用如下格式：
|
||
**已完成子任务：**
|
||
1. <team> : <subtask>
|
||
**Next sub-task:**
|
||
n. <team> : <subtask>
|
||
|
||
You can end with "HUMAN" if you need to, which means you need human approval or other advice or instructions;
|
||
After plan and delegate tasks are complete, end with "START";
|
||
Determine if all sub-teams have completed their tasks, and if so, summarize the findings and end with "TERMINATE".
|
||
Planning with English:
|
||
""",
|
||
reflect_on_tool_use=False
|
||
)
|
||
|
||
# The termination condition is a combination of text mention termination and max message termination.
|
||
text_mention_termination = TextMentionTermination("TERMINATE")
|
||
max_messages_termination = MaxMessageTermination(max_messages=200)
|
||
termination = text_mention_termination | max_messages_termination
|
||
|
||
# The selector function is a function that takes the current message thread of the group chat
|
||
# and returns the next speaker's name. If None is returned, the LLM-based selection method will be used.
|
||
def selector_func(messages: Sequence[AgentEvent | ChatMessage]) -> str | None:
|
||
if messages[-1].source != planning_agent.name:
|
||
return planning_agent.name # Always return to the planning agent after the other agents have spoken.
|
||
elif "HUMAN" in messages[-1].content:
|
||
return user.name
|
||
return None
|
||
|
||
team = SelectorGroupChat(
|
||
[planning_agent, user, scientist_team, engineer_team, robot_platform, analyst_team],
|
||
model_client=model_client, # Use a smaller model for the selector.
|
||
termination_condition=termination,
|
||
selector_func=selector_func,
|
||
)
|
||
# Load state from file.
|
||
state_path = os.path.join(session_dir, "team_state.json")
|
||
if not os.path.exists(state_path):
|
||
return team
|
||
async with aiofiles.open(state_path, "r") as file:
|
||
state = json.loads(await file.read())
|
||
await team.load_state(state)
|
||
return team
|
||
|
||
|
||
|
||
app = FastAPI()
|
||
current_task = None # 用于跟踪当前任务
|
||
# Add CORS middleware
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"], # Allows all origins
|
||
allow_credentials=True,
|
||
allow_methods=["*"], # Allows all methods
|
||
allow_headers=["*"], # Allows all headers
|
||
)
|
||
|
||
# model_config_path = "model_config.yaml"
|
||
# state_path = "team_state.json"
|
||
# history_path = "team_history.json"
|
||
|
||
# Serve static files
|
||
app.mount("/static", StaticFiles(directory="."), name="static")
|
||
|
||
|
||
async def get_session_history(session_dir: str) -> list[dict[str, Any]]:
|
||
"""Get chat history from file using UUID."""
|
||
session_history_path = os.path.join(session_dir, "team_history.json")
|
||
if not os.path.exists(session_history_path):
|
||
return []
|
||
async with aiofiles.open(session_history_path, "r") as file:
|
||
content = await file.read()
|
||
if content:
|
||
return json.loads(content)
|
||
else:
|
||
return []
|
||
|
||
|
||
@app.websocket("/history/{session_uuid}")
|
||
async def history(websocket: WebSocket) -> list[dict[str, Any]]:
|
||
await websocket.accept()
|
||
data = await websocket.receive_json()
|
||
session_uuid = data["uuid"]
|
||
try:
|
||
session_history = await get_session_history(session_uuid)
|
||
await websocket.send_json(session_history)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=str(e)) from e
|
||
|
||
@app.websocket("/sessions")
|
||
async def sessions(websocket: WebSocket) -> list[dict[str, str]]:
|
||
"""Get all history UUIDs and their main content."""
|
||
await websocket.accept()
|
||
cur_path = os.path.dirname(os.path.abspath(__file__))
|
||
history_dir = os.path.join(cur_path, "history")
|
||
session_data = []
|
||
for dir_name in os.listdir(history_dir):
|
||
session_dir = os.path.join(history_dir, dir_name)
|
||
if os.path.isdir(session_dir): # Check if it's a directory
|
||
try:
|
||
history = await get_session_history(session_dir)
|
||
main_content = history[0]["content"] if history and "content" in history[0] else ""
|
||
session_data.append({"uuid": dir_name, "content": main_content})
|
||
|
||
except Exception as e:
|
||
print(f"Error reading history for {dir_name}: {e}") # Log the error but continue
|
||
|
||
await websocket.send_json(session_data)
|
||
|
||
|
||
@app.websocket("/ws/chat")
|
||
async def chat(websocket: WebSocket):
|
||
await websocket.accept()
|
||
|
||
# User input function used by the team.
|
||
async def _user_input(prompt: str, cancellation_token: CancellationToken | None) -> str:
|
||
data = await websocket.receive_json()
|
||
return data['content']
|
||
|
||
try:
|
||
while True:
|
||
# Get user message.
|
||
data = await websocket.receive_json()
|
||
|
||
if 'session_uuid' not in data:
|
||
# New session
|
||
request = TextMessage.model_validate(data)
|
||
# request = data['content']
|
||
session_uuid = str(uuid.uuid4()) # Initialize a unique UUID for each session
|
||
cur_path = os.path.dirname(os.path.abspath(__file__))
|
||
session_dir = os.path.join(cur_path, "history", session_uuid) # Directory for session states
|
||
os.makedirs(session_dir, exist_ok=True) # ensure the directory is created.
|
||
history = []
|
||
else:
|
||
session_uuid = data['session_uuid']
|
||
cur_path = os.path.dirname(os.path.abspath(__file__))
|
||
session_dir = os.path.join(cur_path, "history", session_uuid) # Directory for session states
|
||
history = await get_session_history(session_dir)
|
||
new_data = {k: v for k, v in data.items() if k != "session_uuid"}
|
||
request = TextMessage.model_validate(new_data)
|
||
request = history + request
|
||
|
||
try:
|
||
# Get the team and respond to the message.
|
||
team = await get_team(_user_input, session_dir)
|
||
|
||
stream = team.run_stream(task=request)
|
||
async for message in stream:
|
||
if isinstance(message, TaskResult):
|
||
continue
|
||
print(f"----------------{message.source}----------------\n {message.content}")
|
||
if message.type == 'TextMessage' or message.type == 'HandoffMessage' \
|
||
or message.type == 'UserInputRequestedEvent':
|
||
if isinstance(message.content, str):
|
||
await websocket.send_json(message.model_dump())
|
||
else:
|
||
_message = message
|
||
_message.content = str(_message.content)
|
||
await websocket.send_json(_message.model_dump())
|
||
if not isinstance(message, UserInputRequestedEvent):
|
||
history.append(message.model_dump())
|
||
|
||
# Save chat history to file.
|
||
session_history_path = os.path.join(session_dir, "team_history.json")
|
||
async with aiofiles.open(session_history_path, "w") as file:
|
||
await file.write(json.dumps(history))
|
||
|
||
# # Save team state to file.
|
||
# session_state_path = os.path.join(session_dir, "team_state.json")
|
||
# async with aiofiles.open(session_state_path, "w") as file:
|
||
# state = await team.save_state()
|
||
# await file.write(json.dumps(state))
|
||
|
||
except Exception as e:
|
||
# Send error message to client
|
||
error_message = {
|
||
"type": "error",
|
||
"content": f"Error: {str(e)}",
|
||
"source": "system"
|
||
}
|
||
await websocket.send_json(error_message)
|
||
# Re-enable input after error
|
||
await websocket.send_json({
|
||
"type": "UserInputRequestedEvent",
|
||
"content": "An error occurred. Please try again.",
|
||
"source": "system"
|
||
})
|
||
|
||
except WebSocketDisconnect:
|
||
logger.info("Client disconnected")
|
||
except Exception as e:
|
||
logger.error(f"Unexpected error: {str(e)}")
|
||
try:
|
||
await websocket.send_json({
|
||
"type": "error",
|
||
"content": f"Unexpected error: {str(e)}",
|
||
"source": "system"
|
||
})
|
||
except:
|
||
pass
|
||
|
||
|
||
|
||
RTSP_STREAMS = {
|
||
"camera1": "rtsp://admin:@192.168.1.13:554/live",
|
||
"camera2": "rtsp://admin:@192.168.1.10:554/live",
|
||
}
|
||
|
||
@app.websocket("/video_stream/{camera_id}")
|
||
async def websocket_endpoint(websocket: WebSocket, camera_id: str):
|
||
await websocket.accept()
|
||
|
||
rtsp_url = RTSP_STREAMS.get(camera_id)
|
||
if not rtsp_url:
|
||
await websocket.close(code=4001)
|
||
return
|
||
|
||
fps = 15
|
||
frame_interval = 1 / fps
|
||
frame_buffer = deque() # 每个连接独立的队列
|
||
|
||
# 启动FFmpeg进程
|
||
process = (
|
||
ffmpeg
|
||
.input(rtsp_url, rtsp_transport='tcp')
|
||
.output('pipe:', format='image2pipe', vcodec='mjpeg', r=fps)
|
||
.run_async(pipe_stdout=True, pipe_stderr=True, quiet=True)
|
||
)
|
||
|
||
async def capture_frames():
|
||
"""读取并分割完整的JPEG帧"""
|
||
buffer = b''
|
||
loop = asyncio.get_event_loop()
|
||
while True:
|
||
# 异步读取防止阻塞
|
||
data = await loop.run_in_executor(None, process.stdout.read, 4096)
|
||
if not data:
|
||
break
|
||
buffer += data
|
||
# 提取完整帧
|
||
while True:
|
||
start = buffer.find(b'\xff\xd8') # JPEG开始标记
|
||
if start == -1:
|
||
break
|
||
end = buffer.find(b'\xff\xd9', start + 2) # JPEG结束标记
|
||
if end == -1:
|
||
break
|
||
# 提取帧数据并编码
|
||
frame_data = buffer[start:end+2]
|
||
frame = base64.b64encode(frame_data).decode('utf-8')
|
||
frame_buffer.append(frame)
|
||
buffer = buffer[end+2:] # 移除已处理数据
|
||
|
||
# 清理FFmpeg进程
|
||
process.kill()
|
||
|
||
async def send_frames():
|
||
"""发送帧到客户端"""
|
||
while True:
|
||
if frame_buffer:
|
||
frame = frame_buffer.popleft()
|
||
await websocket.send_text(frame)
|
||
await asyncio.sleep(frame_interval)
|
||
|
||
# 启动任务并处理退出
|
||
capture_task = asyncio.create_task(capture_frames())
|
||
send_task = asyncio.create_task(send_frames())
|
||
try:
|
||
await asyncio.gather(capture_task, send_task)
|
||
except asyncio.CancelledError:
|
||
capture_task.cancel()
|
||
send_task.cancel()
|
||
await asyncio.gather(capture_task, send_task, return_exceptions=True)
|
||
finally:
|
||
if process and process.poll() is None: # Check if the process is still running
|
||
process.kill()
|
||
await websocket.close()
|
||
|
||
|
||
@app.get("/")
|
||
async def debug_page():
|
||
return HTMLResponse("""
|
||
<!DOCTYPE html>
|
||
<html lang="en">
|
||
<head>
|
||
<meta charset="UTF-8">
|
||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||
<title>RTSP Stream</title>
|
||
<style>
|
||
#video {
|
||
width: 100%;
|
||
height: auto;
|
||
}
|
||
</style>
|
||
</head>
|
||
<body>
|
||
<!-- 视频流 1 的画布 -->
|
||
<canvas id="canvas1"></canvas>
|
||
<script>
|
||
const canvas1 = document.getElementById('canvas1');
|
||
const ctx1 = canvas1.getContext('2d');
|
||
const img1 = new Image();
|
||
|
||
img1.onload = () => {
|
||
ctx1.drawImage(img1, 0, 0, canvas1.width, canvas1.height);
|
||
};
|
||
|
||
const ws1 = new WebSocket('ws://100.85.52.31:8000/video_stream/camera1');
|
||
|
||
ws1.onopen = () => {
|
||
console.log('WebSocket connection established for camera1'); // Log connection status.
|
||
};
|
||
|
||
ws1.onmessage = (event) => {
|
||
img1.src = 'data:image/jpeg;base64,' + event.data;
|
||
};
|
||
|
||
ws1.onclose = () => {
|
||
console.log('WebSocket connection closed for camera1');
|
||
};
|
||
|
||
// 动态调整 Canvas 尺寸
|
||
function resizeCanvas1() {
|
||
canvas1.width = 720; // 设置宽度为屏幕的一半
|
||
canvas1.height = 480; // 设置高度为屏幕的一半
|
||
}
|
||
resizeCanvas1();
|
||
window.addEventListener('resize', resizeCanvas1);
|
||
|
||
|
||
</script>
|
||
|
||
<!-- 视频流 2 的画布 -->
|
||
<canvas id="canvas2"></canvas>
|
||
<script>
|
||
const canvas2 = document.getElementById('canvas2');
|
||
const ctx2 = canvas2.getContext('2d');
|
||
const img2 = new Image();
|
||
|
||
img2.onload = () => {
|
||
ctx2.drawImage(img2, 0, 0, canvas2.width, canvas2.height);
|
||
};
|
||
|
||
const ws2 = new WebSocket('ws://100.85.52.31:8000/video_stream/camera2');
|
||
|
||
ws2.onopen = () => {
|
||
console.log('WebSocket connection established for camera2');
|
||
};
|
||
|
||
|
||
ws2.onmessage = (event) => {
|
||
img2.src = 'data:image/jpeg;base64,' + event.data;
|
||
};
|
||
|
||
ws2.onclose = () => {
|
||
console.log('WebSocket connection closed for camera2');
|
||
};
|
||
|
||
// 动态调整 Canvas 尺寸
|
||
function resizeCanvas2() {
|
||
canvas2.width = 720; // 设置宽度为屏幕的一半
|
||
canvas2.height = 480; // 设置高度为屏幕的一半
|
||
}
|
||
resizeCanvas2();
|
||
window.addEventListener('resize', resizeCanvas2);
|
||
</script>
|
||
|
||
</body>
|
||
</html>
|
||
""")
|
||
|
||
# Example usage
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
|
||
uvicorn.run(app, host="0.0.0.0", port=8000)
|