Update app.py

app.py

@@ -1,107 +1,126 @@
+import json
 import os
-import gradio as gr
-from openai import OpenAI
-from typing import List, Tuple
+import fastapi
+from fastapi.responses import StreamingResponse
+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
-client = OpenAI(
+oai_client = AsyncOpenAI(
     base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENROUTER_API_KEY")
+    api_key=OPENROUTER_API_KEY
 )
 
-# Available models
-MODELS = [
-    "nousresearch/hermes-3-llama-3.1-405b",
-    "anthropic/claude-3-opus",
-    "google/gemma-7b",
-    "mistralai/mixtral-8x7b"
-]
+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  # Default to True for ElevenLabs
+    user_id: Optional[str] = None
+    extra_headers: Optional[Dict[str, str]] = None
+    extra_body: Optional[Dict[str, Any]] = None
 
-def respond(
-    message: str,
-    history: List[Tuple[str, str]],
-    system_message: str,
-    model: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-) -> str:
-    # Construct messages array with system message and history
-    messages = [{"role": "system", "content": system_message}]
-    
-    # Add conversation history
-    for user_msg, assistant_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
-    
-    # Add current message
-    messages.append({"role": "user", "content": message})
-    
+@app.post("/v1/chat/completions")
+async def create_chat_completion(request: ChatCompletionRequest) -> StreamingResponse:
     try:
-        completion = client.chat.completions.create(
-            extra_headers={
-                "HTTP-Referer": os.getenv("SITE_URL", "https://huggingface.co/spaces"),
-                "X-Title": os.getenv("SITE_NAME", "OpenRouter Gradio Interface"),
-            },
-            model=model,
-            messages=messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            stream=True  # Enable streaming
-        )
+        # Prepare the request for OpenRouter
+        oai_request = request.dict(exclude_none=True)
+        
+        # Remove fields that OpenRouter doesn't expect
+        if "user_id" in oai_request:
+            oai_request["user"] = oai_request.pop("user_id")
         
-        # Stream the response
-        response = ""
-        for chunk in completion:
-            if chunk.choices[0].delta.content is not None:
-                response += chunk.choices[0].delta.content
-                yield response
+        # 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"
+        
+        # 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:
-        yield f"Error: {str(e)}"
+        logging.error(f"Request error: {str(e)}")
+        raise fastapi.HTTPException(status_code=500, detail=str(e))
 
-# Create Gradio interface
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(
-            value="You are a helpful and friendly AI assistant.",
-            label="System message"
-        ),
-        gr.Dropdown(
-            choices=MODELS,
-            value=MODELS[0],
-            label="Model"
-        ),
-        gr.Slider(
-            minimum=1,
-            maximum=2048,
-            value=512,
-            step=1,
-            label="Max new tokens"
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=2.0,
-            value=0.7,
-            step=0.1,
-            label="Temperature"
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)"
-        ),
-    ],
-    title="OpenRouter Chat Interface",
-    description="Chat with various AI models through OpenRouter API"
-)
+# 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",
+            }
+        ]
+    }
 
-# Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    uvicorn.run(app, host="0.0.0.0", port=8013)