Exposing an API call that streams an LLM response in Foundry

Ask the Community
,
1

Hi There,

I am building a chat bot that has to do some logic and use data in Foundry before responding. I have it working with a Foundry function that hits an LLM and returns the response, but I would strongly prefer to have the response streamed back, so the chat experience feels faster. The end goal will be to call this function from an OSDK app running on mobile. What is the best way to do this?

Thanks,
Andrew Girvin

3

Here is how to do it:

Code on the Foundry side (Code repository of python function)

from anthropic import Anthropic
from openai import OpenAI
from typing import Iterable

from functions.api import function
from functions.aliases import model
from foundry_sdk.v2.language_models.utils import (
    get_anthropic_base_url,
    get_foundry_token,
    get_http_client,
    get_openai_base_url,
)

MODEL_PROVIDERS = {
    "poemOpenAi": "openai",
    "poemClaude": "anthropic",
}



@function
def get_chat_completion(prompt: str, model_key: str ="gpt56Sol") -> str:
    client = OpenAI(
        api_key=get_foundry_token(preview=True),
        base_url=get_openai_base_url(preview=True),
        http_client=get_http_client(preview=True),
    )

    completion = client.chat.completions.create(
        model=model(model_key).rid,
        messages=[
            {
                "role": "user",
                "content": prompt,
            },
        ],
    )

    return str(completion.choices[0].message.content)


@function(api_name="createChatCompletionAsStream")
def create_chat_completion_as_stream(prompt: str, model_key: str ="gpt56Sol") -> Iterable[str]:
    client = OpenAI(
        api_key=get_foundry_token(preview=True),
        base_url=get_openai_base_url(preview=True),
        http_client=get_http_client(preview=True),
    )

    stream = client.chat.completions.create(
        model=model(model_key).rid,
        messages=[{"role": "user", "content": prompt}],
        stream=True,
    )

    for event in stream:
        if event.choices:
            content = event.choices[0].delta.content
            if content:
                yield content

You can then use the streamingExecute endpoint to fetch the response in a streaming way. See https://www.palantir.com/docs/foundry/api/v2/functions-v2-resources/queries/streaming-execute-query

Below is a code example to test it via a python script. You need to replace the url and the token to fetch the response. The parameters need to be tweak to match the prototype of the query/function exposed on the Foundry side.

import json
import requests

HOSTNAME = "REDACTED.com"
TOKEN = "ey...A"
QUERY_API_NAME = "createChatCompletionAsStream" # API name of the query you expose
ONTOLOGY = "ontology-d2770208-1234-4567-4677-67835767f0d6" # API name of the ontology the function belongs to
VERSION = None  # Set to a published version, or leave None for latest.
PARAMETERS = {
    "prompt": "Tell me a short story about a robot.",
    "model_key": "gpt56Sol",
}

PREVIEW = True
TIMEOUT_SECONDS = 120

url = f"https://{HOSTNAME}/api/v2/functions/queries/{QUERY_API_NAME}/streamingExecute"
body = {"ontology": ONTOLOGY, "parameters": PARAMETERS}
if VERSION is not None:
    body["version"] = VERSION

with requests.post(
    url,
    headers={
        "Authorization": f"Bearer {TOKEN}",
        "Accept": "text/event-stream",
    },
    params={"preview": str(PREVIEW).lower()},
    json=body,
    stream=True,
    timeout=TIMEOUT_SECONDS,
) as response:
    response.raise_for_status()

    for line in response.iter_lines(chunk_size=1):
        if not line.startswith(b"data:"):
            continue

        event = json.loads(line[5:])
        if event["type"] == "error":
            raise RuntimeError(event)

        print(event["value"], flush=True)