Compare commits

...

3 Commits

Author SHA1 Message Date
OrbisAI Security 3466496f34
Merge 8607169546 into 4497610ec0 2026-05-23 17:19:18 +08:00
jingyaogong 4497610ec0 [fix] issue#771 2026-05-19 17:40:03 +08:00
orbisai0security 8607169546 fix: V-001 security vulnerability
Automated security fix generated by Orbis Security AI
2026-05-05 07:51:07 +00:00
9 changed files with 19 additions and 2 deletions

View File

@ -13,8 +13,9 @@ import uvicorn
from threading import Thread
from queue import Queue
from fastapi import FastAPI, HTTPException
from fastapi import FastAPI, HTTPException, Depends
from fastapi.responses import StreamingResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
@ -24,6 +25,14 @@ warnings.filterwarnings('ignore')
app = FastAPI()
_security = HTTPBearer(auto_error=False)
def _verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(_security)):
api_key = os.environ.get("MINIMIND_API_KEY")
if api_key and (credentials is None or credentials.credentials != api_key):
raise HTTPException(status_code=401, detail="Invalid or missing API key")
def init_model(args):
tokenizer = AutoTokenizer.from_pretrained(args.load_from)
@ -169,7 +178,7 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
@app.post("/v1/chat/completions")
async def chat_completions(request: ChatRequest):
async def chat_completions(request: ChatRequest, _: None = Depends(_verify_api_key)):
try:
if request.stream:
return StreamingResponse(

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771)
import re
import gc
import json

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771)
import argparse
import time
import warnings

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771)
import argparse
import time
import warnings

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771)
import argparse
import time
import warnings

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771)
import argparse
import math
import re

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771)
import argparse
import time
import warnings

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771)
import argparse
import math
import re

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771)
import argparse
import time
import warnings