LLM Mesh

The LLM Mesh is the common backbone for Enterprise Generative AI Applications. For more details on the LLM Mesh features of Dataiku, please visit Generative AI and LLM Mesh.

The LLM Mesh API allows you to:

• Send completion and embedding queries to all LLMs supported by the LLM Mesh

• Stream responses from LLMs that support it

• Create and manage persisted multi-turn conversations

• Query LLMs using multimodal inputs (image and text)

• Query the LLM Mesh from LangChain code

• Interact with knowledge banks, and perform semantic search

• Create a fine-tuned saved model

Read LLM Mesh metadata

List and get LLMs

By default, list_llms() returns a list of DSSLLMListItem.

List and get LLMs

import dataiku

client = dataiku.api_client()
project = client.get_default_project()
llm_list = project.list_llms()
for llm in llm_list:
    print(f"- {llm.description} (id: {llm.id})")

List LLMs with a purpose

In addition, you can list LLM with a defined purpose.

List LLMs with a purpose

import dataiku

client = dataiku.api_client()
project = client.get_default_project()
llm_list = project.list_llms(purpose="TEXT_EMBEDDING_EXTRACTION")
for llm in llm_list:
    print(f"- {llm.description} (id: {llm.id})")

Perform completion queries on LLMs

Your first simple completion query

This sample receives an LLM and uses a completion query to ask the LLM to “write a haiku on GPT models.”

import dataiku

# Fill with your LLM id. For example, if you have an OpenAI connection called "myopenai", LLM_ID can be "openai:myopenai:gpt-4o"
# To get the list of LLM ids, you can use project.list_llms() (see above)
LLM_ID = ""

# Create a handle for the LLM of your choice
client = dataiku.api_client()
project = client.get_default_project()
llm = project.get_llm(LLM_ID)

# Create and run a completion query
completion = llm.new_completion()
completion.with_message("Write a haiku on GPT models")
resp = completion.execute()

# Display the LLM output
if resp.success:
    print(resp.text)

# GPT, a marvel,
# Deep learning's symphony plays,
# Thoughts dance, words unveil.

Multi-turn and system prompts

You can have multiple messages in the completion object, with roles

completion = llm.new_completion()

# First, put a system prompt
completion.with_message("You are a poetic assistant who always answers in haikus", role="system")

# Then, give an example, or send the conversation history
completion.with_message("What is a transformer", role="user")
completion.with_message("Transformers, marvels\nOf the deep learning research\nAttention, you need", role="assistant")

# Then, the last query of the user
completion.with_message("What's your name", role="user")

resp = completion.execute()

You can follow-up on an LLM response directly from the last completion response

# Create a follow-up completion query with the conversation history pre-filled, including the last assistant response
followup_completion = resp.prepare_followup()

# Add new user messages to the conversation history
followup_completion.with_message("Repeat the last haiku backwards", role="user")

next_resp = followup_completion.execute()

Multimodal input

Multimodal input is supported on a subset of the LLMs in the LLM Mesh:

• OpenAI

• Bedrock Anthropic Claude

• Microsoft Foundry

• Azure OpenAI

• Vertex Gemini

• Snowflake Cortex

• Databricks Mosaic AI

completion = llm.new_completion()

with open("myimage.jpg", "rb") as f:
    image = f.read()

mp_message = completion.new_multipart_message()
mp_message.with_text("The image represents an artwork. Describe it as it would be described by art critics")
mp_message.with_inline_image(image)

# Add it to the completion request
mp_message.add()

resp = completion.execute()

Completion settings

You can set settings on the completion query

completion = llm.new_completion()
completion.with_message("Write a haiku on GPT models")

completion.settings["temperature"] = 0.7
completion.settings["topK"] = 10
completion.settings["topP"] = 0.3
completion.settings["maxOutputTokens"] = 2048
completion.settings["stopSequences"] = [".", "\n"]
completion.settings["presencePenalty"] = 0.6
completion.settings["frequencyPenalty"] = 0.9
completion.settings["logitBias"] = {
  1489: 60,  # apply a logit bias of 60 on token value "1489"
}
completion.settings["logProbs"] = True
completion.settings["topLogProbs"] = 3

resp = completion.execute()

Response streaming

from dataikuapi.dss.llm import DSSLLMStreamedCompletionChunk, DSSLLMStreamedCompletionFooter

completion = llm.new_completion()
completion.with_message("Please explain special relativity")

for chunk in completion.execute_streamed():
    if isinstance(chunk, DSSLLMStreamedCompletionChunk):
        print("Received text: %s" % chunk.data["text"])
    elif isinstance(chunk, DSSLLMStreamedCompletionFooter):
        print("Completion is complete: %s" % chunk.data)

You can optionally retrieve the full response at the end of streaming, or create a followup completion query.

# Use collect_response=True to build the full response as the chunks are streamed
streamer = completion.execute_streamed(collect_response=True)
for chunk in streamer.iter_chunks():
    if isinstance(chunk, DSSLLMStreamedCompletionChunk):
        print("Received text: %s" % chunk.data["text"])
    elif isinstance(chunk, DSSLLMStreamedCompletionFooter):
        print("Completion is complete: %s" % chunk.data)

# The full response is available once completion is complete
resp = streamer.response

# Create a follow-up completion query directly from the streamer
followup_completion = streamer.prepare_followup()

Persisted conversations

Regular completion queries require the client to provide the conversation history with every request and to maintain the state of the context. Persisted conversations instead keep the active conversation history and context in DSS. The client only sends the input for the next turn.

Persisted conversations are project-scoped. They can use completion models and agents available through the LLM Mesh.

Note

Persisted conversations use project-level permissions. Users with Read project content can list and retrieve every conversation in the project. The API client can use end_user_id to filter conversation lists, but it is not an access-control boundary. When a conversation is created without an end_user_id, it defaults to the authenticated Dataiku user. Client applications serving multiple end users must enforce authorization before accessing a conversation.

Each turn accepts a new user message, the outputs for all pending tool calls, or the responses for all pending tool validation requests. The conversation history and context are managed by DSS.

Create and continue a conversation

Create a conversation from an LLM handle, then execute its first turn:

import dataiku

client = dataiku.api_client()
project = client.get_default_project()
llm = project.get_llm(LLM_ID)

conversation = llm.create_conversation(
    end_user_id="customer-123",
    metadata={"channel": "support"},
)

response = conversation.new_completion() \
    .with_message("How do I reset my password?") \
    .execute()

print(response.text)

The conversation uses the LLM from which it was created as its default LLM. DSS generates a conversation ID when conversation_id is omitted.

You can also provide the first message to create_conversation(). In this case, the method returns the first completion response, and the created conversation is available from the response:

response = llm.create_conversation(
    end_user_id="customer-123",
    metadata={"channel": "support"},
    message="How do I reset my password?",
)
conversation = response.conversation

Use prepare_followup() to create a query in reply to the persisted response. DSS rebuilds the history and context from the selected message and its parent thread, which is the history of ancestor messages:

followup = response.prepare_followup()
followup.with_message("Can an administrator do it for me?")
next_response = followup.execute()

Alternatively, use new_completion() to append a turn to the active thread, defined as the thread of the latest message in the conversation.

If a turn is accepted but LLM execution fails, the response has success=False. DSS keeps both the input message and the failed response in the conversation history, so the client application can inspect the failure or continue from an earlier message.

Stream a persisted turn

Persisted conversations support streaming responses like regular completion queries.

from dataikuapi.dss.llm import (
    DSSLLMStreamedCompletionChunk,
    DSSLLMStreamedCompletionFooter,
)

streamer = conversation.new_completion() \
    .with_message("Summarize the steps.") \
    .execute_streamed(collect_response=True)

for chunk in streamer.iter_chunks():
    if isinstance(chunk, DSSLLMStreamedCompletionChunk):
        print(chunk.data.get("text", ""), end="")
    elif isinstance(chunk, DSSLLMStreamedCompletionFooter):
        print()

response = streamer.response
followup = streamer.prepare_followup()

Retrieve and manage conversations

Use the project-level methods to retrieve conversations independently from an LLM handle:

conversations = project.list_llm_conversations(
    end_user_id="customer-123",
    include_archived=False,
    as_type="objects",
)

conversation = project.get_llm_conversation(conversations[0].conversation_id)
messages = conversation.get_messages()

By default, get_messages() returns the active thread.

You can update conversation metadata or archive a conversation. Archived conversations remain available for retrieval and for listing when explicitly requested, but do not accept new turns.

conversation.update(
    archived=True,
    metadata={"channel": "support", "status": "resolved"},
)

archived_conversations = project.list_llm_conversations(
    include_archived=True,
    as_type="objects",
)

conversation.update(archived=False)

Set archived=False to unarchive the conversation. Calling delete() permanently deletes the conversation and its messages.

Threads and branching

Each persisted message has a messageId and a parentMessageId. To continue from an earlier point, pass the desired parent message ID to new_completion():

messages = conversation.get_messages()
parent_message_id = messages[0]["messageId"]

branched_response = conversation.new_completion(
    parent_message_id=parent_message_id,
).with_message(
    "Explain the process for an administrator instead."
).execute()

The new thread becomes the conversation’s active thread since it contains the most recent messages. Call conversation.get_messages(with_threads=True) to retrieve messages from all threads.

Respond to tool calls and validation requests

When a persisted response contains tool calls, use prepare_followup() and call with_tool_output() for each pending tool call. When an agent returns human-in-the-loop tool validation requests, add a response for each request before executing the follow-up:

followup = response.prepare_followup()

for validation_request in response.tool_validation_requests:
    followup.with_tool_validation_response(
        validation_request["id"],
        validated=True,
    )

response = followup.execute()

See Handling tool validation requests for more information about validation requests. Persisted conversations reconstruct the earlier messages and memory fragments from DSS, so the client only supplies the validation responses.

Text embedding

import dataiku

EMBEDDING_MODEL_ID = "" # Fill with your embedding model id, for example: openai:myopenai:text-embedding-3-small

# Create a handle for the embedding model of your choice
client = dataiku.api_client()
project = client.get_default_project()
emb_model = project.get_llm(EMBEDDING_MODEL_ID)

# Create and run an embedding query
txt = "The quick brown fox jumps over the lazy dog."
emb_query = emb_model.new_embeddings()
emb_query.add_text(txt)
emb_resp = emb_query.execute()

# Display the embedding output
print(emb_resp.get_embeddings())

# [[0.000237455,
#   -0.103262354,
#   ...
# ]]

Reranking

import dataiku

RERANKING_MODEL_ID = "" # Fill with your reranking model id, for example: azureaifoundry:myazureaifoundry:Cohere-rerank-v4.0-fast

# Create a handle for the reranking model of your choice
client = dataiku.api_client()
project = client.get_default_project()
rerank_model = project.get_llm(RERANKING_MODEL_ID)

# Create and run a reranking query
reranking_query = rerank_model.new_reranking()
reranking_query.with_query("What is AI?")
documents = [
    "Chocolate is a sweet treat made from cocoa beans.",
    "Artificial Intelligence (AI) is the simulation of human intelligence processes by machines, especially computer systems.",
    "An almost intelligent idea is an idea that seems smart but lacks depth.",
]
for document in documents:
    reranking_query.with_document(document)
response = reranking_query.execute()

# display the reranking results
print([(doc.index, doc.relevance_score) for doc in response.documents])
# [(1, 0.88377875), (2, 0.42403036), (0, 0.18335953)]

# Print the reranked documents
print("Reranked documents:")
for doc in response.documents:
    print(" - " + documents[doc.index])

Tool calls

Tool calls (sometimes referred to as “function calling”) allow you to augment a LLM with “tools”, functions that it can call and provide the arguments. Your client code can then perform those calls, and provide the output back to the LLM so that it can generate the next response.

Tool calls are supported on the compatible completion models of some LLM connections:

• OpenAI

• Azure OpenAI

• Azure LLM

• Anthropic Claude

• Anthropic Claude models on AWS Bedrock connections

• MistralAI

Define tools

You can define tools as settings in the completion query. Tool parameters are defined as JSON Schema objects. See the JSON Schema reference for documentation about the format.

Tools can also be automatically prepared and invoked from Python code, e.g. using Langchain.

completion = llm.new_completion()
completion.settings["tools"] = [
  {
    "type": "function",
    "function": {
      "name": "multiply",
      "description": "Multiply integers",
      "parameters": {
        "type": "object",
        "properties": {
          "a": {
            "type": "integer",
            "description": "The first integer to multiply",
          },
          "b": {
            "type": "integer",
            "description": "The other integer to multiply",
          },
        },
        "required": ["a", "b"],
      }
    }
  }
]

completion.with_message("What is 3 * 6 ?")
resp = completion.execute()

print(resp.tool_calls)

# [{'type': 'function',
# 'function': {'name': 'multiply', 'arguments': '{"a":3,"b":6}'},
#   'id': 'call_gEB9fOdroydyxYuRs0Ge6Izg'}]

Response streaming with tool calls

LLM responses which include tool calls can also leverage streaming. Depending on the LLM, response chunks may include either complete tool calls or partial tool calls. When the LLM sends partial tool calls, the streamed chunk contains an extra field index allowing to reconstruct the whole LLM response.

for chunk in completion.execute_streamed():
    if isinstance(chunk, DSSLLMStreamedCompletionChunk):
        if "text" in chunk.data:
            print("Received text: %s" % chunk.data["text"])
        if "toolCalls" in chunk.data:
            print("Received tool call: %s" % chunk.data["toolCalls"])

    elif isinstance(chunk, DSSLLMStreamedCompletionFooter):
        print("Completion is complete: %s" % chunk.data)

Provide tool outputs

Tool calls can then be parsed and executed. In order to provide the tool response in the chat messages, use the following methods:

import json

# Function to handle the tool call
def multiply(llm_arguments):
    try:
        json_arguments = json.loads(llm_arguments)
        a = json_arguments["a"]
        b = json_arguments["b"]
        return str(a * b)
    except Exception as e:
        return f"Cannot call the 'multiply' tool: {str(e)}"

tool_calls = resp.tool_calls
call_id = tool_calls[0]["id"]
llm_arguments = tool_calls[0]["function"]["arguments"]
result = multiply(llm_arguments)

# Append the tool calls and the tool outputs to the conversation history
completion.with_tool_calls(tool_calls)
completion.with_tool_output(result, tool_call_id=call_id)

next_resp = completion.execute()

print(next_resp.text)

# 3 multiplied by 6 is 18.

Alternatively, creating a pre-filled follow-up completion query from the last LLM response automatically adds the tool calls to the conversation history.

followup = resp.prepare_followup()
followup.with_tool_output(result, tool_call_id=call_id)

next_resp = followup.execute()

Control tool usage

Tool usage can be constrained in the completion settings:

completion = llm.new_completion()

# Let the LLM decide whether to call a tool
completion.settings["toolChoice"] = {"type": "auto"}

# The LLM must call at least one tool
completion.settings["toolChoice"] = {"type": "required"}

# The LLM must not call any tool
completion.settings["toolChoice"] = {"type": "none"}

# The LLM must call the tool with name 'multiply'
completion.settings["toolChoice"] = {"type": "tool_name", "name": "multiply"}

Knowledge Banks (KB)

List and get KBs

To list the KB present in a project:

List and get KBs

import dataiku
client = dataiku.api_client()
project = client.get_default_project()
kb_list = project.list_knowledge_banks()

By default, list_knowledge_banks() returns a list of DSSKnowledgeBankListItem. To get more details:

for kb in kb_list:
    print(f"{kb.name} (id: {kb.id})")

To get a “core handle” on the KB (i.e. to retrieve a KnowledgeBank object) :

KB_ID = "" # Fill with your KB id
kb_public_api = project.get_knowledge_bank(KB_ID)
kb_core = kb_public_api.as_core_knowledge_bank()

LangChain integration

Dataiku LLM model objects can be turned into langchain-compatible objects, making it easy to:

• stream responses

• run asynchronous queries

• batch queries

• chain several models and adapters

• integrate with the wider langchain ecosystem

Transforming LLM handles to LangChain model

# In this sample, llm is the result of calling project.get_llm() (see above)

# Turn a regular LLM handle into a langchain-compatible one
langchain_llm = llm.as_langchain_llm()

# Run a single completion query
langchain_llm.invoke("Write a haiku on GPT models")

# Run a batch of completion queries
langchain_llm.batch(["Write a haiku on GPT models", "Write a haiku on GPT models in German"])

# Run a completion query and stream the response
for chunk in langchain_llm.stream("Write a haiku on GPT models"):
    print(chunk, end="", flush=True)

See the langchain documentation for more details.

You can also turn it into a langchain “chat model”, a specific type of LLM geared towards conversation:

# In this sample, llm is the result of calling project.get_llm() (see above)

# Turn a regular LLM handle into a langchain-compatible one
langchain_llm = llm.as_langchain_chat_model()

# Run a simple query
langchain_llm.invoke("Write a haiku on GPT models")

# Run a chat query
from langchain_core.messages import HumanMessage, SystemMessage

messages = [
    SystemMessage(content="You're a helpful assistant"),
    HumanMessage(content="What is the purpose of model regularization?"),
]
langchain_llm.invoke(messages)

# Streaming and chaining
from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_template("Tell me a joke about {topic}")
chain = prompt | langchain_llm
for chunk in chain.stream({"topic": "parrot"}):
    print(chunk.content, end="", flush=True)

See the langchain documentation for more details.

Creating Langchain models directly

If running from inside DSS, you can also directly create the Langchain model:

from dataiku.langchain.dku_llm import DKULLM, DKUChatModel

langchain_llm = DKUChatModel(llm_id="your llm id") # For example: openai:myopenai:gpt-4o

Response streaming

The LangChain adapter DKUChatModel also support streaming of answer:

from dataiku.langchain.dku_llm import DKULLM, DKUChatModel
from langchain_core.messages import HumanMessage, SystemMessage

langchain_llm = DKUChatModel(llm_id="your llm id") # For example: openai:myopenai:gpt-4o

messages = [
    SystemMessage(content="You're a helpful assistant"),
    HumanMessage(content="What is the purpose of model regularization?"),
]

for gen in langchain_llm.stream(messages):
    print(gen)

Using knowledge banks as LangChain objects

Core handles allow users to leverage the Langchain library and, through it:

• query the KB for semantic similarity search

• combine the KB with an LLM to form a chain and perform complex workflows such as retrieval-augmented generation (RAG).

In practice, core handles expose KBs as a Langchain-native vector store through two different methods:

• as_langchain_retriever() returns a generic VectorStoreRetriever object

• as_langchain_vectorstore() returns an object whose class corresponds to the KB type. For example, for a FAISS-based KB, you will get a langchain.vectorstores.faiss.FAISS object.

import dataiku
client = dataiku.api_client()
project = client.get_default_project()
kb_core = project.get_knowledge_bank(KB_ID).as_core_knowledge_bank()

# Return a langchain.vectorstores.base.VectorStoreRetriever
lc_generic_vs= kb_core.as_langchain_retriever()

# Return an object which type depends on the KB type
lc_vs = kb_core.as_langchain_vectorstore()

# [...] Move forward with similarity search or RAG 

Writing documents to a knowledge bank

Core handles allow users to leverage the LangChain library to write documents to the underlying vector store.

In practice, core handles expose the method get_writer() which allows to get a writer on the said knowledge bank, as a context manager. Such a writer can be used to build a LangChain vector store that inserts documents into the knowledge bank. The writer will synchronize the vector store content to the knowledge bank automatically upon closing. Make sure you have the langchain_community package in your code environment.

import dataiku
from langchain_core.documents import Document

client = dataiku.api_client()
project = client.get_default_project()
dss_kb = project.get_knowledge_bank(KB_ID)
kb_core = dss_kb.as_core_knowledge_bank()

document = Document(page_content="I can write to a knowledge bank!")

with kb_core.get_writer() as writer:
    print(f"Start from folder {writer.folder_path}")
    # writer.clear()  # uncomment to clear the knowledge bank

    langchain_vs = writer.as_langchain_vectorstore()
    langchain_vs.add_documents([document])

Saving knowledge bank metadata

When inserting documents into a vector store, it is possible to specify metadata that can be retrieved later on. To leverage this metadata during retrieval in Dataiku, it is necessary to set the metadata schema in the knowledge bank settings.

kb_settings = dss_kb.get_settings()
kb_settings.set_metadata_schema({
    "source": "string",
    "start_index": "int"
})

kb_settings.save()

document = Document(
    page_content="I can set metadata",
    metadata={
      "source": "developer-guide",
      "start_index": 1
    }
)

with kb_core.get_writer() as writer:
    langchain_vs = writer.as_langchain_vectorstore()
    langchain_vs.add_documents([document])

Setting retrieval content in the document metadata

The write API allows to format retrieval content in the document metadata, so that the knowledge bank can be used for multimodal retrieval.

The image folder of the knowledge bank can be configured in the knowledge bank settings:

images_folder = dataiku.Folder(IMAGE_FOLDER_ID)
images_folder_full_id = f"{images_folder.project_key}.{images_folder.get_id()}"

kb_settings = dss_kb.get_settings()
kb_settings.set_images_folder(images_folder_full_id)
kb_settings.save()

The retrieval content can be formatted using the knowledge bank writer:

from dataikuapi.dss.document_extractor import ManagedFolderDocumentRef

# assumption: the document has only one page
original_document_ref = ManagedFolderDocumentRef("/path/to/file.pdf", ORIGINAL_DOCUMENT_FOLDER_ID)

with kb_core.get_writer() as writer:
    document = Document(page_content="Summarized text from VLM extraction")

    document = (
        writer.get_metadata_formatter()
            .with_original_document_ref(original_document_ref)
            .with_original_document_page_range(1, 1)  # one page
            .with_retrieval_content(image_paths=[
                "/path/to/screenshot/in/image/folder.jpg"
            ])
            .format_metadata(document)
    )

    # writer.clear()  # uncomment to clear the knowledge bank
    langchain_vs = writer.as_langchain_vectorstore()
    langchain_vs.add_documents([document])

Hybrid Search

Combines both similarity search (default behaviour) and keyword search to retrieve more relevant documents. Only supported by Azure AI Search and Elasticsearch; and not compatible with the diversity option.

Additionally, both vector store offer advanced reranking capabilities, to enhance the mix of documents retrieved. Each has its own specific configuration.

Azure AI Search

import dataiku
client = dataiku.api_client()
project = client.get_default_project()
kb_core = project.get_knowledge_bank(KB_ID).as_core_knowledge_bank()

# 1 using as_langchain_retriever
azure_classic_retriever = kb_core.as_langchain_retriever(search_type="similarity")
azure_hybrid_retriever = kb_core.as_langchain_retriever(search_type="hybrid")
azure_hybrid_advanced_retriever = kb_core.as_langchain_retriever(search_type="semantic_hybrid")

# 2 using as_langchain_vectorstore to get retriever
azure_classic_retriever = kb_core.as_langchain_vectorstore().as_retriever(
  search_type="similarity")
azure_hybrid_retriever = kb_core.as_langchain_vectorstore().as_retriever(
  search_type="hybrid")
azure_hybrid_advanced_retriever = kb_core.as_langchain_vectorstore().as_retriever(
  search_type="semantic_hybrid")

# 3 using as_langchain_vectorstore to perform query
query = "A text to match some doccuments"
azure_classic_result = kb_core.as_langchain_vectorstore().similarity_search(query)
azure_hybrid_result = kb_core.as_langchain_vectorstore().hybrid_search(query)
azure_hybrid_advanced_result = kb_core.as_langchain_vectorstore().semantic_hybrid_search(query)

ElasticSearch

For elastic search, since we need the info at db instantiation time, thats why we need to use vectorstore_kwargs.

Only similarity search type is allowed when using a hybrid strategy.

import dataiku
from elasticsearch.helpers.vectorstore import DenseVectorStrategy

client = dataiku.api_client()
project = client.get_default_project()
kb_core = project.get_knowledge_bank(KB_ID).as_core_knowledge_bank()

hybrid_strategy = DenseVectorStrategy(hybrid=True)
hybrid_advanced_strategy = DenseVectorStrategy(hybrid=True, rrf=True)

# 1 using as_langchain_retriever
elastic_classic = kb_core.as_langchain_retriever()
elastic_hybrid = kb_core.as_langchain_retriever(
  vectorstore_kwargs={"strategy": hybrid_strategy})
elastic_hybrid_advanced = kb_core.as_langchain_retriever(
  vectorstore_kwargs={"strategy": hybrid_advanced_strategy})

# 2 using as_langchain_vectorstore
elastic_classic = kb_core.as_langchain_vectorstore().as_retriever()
elastic_hybrid = kb_core.as_langchain_vectorstore(
  strategy=hybrid_strategy).as_retriever(search_type="similarity")
elastic_hybrid_advanced = kb_core.as_langchain_vectorstore(
  strategy=hybrid_advanced_strategy).as_retriever(search_type="similarity")

Using tool calls

The LangChain chat model adapter supports tool calling, assuming that the underlying LLM supports it too.

The first tab shows how to define tools and bind them to your LLM. The second tab shows how to have a more advanced tool output. In this case, the tool returns a tuple containing a content value that can be used to enhance an LLM’s context and an artifact value for advanced programmatic use.

Simple toolAdvanced tool output

import dataiku

from langchain_core.tools import tool
from langchain_core.messages import HumanMessage

# Define tools

@tool
def add(a: int, b: int) -> int:
    """Adds a and b."""
    return a + b

@tool
def multiply(a: int, b: int) -> int:
    """Multiplies a and b."""
    return a * b

tools_by_name = {"add": add, "multiply": multiply}
tools = [add, multiply]
tool_choice = {"type": "auto"}

# Get the LangChain chat model, bind it to the tools
client = dataiku.api_client()
project = client.get_default_project()
llm_id = "<your llm id>"  # For example: "openai:myopenai:gpt-4o"
llm = project.get_llm(llm_id).as_langchain_chat_model()
llm_with_tools = llm.bind_tools(tools, tool_choice=tool_choice)

# Ask your question
messages = [HumanMessage("What is 3 * 12? and 6 + 4?")]
ai_msg = llm_with_tools.invoke(messages)
messages.append(ai_msg)

# Retrieve tool calls, run them and put the results in the chat messages
for tool_call in ai_msg.tool_calls:
    tool_name = tool_call["name"]
    selected_tool = tools_by_name[tool_name]
    tool_msg = selected_tool.invoke(tool_call)
    messages.append(tool_msg)

# Get the final response
ai_msg = llm.invoke(messages)
ai_msg.content
# '3 * 12 is 36, and 6 + 4 is 10.'

Fine-tuning

Create a Fine-tuned LLM Saved Model version

Note

Visual model fine-tuning is also available to customers with the Advanced LLM Mesh add-on.

With a Python recipe or notebook, it is possible to fine-tune an LLM from the HuggingFace Hub and save it as a Fine-tuned LLM Saved Model version. This is done with the create_finetuned_llm_version() method, which takes an LLM Mesh connection name as input. Settings on this connection like usage permission, guardrails, code environment, or container configuration, will apply at inference time.

The above method must be called on an existing Saved Model. Create one either programmatically (if you are in a notebook and don’t have one yet) with create_finetuned_llm_saved_model() or visually from the Agents & GenAI models list via +New GenAI Model > Create Fine-tuned LLM (if you want to do this in a python recipe, its output Saved Model must exist to create the recipe).

Here we fine-tune using several open-source frameworks from HuggingFace: transformers, trl & peft.

Attention

Note that fine-tuning a local LLM requires significant computational resources (GPU). The code samples below show state-of-the-art techniques to optimize memory usage and processing time, but this depends on your setup and might not always work. Also, beware that the size of your training (and optionally validation) dataset(s) greatly impacts the memory use and storage during fine-tuning.

Microsoft Phi3 Mini 4kMistral 7B V0.2DPO

One can fine-tune a smaller LLM with a small GPU available. Phi3 Mini is a good example, with “only” 3.8B parameters.

There are many techniques available to reduce memory usage and speed up computation. One of them is called Low-Rank Adaptation. It consists in freezing the weights from the base model and adding new, trainable matrices to the Transformer architecture. It drastically reduces the number of trainable parameters and, hence, the GPU memory requirement.

import datasets
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer

from dataiku import recipe
from dataiku.llm.finetuning import formatters

base_model_name = "microsoft/Phi-3-mini-4k-instruct"
assert base_model_name, ("please specify a base LLM, it must be available"
                         " on HuggingFace hub")

connection_name = "a_huggingface_connection_name"
assert connection_name, ("please specify a connection name, the fine-tuned "
                         "LLM will be available from this connection")

##################
# Initial setup
##################
# Here, we're assuming that your training dataset is composed of 2 columns:
# the input (user message) and expected output (assistant message).
# If using a validation dataset, format should be the same.
user_message_column = "input"
assistant_message_column = "output"
columns = [user_message_column, assistant_message_column]

system_message_column = ""  # optional
static_system_message = ""  # optional
if system_message_column:
    columns.append(system_message_column)

# Turn Dataiku datasets into SFTTrainer datasets. 
training_dataset = recipe.get_inputs()[0]
df = training_dataset.get_dataframe(columns=columns)
train_dataset = datasets.Dataset.from_pandas(df)

validation_dataset = None
eval_dataset = None
if len(recipe.get_inputs()) > 1:
    validation_dataset = recipe.get_inputs()[1]
    df = validation_dataset.get_dataframe(columns=columns)
    eval_dataset = datasets.Dataset.from_pandas(df)

saved_model = recipe.get_outputs()[0]

##################
# Model loading
##################
model = AutoModelForCausalLM.from_pretrained(base_model_name)
tokenizer = AutoTokenizer.from_pretrained(base_model_name)

# It is mandatory to define a formatting function for fine-tuning,
# because ultimately, the model is fed with only one string:
# the concatenation of your input columns, in a specific format.
# Here, we leverage the apply_chat_template method, which depends on
# the tokenizer. For more information, see
# https://huggingface.co/docs/transformers/v4.43.3/chat_templating
formatting_func = formatters.ConversationalPromptFormatter(tokenizer.apply_chat_template,
                                                           *columns)

##################
# Fine-tune using SFTTrainer
##################
with saved_model.create_finetuned_llm_version(connection_name) as finetuned_llm_version:
    # feel free to customize, the only requirement is for a transformers model
    # to be created in finetuned_model_version.working_directory

    # TRL package offers many possibilities to configure the training job. 
    # For the full list,
    # see https://huggingface.co/docs/transformers/v4.43.3/en/main_classes/trainer#transformers.TrainingArguments
    train_conf = SFTConfig(
        output_dir=finetuned_llm_version.working_directory,
        save_safetensors=True,
        gradient_checkpointing=True,
        num_train_epochs=1,
        logging_steps=5,
        eval_strategy="steps" if eval_dataset else "no",
    )

    # LoRA is one of the most popular adapter-based methods to reduce memory-usage
    # and speed up fine-tuning
    peft_conf = LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        task_type="CAUSAL_LM",
        target_modules="all-linear",
    )

    trainer = SFTTrainer(
        model=model,
        processing_class=tokenizer,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        formatting_func=formatting_func,
        args=train_conf,
        peft_config=peft_conf,
    )
    trainer.train()
    trainer.save_model()

    # Finally, we are logging training information to the Saved Model version
    config = finetuned_llm_version.config
    config["trainingDataset"] = training_dataset.short_name
    if validation_dataset:
        config["validationDataset"] = validation_dataset.short_name
    config["userMessageColumn"] = user_message_column
    config["assistantMessageColumn"] = assistant_message_column
    config["systemMessageColumn"] = system_message_column
    config["staticSystemMessage"] = static_system_message
    config["batchSize"] = trainer.state.train_batch_size
    config["eventLog"] = trainer.state.log_history

In these examples, we used popular techniques to optimize memory usage and processing time, like LoRA, quantization or gradient checkpointing. Note that the research and open source community is constantly coming up with new ways to make fine-tuning more accessible, while trying to avoid too much performance loss. For more information on other techniques you could try, see for instance the Transformers or PEFT documentations.

OpenAI-compatible API

The OpenAI-compatible API provides an easy way to query the LLM Mesh as it is built on top of the LLM Mesh API and implements the most used parts of OpenAI’s APIs for text generation.

The OpenAI-compatible API allows you to send both Chat Completions API and Responses API queries to all LLMs supported by the LLM Mesh, using the standard OpenAI Python client. This includes, for models that support it:

• Streamed Chat Completions API responses

• Streamed Responses API events

• Multimodal inputs (image and text)

• Tool calls

• JSON output mode

The Dataiku OpenAI-compatible public API URL will have the following form:
http://<DATAIKU_HOST>/public/api/projects/<PROJECT_KEY>/llms/openai/v1/

• <DATAIKU_HOST> is the qualifier of your Dataiku instance.

• <PROJECT_KEY> is the identifier of the project containing the LLM Mesh you want to expose. You can retrieve it from your Dataiku designer URL.

Note

For example, let’s say you are working on a Dataiku project using the URL:
https://dataiku.mycompany.io/projects/AGENTCONCEPTION/

Your <DATAIKU_HOST> will be dataiku.mycompany.io and your <PROJECT_KEY> will be AGENTCONCEPTION
The Dataiku OpenAI-compatible public API URL in this case will be:
http://dataiku.mycompany.io/public/api/projects/AGENTCONCEPTION/llms/openai/v1/

Use the same base URL for both APIs:

• openai_client.chat.completions.create(...) maps to the LLM Mesh endpoint for the Chat Completions API

• openai_client.responses.create(...) maps to the LLM Mesh endpoint for the Responses API

Attention

Some arguments from the OpenAI’s API reference are not supported.

Chat Completions API request:

• n

• response_format

• seed

• service_tier

• parallel_tool_calls

• user

• function_call (deprecated)

• functions (deprecated)

Chat Completions API response:

• choices.message.refusal

• choices.logprobs.refusal

• created

• service_tier

• system_fingerprint

• usage.completion_tokens_details

Responses API request:

Features that rely on OpenAI server-side conversation state are not supported through the LLM Mesh. In particular, you can’t use previous_response_id, so you need to send the conversation history back with each request.

Your first OpenAI-compatible query

Chat Completions APIResponses APIStreaming with the Chat Completions APIStreaming with the Responses API

from openai import OpenAI

# Specify the DSS OpenAI-compatible public API URL, e.g. http://<DATAIKU_HOST>/public/api/projects/<PROJECT_KEY>/llms/openai/v1/
BASE_URL = ""
# Fill with your DSS API Key
API_KEY = ""

# Fill with your LLM id. For example, if you have a HuggingFace connection called "myhf", LLM_ID can be "huggingfacelocal:myhf:meta-llama/Meta-Llama-3.1-8B-Instruct:TEXT_GENERATION_LLAMA_2:promptDriven=true"
# To get the list of LLM ids, you can use openai_client.models.list() or project.list_llms() through the dataiku client 
LLM_ID = ""

# Create an OpenAI client
openai_client = OpenAI(
  base_url=BASE_URL,
  api_key=API_KEY
)

resp = openai_client.chat.completions.create(
  model=LLM_ID,
  messages=[{"role": "user", "content": "Write a haiku on GPT models" }],
)

if resp and resp.choices:
  print(resp.choices[0].message.content)

# GPT, a marvel,
# Deep learning's symphony plays,
# Thoughts dance, words unveil.

Image generation using the LLM Mesh

Your first image-generation query

This sample shows how to send an image generation query with the LLM Mesh to ask the image generation model to generate an image of a blue bird.

import dataiku

client = dataiku.api_client()
project = client.get_default_project()

# To list the image generation model ids, you can use project.list_llms(purpose="IMAGE_GENERATION")
IMAGE_GENERATION_MODELS = project.list_llms(purpose="IMAGE_GENERATION")

IMAGE_GENERATION_MODEL_ID = "" # Fill with your image generation model id, for example: openai:my_openai_connection:dall-e-3

# Create a handle for the image generation model of your choice
img_gen_model = project.get_llm(IMAGE_GENERATION_MODEL_ID)

prompt_text = "Vibrant blue bird in a serene scene on a blooming cherry blossom branch. Tranquil morning sky background with soft pastel colors of dawn, gently blending pinks, purples, and soft oranges. Distant view of a calm lake reflecting the colors of the sky and surrounded by lush greenery."

img_gen_query = img_gen_model.new_images_generation()
img_gen_query.with_prompt(prompt_text)
img_gen_resp = img_gen_query.execute()
image_data = img_gen_resp.first_image()

# You can display the image in your notebook
from IPython.display import Image, display
if img_gen_resp.success:
    display(Image(image_data))

# Or you can save the image to a managed folder
FOLDER_ID = ""  # Enter your managed folder id here
my_images_folder = dataiku.Folder(FOLDER_ID)
with my_images_folder.get_writer("blue_bird.png") as writer:
    writer.write(image_data)

You can parameterize the query to impact the resulting image or generate more images.

The LLM Mesh maps each parameter to the corresponding parameter for the underlying model provider. Support varies across models/providers, and in particular not all models can generate more than one image.

If you want to generate multiple images with different prompts, you must query the LLM Mesh multiple times.

import dataiku

IMAGE_GENERATION_MODEL_ID = "" # Fill with your image generation model id

# Create a handle for the image generation model of your choice
client = dataiku.api_client()
project = client.get_default_project()
img_gen_model = project.get_llm(IMAGE_GENERATION_MODEL_ID)
generation = img_gen_model.new_images_generation()
generation.height = 1024
generation.width = 1024
generation.seed = 3
# If the underlying model supports weighted prompts they will be passed with
# their specified weight, otherwise they will just be merged and sent as a single prompt.
generation.with_prompt("meat pizza", weight=0.8).with_prompt("rustic wooden table", weight=0.6)

# Not all models or providers support more than one
generation.images_to_generate = 1

# Regardless of what parameter the underlying provider expects for the image dimensions,
# when using the LLM Mesh API you can specify either the height and width or the aspect_ratio.
# The LLM Mesh will do the translation between its API and the underlying provider.
# Not all models support the same dimensions.
generation.aspect_ratio = 21 / 9

# The following parameters are not relevant for all models
generation.with_negative_prompt("tomatoes, basil, green leaf", weight=1)
generation.fidelity = 0.5 # from 0.1 to 1, how strongly to adhere to prompt
# valid values depend on the targeted model
generation.quality = "hd"
generation.style = "anime"

resp = generation.execute()

Image-to-image query

Some models can generate an image from another image, see this documentation.

• Mask-free variation generates another image guided by a prompt

• Some models can generate unprompted variations

• Inpainting uses a mask (either black pixels in a second input image, or transparent pixels on the original image) to fill the corresponding pixels of the input image

Prompted variationUnprompted variationInpainting: mask imageInpainting: mask pixels

In this example, we ask the model for an image variation by passing an image and a prompt using the MASK_FREE mode.

import dataiku

IMAGE_GENERATION_MODEL_ID = "" # Fill with your image generation model id

img_gen_model = dataiku.api_client().get_default_project().get_llm(IMAGE_GENERATION_MODEL_ID)

# Your image to use as an input.
# Here we're retrieving it from a managed folder but it could also be an image from a previous generation
my_images_folder = dataiku.Folder("my_folder_id")
with my_images_folder.get_download_stream("cat_on_the_beach.png") as img_file:
    input_img_data = img_file.read()

# Create the generation query
generation = img_gen_model.new_images_generation()
generation.with_original_image(input_img_data, mode="MASK_FREE", weight=0.3)
generation.with_prompt("dog on the beach")
resp = generation.execute()

Image-to-image generation with a prompt can also be used with the CONTROLNET_STRUCTURE and CONTROLNET_SKETCH modes.

Reference documentation

Classes

dataiku.KnowledgeBank(id[, project_key, ...])	This is a handle to interact with a Dataiku Knowledge Bank flow object
dataiku.core.vector_stores.data.metadata.DocumentMetadataFormatter(...)	Helper class to format vector store documents metadata for usage within Dataiku.
dataiku.core.vector_stores.data.writer.VectorStoreWriter(...)	A helper class to write vector store data to the underlying knowledge bank folder.
dataikuapi.dss.document_extractor.ManagedFolderDocumentRef(...)	A reference to a file in a DSS-managed folder.
dataikuapi.dss.llm.DSSLLM(client, ...)	A handle to interact with a DSS-managed LLM.
dataikuapi.dss.llm.DSSLLMListItem(client, ...)	An item in a list of llms
dataikuapi.dss.llm.DSSLLMCompletionQuery(llm)	A handle to interact with a completion query.
dataikuapi.dss.llm.DSSLLMCompletionsQuery(llm)	A handle to interact with a multi-completion query.
dataikuapi.dss.llm.DSSLLMCompletionsQuerySingleQuery()
dataikuapi.dss.llm.DSSLLMCompletionQueryMultipartMessage(q, role)
dataikuapi.dss.llm.DSSLLMCompletionQueryMultipartToolOutput(q, ...)
dataikuapi.dss.llm.DSSLLMCompletionResponse([...])	A handle to interact with a completion query result.
dataikuapi.dss.llm.DSSLLMConversationListItem(...)	An item in a list of persisted LLM conversations.
dataikuapi.dss.llm.DSSLLMConversation(...[, ...])	A handle to interact with a persisted LLM conversation.
dataikuapi.dss.llm.DSSLLMConversationCompletionQuery(...)	A query that appends a turn to an existing persisted conversation.
dataikuapi.dss.llm.DSSLLMConversationCompletionResponse(...)	Response to a persisted conversation turn.
dataikuapi.dss.llm.DSSLLMConversationStreamedCompletionChunks(query)	Streamed chunks for a persisted conversation turn.
dataikuapi.dss.llm.DSSLLMEmbeddingsQuery(...)	A handle to interact with an embedding query.
dataikuapi.dss.llm.DSSLLMEmbeddingsResponse(...)	A handle to interact with an embedding query result.
dataikuapi.dss.knowledgebank.DSSKnowledgeBank(...)	A handle to interact with a DSS-managed knowledge bank.
dataikuapi.dss.knowledgebank.DSSKnowledgeBankListItem(...)	An item in a list of knowledge banks
dataikuapi.dss.knowledgebank.DSSKnowledgeBankSettings(...)	Settings for a knowledge bank
dataikuapi.dss.langchain.DKUChatModel(*args, ...)	Langchain-compatible wrapper around Dataiku-mediated chat LLMs
dataikuapi.dss.langchain.DKULLM(*args, **kwargs)	Langchain-compatible wrapper around Dataiku-mediated LLMs
dataikuapi.dss.langchain.DKUEmbeddings(...)	Langchain-compatible wrapper around Dataiku-mediated embedding LLMs
dataikuapi.dss.project.DSSProject(client, ...)	A handle to interact with a project on the DSS instance.
dataikuapi.dss.utils.DSSSimpleFilter(operator)	A simplified representation of a DSS filter.
dataikuapi.dss.utils.DSSSimpleFilterOperator(*values)	Operators for the DSSSimpleFilter.

Functions

add()	Add this message to the completion query
create_conversation([conversation_id, ...])	Create a persisted conversation bound to this LLM.
create_llm_conversation([conversation_id, ...])	Create a persisted LLM conversation.
delete()	Hard-delete the conversation.
get_messages([message_id, with_threads])	Retrieve persisted conversation messages.
get_llm_conversation(conversation_id)	Get a handle to interact with a specific persisted LLM conversation.
add_text(text)	Add text to the embedding query.
as_core_knowledge_bank()	Get the dataiku.KnowledgeBank object corresponding to this knowledge bank
aspect_ratio
bind_tools(tools[, tool_choice, strict, ...])	Bind tool-like objects to this chat model.
as_langchain_chat_model(**data)	Create a langchain-compatible chat LLM object for this LLM.
as_langchain_llm(**data)	Create a langchain-compatible LLM object for this LLM.
description
execute()	Run the completion query and retrieve the LLM response.
execute()	Run the embedding query.
dataikuapi.dss.llm.DSSLLMImageGenerationQuery.execute()	Executes the image generation
execute_streamed([collect_response])	Run the completion query and retrieve the LLM response as streamed chunks.
fidelity
first_image([as_type])
format_metadata(document)	Formats the metadata in the provided document, so that it can be used for retrieval in Dataiku.
get_embeddings()	Retrieve vectors resulting from the embeddings query.
get_knowledge_bank(id)	Get a handle to interact with a specific knowledge bank
get_llm(llm_id)	Get a handle to interact with a specific LLM
get_metadata_formatter()	Gets the metadata formatter to help writing documents to this vector store.
get_settings()	Get the knowledge bank's definition
get_writer()	Gets a writer on the latest vector store files on disk.
id
images_to_generate
list_knowledge_banks([as_type])	List the knowledge banks of this project
list_llms([purpose, as_type])	List the LLM usable in this project
list_llm_conversations([end_user_id, ...])	List persisted LLM conversations in this project.
new_completion()	Create a new completion query.
new_completion([parent_message_id, llm_id])	Prepare a new turn on this persisted conversation.
new_embeddings([text_overflow_mode])	Create a new embedding query.
new_images_generation()
new_multipart_message([role])	Start adding a multipart-message to the completion query.
quality
save()	Saves the settings on the knowledge bank
set_images_folder(managed_folder_id[, ...])	Sets the images folder to use with this knowledge bank.
set_metadata_schema(schema)	Sets the schema for metadata fields.
settings
style
success
text
tool_calls
prepare_followup()	Prepare a follow-up turn pinned to this persisted response.
refresh()	Refresh the conversation metadata snapshot.
update([end_user_id, archived, ...])	Update persisted conversation metadata.
with_inline_image(image[, mime_type])	Add an image part to the multipart message
with_message(message[, role])	Add a message to the completion query.
with_original_document_page_range(...)	Adds the page range in the original document.
with_original_document_ref(document_ref[, ...])	Adds the original document information in the metadata.
with_original_image(image[, mode, weight])	Add an image to the generation query.
with_mask(mode[, image])	Add a mask for edition to the generation query.
with_negative_prompt(prompt[, weight])	Add a negative prompt to the image generation query.
with_prompt(prompt[, weight])	Add a prompt to the image generation query.
with_retrieval_content([text, image_paths, ...])	Adds the retrieval content in the metadata. Accepts either
with_text(text)	Add a text part to the multipart message
with_tool_calls(tool_calls[, role])	Add tool calls to the completion query.
with_tool_output(tool_output, tool_call_id)	Add a tool message to the completion query.