5702dd8a9f
将原先独立的记忆可视化工具（Memory Visualizer）和LLM使用统计逻辑深度整合到项目主服务中。 主要变更包括： - **移除独立的可视化工具**: 删除了 `tools/memory_visualizer` 目录下的所有独立服务器、脚本和文档，清理了项目结构。 - **API路由整合**: 在主 FastAPI 应用中注册了记忆可视化工具的路由，使其成为核心功能的一部分，可通过 `/visualizer` 访问。 - **统计逻辑重构**: 将LLM使用统计的计算逻辑从API路由层 `statistic_router.py` 中剥离，迁移到 `src/chat/utils/statistic.py` 中，实现了逻辑的解耦和复用。API路由现在直接调用重构后的统计任务。 - **依赖清理与添加**: 添加了 `jinja2` 作为模板渲染的依赖，并清除了与独立可视化工具相关的旧依赖。 此次重构简化了项目的维护和部署，将原本分散的功能统一管理，提升了代码的内聚性和可维护性。
127 lines
3.9 KiB
Python
127 lines
3.9 KiB
Python
import os
|
|
import socket
|
|
|
|
from fastapi import APIRouter, FastAPI
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.staticfiles import StaticFiles
|
|
from rich.traceback import install
|
|
from uvicorn import Config
|
|
from uvicorn import Server as UvicornServer
|
|
|
|
from src.common.logger import get_logger
|
|
|
|
install(extra_lines=3)
|
|
|
|
logger = get_logger("Server")
|
|
|
|
|
|
class Server:
|
|
def __init__(self, host: str | None = None, port: int | None = None, app_name: str = "MaiMCore"):
|
|
self.app = FastAPI(title=app_name)
|
|
self.host: str = "127.0.0.1"
|
|
self.port: int = 8080
|
|
self._server: UvicornServer | None = None
|
|
self.set_address(host, port)
|
|
|
|
# 配置 CORS
|
|
origins = [
|
|
"http://localhost:3000", # 允许的前端源
|
|
"http://127.0.0.1:3000",
|
|
"http://127.0.0.1:3000",
|
|
# 在生产环境中，您应该添加实际的前端域名
|
|
]
|
|
|
|
self.app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=origins,
|
|
allow_credentials=True, # 是否支持 cookie
|
|
allow_methods=["*"], # 允许所有 HTTP 方法
|
|
allow_headers=["*"], # 允许所有 HTTP 请求头
|
|
)
|
|
|
|
def register_router(self, router: APIRouter, prefix: str = ""):
|
|
"""注册路由
|
|
|
|
APIRouter 用于对相关的路由端点进行分组和模块化管理：
|
|
1. 可以将相关的端点组织在一起，便于管理
|
|
2. 支持添加统一的路由前缀
|
|
3. 可以为一组路由添加共同的依赖项、标签等
|
|
|
|
示例:
|
|
router = APIRouter()
|
|
|
|
@router.get("/users")
|
|
def get_users():
|
|
return {"users": [...]}
|
|
|
|
@router.post("/users")
|
|
def create_user():
|
|
return {"msg": "user created"}
|
|
|
|
# 注册路由，添加前缀 "/api/v1"
|
|
server.register_router(router, prefix="/api/v1")
|
|
"""
|
|
self.app.include_router(router, prefix=prefix)
|
|
|
|
def set_address(self, host: str | None = None, port: int | None = None):
|
|
"""设置服务器地址和端口"""
|
|
if host:
|
|
self.host = host
|
|
if port:
|
|
self.port = port
|
|
|
|
def _is_port_in_use(self, port: int):
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
return s.connect_ex(("127.0.0.1", port)) == 0
|
|
|
|
async def run(self):
|
|
"""启动服务器"""
|
|
while self._is_port_in_use(self.port):
|
|
logger.warning(f"端口 {self.port} 已被占用，正在尝试下一个端口...")
|
|
self.port += 1
|
|
|
|
logger.info(f"将在 {self.host}:{self.port} 上启动服务器")
|
|
# 禁用 uvicorn 默认日志和访问日志
|
|
config = Config(app=self.app, host=self.host, port=self.port, log_config=None, access_log=False)
|
|
self._server = UvicornServer(config=config)
|
|
try:
|
|
await self._server.serve()
|
|
except KeyboardInterrupt:
|
|
await self.shutdown()
|
|
raise
|
|
except Exception as e:
|
|
await self.shutdown()
|
|
raise RuntimeError(f"服务器运行错误: {e!s}") from e
|
|
finally:
|
|
await self.shutdown()
|
|
|
|
async def shutdown(self):
|
|
"""安全关闭服务器"""
|
|
if self._server:
|
|
self._server.should_exit = True
|
|
await self._server.shutdown()
|
|
self._server = None
|
|
|
|
def get_app(self) -> FastAPI:
|
|
"""获取 FastAPI 实例"""
|
|
return self.app
|
|
|
|
|
|
global_server = None
|
|
|
|
|
|
def get_global_server() -> Server:
|
|
"""获取全局服务器实例"""
|
|
global global_server
|
|
if global_server is None:
|
|
host = os.getenv("HOST", "127.0.0.1")
|
|
port_str = os.getenv("PORT", "8000")
|
|
|
|
try:
|
|
port = int(port_str)
|
|
except ValueError:
|
|
port = 8000
|
|
|
|
global_server = Server(host=host, port=port)
|
|
return global_server
|