[feat] Add real-time access to model data method

This commit is contained in:
LittleMouse
2025-03-17 19:55:00 +08:00
parent b12b845702
commit 5a5fbf2870
4 changed files with 152 additions and 15 deletions
+38 -6
View File
@@ -22,6 +22,8 @@ from backend import (
Message,
)
from services.model_list import GetModelList
logging.basicConfig(
level=logging.DEBUG,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
@@ -45,11 +47,11 @@ config = Config()
async def auth_middleware(request: Request, call_next):
if request.url.path.startswith("/v1"):
api_key = request.headers.get("Authorization", "").replace("Bearer ", "")
if api_key != os.getenv("API_KEY"):
return JSONResponse(
status_code=401,
content={"error": "Invalid authentication credentials"}
)
# if api_key != os.getenv("API_KEY"):
# return JSONResponse(
# status_code=401,
# content={"error": "Invalid authentication credentials"}
# )
return await call_next(request)
class ModelDispatcher:
@@ -75,7 +77,18 @@ class ModelDispatcher:
def get_backend(self, model_name):
return self.backends.get(model_name)
_dispatcher = ModelDispatcher()
async def initialize():
global config
model_list = GetModelList(
host=config.data["server"]["host"],
port=config.data["server"]["port"]
)
await model_list.get_model_list(required_mem=0)
config = Config()
dispatcher = ModelDispatcher()
return dispatcher
_dispatcher = asyncio.run(initialize())
@app.post("/v1/chat/completions")
async def chat_completions(request: Request, body: ChatCompletionRequest):
@@ -280,6 +293,25 @@ async def create_translation(
logger.error(f"Translation error: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/v1/models")
async def list_models():
models_info = []
for model_name in _dispatcher.backends.keys():
model_config = config.data["models"].get(model_name, {})
models_info.append({
"id": model_name,
"object": "model",
"created": model_config.get("created", 0),
"owned_by": model_config.get("owner", "user"),
"permission": [],
"root": model_config.get("root", "")
})
return {
"data": models_info,
"object": "list"
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
+4
View File
@@ -10,6 +10,10 @@ logger = logging.getLogger("llm_client")
logger.setLevel(logging.DEBUG)
class LLMClient:
def __repr__(self):
attrs = ", ".join(f"{k}={v}" for k, v in self.__dict__.items() if not k.startswith("_"))
return f"LLMClient({attrs})"
def __init__(self, host: str = "localhost", port: int = 10001):
self._lock = threading.Lock()
self.host = host
+23 -9
View File
@@ -94,16 +94,28 @@ class SYSClient:
request_id = self._send_request("hwinfo", "", {})
return self._wait_response(request_id)
def model_list(self) -> dict:
request_id = self._send_request("lsmode", "", {})
return self._wait_response(request_id)
def _wait_response(self, request_id: str) -> dict:
start_time = time.time()
buffer = b""
while time.time() - start_time < 10:
response = json.loads(self.sock.recv(4096).decode())
if response["request_id"] == request_id:
if response["error"]["code"] != 0:
raise RuntimeError(f"Server error: {response['error']['message']}")
self.work_id = response["work_id"]
return response
raise TimeoutError("No response from server")
chunk = self.sock.recv(4096)
if not chunk:
break
buffer += chunk
try:
response = json.loads(buffer.decode('utf-8'))
if response["request_id"] == request_id:
if response["error"]["code"] != 0:
raise RuntimeError(f"Server error: {response['error']['message']}")
self.work_id = response["work_id"]
return response
except json.JSONDecodeError:
continue
raise TimeoutError("No valid response from server")
def connect(self):
with self._lock:
@@ -128,8 +140,10 @@ class SYSClient:
return full_text
if __name__ == "__main__":
with SYSClient(host='192.168.20.65') as client:
with SYSClient(host='192.168.20.48') as client:
hw_response = client.hwinfo()
print("hwinfo response:", hw_response)
cmm_response = client.cmminfo()
print("cmm response:", cmm_response)
print("cmm response:", cmm_response)
model_list_response = client.model_list()
print("model_list_response:", model_list_response)
+87
View File
@@ -0,0 +1,87 @@
import logging
import asyncio
import yaml
from typing import Optional
from client.sys_client import SYSClient
class GetModelList:
def __init__(self, host: str, port: int):
self.host = host
self.port = port
self.logger = logging.getLogger("get_model_list")
self._sys_client: Optional[SYSClient] = None
async def get_model_list(self, required_mem: int) -> None:
try:
if not self._sys_client:
self._sys_client = SYSClient(host=self.host, port=self.port)
with open('config/config.yaml', 'r') as f:
config = yaml.safe_load(f)
models_config = config.get('models', {})
model_list = await self._get_model_list()
for model_data in model_list["data"]:
mode = model_data.get("mode")
model_type = model_data.get("type")
if not mode or not model_type:
continue
if model_type not in ['llm', 'vlm', 'tts', 'asr']:
continue
if mode not in models_config:
new_entry = {
"host": self.host,
"port": self.port,
"type": model_type,
"input": f"{model_type}.utf-8",
"model_name": mode,
}
if model_type in ['llm', 'vlm']:
new_entry.update({
"response_format": f"{model_type}.utf-8.stream",
"object": f"{model_type}.setup",
"system_prompt": "You are a helpful assistant."
})
elif model_type == 'tts':
if 'melotts' in mode.lower():
obj = 'melotts.setup'
else:
obj = 'tts.setup'
new_entry.update({
"response_format": "wav.base64",
"object": "melotts.setup",
"object": obj
})
elif model_type == 'asr':
if 'whisper' in mode.lower():
obj = 'whisper.setup'
else:
obj = 'asr.setup'
new_entry.update({
"input": "pcm.base64",
"response_format": "asr.utf-8",
"object": "whisper.setup"
})
else:
continue
models_config[mode] = new_entry
config['models'] = models_config
with open('config/config.yaml', 'w') as f:
yaml.safe_dump(config, f, default_flow_style=False, sort_keys=False)
except Exception as e:
self.logger.error(f"Get model failed: {str(e)}")
raise
async def _get_model_list(self):
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
None,
self._sys_client.model_list
)