tjdify/api/tests/integration_tests/model_runtime/__mock/huggingface_chat.py

import re
from collections.abc import Generator
from typing import Any, Literal, Optional, Union

from _pytest.monkeypatch import MonkeyPatch
from huggingface_hub import InferenceClient
from huggingface_hub.inference._text_generation import (
    Details,
    StreamDetails,
    TextGenerationResponse,
    TextGenerationStreamResponse,
    Token,
)
from huggingface_hub.utils import BadRequestError


class MockHuggingfaceChatClass:
    @staticmethod
    def generate_create_sync(model: str) -> TextGenerationResponse:
        response = TextGenerationResponse(
            generated_text="You can call me Miku Miku o~e~o~",
            details=Details(
                finish_reason="length",
                generated_tokens=6,
                tokens=[Token(id=0, text="You", logprob=0.0, special=False) for i in range(0, 6)],
            ),
        )

        return response

    @staticmethod
    def generate_create_stream(model: str) -> Generator[TextGenerationStreamResponse, None, None]:
        full_text = "You can call me Miku Miku o~e~o~"

        for i in range(0, len(full_text)):
            response = TextGenerationStreamResponse(
                token=Token(id=i, text=full_text[i], logprob=0.0, special=False),
            )
            response.generated_text = full_text[i]
            response.details = StreamDetails(finish_reason="stop_sequence", generated_tokens=1)

            yield response

    def text_generation(
        self: InferenceClient, prompt: str, *, stream: Literal[False] = ..., model: Optional[str] = None, **kwargs: Any
    ) -> Union[TextGenerationResponse, Generator[TextGenerationStreamResponse, None, None]]:
        # check if key is valid
        if not re.match(r"Bearer\shf\-[a-zA-Z0-9]{16,}", self.headers["authorization"]):
            raise BadRequestError("Invalid API key")

        if model is None:
            raise BadRequestError("Invalid model")

        if stream:
            return MockHuggingfaceChatClass.generate_create_stream(model)
        return MockHuggingfaceChatClass.generate_create_sync(model)