File size: 5,018 Bytes
27255de 2a9bd55 27255de ce727ee 27255de 79fa11b 2a9bd55 27255de 2a9bd55 27255de 2a9bd55 79fa11b 27255de ce727ee 27255de 79fa11b ce727ee a7c8fda ce727ee a7c8fda ce727ee a7c8fda ce727ee 27255de 5e3f932 2a9bd55 27255de 420388d 27255de 2a9bd55 27255de ce727ee 27255de 2a9bd55 27255de 2a9bd55 27255de 79fa11b 27255de 79fa11b ce727ee 79fa11b 95e4375 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
import json
import os
import fastapi
from fastapi.responses import StreamingResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from openai import AsyncOpenAI
import uvicorn
import logging
from dotenv import load_dotenv
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
# Load environment variables
load_dotenv()
# Retrieve API key from environment
OPENROUTER_API_KEY = os.getenv('OPENROUTER_API_KEY')
if not OPENROUTER_API_KEY:
raise ValueError("OPENROUTER_API_KEY not found in environment variables")
# Setup FastAPI app
app = fastapi.FastAPI()
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Initialize OpenRouter client
oai_client = AsyncOpenAI(
base_url="https://openrouter.ai/api/v1",
api_key=OPENROUTER_API_KEY
)
class Message(BaseModel):
role: str
content: str
class ChatCompletionRequest(BaseModel):
messages: List[Message]
model: str
temperature: Optional[float] = 0.7
max_tokens: Optional[int] = None
stream: Optional[bool] = True
user_id: Optional[str] = None
extra_headers: Optional[Dict[str, str]] = None
extra_body: Optional[Dict[str, Any]] = None
# Root path handler
@app.get("/", response_class=HTMLResponse)
async def root():
return """
<html>
<head>
<title>Proxy Server</title>
</head>
<body>
<h1>OpenRouter Proxy Server</h1>
<p>V0.1</p>
<ul>
<li>Chat completions endpoint</li>
<li>List available models</li>
<li>Health check</li>
</ul>
</body>
</html>
"""
# Robots.txt handler
@app.get("/robots.txt")
async def robots():
return "User-agent: *\nDisallow: /"
@app.post("/v1/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest) -> StreamingResponse:
try:
# Prepare the request for OpenRouter
oai_request = request.dict(exclude_none=True)
# Add system message at the beginning of messages array
system_message = {
"role": "system",
"content": "You are a very strict girlfriend. Almost like a mistress."
}
oai_request["messages"].insert(0, system_message)
# Remove fields that OpenRouter doesn't expect
if "user_id" in oai_request:
oai_request["user"] = oai_request.pop("user_id")
# Add OpenRouter specific headers
extra_headers = {
"HTTP-Referer": os.getenv("SITE_URL", "https://huggingface.co/spaces"),
"X-Title": os.getenv("SITE_NAME", "ElevenLabs-OpenRouter Bridge"),
}
oai_request["extra_headers"] = extra_headers
# Ensure the model is an OpenRouter model
if not oai_request["model"].startswith("openai/") and "/" not in oai_request["model"]:
oai_request["model"] = "nousresearch/hermes-3-llama-3.1-405b"
logging.info(f"Sending request to OpenRouter with model: {oai_request['model']}")
# Create the chat completion
chat_completion_coroutine = await oai_client.chat.completions.create(**oai_request)
async def event_stream():
try:
async for chunk in chat_completion_coroutine:
# Convert the ChatCompletionChunk to a dictionary
chunk_dict = chunk.model_dump()
yield f"data: {json.dumps(chunk_dict)}\n\n"
yield "data: [DONE]\n\n"
except Exception as e:
logging.error(f"Streaming error: {str(e)}")
yield f"data: {json.dumps({'error': str(e)})}\n\n"
return StreamingResponse(event_stream(), media_type="text/event-stream")
except Exception as e:
logging.error(f"Request error: {str(e)}")
raise fastapi.HTTPException(status_code=500, detail=str(e))
# Health check endpoint
@app.get("/health")
async def health_check():
return {"status": "healthy"}
# Models endpoint
@app.get("/v1/models")
async def list_models():
return {
"data": [
{
"id": "nousresearch/hermes-3-llama-3.1-405b",
"object": "model",
"created": 1677610602,
"owned_by": "openrouter",
},
{
"id": "anthropic/claude-3-opus",
"object": "model",
"created": 1677610602,
"owned_by": "openrouter",
},
{
"id": "mistralai/mixtral-8x7b",
"object": "model",
"created": 1677610602,
"owned_by": "openrouter",
}
]
}
# Configure logging
logging.basicConfig(level=logging.INFO)
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860) |