app.py

import os
import asyncio
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import StreamingResponse
from openai import AsyncOpenAI
from collections import defaultdict

app = FastAPI()

# Define available models
AVAILABLE_MODELS = {
    "openai/gpt-4.1": "OpenAI GPT-4.1",
    "openai/gpt-4.1-mini": "OpenAI GPT-4.1-mini",
    "openai/gpt-4.1-nano": "OpenAI GPT-4.1-nano",
    "openai/gpt-4o": "OpenAI GPT-4o",
    "openai/gpt-4o-mini": "OpenAI GPT-4o mini",
    "openai/o4-mini": "OpenAI o4-mini",
    "microsoft/MAI-DS-R1": "MAI-DS-R1",
    "microsoft/Phi-3.5-MoE-instruct": "Phi-3.5-MoE instruct (128k)",
    "microsoft/Phi-3.5-mini-instruct": "Phi-3.5-mini instruct (128k)",
    "microsoft/Phi-3.5-vision-instruct": "Phi-3.5-vision instruct (128k)",
    "microsoft/Phi-3-medium-128k-instruct": "Phi-3-medium instruct (128k)",
    "microsoft/Phi-3-medium-4k-instruct": "Phi-3-medium instruct (4k)",
    "microsoft/Phi-3-mini-128k-instruct": "Phi-3-mini instruct (128k)",
    "microsoft/Phi-3-small-128k-instruct": "Phi-3-small instruct (128k)",
    "microsoft/Phi-3-small-8k-instruct": "Phi-3-small instruct (8k)",
    "microsoft/Phi-4": "Phi-4",
    "microsoft/Phi-4-mini-instruct": "Phi-4-mini-instruct",
    "microsoft/Phi-4-multimodal-instruct": "Phi-4-multimodal-instruct",
    "ai21-labs/AI21-Jamba-1.5-Large": "AI21 Jamba 1.5 Large",
    "ai21-labs/AI21-Jamba-1.5-Mini": "AI21 Jamba 1.5 Mini",
    "mistral-ai/Codestral-2501": "Codestral 25.01",
    "cohere/Cohere-command-r": "Cohere Command R",
    "cohere/Cohere-command-r-08-2024": "Cohere Command R 08-2024",
    "cohere/Cohere-command-r-plus": "Cohere Command R+",
    "cohere/Cohere-command-r-plus-08-2024": "Cohere Command R+ 08-2024",
    "deepseek/DeepSeek-R1": "DeepSeek-R1",
    "deepseek/DeepSeek-V3-0324": "DeepSeek-V3-0324",
    "meta/Llama-3.2-11B-Vision-Instruct": "Llama-3.2-11B-Vision-Instruct",
    "meta/Llama-3.2-90B-Vision-Instruct": "Llama-3.2-90B-Vision-Instruct",
    "meta/Llama-3.3-70B-Instruct": "Llama-3.3-70B-Instruct",
    "meta/Llama-4-Maverick-17B-128E-Instruct-FP8": "Llama 4 Maverick 17B 128E Instruct FP8",
    "meta/Llama-4-Scout-17B-16E-Instruct": "Llama 4 Scout 17B 16E Instruct",
    "meta/Meta-Llama-3.1-405B-Instruct": "Meta-Llama-3.1-405B-Instruct",
    "meta/Meta-Llama-3.1-70B-Instruct": "Meta-Llama-3.1-70B-Instruct",
    "meta/Meta-Llama-3.1-8B-Instruct": "Meta-Llama-3.1-8B-Instruct",
    "meta/Meta-Llama-3-70B-Instruct": "Meta-Llama-3-70B-Instruct",
    "meta/Meta-Llama-3-8B-Instruct": "Meta-Llama-3-8B-Instruct",
    "mistral-ai/Ministral-3B": "Ministral 3B",
    "mistral-ai/Mistral-Large-2411": "Mistral Large 24.11",
    "mistral-ai/Mistral-Nemo": "Mistral Nemo",
    "mistral-ai/Mistral-large-2407": "Mistral Large (2407)",
    "mistral-ai/Mistral-small": "Mistral Small",
    "cohere/cohere-command-a": "Cohere Command A",
    "core42/jais-30b-chat": "JAIS 30b Chat",
    "mistral-ai/mistral-small-2503": "Mistral Small 3.1"
}

# In-memory chat history and locks
chat_histories = defaultdict(list)
chat_locks = defaultdict(asyncio.Lock)
MAX_HISTORY = 100

# Streaming AI generation
async def generate_ai_response(chat_id: str, model: str):
    token = os.getenv("GITHUB_TOKEN")
    if not token:
        yield "Error: GitHub token not configured"
        raise HTTPException(status_code=500, detail="GitHub token not configured")

    if model not in AVAILABLE_MODELS:
        yield f"Error: Invalid model {model}"
        raise HTTPException(status_code=400, detail="Invalid model")

    client = AsyncOpenAI(
        base_url="https://models.github.ai/inference",
        api_key=token
    )

    try:
        async with chat_locks[chat_id]:
            stream = await asyncio.wait_for(
                client.chat.completions.create(
                    messages=chat_histories[chat_id],
                    model=model,
                    temperature=1.0,
                    top_p=1.0,
                    stream=True
                ),
                timeout=60
            )

        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                yield content
                async with chat_locks[chat_id]:
                    chat_histories[chat_id].append({"role": "assistant", "content": content})
                    chat_histories[chat_id] = chat_histories[chat_id][-MAX_HISTORY:]

    except asyncio.TimeoutError:
        yield "Error: Response timed out."
        raise HTTPException(status_code=504, detail="Timeout")
    except Exception as e:
        yield f"Error: {str(e)}"
        raise HTTPException(status_code=500, detail="AI generation failed")

# POST /generate
@app.post("/generate")
async def generate_response(
    chat_id: str = Query(..., description="Unique chat ID"),
    prompt: str = Query(..., description="User prompt"),
    model: str = Query("openai/gpt-4.1-mini", description="Model to use")
):
    if not prompt.strip():
        raise HTTPException(status_code=400, detail="Prompt is required")

    async with chat_locks[chat_id]:
        chat_histories[chat_id].append({"role": "user", "content": prompt})
        chat_histories[chat_id] = chat_histories[chat_id][-MAX_HISTORY:]

    return StreamingResponse(
        generate_ai_response(chat_id, model),
        media_type="text/event-stream"
    )

# POST /reset
@app.post("/reset")
async def reset_chat(chat_id: str = Query(..., description="Chat ID to reset")):
    async with chat_locks[chat_id]:
        chat_histories[chat_id].clear()
    return {"message": f"Chat history for {chat_id} cleared."}

# For ASGI hosting
def get_app():
    return app