How to Add a New LLM Model to CreativeMath: A Complete Integration Guide
Adding a new LLM to CreativeMath requires updating four key files: declaring the model type in model_loader.py, defining prompt templates in prompt_utils.py, implementing the loader in api_models.py or local_models.py, and optionally exposing it in config.py.
CreativeMath is an open-source mathematical reasoning framework that abstracts LLM interactions through a unified ModelWrapper class. To add a new LLM model to CreativeMath, you must integrate it into the model loading pipeline and define how the system formats prompts for that specific provider. The architecture cleanly separates API-based models from local HuggingFace transformers, making extension straightforward once you understand the entry points.
Understanding the Model Architecture in CreativeMath
The ModelWrapper class in src/models/model_loader.py serves as the central dispatcher. It checks self.is_api_model to determine whether to route requests through api_models.py or local_models.py. This boolean is calculated by checking if the model name exists in a predefined list of API models.
When you add a new LLM model to CreativeMath, you are essentially teaching this dispatcher three things: whether the model is remote or local, how to structure prompts for it, and how to initialize its client or tokenizer.
Step-by-Step Guide to Add a New LLM Model to CreativeMath
Step 1: Declare the Model Type in model_loader.py
First, indicate whether your model is API-based or local. Open src/models/model_loader.py and locate the is_api_model check. Add your model identifier to the list if it is an API model:
# src/models/model_loader.py
self.is_api_model = model_name in [
"claude-3-opus",
"claude-3-5-sonnet",
"deepseek-v2",
"gemini-1.5-pro",
"gpt-4",
"gpt-4o",
"gpt-4o-mini",
# ← add your new model here
"my-new-llm",
]
If you do not add the name here, ModelWrapper will treat it as a local model and attempt to load it via local_models.py.
Step 2: Define the Prompt Template in prompt_utils.py
Different LLMs expect different input formats. Open src/models/prompt_utils.py and add a template for your model in the templates dictionary. For chat-style APIs that expect a list of messages:
# src/models/prompt_utils.py
templates = {
# … existing entries …
"my-new-llm": [
{"role": "system", "content": "You are a helpful math assistant."},
{"role": "user", "content": prompt},
],
}
If your provider expects a raw string (like Gemini), simply map to the string directly:
"my-new-llm": prompt,
Step 3: Implement the Model Loader (API or Local)
For API Models: Edit api_models.py
Open src/models/api_models.py and create two functions: one to load the client and one to generate responses. For example, using a hypothetical provider:
# src/models/api_models.py
from openai import OpenAI # example; replace with the provider's SDK
def load_my_new_llm_api(model_name: str):
# Create and return the client object; authentication is read from env vars.
client = OpenAI() # <-- adjust to the provider's client
return client # the wrapper stores this as `self.model`
def generate_my_new_llm_response(model_name, client, messages):
# Most providers expose a `chat.completions.create`‑like method.
resp = client.chat.completions.create(
model=model_name,
messages=messages,
temperature=0.7,
)
return resp.choices[0].message.content
Then wire these into the public interface functions load_api_model and generate_api_response:
def load_api_model(name):
if name == "my-new-llm":
return load_my_new_llm_api(name)
# existing branches …
def generate_api_response(name, model, messages):
if name == "my-new-llm":
return generate_my_new_llm_response(name, model, messages)
# existing branches …
For Local Models: Edit local_models.py
If your model runs locally via HuggingFace transformers, open src/models/local_models.py and add a loader function:
def load_my_new_llm_local(name):
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("myorg/my-new-llm")
tokenizer = AutoTokenizer.from_pretrained("myorg/my-new-llm")
return model, tokenizer
Then integrate this into the main loading dispatcher within the same file.
Step 4: Expose the Model in Configuration (Optional)
To make your model selectable via the config.json UI or CLI, open src/config.py and add an entry under "model_config":
# src/config.py
"model_config": {
"my-new-llm": {
"temperature": 0.7,
"max_tokens": 2048,
# other provider-specific settings
}
}
Complete Code Example: Integrating a Custom API Model
Here is the full integration path for a hypothetical API provider called mathllm:
# src/models/model_loader.py
self.is_api_model = model_name in [
"gpt-4",
"claude-3-opus",
"mathllm-v1", # ← added
]
# src/models/prompt_utils.py
templates = {
"mathllm-v1": [
{"role": "system", "content": "You are an expert mathematical reasoning assistant."},
{"role": "user", "content": prompt},
],
}
# src/models/api_models.py
import os
from openai import OpenAI # assuming OpenAI-compatible endpoint
def load_mathllm_v1(model_name: str):
client = OpenAI(
base_url=os.getenv("MATHLLM_BASE_URL"),
api_key=os.getenv("MATHLLM_API_KEY")
)
return client
def generate_mathllm_v1_response(model_name, client, messages):
resp = client.chat.completions.create(
model=model_name,
messages=messages,
temperature=0.7,
max_tokens=2048,
)
return resp.choices[0].message.content
# Wire into dispatchers
def load_api_model(name):
if name == "mathllm-v1":
return load_mathllm_v1(name)
# ... existing models ...
def generate_api_response(name, model, messages):
if name == "mathllm-v1":
return generate_mathllm_v1_response(name, model, messages)
# ... existing models ...
Testing Your New LLM Integration
Once the files are updated, test the integration by instantiating ModelWrapper directly:
from src.models.model_loader import ModelWrapper
# Initialize the wrapper
model = ModelWrapper("mathllm-v1")
# Generate a response
response = model.generate_response("Solve the equation: 2x + 5 = 13")
print(response)
If the model loads without errors and returns a response, the integration is complete. The rest of the CreativeMath pipeline—including generation.py and evaluation.py—will automatically recognize the new model because they all rely on the ModelWrapper abstraction.
Summary
- Declare the model type in
src/models/model_loader.pyby adding the name to theis_api_modellist if it is an API-based model. - Define the prompt template in
src/models/prompt_utils.pyto specify whether the model expects a list of messages or a raw string. - Implement the loader in
src/models/api_models.py(for APIs) orsrc/models/local_models.py(for local HuggingFace models) withload_<model>andgenerate_<model>_responsefunctions. - Wire the functions into the dispatchers
load_api_modelandgenerate_api_response(or their local equivalents) soModelWrappercan route requests correctly. - Optionally expose configuration in
src/config.pyunder"model_config"to make the model selectable via UI or CLI tools.
Frequently Asked Questions
What is the difference between API and local models in CreativeMath?
API models are hosted remotely by third-party providers (such as OpenAI, Anthropic, or Google) and require network requests to generate responses. Local models are loaded directly into memory using the HuggingFace transformers library and run on your own hardware. The ModelWrapper class checks the is_api_model list in model_loader.py to determine which loading path to use.
How do I handle authentication for new API models?
Authentication credentials should be read from environment variables within your loader function in src/models/api_models.py. For example, use os.getenv("MY_API_KEY") to retrieve keys and pass them to the client constructor. Never hardcode credentials in the source files. The existing implementations for OpenAI and Anthropic follow this pattern.
Can I add a local HuggingFace model that requires specific loading arguments?
Yes. When implementing your loader function in src/models/local_models.py, you can pass additional arguments to AutoModelForCausalLM.from_pretrained(). For example, you might specify device_map="auto", torch_dtype=torch.float16, or trust_remote_code=True depending on the model's requirements. Return both the model and tokenizer objects so the generation functions can access them.
Where should I define custom generation parameters for my new model?
Default generation parameters such as temperature, max_tokens, and top_p should be defined in src/config.py under the "model_config" dictionary using your model name as the key. These values are typically read by the generation functions in api_models.py or local_models.py when constructing the API call or the model's generate() method. This centralizes configuration and allows easy adjustment without modifying the loading logic.
Have a question about this repo?
These articles cover the highlights, but your codebase questions are specific. Give your agent direct access to the source. Share this with your agent to get started:
curl -s "https://instagit.com/install.md" Maintain an open-source project? Get it listed too →