We have some upcoming features coming out of our research team that will require Pipelines that contain multiple LLMBlocks, with each LLMBlock potentially using a different model id deployed on the same inference server, different model family (ie granite vs mistral and so on), or even different inference endpoints entirely for each.

Today users specify the inference endpoint via passing an OpenAI Client into the PipelineContext. This client gets used for every LLMBlock, with no ability to map the client to each LLMBlock.

What we need is some way to pass in multiple OpenAI clients, and map each to the relevant LLMBlock. One example of what this could look like:

Python code

clients = {
  "default": OpenAI(base_url="http://localhost:8000/v1", api_key="foo"),
  "granite-endpoint": OpenAI(base_url="http://localhost:8001/v1", api_key="bar"),
}
ctx = PipelineContext(clients=clients)

We pass a map of OpenAI clients to PipelineContext in this example instead of a single one. If a single one was passed (ie for backwards compatibility), we could turn this internally into a map with that single client as the "default" value. Users would be able to control the full array of client parameters here, including SSL cert handling, timeouts, retries, and anything else that can be configured on the OpenAI client or its underlying httpx.Client.

pipeline.yaml

version: "1.0"
blocks:
  - name: model_one
    type: LLMBlock
    config:
      model_family: mixtral
      model_id: Mixtral-8x7B-Instruct-v0.1
      config_path: model_one_config.yaml
      output_cols:
        - column_one
      gen_kwargs:
        max_tokens: 2048

  - name: model_two
    type: LLMBlock
    config:
      client: granite-endpoint
      model_family: granite
      model_id: granite-7b-lab
      config_path: model_two_config.yaml
      output_cols:
        - column_two
      gen_kwargs:
        max_tokens: 512

  - name: model_three
    type: LLMBlock
    config:
      client: granite-endpoint
      model_id: granite-7b-lab
      model_prompt: <s> [INST] {prompt} [/INST]
      config_path: model_three_config.yaml
      output_cols:
        - column_three
      gen_kwargs:
        max_tokens: 5

Then, in your pipeline.yaml, we map each LLMBlock to a client. Any block that does not specify a client gets the "default" client. Otherwise, they can choose a string value that picks one of the clients passed into the PipelineContext out of that map. This allows us to create a re-usable Pipeline that expects N clients, and provides users running that Pipeline a way to configure those clients for their specific environment without modifying the Pipeline yaml itself.