Update app.py

app.py

@@ -6,10 +6,9 @@ from pydantic import BaseModel
 import httpx
 import secrets
 import os
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Optional
 
-# Initialize FastAPI app
-app = FastAPI(title="AI Chat Application")
+# Initialize FastAPI app\ napp = FastAPI(title="AI Chat Application")
 
 # Mount static files
 app.mount("/static", StaticFiles(directory="static"), name="static")
@@ -19,18 +18,28 @@ security = HTTPBasic()
 ADMIN_USERNAME = os.environ.get("ADMIN_USERNAME", "admin")
 ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "password")  # Change in production!
 
-# Store active chats and settings in memory
-active_chats = {}
-chat_history = {}
-rules = [
+def get_admin_user(credentials: HTTPBasicCredentials = Depends(security)):
+    valid_user = secrets.compare_digest(credentials.username, ADMIN_USERNAME)
+    valid_pass = secrets.compare_digest(credentials.password, ADMIN_PASSWORD)
+    if not (valid_user and valid_pass):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect username or password",
+            headers={"WWW-Authenticate": "Basic"},
+        )
+    return credentials.username
+
+# In-memory storage
+chat_history: Dict[str, List[Dict[str, str]]] = {}
+rules: List[str] = [
     "Be respectful and helpful",
     "Do not generate harmful content",
     "Provide accurate information"
 ]
-available_models = ["mistral-large-latest", "gemini", "openai-xlarge"]
-current_model = "mistral-large-latest"
+available_models: List[str] = ["mistral-large-latest", "gemini", "openai-xlarge"]
+current_model: str = available_models[0]
 
-# API endpoint for the chat
+# Request models
 class ChatMessage(BaseModel):
     message: str
     chat_id: Optional[str] = None
@@ -42,51 +51,38 @@ class Rule(BaseModel):
 class ModelRequest(BaseModel):
     model: str
 
-def get_admin_user(credentials: HTTPBasicCredentials = Depends(security)):
-    is_username_correct = secrets.compare_digest(credentials.username, ADMIN_USERNAME)
-    is_password_correct = secrets.compare_digest(credentials.password, ADMIN_PASSWORD)
-    if not (is_username_correct and is_password_correct):
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Incorrect username or password",
-            headers={"WWW-Authenticate": "Basic"},
-        )
-    return credentials.username
-
+# Root endpoint to serve index.html
 @app.get("/", response_class=HTMLResponse)
 async def read_root():
-    with open("index.html", "r") as f:
-        return f.read()
-
-@app.post("/api/chat")
+    try:
+        with open("index.html", "r") as f:
+            content = f.read()
+        return HTMLResponse(content=content)
+    except FileNotFoundError:
+        raise HTTPException(status_code=404, detail="index.html not found")
+
+# Chat API endpoint
+@app.post("/api/chat", response_class=JSONResponse)
 async def process_chat(chat_request: ChatMessage):
-    # Generate a new chat ID if not provided
     chat_id = chat_request.chat_id or secrets.token_hex(8)
-    
-    # Initialize chat history if it doesn't exist
     if chat_id not in chat_history:
         chat_history[chat_id] = []
-    
-    # Store the message in chat history
     chat_history[chat_id].append({"role": "user", "content": chat_request.message})
-    
-    # Apply rules to message (simplified implementation)
+
+    # Content rules
     for rule in rules:
         if rule.startswith("Block:") and rule[6:].lower() in chat_request.message.lower():
-            return {"response": "I cannot respond to that due to content rules.", "chat_id": chat_id}
-    
-    # Select model to use
+            return JSONResponse({"response": "I cannot respond to that due to content rules.", "chat_id": chat_id})
+
     model_to_use = chat_request.model if chat_request.model in available_models else current_model
-    
-    # Prepare request to external API
+
     api_request = {
         "model": model_to_use,
         "messages": chat_history[chat_id],
         "temperature": 0.7,
         "max_tokens": 1000
     }
-    
-    # Call external API
+
     try:
         async with httpx.AsyncClient(timeout=60.0) as client:
             response = await client.post(
@@ -95,23 +91,36 @@ async def process_chat(chat_request: ChatMessage):
                 headers={"Authorization": "sigma"}
             )
             response_data = response.json()
-            
-            if "choices" in response_data and len(response_data["choices"]) > 0:
-                ai_message = response_data["choices"][0]["message"]["content"]
+            if choices := response_data.get("choices"):
+                ai_message = choices[0]["message"]["content"]
                 chat_history[chat_id].append({"role": "assistant", "content": ai_message})
                 return {"response": ai_message, "chat_id": chat_id, "model": model_to_use}
             else:
                 error_msg = "No response from AI model"
                 chat_history[chat_id].append({"role": "assistant", "content": error_msg})
                 return {"response": error_msg, "chat_id": chat_id}
-    
-    except Exception as e:
+    except httpx.RequestError as e:
         error_msg = f"Error communicating with AI service: {str(e)}"
         chat_history[chat_id].append({"role": "assistant", "content": error_msg})
         return {"response": error_msg, "chat_id": chat_id}
 
+# Admin panel
 @app.get("/admin", response_class=HTMLResponse)
 async def admin_panel(username: str = Depends(get_admin_user)):
+    # Precompute HTML fragments
+    model_options = "".join(
+        f'<option value="{model}"{" selected" if model == current_model else ""}>{model}</option>'
+        for model in available_models
+    )
+    rules_list = "".join(
+        f'<li>{rule} <button onclick="deleteRule({i})">Delete</button></li>'
+        for i, rule in enumerate(rules)
+    )
+    chats_list = "".join(
+        f'''<li><a href="#" onclick="viewChat('{chat_id}')">{chat_id}</a></li>'''
+        for chat_id in chat_history
+    )
+
     html_content = f"""
     <!DOCTYPE html>
     <html>
@@ -122,30 +131,27 @@ async def admin_panel(username: str = Depends(get_admin_user)):
     </head>
     <body class="admin-page">
         <h1>Admin Panel</h1>
-        
         <section class="admin-section">
             <h2>Active Model: <span id="current-model">{current_model}</span></h2>
             <select id="model-selector">
-                {"".join(f'<option value="{model}">{model}</option>' for model in available_models)}
+                {model_options}
             </select>
             <button onclick="updateModel()">Update Model</button>
         </section>
-        
         <section class="admin-section">
             <h2>Rules</h2>
             <ul id="rules-list">
-                {"".join(f'<li>{rule} <button onclick="deleteRule({i})">Delete</button></li>' for i, rule in enumerate(rules))}
+                {rules_list}
             </ul>
             <div>
                 <input type="text" id="new-rule" placeholder="Add new rule">
                 <button onclick="addRule()">Add Rule</button>
             </div>
         </section>
-        
         <section class="admin-section">
             <h2>Active Chats</h2>
             <ul id="active-chats">
-                {"".join(f'<li><a href="#" onclick="viewChat(\'{chat_id}\')">{chat_id}</a></li>' for chat_id in chat_history)}
+                {chats_list}
             </ul>
             <div id="chat-viewer"></div>
         </section>
@@ -154,23 +160,25 @@ async def admin_panel(username: str = Depends(get_admin_user)):
     """
     return HTMLResponse(content=html_content)
 
-@app.get("/api/rules")
+# Rules management
+@app.get("/api/rules", response_class=JSONResponse)
 async def get_rules(username: str = Depends(get_admin_user)):
     return {"rules": rules}
 
-@app.post("/api/rules")
+@app.post("/api/rules", response_class=JSONResponse)
 async def add_rule(rule_request: Rule, username: str = Depends(get_admin_user)):
     rules.append(rule_request.rule)
     return {"status": "success", "rules": rules}
 
-@app.delete("/api/rules/{rule_id}")
+@app.delete("/api/rules/{rule_id}", response_class=JSONResponse)
 async def delete_rule(rule_id: int, username: str = Depends(get_admin_user)):
     if 0 <= rule_id < len(rules):
         rules.pop(rule_id)
         return {"status": "success", "rules": rules}
     raise HTTPException(status_code=404, detail="Rule not found")
 
-@app.post("/api/model")
+# Model selection endpoint
+@app.post("/api/model", response_class=JSONResponse)
 async def set_model(model_request: ModelRequest, username: str = Depends(get_admin_user)):
     global current_model
     if model_request.model in available_models:
@@ -178,11 +186,12 @@ async def set_model(model_request: ModelRequest, username: str = Depends(get_adm
         return {"status": "success", "model": current_model}
     raise HTTPException(status_code=400, detail="Invalid model")
 
-@app.get("/api/chats")
+# Chat history retrieval
+@app.get("/api/chats", response_class=JSONResponse)
 async def get_chats(username: str = Depends(get_admin_user)):
     return {"chats": list(chat_history.keys())}
 
-@app.get("/api/chats/{chat_id}")
+@app.get("/api/chats/{chat_id}", response_class=JSONResponse)
 async def get_chat_history(chat_id: str, username: str = Depends(get_admin_user)):
     if chat_id in chat_history:
         return {"history": chat_history[chat_id]}
@@ -190,5 +199,6 @@ async def get_chat_history(chat_id: str, username: str = Depends(get_admin_user)
 
 # Run the app (for local development)
 if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=port)