Create fastapi_app.py

fastapi_app.py

@@ -0,0 +1,132 @@
+from fastapi import FastAPI, Request, Response
+from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+import uvicorn
+import json
+
+from typegpt_api import generate, model_mapping, simplified_models
+from api_info import developer_info, model_providers
+
+app = FastAPI()
+
+# Set up CORS middleware if needed
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.get("/health_check")
+async def health_check():
+    return {"status": "OK"}
+
+@app.get("/models")
+async def get_models():
+    try:
+        response = {
+            "object": "list",
+            "data": []
+        }
+        for provider, info in model_providers.items():
+            for model in info["models"]:
+                response["data"].append({
+                    "id": model,
+                    "object": "model",
+                    "provider": provider,
+                    "description": info["description"]
+                })
+                
+        return JSONResponse(content=response)
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=500)
+
+@app.post("/chat/completions")
+async def chat_completions(request: Request):
+    # Receive the JSON payload
+    try:
+        body = await request.json()
+    except Exception as e:
+        return JSONResponse(content={"error": "Invalid JSON payload"}, status_code=400)
+
+    # Extract parameters
+    model = body.get("model")
+    messages = body.get("messages")
+    temperature = body.get("temperature", 0.7)
+    top_p = body.get("top_p", 1.0)
+    n = body.get("n", 1)
+    stream = body.get("stream", False)
+    stop = body.get("stop")
+    max_tokens = body.get("max_tokens")
+    presence_penalty = body.get("presence_penalty", 0.0)
+    frequency_penalty = body.get("frequency_penalty", 0.0)
+    logit_bias = body.get("logit_bias")
+    user = body.get("user")
+    timeout = 30  # or set based on your preference
+
+    # Validate required parameters
+    if not model:
+        return JSONResponse(content={"error": "The 'model' parameter is required."}, status_code=400)
+    if not messages:
+        return JSONResponse(content={"error": "The 'messages' parameter is required."}, status_code=400)
+
+    # Call the generate function
+    try:
+        if stream:
+            async def generate_stream():
+                response = generate(
+                    model=model,
+                    messages=messages,
+                    temperature=temperature,
+                    top_p=top_p,
+                    n=n,
+                    stream=True,
+                    stop=stop,
+                    max_tokens=max_tokens,
+                    presence_penalty=presence_penalty,
+                    frequency_penalty=frequency_penalty,
+                    logit_bias=logit_bias,
+                    user=user,
+                    timeout=timeout,
+                )
+                
+                for chunk in response:
+                    yield f"data: {json.dumps(chunk)}\n\n"
+                yield "data: [DONE]\n\n"
+
+            return StreamingResponse(
+                generate_stream(),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "Transfer-Encoding": "chunked"
+                }
+            )
+        else:
+            response = generate(
+                model=model,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                n=n,
+                stream=False,
+                stop=stop,
+                max_tokens=max_tokens,
+                presence_penalty=presence_penalty,
+                frequency_penalty=frequency_penalty,
+                logit_bias=logit_bias,
+                user=user,
+                timeout=timeout,
+            )
+            return JSONResponse(content=response)
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=500)
+
+@app.get("/developer_info")
+async def get_developer_info():
+    return JSONResponse(content=developer_info)
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)