demo/BackendApi.py

import json
import re
import time
from urllib.parse import quote_plus

import g4f.api
import g4f.Provider

from g4f.Provider.base_provider import AsyncGeneratorProvider, ProviderModelMixin
from g4f.typing import AsyncResult, Messages
from g4f.requests import StreamSession
from g4f.providers.response import *
from g4f.errors import ModelNotSupportedError
from g4f import debug
from . import headers

class BackendApi(AsyncGeneratorProvider, ProviderModelMixin):
    url = "https://ahe.hopto.org"
    working = True
    ssl = False

    models = [
        *g4f.Provider.OpenaiAccount.get_models(),
        *g4f.Provider.PerplexityLabs.get_models(),
        "flux",
        "flux-pro",
        "MiniMax-01",
        "Microsoft Copilot",
    ]

    @classmethod
    def get_model(cls, model):
        if "MiniMax" in model:
            model = "MiniMax"
        elif "Copilot" in model:
            model = "Copilot"
        elif "FLUX" in model:
            model = f"flux-{model.split('-')[-1]}"
        elif "flux" in model:
            model = model.split(' ')[-1]
        elif model in g4f.Provider.OpenaiAccount.get_models():
            pass
        elif model in g4f.Provider.PerplexityLabs.get_models():
            pass
        else:
            raise ModelNotSupportedError(f"Model: {model}")
        return model

    @classmethod
    def get_provider(cls, model):
        if model.startswith("MiniMax"):
            return "HailuoAI"
        elif model == "Copilot" or "dall-e" in model:
            return "CopilotAccount"
        elif model in g4f.Provider.OpenaiAccount.get_models():
            return "OpenaiAccount"
        elif model in g4f.Provider.PerplexityLabs.get_models():
            return "PerplexityLabs"
        return None

    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: Messages,
        api_key: str = None,
        proxy: str = None,
        timeout: int = 0,
        **kwargs
    ) -> AsyncResult:
        debug.log(f"{cls.__name__}: {api_key}")
    
        if "dall-e" in model and "prompt" not in kwargs:
            kwargs["prompt"] = messages[-1]["content"]
            messages[-1]["content"] = f"Generate a image: {kwargs['prompt']}"

        async with StreamSession(
            proxy=proxy,
            headers={"Accept": "text/event-stream", **headers},
            timeout=timeout
        ) as session:
            model = cls.get_model(model)
            provider = cls.get_provider(model)
            async with session.post(f"{cls.url}/backend-api/v2/conversation", json={
                **kwargs,
                "model": model,
                "messages": messages,
                "provider": provider
            }, ssl=cls.ssl) as response:
                is_thinking = 0
                async for line in response.iter_lines():
                    response.raise_for_status()
                    data = json.loads(line)
                    data_type = data.pop("type")
                    if data_type == "provider":
                        yield ProviderInfo(**data[data_type])
                        provider = data[data_type]["name"]
                    elif data_type == "conversation":
                        yield JsonConversation(**data[data_type][provider] if provider in data[data_type] else data[data_type][""])
                    elif data_type == "conversation_id":
                        pass
                    elif data_type == "message":
                        yield Exception(data)
                    elif data_type == "preview":
                        yield PreviewResponse(data[data_type])
                    elif data_type == "content":
                        def on_image(match):
                            extension = match.group(3).split(".")[-1].replace("%3F", "?").split("?")[0]
                            extension = ".jpg" if not extension or len(extension) > 4 else f".{extension}"
                            filename = f"{int(time.time())}_{quote_plus(model, '')}.{quote_plus(match.group(1)[:100], '')}{extension}"
                            download_url = f"/download/{filename}?url={cls.url}{match.group(3)}"
                            return f"[![{match.group(1)}]({download_url})](/images/{filename})"
                        if "<think>" in data[data_type]:
                            data[data_type] = data[data_type].split("<think>", 1)
                            yield data[data_type][0]
                            yield Reasoning(data[data_type][1])
                            yield Reasoning(None, "Is thinking...")
                            is_thinking = time.time()
                        if "</think>" in data[data_type]:
                            data[data_type][1] = data[data_type].split("</think>", 1)
                            yield Reasoning(data[data_type][0])
                            yield Reasoning(None, f"Finished in {round(time.time()-is_thinking, 2)} seconds")
                            yield data[data_type][1]
                            is_thinking = 0
                        elif is_thinking:
                            yield Reasoning(data[data_type])
                        else:
                            yield re.sub(r'\[\!\[(.+?)\]\(([^)]+?)\)\]\(([^)]+?)\)', on_image, data[data_type])
                    elif data_type =="synthesize":
                        yield SynthesizeData(**data[data_type])
                    elif data_type == "parameters":
                        yield Parameters(**data[data_type])
                    elif data_type == "usage":
                        yield Usage(**data[data_type])
                    elif data_type == "reasoning":
                        yield Reasoning(**data)
                    elif data_type == "login":
                        pass
                    elif data_type == "title":
                        yield TitleGeneration(data[data_type])
                    elif data_type == "finish":
                        yield FinishReason(data[data_type]["reason"])
                    elif data_type == "log":
                        debug.log(data[data_type])
                    else:
                        debug.log(f"Unknown data: ({data_type}) {data}")