[FEATURE] Support predibase LLM serving a base model with optional fine-tuned adapter.

[alexsherstinsky]

Title: Support predibase LLM serving a base model with an optional fine-tuned adapter.

• Brief Description of Changes

• We add "predibase" as the provider using the OpenAI-compliant API that Predibase has.
• Format for supplying a fine-tuned adapter is "<base_model>[:adapter_id]", where adapter_id format is
  "<adapter_repository_reference/version_number" (version_number is required).

client = OpenAI(
    api_key=os.environ["PREDIBASE_API_TOKEN"],
    base_url=PORTKEY_GATEWAY_URL,
    default_headers=createHeaders(
        provider="predibase",
    )
)

chat_complete = client.chat.completions.create(
    user=os.environ["PREDIBASE_TENANT_ID"],
    # model=os.environ["PREDIBASE_DEPLOYMENT"],
    model=f'{os.environ["PREDIBASE_DEPLOYMENT"]}:test-phi-3/4',
    stream=False,  # True is also supported.
    max_tokens=128,
    temperature=0.2,
    messages=[
        {
            "role": "user",
            "content": "How fast can a horse run?",
        },
    ],
)

if isinstance(chat_complete, Stream):
    completion_stream: Stream = chat_complete
    text: list[str] = []
    for message in completion_stream:
        print(message)
        delta_content: str | None = message.choices[0].delta.content
        if delta_content:
            text.append(delta_content)
    print("".join(text))
else:
    print(chat_complete.choices[0])
    print(chat_complete.choices[0].message)
    print(chat_complete.choices[0].message.content)


Supported body parameters:

  model,
  messages,
  max_tokens,
  temperature,
  top_p,
  stream,
  n,
  stop,

Description: (optional)

• Detailed change 1
• Detailed change 2

Motivation: (optional)

• Predibase becomes an option as an LLM provider for the users of Gateway. Predibase is the most robust platform for LLM fine-tuning and serving for a wide variety of pre-trained models, including the open-source ones.

Related Issues: (optional)

• #issue-number

[roh26it]

Looks good to me and excited to merge this.

@VisargD will run some tests at our end on this and come back if there's any changes.

[VisargD]

Hey @alexsherstinsky - Thanks for the PR! In the Predibase docs, I can see that streaming (/generate_stream) is also supported. Are you planning to add it to this PR as well? Here is the doc that I am referring to: https://docs.predibase.com/user-guide/inference/rest_api#notes

If its not planned for this PR then I can merge this and raise a new one with streaming support for Predibase.

You can also use the /generate_stream endpoint to have the tokens be streamed from the deployment. The parameters also follow the same format as the [LoRAX generate endpoints](https://github.com/predibase/lorax/tree/main/clients/python).

[alexsherstinsky]

Hey @alexsherstinsky - Thanks for the PR! In the Predibase docs, I can see that streaming (/generate_stream) is also supported. Are you planning to add it to this PR as well? Here is the doc that I am referring to: https://docs.predibase.com/user-guide/inference/rest_api#notes

If its not planned for this PR then I can merge this and raise a new one with streaming support for Predibase.

You can also use the /generate_stream endpoint to have the tokens be streamed from the deployment. The parameters also follow the same format as the [LoRAX generate endpoints](https://github.com/predibase/lorax/tree/main/clients/python).

@VisargD Streaming is already supported! If you see in my example above, there is a stream flag -- I test it, and it works! Thank you!

[VisargD]

Gateway expects separate responseTransforms for stream and non-stream mode. In the index.ts of a provider, we define chatComplete and stream-chatComplete responseTransforms. Just like you defined PredibaseChatCompleteResponseTransform function, there should be one more function like PredibaseChatCompleteStreamChunkTransform which maps provider's stream chunk to Gateway compatible stream chunk. You can check other providers like perplexity-ai, mistral-ai to get an idea of this.

The reason why stream is working fine currently is because it is not able to find a stream chunk transform function and so it is doing a passthrough of all the chunks as it is. Even if predibase is sending OpenAI compatible chunks, its preferred to atleast add this function and map the chunk data as it is so that in future it does not break whenever predibase makes any change.

Please let me know if you need any help with this. I can provide more details if required.

[VisargD]

Here is what I am suggesting:

• Add a new stream chunk transform function in chatComplete.ts:
  Predibase does not send id in chunks. So you can use fallbackId which is passed to all transform function calls.

export const PredibaseChatCompleteStreamChunkTransform: (
  response: string,
  fallbackId: string,
) => string | string[] = (responseChunk, fallbackId) => {
  let chunk = responseChunk.trim();
  chunk = chunk.replace(/^data:/, '');
  chunk = chunk.trim();

  const parsedChunk: PredibaseChatCompletionStreamChunk = JSON.parse(chunk);

  return `data: ${JSON.stringify({
      id: fallbackId,
      object: parsedChunk.object,
      created: Math.floor(Date.now() / 1000),
      model: parsedChunk.model,
      provider: PREDIBASE,
      choices: [
        {
          delta: {
            role: parsedChunk.choices[0]?.delta?.role,
            content: parsedChunk.choices[0]?.delta?.content,
          },
          index: 0,
          finish_reason: parsedChunk.choices[0]?.finish_reason,
        },
      ],
    })}` + '\n\n';
};

• Add stream-chatComplete responseTransform in predibase index.ts:

const PredibaseConfig: ProviderConfigs = {
  chatComplete: PredibaseChatCompleteConfig,
  api: PredibaseAPIConfig,
  responseTransforms: {
    chatComplete: PredibaseChatCompleteResponseTransform,
    'stream-chatComplete': PredibaseChatCompleteStreamChunkTransform,
  },
};

[alexsherstinsky]

Here is what I am suggesting:

• Add a new stream chunk transform function in chatComplete.ts:
  Predibase does not send id in chunks. So you can use fallbackId which is passed to all transform function calls.

export const PredibaseChatCompleteStreamChunkTransform: (
  response: string,
  fallbackId: string,
) => string | string[] = (responseChunk, fallbackId) => {
  let chunk = responseChunk.trim();
  chunk = chunk.replace(/^data:/, '');
  chunk = chunk.trim();

  const parsedChunk: PredibaseChatCompletionStreamChunk = JSON.parse(chunk);

  return `data: ${JSON.stringify({
      id: fallbackId,
      object: parsedChunk.object,
      created: Math.floor(Date.now() / 1000),
      model: parsedChunk.model,
      provider: PREDIBASE,
      choices: [
        {
          delta: {
            role: parsedChunk.choices[0]?.delta?.role,
            content: parsedChunk.choices[0]?.delta?.content,
          },
          index: 0,
          finish_reason: parsedChunk.choices[0]?.finish_reason,
        },
      ],
    })}` + '\n\n';
};

• Add stream-chatComplete responseTransform in predibase index.ts:

const PredibaseConfig: ProviderConfigs = {
  chatComplete: PredibaseChatCompleteConfig,
  api: PredibaseAPIConfig,
  responseTransforms: {
    chatComplete: PredibaseChatCompleteResponseTransform,
    'stream-chatComplete': PredibaseChatCompleteStreamChunkTransform,
  },
};

@VisargD Thank you very much for this -- it was extremely helpful! I incorporated your suggestions and looked up how perplexity-ai does it as well. Thanks to your suggestion, I already found one error (one of my tests is failing, which is a good thing, because it is happening now, while we are still developing it!). I will ping you again once I have figured it out and made the fix. Thanks again!

[alexsherstinsky]

@VisargD Please re-review; I incorporated your suggestion and also added error handling. The error handling this way enables the client to see the actual error; otherwise, the error response does not work, because the HTTP response is 200 OK. Thank you.

[VisargD]

Thanks for the quick changes. Looks good to me! I will merge this PR and make it a part of the next gateway release.

[VisargD]

Closes #126