Compare commits

...

2 Commits

Author SHA1 Message Date
orbisai0security e091e6c13d
Merge 8607169546 into 802c15b2b4 2026-05-06 16:39:54 +05:30
orbisai0security 8607169546 fix: V-001 security vulnerability
Automated security fix generated by Orbis Security AI
2026-05-05 07:51:07 +00:00
1 changed files with 11 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(