Upload 2 files

Dockerfile

@@ -0,0 +1,6 @@
+FROM python:3.11
+
+COPY main.py /app/main.py
+WORKDIR /app
+RUN pip install --break-system-packages fastapi uvicorn uvloop fastapi_poe
+

main.py

@@ -0,0 +1,99 @@
+import os
+
+DEFAULT_MODEL = os.getenv("BOT", default="Claude-3-Sonnet")
+LISTEN_PORT = int(os.getenv("PORT", default=7860))
+BASE_URL = os.getenv("BASE", default="https://api.poe.com/bot/")
+POE_API_KEY = os.getenv("POE_API_KEY","")
+AUTHORIZATION_API_KEY = os.getenv("AUTHORIZATION_API_KEY","")
+if not POE_API_KEY or not AUTHORIZATION_API_KEY:
+    raise ValueError("POE_API_KEY and AUTHORIZATION_API_KEY must be set in the environment variables")
+
+from fastapi import FastAPI, Request, Header, HTTPException
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+
+import uvicorn
+from typing import AsyncGenerator
+import json
+
+from fastapi_poe.types import ProtocolMessage
+from fastapi_poe.client import get_bot_response
+
+
+app = FastAPI()
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
+
+async def generate_responses(api_key: str, formatted_messages: list, bot_name: str) -> AsyncGenerator[str, None]:
+    """An async generator to stream responses from the POE API."""
+
+    # Create a base response template
+    response_template = {
+        "id": "chatcmpl-123",
+        "object": "chat.completion.chunk",
+        "created": 1694268190,
+        "model": bot_name,
+        "choices": [{
+            "index": 0,
+            "delta": {
+                "content": "",  # Placeholder, to be filled for each partial response
+                "logprobs": None,
+                "finish_reason": None
+            }
+        }]
+    }
+
+    async for partial in get_bot_response(messages=formatted_messages, bot_name=bot_name, api_key=api_key, 
+                                          base_url=BASE_URL,
+                                          skip_system_prompt=False,
+                                          logit_bias={'24383':-100}):
+
+        # Fill the required field for this partial response
+        response_template["choices"][0]["delta"]["content"] = partial.text
+
+        # Create the SSE formatted string, and then yield
+        yield f"data: {json.dumps(response_template)}\n\n"
+
+    # Send termination sequence
+    response_template["choices"][0]["delta"] = {}  # Empty 'delta' field
+    response_template["choices"][0]["finish_reason"] = "stop"  # Set 'finish_reason' to 'stop'
+
+    yield f"data: {json.dumps(response_template)}\n\ndata: [DONE]\n\n"
+
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request, authorization: str = Header(None)):
+    if not authorization:
+        raise HTTPException(status_code=401, detail="Authorization header is missing")
+
+    api_key = authorization.split(" ")[1]  # Assuming the header follows the standard format: "Bearer $API_KEY"
+    if  api_key != AUTHORIZATION_API_KEY:
+        return HTTPException(status_code=401, detail="Invalid API Key")
+    body = await request.json()
+    
+    # Extract bot_name (model) and messages from the request body
+    bot_name = body.get("model", DEFAULT_MODEL)  # Defaulting to a specific bot if not provided
+    messages = body.get("messages", [])
+    
+    formatted_messages = [ProtocolMessage(role=msg["role"].lower().replace("assistant", "bot"),
+                                      content=msg["content"],
+                                      temperature=msg.get("temperature", 0.95)) 
+                      for msg in messages]
+
+    async def response_stream() -> AsyncGenerator[str, None]:
+        async for response_content in generate_responses(api_key, formatted_messages, bot_name):
+            # Assuming each response_content is a complete "message" response from the bot.
+            # Adjust according to actual response pattern if needed.
+            yield response_content
+        
+    # Stream responses back to the client
+    # Wrap the streamed content to fit the desired response format
+    return StreamingResponse(response_stream(), media_type="application/json")
+
+if __name__ == '__main__':
+    try:
+        import uvloop
+    except ImportError:
+        uvloop = None
+    if uvloop:
+        uvloop.install()
+    uvicorn.run(app, host="0.0.0.0", port=LISTEN_PORT)