main.py

"""
AgentScope AI App — FastAPI Backend
Hugging Face Spaces compatible
"""

import os
import json
import base64
import asyncio
import httpx
from typing import AsyncGenerator, Optional
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel

app = FastAPI(title="AgentScope AI Backend")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ─── Models ───────────────────────────────────────────────────────────────────

class ChatRequest(BaseModel):
    provider: str
    model: str
    messages: list
    api_key: str
    temperature: float = 0.7
    max_tokens: int = 4096
    top_p: float = 0.9
    stream: bool = True
    system_prompt: str = ""

class ModelsRequest(BaseModel):
    provider: str
    api_key: str

# ─── Provider configs ──────────────────────────────────────────────────────────

PROVIDER_CONFIGS = {
    "groq": {
        "base_url": "https://api.groq.com/openai/v1",
        "models_endpoint": "/models",
        "chat_endpoint": "/chat/completions",
        "auth_header": "Bearer",
    },
    "mistral": {
        "base_url": "https://api.mistral.ai/v1",
        "models_endpoint": "/models",
        "chat_endpoint": "/chat/completions",
        "auth_header": "Bearer",
    },
    "sambanova": {
        "base_url": "https://api.sambanova.ai/v1",
        "models_endpoint": "/models",
        "chat_endpoint": "/chat/completions",
        "auth_header": "Bearer",
    },
    "nvidia": {
        "base_url": "https://integrate.api.nvidia.com/v1",
        "models_endpoint": "/models",
        "chat_endpoint": "/chat/completions",
        "auth_header": "Bearer",
    },
    "cohere": {
        "base_url": "https://api.cohere.com/v2",
        "models_endpoint": "/models",
        "chat_endpoint": "/chat",
        "auth_header": "Bearer",
    },
    "gemini": {
        "base_url": "https://generativelanguage.googleapis.com/v1beta",
        "models_endpoint": "/models",
        "chat_endpoint": None,  # Special handling
        "auth_header": "key",
    },
    "openai": {
        "base_url": "https://api.openai.com/v1",
        "models_endpoint": "/models",
        "chat_endpoint": "/chat/completions",
        "auth_header": "Bearer",
    },
    "anthropic": {
        "base_url": "https://api.anthropic.com/v1",
        "models_endpoint": None,
        "chat_endpoint": "/messages",
        "auth_header": "x-api-key",
    },
}

# ─── Fetch models ──────────────────────────────────────────────────────────────

@app.post("/api/models")
async def fetch_models(req: ModelsRequest):
    config = PROVIDER_CONFIGS.get(req.provider)
    if not config:
        raise HTTPException(400, f"Unknown provider: {req.provider}")

    if req.provider == "anthropic":
        return {"models": [
            {"id": "claude-opus-4-5", "name": "Claude Opus 4.5", "context_window": 200000},
            {"id": "claude-sonnet-4-5", "name": "Claude Sonnet 4.5", "context_window": 200000},
            {"id": "claude-haiku-4-5", "name": "Claude Haiku 4.5", "context_window": 200000},
        ]}

    if not config["models_endpoint"]:
        raise HTTPException(400, "No models endpoint for this provider")

    url = config["base_url"] + config["models_endpoint"]
    headers = {}

    if req.provider == "gemini":
        url += f"?key={req.api_key}"
    elif config["auth_header"] == "x-api-key":
        headers["x-api-key"] = req.api_key
        headers["anthropic-version"] = "2023-06-01"
    else:
        headers["Authorization"] = f"Bearer {req.api_key}"

    async with httpx.AsyncClient(timeout=15) as client:
        try:
            resp = await client.get(url, headers=headers)
            if resp.status_code != 200:
                raise HTTPException(resp.status_code, f"Provider error: {resp.text}")
            data = resp.json()
        except httpx.TimeoutException:
            raise HTTPException(504, "Provider timeout")

    # Normalize response across providers
    models = []
    if req.provider == "gemini":
        for m in data.get("models", []):
            name = m.get("name", "").replace("models/", "")
            if "generateContent" in m.get("supportedGenerationMethods", []):
                models.append({
                    "id": name,
                    "name": m.get("displayName", name),
                    "context_window": m.get("inputTokenLimit", 0),
                })
    elif req.provider == "cohere":
        for m in data.get("models", []):
            models.append({
                "id": m.get("name", ""),
                "name": m.get("name", ""),
                "context_window": m.get("context_length", 0),
            })
    else:
        for m in data.get("data", data.get("models", [])):
            models.append({
                "id": m.get("id", m.get("name", "")),
                "name": m.get("id", m.get("name", "")),
                "context_window": m.get("context_window", m.get("max_tokens", 0)),
            })

    return {"models": models}


# ─── Chat / Streaming ──────────────────────────────────────────────────────────

def build_openai_messages(req: ChatRequest) -> list:
    msgs = []
    if req.system_prompt:
        msgs.append({"role": "system", "content": req.system_prompt})
    msgs.extend(req.messages)
    return msgs


async def stream_openai(req: ChatRequest, base_url: str) -> AsyncGenerator[str, None]:
    headers = {"Authorization": f"Bearer {req.api_key}", "Content-Type": "application/json"}
    payload = {
        "model": req.model,
        "messages": build_openai_messages(req),
        "temperature": req.temperature,
        "max_tokens": req.max_tokens,
        "top_p": req.top_p,
        "stream": True,
    }

    async with httpx.AsyncClient(timeout=60) as client:
        async with client.stream("POST", base_url + "/chat/completions", headers=headers, json=payload) as resp:
            if resp.status_code != 200:
                body = await resp.aread()
                yield f"data: {json.dumps({'error': body.decode()})}\n\n"
                return
            async for line in resp.aiter_lines():
                if line.startswith("data: "):
                    chunk = line[6:]
                    if chunk == "[DONE]":
                        yield "data: [DONE]\n\n"
                        return
                    try:
                        j = json.loads(chunk)
                        delta = j["choices"][0]["delta"].get("content", "")
                        if delta:
                            yield f"data: {json.dumps({'content': delta})}\n\n"
                    except Exception:
                        pass


async def stream_anthropic(req: ChatRequest) -> AsyncGenerator[str, None]:
    headers = {
        "x-api-key": req.api_key,
        "anthropic-version": "2023-06-01",
        "Content-Type": "application/json",
    }
    payload = {
        "model": req.model,
        "max_tokens": req.max_tokens,
        "stream": True,
        "messages": req.messages,
    }
    if req.system_prompt:
        payload["system"] = req.system_prompt

    async with httpx.AsyncClient(timeout=60) as client:
        async with client.stream("POST", "https://api.anthropic.com/v1/messages", headers=headers, json=payload) as resp:
            async for line in resp.aiter_lines():
                if line.startswith("data: "):
                    try:
                        j = json.loads(line[6:])
                        if j.get("type") == "content_block_delta":
                            delta = j.get("delta", {}).get("text", "")
                            if delta:
                                yield f"data: {json.dumps({'content': delta})}\n\n"
                        elif j.get("type") == "message_stop":
                            yield "data: [DONE]\n\n"
                            return
                    except Exception:
                        pass


async def stream_gemini(req: ChatRequest) -> AsyncGenerator[str, None]:
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{req.model}:streamGenerateContent?key={req.api_key}&alt=sse"
    contents = []
    for m in req.messages:
        role = "user" if m["role"] == "user" else "model"
        contents.append({"role": role, "parts": [{"text": m["content"]}]})

    payload = {"contents": contents, "generationConfig": {"temperature": req.temperature, "maxOutputTokens": req.max_tokens}}

    async with httpx.AsyncClient(timeout=60) as client:
        async with client.stream("POST", url, json=payload) as resp:
            async for line in resp.aiter_lines():
                if line.startswith("data: "):
                    try:
                        j = json.loads(line[6:])
                        text = j["candidates"][0]["content"]["parts"][0]["text"]
                        if text:
                            yield f"data: {json.dumps({'content': text})}\n\n"
                    except Exception:
                        pass
    yield "data: [DONE]\n\n"


async def stream_cohere(req: ChatRequest) -> AsyncGenerator[str, None]:
    headers = {"Authorization": f"Bearer {req.api_key}", "Content-Type": "application/json"}
    msgs = []
    system_txt = req.system_prompt or ""
    for m in req.messages:
        role = "user" if m["role"] == "user" else "assistant"
        msgs.append({"role": role, "content": m["content"]})

    payload = {"model": req.model, "messages": msgs, "stream": True, "temperature": req.temperature, "max_tokens": req.max_tokens}
    if system_txt:
        payload["system"] = system_txt

    async with httpx.AsyncClient(timeout=60) as client:
        async with client.stream("POST", "https://api.cohere.com/v2/chat", headers=headers, json=payload) as resp:
            async for line in resp.aiter_lines():
                if not line:
                    continue
                try:
                    j = json.loads(line)
                    if j.get("type") == "content-delta":
                        text = j.get("delta", {}).get("message", {}).get("content", {}).get("text", "")
                        if text:
                            yield f"data: {json.dumps({'content': text})}\n\n"
                    elif j.get("type") == "message-end":
                        yield "data: [DONE]\n\n"
                        return
                except Exception:
                    pass


@app.post("/api/chat/stream")
async def chat_stream(req: ChatRequest):
    async def generate():
        try:
            if req.provider == "anthropic":
                async for chunk in stream_anthropic(req):
                    yield chunk
            elif req.provider == "gemini":
                async for chunk in stream_gemini(req):
                    yield chunk
            elif req.provider == "cohere":
                async for chunk in stream_cohere(req):
                    yield chunk
            else:
                config = PROVIDER_CONFIGS.get(req.provider)
                if not config:
                    yield f"data: {json.dumps({'error': 'Unknown provider'})}\n\n"
                    return
                async for chunk in stream_openai(req, config["base_url"]):
                    yield chunk
        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")


# ─── Web search proxy (DuckDuckGo) ────────────────────────────────────────────

@app.get("/api/search")
async def web_search(q: str):
    url = f"https://api.duckduckgo.com/?q={q}&format=json&no_html=1&skip_disambig=1"
    async with httpx.AsyncClient(timeout=10) as client:
        try:
            resp = await client.get(url)
            data = resp.json()
            results = []
            for r in data.get("RelatedTopics", [])[:8]:
                if "Text" in r and "FirstURL" in r:
                    results.append({"title": r["Text"][:120], "url": r["FirstURL"]})
            if data.get("AbstractText"):
                results.insert(0, {"title": data["AbstractText"][:200], "url": data.get("AbstractURL", "")})
            return {"results": results, "query": q}
        except Exception as e:
            raise HTTPException(500, str(e))


# ─── Health ────────────────────────────────────────────────────────────────────

@app.get("/api/health")
async def health():
    return {"status": "ok"}


# ─── Serve static frontend ─────────────────────────────────────────────────────

app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
async def root():
    return FileResponse("static/index.html")

@app.get("/{path:path}")
async def spa(path: str):
    return FileResponse("static/index.html")