client.py

# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT

import logging
import os
from typing import Optional, Dict, Union, IO, Type, Tuple

import httpx
from pydantic import BaseModel

from cozeloop.internal import consts
from cozeloop.internal.httpclient.auth import Auth
from cozeloop.internal.httpclient.http_client import T, parse_response, HTTPClient
from cozeloop.internal.httpclient.user_agent import user_agent_header

logger = logging.getLogger(__name__)

FileContent = Union[IO[bytes], bytes]
FileType = Tuple[str, FileContent]


class Client:
    def __init__(
            self,
            api_base_url: str,
            http_client: HTTPClient,
            auth: Auth,
            timeout: int = consts.DEFAULT_TIMEOUT,
            upload_timeout: int = consts.DEFAULT_UPLOAD_TIMEOUT,
    ):
        self.api_base_url = api_base_url
        self.http_client = http_client
        self.auth = auth
        self.timeout = timeout
        self.upload_timeout = upload_timeout

    def _build_url(self, path: str) -> str:
        return f"{self.api_base_url}{path}"

    def _set_headers(self, headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
        res = user_agent_header()
        if headers:
            res.update(headers)
        res[consts.AUTHORIZE_HEADER] = f"Bearer {self.auth.token}"

        tt_env = os.getenv("x_tt_env")
        if tt_env:
            res["x-tt-env"] = tt_env
        ppe_env = os.getenv("x_use_ppe")
        if ppe_env:
            res["x-use-ppe"] = "1"

        return res

    def request(
            self,
            path: str,
            method: str,
            response_model: Type[T],
            *,
            params: Optional[Dict[str, str]] = None,
            form: Optional[Dict[str, str]] = None,
            json: Optional[Union[BaseModel, Dict]] = None,
            files: Optional[Dict[str, FileType]] = None,
            headers: Optional[Dict[str, str]] = None,
            timeout: Optional[int] = None,
    ) -> T:
        url = self._build_url(path)
        _headers = self._set_headers(headers)

        _timeout = timeout if timeout is not None else self.timeout

        if isinstance(json, BaseModel):
            json = json.model_dump(by_alias=True)

        try:
            response = self.http_client.request(
                method,
                url,
                params=params,
                data=form,
                json=json,
                files=files,
                headers=_headers,
                timeout=_timeout
            )
        except httpx.HTTPError as e:
            logger.error(f"Http client request failed, path: {path}, err: {e}.")
            raise consts.NetworkError from e

        return parse_response(url, response, response_model)

    def get(
            self,
            path: str,
            response_model: Type[T],
            params: Optional[Dict[str, str]] = None,
    ) -> T:
        return self.request(path, "GET", response_model, params=params,
                            headers={"Content-Type": "application/json"})

    def post(
            self,
            path: str,
            response_model: Type[T],
            json: Union[BaseModel, Dict] = None,
    ) -> T:
        return self.request(path, "POST", response_model, json=json,
                            headers={"Content-Type": "application/json"})

    def upload_file(
            self,
            path: str,
            response_model: Type[T],
            file: FileContent,
            file_name: str,
            form: Dict[str, str],
    ) -> T:
        _file = {"file": (file_name, file)}
        return self.request(path, "POST", response_model, form=form, files=_file, timeout=self.upload_timeout)