Functions
Functions make your code reusable, readable, and testable. Every AI engineering project is built from functions.
Learning Objectives
- Define and call functions with parameters
- Use default and keyword arguments
- Return multiple values with tuples
- Write docstrings for documentation
- Use type hints for clarity
- Understand *args and **kwargs
- Use lambda for simple inline functions
Defining Functions
python
def calculate_cost(tokens: int, price_per_1k: float) -> float:
"""Calculate API cost given token count and price per 1K tokens."""
return (tokens / 1000) * price_per_1k
# Call it
cost = calculate_cost(1500, 0.00015)
print(f"Cost: ${cost:.6f}") # Cost: $0.000225Default & Keyword Arguments
python
def call_llm(
prompt: str,
model: str = "gpt-4o-mini",
temperature: float = 0.7,
max_tokens: int = 1000,
) -> dict:
"""Call an LLM with configurable parameters."""
return {
"model": model,
"prompt": prompt,
"temperature": temperature,
"max_tokens": max_tokens,
}
# Call with defaults
result = call_llm("Explain RAG")
# Override specific args with keywords
result = call_llm("Explain RAG", model="gpt-4o", max_tokens=500)
result = call_llm("Explain RAG", temperature=0.3) # named args!Multiple Return Values
python
def parse_api_response(response: dict) -> tuple[str, int, float]:
"""Extract reply, tokens used, and cost from API response."""
reply = response["choices"][0]["message"]["content"]
tokens = response["usage"]["total_tokens"]
cost = (tokens / 1000) * 0.00015
return reply, tokens, cost
# Unpack the tuple
reply, tokens, cost = parse_api_response(api_response)
print(f"Reply: {reply}")
print(f"Tokens: {tokens}, Cost: ${cost:.6f}")*args and **kwargs
python
# *args = variable number of positional arguments
def total_cost(*prices: float) -> float:
"""Sum any number of API call costs."""
return sum(prices)
total_cost(0.01, 0.02, 0.003) # 0.033
# **kwargs = variable keyword arguments
def build_request(url: str, **headers: str) -> dict:
"""Build an API request with any headers."""
return {"url": url, "headers": headers}
build_request(
"https://api.openai.com/v1/chat",
Authorization="Bearer sk-xxx",
X_Custom="value"
)
# {"url": "...", "headers": {"Authorization": "Bearer sk-xxx", "X_Custom": "value"}}Docstrings — Document Your Functions
python
def retrieve_context(query: str, top_k: int = 5) -> list[dict]:
"""Retrieve relevant document chunks for a RAG query.
Args:
query: The user's search query.
top_k: Number of top results to return. Defaults to 5.
Returns:
List of dicts with keys: text, score, source.
Raises:
ValueError: If query is empty.
Example:
>>> retrieve_context("What is RAG?", top_k=3)
[{"text": "...", "score": 0.92, "source": "doc1.pdf"}]
"""
if not query.strip():
raise ValueError("Query cannot be empty")
# ... actual implementation ...
return []Lambda — Inline Functions
python
# Lambda = one-line anonymous function
# Use sparingly — only for simple transforms
# Sort models by price
models = [("gpt-4o", 0.005), ("gpt-4o-mini", 0.00015), ("claude-opus-4-7", 0.015)]
sorted_models = sorted(models, key=lambda m: m[1])
# [("gpt-4o-mini", 0.00015), ("gpt-4o", 0.005), ("claude-opus-4-7", 0.015)]
# Filter with lambda
cheap = list(filter(lambda m: m[1] < 0.001, models))
# [("gpt-4o-mini", 0.00015)]
# Map with lambda
names = list(map(lambda m: m[0], models))
# ["gpt-4o", "gpt-4o-mini", "claude-opus-4-7"]Exercise — Cost Calculator
Write a function estimate_monthly_cost(daily_requests, avg_tokens, price_per_1k) that returns the daily and monthly cost. Add a docstring, type hints, and a default price of $0.00015. Test it with 1000 daily requests averaging 500 tokens.