Large language models are transforming how we build applications, but managing prompts across different environments and versions can quickly become chaotic. What if you could version your prompts the same way you version your ML models? In this article, I'll show you how to leverage Snowflake's Model Registry to create a robust, versioned prompt management system that treats your prompts as first-class artifacts.
Setting Up Your Prompt Configuration
Before we can register anything, we need to define our prompt and its parameters. Here, we focus on structured output prompts which return a predefined JSON schema given by a Pydantic BaseModel.
import pydantic
from typing import Literal
# Define the prompt template with placeholder for dynamic content
prompt_template = """Classify the type of text given as follows: {given_text}"""
# List of parameters that will be injected into the prompt
prompt_parameters = ["given_text"]
# Version control for your prompt
prompt_name = "text_classification"
prompt_version = "v1"
# LLM configuration parameters
temperature = 0.0
max_tokens = 15_000
model = "claude-4-0-sonnet"
# Define the expected output structure using Pydantic
class StructuredOutput(pydantic.BaseModel):
text_category: Literal["scientific_paper", "newspaper_article"] = pydantic.Field(
description="Is the given text a newspaper article ('newspaper_article') or a scientific paper ('scientific_paper')?",
)
Sources:
- AI_COMPLETE structured outputs | Snowflake Documentation
- LLM Coding Concepts: Static Typing, Structured Output, and AsyncIO | by Dr. Johann Hagerer | Medium
Serialize the Prompt Configuration
The model context is where we serialize our prompt configuration into files that Snowflake can track. This step is crucial because it transforms your prompt from ephemeral code into persistent artifacts.
import json
from snowflake.ml.model import custom_model
config_file: str = "config.json"
prompt_template_file: str = "prompt_template.txt"
json_schema_file: str = "json_schema.json"
LlmConfig = dict[str, str | float | int | bool] # stupid pandas limitation
# Bundle all LLM parameters into a single configuration dictionary
llm_params: LlmConfig = dict(
prompt_name=prompt_name,
prompt_version=prompt_version,
temperature=temperature,
max_tokens=max_tokens,
model=model,
prompt_parameters=json.dumps(prompt_parameters),
json_schema_file=json_schema_file,
prompt_template_file=prompt_template_file,
)
# Write configuration to disk for model context
with open(config_file, "w") as f:
json.dump(llm_params, f)
# Save the prompt template as a separate file for easy editing
with open(prompt_template_file, "w") as f:
f.write(prompt_template)
# Export the Pydantic schema for structured output validation
with open(json_schema_file, "w") as f:
json.dump(StructuredOutput.model_json_schema(), f, indent=2)
# Create model context that bundles all files together
model_context = custom_model.ModelContext(
config_file=config_file,
prompt_template_file=prompt_template_file,
json_schema_file=json_schema_file,
)
Sources:
Building the Custom Prompt Model
Now we implement the actual model class that will handle prompt generation. This class reads your configuration and creates properly formatted prompts with all the necessary metadata attached.
import pandas as pd
import polars as pl
import json
from snowflake.ml.model import custom_model
class PromptModel(custom_model.CustomModel):
def __init__(self, context: custom_model.ModelContext) -> None:
super().__init__(context)
# Load all configuration files from the model context
config_file = self.context.path("config_file")
json_schema_file = self.context.path("json_schema_file")
prompt_template_file = self.context.path("prompt_template_file")
with open(config_file) as f:
self.config: LlmConfig = json.load(f)
with open(json_schema_file) as f:
self.json_schema = f.read()
with open(prompt_template_file) as f:
self.prompt_template = f.read()
def has_all_params(self, df: pl.DataFrame) -> bool:
"""Check if DataFrame contains all required prompt parameters"""
params: list[str] = json.loads(self.config['prompt_parameters'])
has_all_columns = all(x in df.columns for x in params)
if not has_all_columns:
return False
# Ensure no null values in required columns
return sum(df.null_count().row(0)) == 0
@custom_model.inference_api
def create_prompts(self, df: pd.DataFrame) -> pd.DataFrame:
"""Main inference API that converts pandas DataFrames to prompts"""
return self._create_prompts(pl.from_pandas(df)).to_pandas()
def _create_prompts(self, df: pl.DataFrame) -> pl.DataFrame:
"""Internal method using Polars for efficient processing"""
# Return empty output if required parameters are missing
if not self.has_all_params(df):
output_dict = {'output': [""] * len(df)}
return pl.DataFrame(output_dict)
# Create a row for each input with formatted prompt and metadata
rows: LlmConfig = [
dict(
**self.config, # Include all config parameters
**row, # Include input data
prompt=self.prompt_template.format(**row), # Format the prompt
json_schema = self.json_schema,
prompt_template = self.prompt_template,
)
for row in df.iter_rows(named=True)
]
return pl.from_dicts(rows)
# Instantiate the model with our context
prompt_model = PromptModel(model_context)
Sources:
Testing Your Prompt Model
Before registering anything, let's verify that our prompt model works correctly. This step ensures you're not committing broken code to your registry.
import polars as pl
# Create sample input data
df = pl.DataFrame({
"given_text": [
"Lorem ipsum ..."
]
})
# Generate prompts from the input data
prompts_df = prompt_model.create_prompts(df.to_pandas())
prompts_df
Defining the Model Signature
The model signature tells Snowflake what inputs your model expects and what outputs it produces. This is essential for type safety and model validation.
from snowflake.ml.model.model_signature import infer_signature
# Automatically infer the signature from sample input/output
sig = infer_signature(
input_data=df.to_pandas(),
output_data=prompts_df
)
sig
Sources:
Registering to Snowflake Model Registry
This is where the magic happens---we're logging our prompt model to the registry, making it discoverable and versionable alongside your other ML artifacts.
import snowflake.snowpark as sp
from snowflake.ml.registry import Registry
# Get the active Snowflake session
session = sp.context.get_active_session()
snowml_registry = Registry(session)
# Log the model to the registry with all metadata
custom_mv = snowml_registry.log_model(
prompt_model,
model_name=prompt_name,
version_name=prompt_version,
conda_dependencies=["polars"], # Specify runtime dependencies
options={"relax_version": False}, # Enforce exact version matching
signatures={"predict": sig},
comment = 'A prompt for KFZ liability completeness check',
)
Adding Custom Metadata
Finally, we attach our full configuration as metadata to the model version. This creates a complete audit trail of all parameters used for this specific prompt version.
import json
# Serialize configuration as JSON string
metadata_dict = json.dumps(llm_params)
# Attach metadata to the model version using SQL
session.sql(f"""
ALTER MODEL {prompt_name} MODIFY VERSION {prompt_version}
SET METADATA = $${metadata_dict}$$
""").collect()
With this approach, you now have a fully versioned, auditable prompt management system built directly into your Snowflake infrastructure. No more hunting through git commits to find which prompt was used in production---it's all tracked in the same place as your models.
In future articles, I'll explore how to integrate this prompt management approach with Snowflake's Cortex LLM service for seamless inference, and how to leverage feature views to establish proper data lineage between your source data, prompt versions, and model predictions. This will complete the picture of end-to-end MLOps for LLM applications within the Snowflake ecosystem.
Top comments (0)