simple_llm_server.py

"""
Simple Local LLM Server Simulator
For testing the 3-Agent code assistant system
"""

from flask import Flask, request, jsonify
import json
import random
import time
from typing import Dict, Any, List

app = Flask(__name__)

# Simulate different LLM responses
LLM_RESPONSES = {
    "coder": {
        "system_prompt": "You are a coding assistant. Generate code solutions.",
        "examples": [
            {
                "prefix": "A simple function to add two numbers",
                "imports": "",
                "code": "def add_numbers(a, b):\n    return a + b"
            },
            {
                "prefix": "A RAG chain implementation using LCEL",
                "imports": "from langchain_core.prompts import ChatPromptTemplate\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.output_parsers import StrOutputParser",
                "code": "# Build RAG chain\nrag_chain = ChatPromptTemplate.from_template('Answer: {question}') | ChatOpenAI() | StrOutputParser()"
            }
        ]
    },
    "checker": {
        "system_prompt": "You validate code.",
        "examples": [
            {"validation": "All checks passed"},
            {"validation": "Import error detected"}
        ]
    },
    "debugger": {
        "system_prompt": "You fix code errors.",
        "examples": [
            {
                "prefix": "Fixed version with error handling",
                "imports": "import math",
                "code": "def fixed_function(x):\n    if x < 0:\n        raise ValueError('x must be non-negative')\n    return math.sqrt(x)"
            }
        ]
    }
}


def generate_coder_response(messages: List[Dict[str, str]]) -> Dict[str, Any]:
    """Generate responses for coder model"""
    # Extract user question from messages
    user_content = ""
    for msg in reversed(messages):
        if msg.get("role") == "user":
            user_content = msg.get("content", "")
            break

    # Select response based on question type
    if "add" in user_content.lower() or "sum" in user_content.lower():
        code_data = {
            "prefix": "A simple function to add two numbers",
            "imports": "",
            "code": "def add_numbers(a, b):\n    \"\"\"Add two numbers\"\"\"\n    return a + b"
        }
    elif "rag" in user_content.lower() or "retrieval" in user_content.lower():
        code_data = {
            "prefix": "A RAG chain implementation using LCEL",
            "imports": "from langchain_core.prompts import ChatPromptTemplate\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.output_parsers import StrOutputParser",
            "code": """# RAG chain implementation
prompt = ChatPromptTemplate.from_template("Answer based on context: {context}\nQuestion: {question}")
model = ChatOpenAI()
parser = StrOutputParser()

rag_chain = prompt | model | parser"""
        }
    elif "error" in user_content.lower() or "fix" in user_content.lower():
        code_data = {
            "prefix": "Fixed version with proper error handling",
            "imports": "import math",
            "code": "def fixed_function(x):\n    \"\"\"Fixed function with validation\"\"\"\n    if not isinstance(x, (int, float)):\n        raise TypeError('x must be a number')\n    return math.sqrt(abs(x))"
        }
    else:
        # Default response
        code_data = {
            "prefix": "A simple Python function",
            "imports": "",
            "code": "def hello_world():\n    print('Hello, World!')"
        }

    return {
        "choices": [{
            "message": {
                "content": json.dumps(code_data),
                "role": "assistant"
            },
            "finish_reason": "stop",
            "index": 0
        }]
    }


def generate_debugger_response(messages: List[Dict[str, str]]) -> Dict[str, Any]:
    """Generate responses for debugger model"""
    # Extract error information
    user_content = ""
    for msg in reversed(messages):
        if msg.get("role") == "user":
            user_content = msg.get("content", "")
            break

    # Generate fixed code
    if "import" in user_content.lower():
        code_data = {
            "prefix": "Fixed imports and code",
            "imports": "import math  # Fixed import",
            "code": "def fixed_function():\n    return math.pi"
        }
    elif "syntax" in user_content.lower():
        code_data = {
            "prefix": "Fixed syntax errors",
            "imports": "",
            "code": "def fixed_function():\n    return 'syntax fixed'"
        }
    else:
        code_data = {
            "prefix": "Fixed version with error handling",
            "imports": "",
            "code": "def fixed_function(x):\n    try:\n        return x / 1\n    except:\n        return 0"
        }

    return {
        "choices": [{
            "message": {
                "content": json.dumps(code_data),
                "role": "assistant"
            },
            "finish_reason": "stop",
            "index": 0
        }]
    }


def generate_checker_response(messages: List[Dict[str, str]]) -> Dict[str, Any]:
    """生成checker模型的响应"""
    # 简单返回成功消息
    return {
        "choices": [{
            "message": {
                "content": "Code validation completed",
                "role": "assistant"
            },
            "finish_reason": "stop",
            "index": 0
        }]
    }


@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    """Handle chat completion requests"""
    try:
        data = request.get_json()
        messages = data.get('messages', [])
        model = data.get('model', 'default')

        # Simulate processing delay
        time.sleep(random.uniform(0.1, 0.5))

        # Generate different responses based on model type
        if 'coder' in model.lower() or 'code' in model.lower():
            response = generate_coder_response(messages)
        elif 'debug' in model.lower() or 'fix' in model.lower():
            response = generate_debugger_response(messages)
        elif 'check' in model.lower() or 'valid' in model.lower():
            response = generate_checker_response(messages)
        else:
            # Default response
            response = generate_coder_response(messages)

        # Add response metadata
        response.update({
            "model": model,
            "object": "chat.completion",
            "created": int(time.time()),
            "usage": {
                "prompt_tokens": len(str(messages)),
                "completion_tokens": len(str(response)),
                "total_tokens": len(str(messages)) + len(str(response))
            }
        })

        return jsonify(response)

    except Exception as e:
        return jsonify({
            "error": {
                "message": str(e),
                "type": "internal_error",
                "code": "internal_error"
            }
        }), 500


@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    return jsonify({
        "status": "healthy",
        "service": "simple-llm-server",
        "timestamp": int(time.time())
    })


@app.route('/models', methods=['GET'])
def list_models():
    """List available models"""
    return jsonify({
        "object": "list",
        "data": [
            {
                "id": "coder-model",
                "object": "model",
                "created": int(time.time()),
                "owned_by": "local-server"
            },
            {
                "id": "checker-model",
                "object": "model",
                "created": int(time.time()),
                "owned_by": "local-server"
            },
            {
                "id": "debugger-model",
                "object": "model",
                "created": int(time.time()),
                "owned_by": "local-server"
            }
        ]
    })


if __name__ == '__main__':
    print("🚀 Starting Simple LLM Server Simulator...")
    print("📡 Endpoints:")
    print("   - POST /v1/chat/completions")
    print("   - GET  /health")
    print("   - GET  /models")
    print("\nSupported Models:")
    print("   - coder-model (Code Generation)")
    print("   - checker-model (Code Validation)")
    print("   - debugger-model (Error Fixing)")
    print("\n📝 Usage Example:")
    print("   curl -X POST http://localhost:8000/v1/chat/completions \\")
    print("     -H 'Content-Type: application/json' \\")
    print("     -d '{\"model\": \"coder-model\", \"messages\": [{\"role\": \"user\", \"content\": \"Generate add function\"}]}'")
    print("\n🔧 Press Ctrl+C to stop service")

    # Start server
    app.run(host='0.0.0.0', port=8000, debug=False, threaded=True)