refactor: 统一使用现代类型注解替代传统类型注解
- 将 Dict、List 等传统类型注解替换为 dict、list 等现代类型注解 - 更新类型注解以更精确地反映变量类型 - 修复部分类型注解与实际使用不匹配的问题 - 优化部分代码逻辑以提高类型安全性
This commit is contained in:
+9
-9
@@ -27,18 +27,18 @@ class NodeConnection:
|
||||
display_name: str = ""
|
||||
serves_users: Set[str] = field(default_factory=set)
|
||||
working_dir: str = ""
|
||||
capabilities: List[str] = field(default_factory=list)
|
||||
capabilities: list[str] = field(default_factory=list)
|
||||
connected_at: float = field(default_factory=time.time)
|
||||
last_heartbeat: float = field(default_factory=time.time)
|
||||
sessions: int = 0
|
||||
active_sessions: List[Dict[str, Any]] = field(default_factory=list)
|
||||
active_sessions: list[dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
@property
|
||||
def is_online(self) -> bool:
|
||||
"""Check if node is still considered online (heartbeat within 60s)."""
|
||||
return time.time() - self.last_heartbeat < 60
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Serialize for API responses."""
|
||||
return {
|
||||
"node_id": self.node_id,
|
||||
@@ -56,9 +56,9 @@ class NodeRegistry:
|
||||
"""Registry of connected host clients."""
|
||||
|
||||
def __init__(self, router_secret: str = ""):
|
||||
self._nodes: Dict[str, NodeConnection] = {}
|
||||
self._user_nodes: Dict[str, Set[str]] = {}
|
||||
self._active_node: Dict[str, str] = {}
|
||||
self._nodes: dict[str, NodeConnection] = {}
|
||||
self._user_nodes: dict[str, Set[str]] = {}
|
||||
self._active_node: dict[str, str] = {}
|
||||
self._secret = router_secret
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
@@ -159,7 +159,7 @@ class NodeRegistry:
|
||||
"""Get a node by ID."""
|
||||
return self._nodes.get(node_id)
|
||||
|
||||
def get_nodes_for_user(self, user_id: str) -> List[NodeConnection]:
|
||||
def get_nodes_for_user(self, user_id: str) -> list[NodeConnection]:
|
||||
"""Get all nodes that serve a user."""
|
||||
node_ids = self._user_nodes.get(user_id, set())
|
||||
return [self._nodes[nid] for nid in node_ids if nid in self._nodes]
|
||||
@@ -186,11 +186,11 @@ class NodeRegistry:
|
||||
logger.info("Active node for user %s set to %s", user_id, node_id)
|
||||
return True
|
||||
|
||||
def list_nodes(self) -> List[Dict[str, Any]]:
|
||||
def list_nodes(self) -> list[dict[str, Any]]:
|
||||
"""List all nodes with their status."""
|
||||
return [node.to_dict() for node in self._nodes.values()]
|
||||
|
||||
def get_affected_users(self, node_id: str) -> List[str]:
|
||||
def get_affected_users(self, node_id: str) -> list[str]:
|
||||
"""Get users affected by a node disconnect."""
|
||||
node = self._nodes.get(node_id)
|
||||
if node:
|
||||
|
||||
@@ -12,6 +12,7 @@ from typing import List, Optional
|
||||
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_openai import ChatOpenAI
|
||||
from pydantic import SecretStr
|
||||
|
||||
from config import OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_MODEL
|
||||
from router.nodes import NodeConnection, get_node_registry
|
||||
@@ -36,7 +37,7 @@ Respond with a JSON object:
|
||||
"""
|
||||
|
||||
|
||||
def _format_nodes_info(nodes: List[NodeConnection], active_node_id: Optional[str] = None) -> str:
|
||||
def _format_nodes_info(nodes: list[NodeConnection], active_node_id: Optional[str] = None) -> str:
|
||||
"""Format node information for the routing prompt."""
|
||||
lines = []
|
||||
for node in nodes:
|
||||
@@ -86,8 +87,8 @@ async def route(user_id: str, chat_id: str, text: str) -> tuple[Optional[str], s
|
||||
try:
|
||||
llm = ChatOpenAI(
|
||||
model=OPENAI_MODEL,
|
||||
openai_api_key=OPENAI_API_KEY,
|
||||
openai_api_base=OPENAI_BASE_URL,
|
||||
api_key=SecretStr(OPENAI_API_KEY),
|
||||
base_url=OPENAI_BASE_URL,
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
@@ -98,7 +99,11 @@ async def route(user_id: str, chat_id: str, text: str) -> tuple[Optional[str], s
|
||||
]
|
||||
|
||||
response = await llm.ainvoke(messages)
|
||||
content = response.content.strip()
|
||||
content = response.content
|
||||
if isinstance(content, str):
|
||||
content = content.strip()
|
||||
else:
|
||||
content = str(content).strip()
|
||||
|
||||
if content.startswith("```"):
|
||||
content = content.split("\n", 1)[1]
|
||||
|
||||
+2
-2
@@ -18,7 +18,7 @@ from router.nodes import get_node_registry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_pending_requests: Dict[str, asyncio.Future] = {}
|
||||
_pending_requests: dict[str, asyncio.Future[str]] = {}
|
||||
_default_timeout = 600.0
|
||||
|
||||
|
||||
@@ -52,7 +52,7 @@ async def forward(
|
||||
raise RuntimeError(f"Node not connected: {node_id}")
|
||||
|
||||
request_id = str(uuid.uuid4())
|
||||
future: asyncio.Future = asyncio.get_event_loop().create_future()
|
||||
future: asyncio.Future[str] = asyncio.get_event_loop().create_future()
|
||||
_pending_requests[request_id] = future
|
||||
|
||||
request = ForwardRequest(
|
||||
|
||||
+1
-1
@@ -47,7 +47,7 @@ async def ws_node_endpoint(websocket: WebSocket) -> None:
|
||||
return
|
||||
|
||||
node_id: Optional[str] = None
|
||||
heartbeat_task: Optional[asyncio.Task] = None
|
||||
heartbeat_task: Optional[asyncio.Task[None]] = None
|
||||
|
||||
async def send_heartbeat():
|
||||
"""Send periodic pings to the host client."""
|
||||
|
||||
Reference in New Issue
Block a user