Querying Underlying LLM-as-a-Service through DSS LLM Mesh

A member from our team looked into the source code of the dataikuapi module in our /usr/lib/python3.x/site-packages/ folder. We were able to extract the core functionality of the dataikuapi call to the LLM mesh with just the requests package. As a result we created our own DSSOpenAI class that allows us to preserve the SDK functionality of the OpenAI library while still routing the request through the LLM Mesh. We currently use llama_index, but replace the import with which ever gateway/session object you use.

We wanted to share our solution here for posterity. Perhaps it may be helpful to some of your other clients, @ClemenceB
.

from llama_index.llms.openai import OpenAI
from typing import (
    List,
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Sequence,
)
from requests.auth import HTTPBasicAuth
from requests import Session
import json
import httpx
from llama_index.core.base.llms.types import (
    ChatMessage,
    ChatResponse,
    CompletionResponse,
    MessageRole,
)
from llama_index.core.callbacks import CallbackManager
from llama_index.core.types import BaseOutputParser, PydanticProgramMode
from llama_index.llms.openai.base import llm_retry_decorator


class DSSOpenAI(OpenAI):
    # TODO Add async and streaming logic as needed

    dss_project: str = ""

    dss_model_deployment: str = ""

    def __init__(
        self,
        model: str = None,
        temperature: float = 1.0,
        max_tokens: Optional[int] = None,
        additional_kwargs: Optional[Dict[str, Any]] = None,
        max_retries: int = 3,
        timeout: float = 60.0,
        reuse_client: bool = True,
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        api_version: Optional[str] = None,
        callback_manager: Optional[CallbackManager] = None,
        default_headers: Optional[Dict[str, str]] = None,
        http_client: Optional[httpx.Client] = None,
        # base class
        system_prompt: Optional[str] = None,
        messages_to_prompt: Optional[Callable[[Sequence[ChatMessage]], str]] = None,
        completion_to_prompt: Optional[Callable[[str], str]] = None,
        pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
        output_parser: Optional[BaseOutputParser] = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            additional_kwargs=additional_kwargs,
            max_retries=max_retries,
            callback_manager=callback_manager,
            api_key=api_key,
            api_version=api_version,
            api_base=api_base,
            timeout=timeout,
            reuse_client=reuse_client,
            default_headers=default_headers,
            system_prompt=system_prompt,
            messages_to_prompt=messages_to_prompt,
            completion_to_prompt=completion_to_prompt,
            pydantic_program_mode=pydantic_program_mode,
            output_parser=output_parser,
            **kwargs,
        )

        self.dss_project = additional_kwargs["dss_project"]
        self.dss_model_deployment = additional_kwargs["dss_model_deployment"]

    def class_name(cls) -> str:
        return "dss_openai_llm"

    def _get_client(self):
        raise NotImplementedError(
            "DSS OpenAI interface does not support this functionality."
        )


    # DSS SPECIFIC FUNCTIONS 

    def _format_dss_payload(self, messages: list, settings: dict) -> dict:
        """ """

        data = {
            "queries": [{"messages": messages}],
            "settings": settings,
            "llmId": self.dss_model_deployment,
        }

        data = json.dumps(data)

        return data

    def _auth_dss_session(self, session: Session) -> Session:
        """ """

        session = HTTPBasicAuth(self.api_key, "")(session)

        return session

    def _define_dss_session(self, verify: bool = False) -> Session:
        """ """

        session = Session()
        session.verify = verify

        session = self._auth_dss_session(session=session)

        return session

    def _get_dss_url(self) -> str:
        """ """

        url = f"{self.api_base}/dip/publicapi/projects/{self.dss_project}/llms/completions"

        return url

    def _parse_messages_to_dss(self, messages: List[ChatMessage]) -> list:
        """
        Converts llama index ChatMessages to a DSS understandable type
        """

        return_list = []

        for chat_message in messages:
            message_dict = {
                "content": chat_message.content,
                "role": chat_message.role.value,
            }

            return_list.append(message_dict)

        return return_list

    def _dss_call_llm(self, messages: List[ChatMessage], settings: dict = {}) -> dict:
        """ """
        # define variables for DSS request
        messages = self._parse_messages_to_dss(messages=messages)

        data = self._format_dss_payload(messages=messages, settings=settings)

        url = self._get_dss_url()

        session = self._define_dss_session()

        # make the requet
        response = session.request(
            method="POST", url=url, params=None, data=data, files=None, stream=False
        )

        # parse to json
        response = response.json()

        return response


    # END OF DSS FUNCTIONS, NOW OVERRIDING LLAMA FUNCTIONS TO INTEGRATE

    @llm_retry_decorator
    def complete(self, prompt, **kwargs: Any) -> CompletionResponse:
        return self._complete(prompt=prompt, **kwargs)

    @llm_retry_decorator
    def _complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        """ """

        if isinstance(prompt, str):
            prompt = [ChatMessage(content=prompt, role=MessageRole.USER)]

        response = self._dss_call_llm(messages=prompt)

        text = response["responses"][0]["text"]

        return CompletionResponse(
            text=text,
            raw=response,
            logprobs=None,
            additional_kwargs=self._get_response_token_counts(response),
        )

    @llm_retry_decorator
    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
        return self._chat(messages=messages, **kwargs)

    @llm_retry_decorator
    def _chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
        """ """

        response = self._dss_call_llm(messages=messages)
        message = response["responses"][0]["text"]
        logprobs = None
        return ChatResponse(
            message=ChatMessage(content=message, role=MessageRole.ASSISTANT),
            raw=response,
            logprobs=logprobs,
            additional_kwargs=self._get_response_token_counts(response),
        )