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 threading import Thread
from queue import Queue from queue import Queue
from fastapi import FastAPI, HTTPException from fastapi import FastAPI, HTTPException, Depends
from fastapi.responses import StreamingResponse from fastapi.responses import StreamingResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
@ -24,6 +25,14 @@ warnings.filterwarnings('ignore')
app = FastAPI() 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): def init_model(args):
tokenizer = AutoTokenizer.from_pretrained(args.load_from) 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") @app.post("/v1/chat/completions")
async def chat_completions(request: ChatRequest): async def chat_completions(request: ChatRequest, _: None = Depends(_verify_api_key)):
try: try:
if request.stream: if request.stream:
return StreamingResponse( return StreamingResponse(

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer" __package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 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 re
import gc import gc
import json import json

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer" __package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 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 argparse
import time import time
import warnings import warnings

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer" __package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 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 argparse
import time import time
import warnings import warnings

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer" __package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 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 argparse
import time import time
import warnings import warnings

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer" __package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 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 argparse
import math import math
import re import re

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer" __package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 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 argparse
import time import time
import warnings import warnings

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer" __package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 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 argparse
import math import math
import re import re

View File

@ -4,6 +4,7 @@ import sys
__package__ = "trainer" __package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 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 argparse
import time import time
import warnings import warnings