litellm/litellm/llms/github_copilot/embedding/transformation.py

"""
GitHub Copilot Embedding API Configuration.

This module provides the configuration for GitHub Copilot's Embedding API.

Implementation based on analysis of the copilot-api project by caozhiyuan:
https://github.com/caozhiyuan/copilot-api
"""

import os
from typing import (
    TYPE_CHECKING,
    Any,
    Final,
    cast,  # noqa: TID251  # narrows untyped request dicts at the session boundary
)

import httpx
from pydantic import ConfigDict, TypeAdapter

from litellm._logging import verbose_logger
from litellm.exceptions import AuthenticationError
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.types.llms.openai import AllEmbeddingInputValues
from litellm.types.utils import EmbeddingResponse
from litellm.utils import convert_to_model_response_object

from ..authenticator import Authenticator
from ..common_utils import (
    DEFAULT_GITHUB_COPILOT_API_BASE,
    GetAPIKeyError,
    get_copilot_default_headers,
    pin_session_authorization,
)
from ..per_user_auth import require_github_copilot_user_session

if TYPE_CHECKING:
    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj

    LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
    LiteLLMLoggingObj = Any

_RESPONSE_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))


class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig):
    """
    Configuration for GitHub Copilot's Embedding API.

    Reference: https://api.githubcopilot.com/embeddings
    """

    def __init__(self) -> None:
        super().__init__()
        self.authenticator = Authenticator()

    def validate_environment(
        self,
        headers: dict,
        model: str,
        messages: list,
        optional_params: dict,
        litellm_params: dict,
        api_key: str | None = None,
        api_base: str | None = None,
    ) -> dict:
        """
        Validate environment and set up headers for GitHub Copilot API.
        """
        user_session: Final = require_github_copilot_user_session(
            cast("dict[str, object]", litellm_params)  # cast-ok: litellm_params arrives as an untyped request dict
        )
        if user_session is not None:
            session_headers: Final = cast(  # cast-ok: both spreads are str-valued header dicts
                "dict[str, str]", {**get_copilot_default_headers(user_session.token), **headers}
            )
            pin_session_authorization(session_headers, user_session.token)
            return session_headers
        try:
            # Get GitHub Copilot API key via OAuth
            api_key = self.authenticator.get_api_key()

            if not api_key:
                raise AuthenticationError(
                    model=model,
                    llm_provider="github_copilot",
                    message="GitHub Copilot API key is required. Please authenticate via OAuth Device Flow.",
                )

            # Get default headers
            default_headers: Final = get_copilot_default_headers(api_key)

            # Merge with existing headers (user's extra_headers take priority)
            merged_headers: Final = {**default_headers, **headers}

            verbose_logger.debug("GitHub Copilot Embedding API: Successfully configured headers for model %s", model)

            return merged_headers

        except GetAPIKeyError as e:
            raise AuthenticationError(
                model=model,
                llm_provider="github_copilot",
                message=str(e),
            )

    def get_complete_url(
        self,
        api_base: str | None,
        api_key: str | None,
        model: str,
        optional_params: dict,
        litellm_params: dict,
        stream: bool | None = None,
    ) -> str:
        """
        Get the complete URL for GitHub Copilot Embedding API endpoint.
        """
        user_session: Final = require_github_copilot_user_session(
            cast("dict[str, object]", litellm_params)  # cast-ok: litellm_params arrives as an untyped request dict
        )
        effective_api_base: Final = (
            (user_session.api_base if user_session is not None else None)
            or api_base
            or self.authenticator.get_api_base()
            or os.getenv("GITHUB_COPILOT_API_BASE")
            or DEFAULT_GITHUB_COPILOT_API_BASE
        )

        return join_api_base(effective_api_base, "embeddings")

    def transform_embedding_request(
        self,
        model: str,
        input: AllEmbeddingInputValues,
        optional_params: dict,
        headers: dict,
    ) -> dict:
        """
        Transform embedding request to GitHub Copilot format.
        """

        # Ensure input is a list
        if isinstance(input, str):
            input = [input]

        # Strip 'github_copilot/' prefix if present
        if model.startswith("github_copilot/"):
            model = model.replace("github_copilot/", "", 1)

        return {
            "model": model,
            "input": input,
            **optional_params,
        }

    def transform_embedding_response(
        self,
        model: str,
        raw_response: httpx.Response,
        model_response: EmbeddingResponse,
        logging_obj: LiteLLMLoggingObj,
        api_key: str | None,
        request_data: dict,
        optional_params: dict,
        litellm_params: dict,
    ) -> EmbeddingResponse:
        """
        Transform embedding response from GitHub Copilot format.
        """
        logging_obj.post_call(original_response=raw_response.text)

        # GitHub Copilot returns standard OpenAI-compatible embedding response
        response_json: Final = _RESPONSE_OBJECT.validate_python(raw_response.json())

        return convert_to_model_response_object(
            response_object=response_json,
            model_response_object=model_response,
            response_type="embedding",
        )

    def get_supported_openai_params(self, model: str) -> list:
        return [
            "timeout",
            "dimensions",
            "encoding_format",
            "user",
        ]

    def map_openai_params(
        self,
        non_default_params: dict,
        optional_params: dict,
        model: str,
        drop_params: bool,
    ) -> dict:
        for param, value in non_default_params.items():
            if param in self.get_supported_openai_params(model):
                optional_params[param] = value
        return optional_params

    def get_error_class(
        self, error_message: str, status_code: int, headers: dict[str, str] | httpx.Headers
    ) -> BaseLLMException:
        from litellm.llms.openai.openai import OpenAIConfig

        return OpenAIConfig().get_error_class(error_message=error_message, status_code=status_code, headers=headers)