Ylem / Mock Openai

import time from typing import List, Optional from pydantic import BaseModel class ChatMessage(BaseModel): role: str content: str class ChatCompletionRequest(BaseModel): model: str = "mock-gpt-model" messages: List[ChatMessage] max_tokens: Optional[int] = 512 temperature: Optional[float] = 0.5 stream: Optional[bool] = False from fastapi import FastAPI, Depends, HTTPException, status from fastapi.security import HTTPBearer app = FastAPI(title="OpenAI-compatible API") bearer_scheme = HTTPBearer(auto_error=False) async def credentials(authorization=Depends(bearer_scheme)): if authorization and authorization.credentials == 'sk-1234': # api key is valid return authorization.credentials # raise http error 401 raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key", ) import asyncio import json async def _resp_async_generator(text_resp: str): # let's pretend every word is a token and return it over time tokens = text_resp.split(" ") for i, token in enumerate(tokens): chunk = { "id": i, "object": "chat.completion.chunk", "created": time.time(), "model": "mock-gpt-model", "choices": [{"delta": {"content": token + " "}}], } yield f"data: {json.dumps(chunk)}\n\n" await asyncio.sleep(0.1) yield "data: [DONE]\n\n" from starlette.responses import StreamingResponse @app.post("/v1/chat/completions", dependencies=[Depends(credentials)]) async def chat_completions(request: ChatCompletionRequest): if request.messages and request.messages[0].role == 'user': resp_content = "As a mock AI Assistant, I only can echo your last message:" + request.messages[-1].content else: resp_content = "As a mock AI Assistant, I only can echo your last message, but there were no messages!" if request.stream: return StreamingResponse(_resp_async_generator(resp_content), media_type="text/event-stream") return { "id": time.gmtime().tm_year, "object": "chat.completion", "created": time.time(), "model": request.model, "choices": [{ "message": ChatMessage(role="assistant", content=resp_content) }] } if __name__ == '__main__': import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)