Create supercoder.py

supercoder.py

@@ -0,0 +1,412 @@
+"""
+SuperCoder - Unified Application
+All-in-one file containing Gradio UI, API server, tunnel support, and AI logic.
+"""
+import os
+import sys
+import time
+import uuid
+import argparse
+import subprocess
+import traceback
+import requests
+import json
+from pathlib import Path
+from typing import Optional, List, Dict, Any, Generator, Tuple
+from collections import defaultdict
+from functools import partial
+from multiprocessing import Process
+
+import gradio as gr
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import uvicorn
+
+# Import config (only external dependency)
+from config import *
+
+# ============================================================================
+# SERVER MANAGER - llama.cpp server lifecycle
+# ============================================================================
+_server_process = None
+_server_info = {}
+
+def check_server_health() -> bool:
+    try:
+        # Check if Ollama is responding
+        response = requests.get(f"{LLAMA_SERVER_URL}/api/tags", timeout=2)
+        return response.status_code == 200 and len(response.json().get("models", [])) > 0
+    except:
+        return False
+
+def start_llama_server() -> bool:
+    global _server_process, _server_info
+    
+    if _server_process and check_server_health():
+        return True
+    
+    print(f"\n🚀 Starting llama.cpp server on {LLAMA_SERVER_URL}")
+    
+    try:
+        cmd = [
+            LLAMA_SERVER_PATH, "-hf", LLAMA_MODEL,
+            "-c", str(MODEL_CONTEXT_WINDOW),
+            "-t", str(MODEL_THREADS),
+            "-ngl", str(MODEL_GPU_LAYERS),
+            "--host", LLAMA_SERVER_HOST, "--port", str(LLAMA_SERVER_PORT)
+        ]
+        
+        _server_process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        _server_info = {'pid': _server_process.pid, 'url': LLAMA_SERVER_URL}
+        
+        # Wait for ready
+        for _ in range(SERVER_STARTUP_TIMEOUT * 2):
+            if check_server_health():
+                print(f"✅ Server ready (PID: {_server_process.pid})")
+                return True
+            time.sleep(0.5)
+        
+        return False
+    except Exception as e:
+        print(f"❌ Server start failed: {e}")
+        return False
+
+def stop_llama_server():
+    global _server_process
+    if _server_process:
+        _server_process.terminate()
+        _server_process.wait()
+        _server_process = None
+
+def get_llm():
+    return True if check_server_health() else None
+
+def get_model_info():
+    return _server_info.copy()
+
+# ============================================================================
+# SESSION MANAGER - Chat history
+# ============================================================================
+SESSION_STORE = {}
+SESSION_METADATA = defaultdict(dict)
+
+def get_session_id(request: gr.Request) -> str:
+    return request.session_hash
+
+def get_history(session_id: str, create_if_missing: bool = False) -> List[Dict]:
+    if session_id not in SESSION_STORE and create_if_missing:
+        SESSION_STORE[session_id] = []
+    return SESSION_STORE.get(session_id, [])
+
+def add_to_history(session_id: str, role: str, text: str):
+    history = get_history(session_id, create_if_missing=True)
+    history.append({"role": role, "text": text, "timestamp": time.time()})
+
+def clear_history(session_id: str):
+    if session_id in SESSION_STORE:
+        SESSION_STORE[session_id] = []
+
+def convert_history_to_gradio_messages(history: List[Dict]) -> List[Dict]:
+    return [{"role": msg["role"], "content": msg["text"]} for msg in history]
+
+def calculate_safe_max_tokens(history: List[Dict], requested: int, max_context: int) -> int:
+    history_chars = sum(len(msg["text"]) for msg in history)
+    estimated_tokens = history_chars // 4
+    available = max_context - estimated_tokens - SYSTEM_OVERHEAD_TOKENS
+    return max(min(requested, available, SAFE_MAX_TOKENS_CAP), MIN_TOKENS)
+
+def get_recent_history(session_id: str, max_messages: int = 10) -> List[Dict]:
+    history = get_history(session_id)
+    return history[-max_messages:] if len(history) > max_messages else history
+
+def update_session_activity(session_id: str):
+    SESSION_METADATA[session_id]['last_activity'] = time.time()
+
+# ============================================================================
+# GENERATION - AI response generation
+# ============================================================================
+def generate_response_stream(session_id: str, user_message: str, max_tokens: int, 
+                            temperature: float, stream: bool = True) -> Generator[str, None, None]:
+    if not get_llm():
+        yield "⚠️ Server not running"
+        return
+
+    update_session_activity(session_id)
+    recent_history = get_recent_history(session_id, max_messages=6)
+    safe_tokens = calculate_safe_max_tokens(recent_history, max_tokens, MODEL_CONTEXT_WINDOW)
+    
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for msg in recent_history:
+        messages.append({"role": msg["role"], "content": msg["text"]})
+    messages.append({"role": "user", "content": user_message})
+    
+    try:
+        payload = {
+            "messages": messages, "max_tokens": safe_tokens,
+            "temperature": max(0.01, temperature),
+            "top_p": DEFAULT_TOP_P, "stream": stream
+        }
+        
+        if stream:
+            response = requests.post(f"{LLAMA_SERVER_URL}/v1/chat/completions", 
+                                   json=payload, stream=True, timeout=300)
+            full_response = ""
+            for line in response.iter_lines():
+                if line:
+                    line_text = line.decode('utf-8')
+                    if line_text.startswith('data: '):
+                        line_text = line_text[6:]
+                    if line_text.strip() == '[DONE]':
+                        break
+                    try:
+                        chunk = json.loads(line_text)
+                        content = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "")
+                        if content:
+                            full_response += content
+                            yield full_response.strip()
+                    except:
+                        continue
+        else:
+            # Use Ollama API format instead of OpenAI format
+            ollama_payload = {
+                "model": LLAMA_MODEL,
+                "messages": messages,
+                "stream": False
+            }
+            response = requests.post(f"{LLAMA_SERVER_URL}/api/chat", 
+                                   json=ollama_payload, timeout=300)
+            yield response.json()["message"]["content"].strip()
+            
+    except Exception as e:
+        yield f"⚠️ Error: {str(e)}"
+
+# ============================================================================
+# GRADIO UI COMPONENTS
+# ============================================================================
+def create_gradio_interface(error_msg: Optional[str] = None):
+    with gr.Blocks(title=APP_TITLE, theme=gr.themes.Soft(primary_hue=PRIMARY_HUE)) as demo:
+        gr.Markdown(f"# 🤖 {APP_TITLE}\n### {APP_DESCRIPTION}\n---")
+        
+        if error_msg:
+            gr.Markdown(f"⚠️ {error_msg}")
+        
+        with gr.Row():
+            with gr.Column(scale=3):
+                chatbot = gr.Chatbot(label="💬 Conversation", height=CHAT_HEIGHT, 
+                                   type="messages", show_copy_button=True)
+                with gr.Row():
+                    txt_input = gr.Textbox(placeholder="Ask me about code...", 
+                                         show_label=False, scale=5, lines=2)
+                    send_btn = gr.Button("Send 🚀", scale=1, variant="primary")
+            
+            with gr.Column(scale=1):
+                gr.Markdown("### ⚙️ Settings")
+                temp_slider = gr.Slider(0.0, 1.0, value=DEFAULT_TEMPERATURE, step=0.05,
+                                      label="🌡️ Temperature")
+                tokens_slider = gr.Slider(MIN_TOKENS, SAFE_MAX_TOKENS_CAP, 
+                                        value=DEFAULT_MAX_TOKENS, step=128, label="📏 Max Tokens")
+                stream_checkbox = gr.Checkbox(label="⚡ Stream", value=True)
+                clear_btn = gr.Button("🗑️ Clear", variant="stop", size="sm")
+        
+        session_state = gr.State()
+        
+        # Event handlers
+        def handle_message(session_id, msg, temp, tokens, stream, request: gr.Request):
+            session_id = session_id or get_session_id(request)
+            if not msg.strip():
+                return session_id, convert_history_to_gradio_messages(get_history(session_id)), ""
+            
+            add_to_history(session_id, "user", msg)
+            yield session_id, convert_history_to_gradio_messages(get_history(session_id)), ""
+            
+            full_response = ""
+            for partial in generate_response_stream(session_id, msg, tokens, temp, stream):
+                full_response = partial
+                temp_hist = get_history(session_id).copy()
+                temp_hist.append({"role": "assistant", "text": full_response})
+                yield session_id, convert_history_to_gradio_messages(temp_hist), ""
+            
+            add_to_history(session_id, "assistant", full_response)
+            yield session_id, convert_history_to_gradio_messages(get_history(session_id)), ""
+        
+        def handle_clear(session_id, request: gr.Request):
+            session_id = session_id or get_session_id(request)
+            clear_history(session_id)
+            return session_id, [], ""
+        
+        txt_input.submit(handle_message, [session_state, txt_input, temp_slider, tokens_slider, stream_checkbox],
+                        [session_state, chatbot, txt_input])
+        send_btn.click(handle_message, [session_state, txt_input, temp_slider, tokens_slider, stream_checkbox],
+                      [session_state, chatbot, txt_input])
+        clear_btn.click(handle_clear, [session_state], [session_state, chatbot, txt_input])
+    
+    return demo
+
+# ============================================================================
+# FASTAPI SERVER
+# ============================================================================
+api_app = FastAPI(title="SuperCoder API")
+api_app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+
+api_sessions = {}
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatRequest(BaseModel):
+    messages: List[ChatMessage]
+    temperature: Optional[float] = 0.1
+    max_tokens: Optional[int] = 512
+
+class ChatResponse(BaseModel):
+    response: str
+    session_id: str
+
+@api_app.get("/health")
+async def health():
+    return {"status": "ok" if get_llm() else "model_not_loaded"}
+
+@api_app.post("/api/chat", response_model=ChatResponse)
+async def chat(request: ChatRequest):
+    if not get_llm():
+        raise HTTPException(503, "Model not loaded")
+    
+    session_id = str(uuid.uuid4())
+    api_sessions[session_id] = []
+    
+    user_message = request.messages[-1].content
+    api_sessions[session_id].append({"role": "user", "text": user_message})
+    
+    full_response = ""
+    for partial in generate_response_stream(session_id, user_message, request.max_tokens, 
+                                          request.temperature, False):
+        full_response = partial
+    
+    api_sessions[session_id].append({"role": "assistant", "text": full_response})
+    return ChatResponse(response=full_response, session_id=session_id)
+
+def run_api_server():
+    uvicorn.run(api_app, host="0.0.0.0", port=8000, log_level="info")
+
+# ============================================================================
+# TUNNEL SUPPORT
+# ============================================================================
+def start_ngrok_tunnel(port: int = 8000) -> Optional[str]:
+    try:
+        subprocess.run(["which", "ngrok"], capture_output=True, check=True)
+        subprocess.Popen(["ngrok", "http", str(port)], stdout=subprocess.PIPE)
+        time.sleep(3)
+        
+        response = requests.get("http://127.0.0.1:4040/api/tunnels", timeout=5)
+        tunnels = response.json()
+        if tunnels.get("tunnels"):
+            url = tunnels["tunnels"][0]["public_url"]
+            print(f"✅ Tunnel: {url}")
+            return url
+    except:
+        print("❌ ngrok not found. Install: brew install ngrok")
+    return None
+
+def start_cloudflare_tunnel(port: int = 8000) -> Optional[str]:
+    try:
+        subprocess.run(["which", "cloudflared"], capture_output=True, check=True)
+        proc = subprocess.Popen(["cloudflared", "tunnel", "--url", f"http://localhost:{port}"],
+                              stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+        time.sleep(3)
+        
+        for _ in range(30):
+            line = proc.stdout.readline()
+            if "trycloudflare.com" in line:
+                import re
+                urls = re.findall(r'https://[^\s]+\.trycloudflare\.com', line)
+                if urls:
+                    print(f"✅ Tunnel: {urls[0]}")
+                    return urls[0]
+            time.sleep(1)
+    except:
+        print("❌ cloudflared not found. Install: brew install cloudflared")
+    return None
+
+# ============================================================================
+# MAIN LAUNCHER
+# ============================================================================
+def main():
+    parser = argparse.ArgumentParser(description="SuperCoder - All-in-One AI Coding Assistant")
+    parser.add_argument("--mode", choices=["gradio", "api", "both"], default="gradio",
+                       help="Run mode: gradio (UI), api (server), or both")
+    parser.add_argument("--tunnel", choices=["ngrok", "cloudflare"], 
+                       help="Start tunnel for public access")
+    parser.add_argument("--no-server", action="store_true",
+                       help="Don't start llama.cpp server (assume already running)")
+    
+    args = parser.parse_args()
+    
+    print("╔════════════════════════════════════════════════╗")
+    print("║         SuperCoder - Unified Launcher          ║")
+    print("╚════════════════════════════════════════════════╝")
+    
+    # Start llama.cpp server
+    if not args.no_server:
+        success = start_llama_server()
+        error_msg = None if success else "Failed to start llama.cpp server"
+    else:
+        error_msg = None
+    
+    # Run selected mode
+    if args.mode == "gradio":
+        print(f"\n📌 Mode: Gradio UI\n🌐 Access: http://localhost:{SERVER_PORT}\n")
+        demo = create_gradio_interface(error_msg)
+        demo.launch(server_name=SERVER_NAME, server_port=SERVER_PORT)
+    
+    elif args.mode == "api":
+        print(f"\n📌 Mode: API Server\n📡 API: http://localhost:8000/api/chat\n")
+        
+        if args.tunnel:
+            api_proc = Process(target=run_api_server)
+            api_proc.start()
+            time.sleep(3)
+            
+            if args.tunnel == "ngrok":
+                start_ngrok_tunnel(8000)
+            else:
+                start_cloudflare_tunnel(8000)
+            
+            try:
+                api_proc.join()
+            except KeyboardInterrupt:
+                api_proc.terminate()
+        else:
+            run_api_server()
+    
+    elif args.mode == "both":
+        print(f"\n📌 Mode: Both Gradio + API\n🎨 UI: http://localhost:{SERVER_PORT}\n📡 API: http://localhost:8000\n")
+        
+        gradio_proc = Process(target=lambda: create_gradio_interface(error_msg).launch(
+            server_name=SERVER_NAME, server_port=SERVER_PORT))
+        api_proc = Process(target=run_api_server)
+        
+        gradio_proc.start()
+        api_proc.start()
+        
+        if args.tunnel:
+            time.sleep(3)
+            if args.tunnel == "ngrok":
+                start_ngrok_tunnel(8000)
+            else:
+                start_cloudflare_tunnel(8000)
+        
+        try:
+            gradio_proc.join()
+            api_proc.join()
+        except KeyboardInterrupt:
+            gradio_proc.terminate()
+            api_proc.terminate()
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\n👋 Shutting down...")
+        stop_llama_server()