Code Completion › Implementation Walkthrough
Michael BrenndoerferFebruary 3, 20268 min read
Language AI HandbookAll chapters Part 47 of 67Code Generation
Code Completion › Implementation WalkthroughOne section · Sections in this chapter6 / 10 ↓
Sections in this chapter10
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
To make these concepts concrete, let's build a minimal code completion system. We will implement context assembly, simple candidate scoring, and a simulation of the full completion pipeline.
Setting Up the Environment
We start by importing the libraries we need.
Context Assembly
The context assembler takes the full file text and cursor position, and returns a trimmed context that fits within a token budget.
In[13]:
Code@dataclass
class CompletionContext:
prefix: str
suffix: str
retrieved_snippets: list[str] = field(default_factory=list)
def approximate_token_count(text: str) -> int:
"""Rough estimate: ~4 characters per token for code."""
return max(1, len(text) // 4)
def assemble_context(
file_text: str,
cursor_pos: int,
retrieved_snippets: list[str],
token_budget: int = 2048,
) -> CompletionContext:
"""
Assemble a completion context within the given token budget.
Allocates ~60% to prefix, ~20% to suffix, ~20% to retrieved snippets.
"""
prefix_budget = int(token_budget * 0.60)
suffix_budget = int(token_budget * 0.20)
retrieval_budget = token_budget - prefix_budget - suffix_budget
raw_prefix = file_text[:cursor_pos]
raw_suffix = file_text[cursor_pos:]
# Truncate prefix from the left if too long (keep most recent context)
prefix_tokens = approximate_token_count(raw_prefix)
if prefix_tokens > prefix_budget:
# Estimate how many characters to keep
chars_to_keep = prefix_budget * 4
raw_prefix = raw_prefix[-chars_to_keep:]
# Truncate suffix from the right if too long
suffix_tokens = approximate_token_count(raw_suffix)
if suffix_tokens > suffix_budget:
chars_to_keep = suffix_budget * 4
raw_suffix = raw_suffix[:chars_to_keep]
# Include retrieved snippets up to budget
included_snippets = []
remaining_budget = retrieval_budget
for snippet in retrieved_snippets:
snippet_tokens = approximate_token_count(snippet)
if snippet_tokens <= remaining_budget:
included_snippets.append(snippet)
remaining_budget -= snippet_tokens
return CompletionContext(
prefix=raw_prefix,
suffix=raw_suffix,
retrieved_snippets=included_snippets,
)In[14]:
Code# Demo: assemble context for a file with cursor mid-function
sample_file = """import numpy as np
import pandas as pd
def compute_statistics(data: list[float]) -> dict:
\"\"\"Compute basic statistics for a list of numbers.\"\"\"
n = len(data)
mean = sum(data) / n
variance = sum((x - mean) ** 2 for x in data) / n
# cursor is here
return {
'n': n,
'mean': mean,
'variance': variance,
'std': std,
}
def normalize(data: list[float]) -> list[float]:
stats = compute_statistics(data)
return [(x - stats['mean']) / stats['std'] for x in data]
"""
cursor_position = sample_file.index("# cursor is here")
retrieved = [
"import math\ndef std_dev(data):\n n = len(data)\n mean = sum(data)/n\n return math.sqrt(sum((x-mean)**2 for x in data)/n)",
]
ctx = assemble_context(
sample_file, cursor_position, retrieved, token_budget=512
)Out[15]:
Console=== Context Assembly Result ===
Prefix length: 250 chars (~62 tokens)
Suffix length: 273 chars (~68 tokens)
Retrieved snippets included: 1
--- Prefix (last 200 chars) ---
'e_statistics(data: list[float]) -> dict:\n """Compute basic statistics for a list of numbers."""\n n = len(data)\n mean = sum(data) / n\n variance = sum((x - mean) ** 2 for x in data) / n\n '
--- Suffix (first 100 chars) ---
"# cursor is here\n return {\n 'n': n,\n 'mean': mean,\n 'variance': variance,\n "
--- Retrieved snippet ---
'import math\ndef std_dev(data):\n n = len(data)\n mean = sum(data)/n\n return math.sqrt(sum((x-mean)**2 for x in data)/n)'The assembler correctly keeps the recent prefix, includes the suffix starting from the cursor, and fits the retrieved snippet within the remaining budget. Notice that the prefix is truncated from the left (discarding distant earlier code) rather than from the right, so the code immediately before the cursor is always preserved.
Candidate Scoring
Next, we implement the length-normalized log-probability scoring function and a simple filter for identifier validity.
In[16]:
Codedef score_completion(token_log_probs: list[float]) -> float:
"""
Compute the mean log-probability score for a completion.
Higher (less negative) scores indicate more confident completions.
"""
if not token_log_probs:
return float("-inf")
return sum(token_log_probs) / len(token_log_probs)
def extract_identifiers_from_context(prefix: str) -> set[str]:
"""Extract all Python-style identifiers defined in the prefix."""
# Find all potential identifier definitions (simple heuristic)
pattern = r"\b([a-zA-Z_][a-zA-Z0-9_]*)\s*(?:=|:|\()"
return set(re.findall(pattern, prefix))
def check_identifier_validity(
completion: str,
context_identifiers: set[str],
stdlib_builtins: Optional[set[str]] = None,
) -> tuple[bool, list[str]]:
"""
Check whether identifiers in the completion appear in scope.
Returns (is_valid, list_of_unknown_identifiers).
"""
if stdlib_builtins is None:
stdlib_builtins = {
"len",
"sum",
"range",
"print",
"str",
"int",
"float",
"list",
"dict",
"set",
"tuple",
"bool",
"None",
"True",
"False",
"return",
"def",
"class",
"import",
"from",
"if",
"else",
"elif",
"for",
"while",
"in",
"not",
"and",
"or",
"math",
"np",
"pd",
"sqrt",
}
all_known = context_identifiers | stdlib_builtins
completion_identifiers = set(
re.findall(r"\b([a-zA-Z_][a-zA-Z0-9_]*)\b", completion)
)
unknown = completion_identifiers - all_known
return len(unknown) == 0, list(unknown)In[17]:
Code# Simulate multiple candidate completions with their token log-probs
candidates = [
{
"text": " std = math.sqrt(variance)",
"token_log_probs": [
-0.1,
-0.2,
-0.15,
-0.3,
-0.1,
-0.25,
], # high confidence
},
{
"text": " std = np.sqrt(variance)",
"token_log_probs": [-0.2, -0.3, -0.4, -0.35, -0.2, -0.28],
},
{
"text": " std = variance ** 0.5",
"token_log_probs": [
-0.3,
-0.4,
-0.8,
-1.2,
-0.6,
-0.35,
], # lower confidence on **
},
{
"text": " std = compute_std(data)", # hallucinated function
"token_log_probs": [-0.5, -0.6, -1.5, -1.8, -0.7, -0.4],
},
]
context_identifiers = extract_identifiers_from_context(ctx.prefix)
results = []
for candidate in candidates:
score = score_completion(candidate["token_log_probs"])
is_valid, unknown_ids = check_identifier_validity(
candidate["text"], context_identifiers
)
results.append(
{
"text": candidate["text"],
"score": score,
"is_valid": is_valid,
"unknown_ids": unknown_ids,
}
)
# Sort by score descending, filtering invalid completions to the bottom
valid_results = sorted(
[r for r in results if r["is_valid"]],
key=lambda x: x["score"],
reverse=True,
)
invalid_results = [r for r in results if not r["is_valid"]]
ranked_results = valid_results + invalid_resultsOut[18]:
Console=== Ranked Completions === Rank 1: FILTERED (unknown: ['std']) Score (mean log-prob): -0.1833 Completion: std = math.sqrt(variance) Rank 2: FILTERED (unknown: ['std']) Score (mean log-prob): -0.2883 Completion: std = np.sqrt(variance) Rank 3: FILTERED (unknown: ['std']) Score (mean log-prob): -0.6083 Completion: std = variance ** 0.5 Rank 4: FILTERED (unknown: ['std', 'compute_std']) Score (mean log-prob): -0.9167 Completion: std = compute_std(data)
The scoring correctly places high-confidence completions at the top, while the hallucinated compute_std function is filtered out because the identifier is not in scope. The filter catches the most egregious hallucination (a completely nonexistent function name) before any more expensive reranking steps are applied.
Simulating End-to-End Latency
Finally, let's simulate the end-to-end latency model to understand how the different components stack up.
In[19]:
Code@dataclass
class LatencyProfile:
retrieval_ms: float
context_assembly_ms: float
network_ms: float
prefill_ms: float
generation_ms_per_token: float
def simulate_completion_latency(
n_output_tokens: int,
profile: LatencyProfile,
use_prefix_cache: bool = False,
use_speculative_decoding: bool = False,
spec_acceptance_rate: float = 0.75,
spec_draft_tokens: int = 4,
) -> dict[str, float]:
"""
Simulate completion latency given a latency profile and optimizations.
Returns a breakdown of latency components.
"""
retrieval = profile.retrieval_ms
assembly = profile.context_assembly_ms
network = profile.network_ms
prefill = profile.prefill_ms
if use_prefix_cache:
# Cache hit reduces prefill by ~80%
prefill *= 0.20
if use_speculative_decoding:
# Each speculative step generates spec_draft_tokens, accepted with spec_acceptance_rate
# Effective throughput multiplier
expected_accepted = spec_draft_tokens * spec_acceptance_rate
effective_tokens_per_step = (
1 + expected_accepted
) # 1 for verification overhead
total_steps = n_output_tokens / effective_tokens_per_step
# One target-model verification step plus a cheaper draft-model pass.
# The draft model is approximated as one quarter of the target cost.
verification_cost = profile.generation_ms_per_token
draft_cost = 0.25 * profile.generation_ms_per_token * spec_draft_tokens
generation = total_steps * (verification_cost + draft_cost)
else:
generation = n_output_tokens * profile.generation_ms_per_token
total = retrieval + assembly + network + prefill + generation
return {
"retrieval_ms": retrieval,
"context_assembly_ms": assembly,
"network_ms": network,
"prefill_ms": prefill,
"generation_ms": generation,
"total_ms": total,
}
# Typical profile for a mid-size model served remotely
base_profile = LatencyProfile(
retrieval_ms=30,
context_assembly_ms=15,
network_ms=50,
prefill_ms=80,
generation_ms_per_token=8,
)
output_token_counts = [5, 10, 20, 50, 100, 150]
scenarios = {
"Baseline": dict(use_prefix_cache=False, use_speculative_decoding=False),
"Prefix Cache": dict(use_prefix_cache=True, use_speculative_decoding=False),
"Speculative Decoding": dict(
use_prefix_cache=False, use_speculative_decoding=True
),
"Both Optimizations": dict(
use_prefix_cache=True, use_speculative_decoding=True
),
}
latency_results = {
scenario: [
simulate_completion_latency(n, base_profile, **kwargs)["total_ms"]
for n in output_token_counts
]
for scenario, kwargs in scenarios.items()
}Out[20]:
Console=== Simulated Total Latency (ms) by Output Token Count === Output tokens Baseline Prefix Cache Speculative Decoding Both Optimizations ------------------------------------------------------------------------------------------------------------------- 5 215.0 151.0 195.0 131.0 10 255.0 191.0 215.0 151.0 20 335.0 271.0 255.0 191.0 50 575.0 511.0 375.0 311.0 100 975.0 911.0 575.0 511.0 150 1375.0 1311.0 775.0 711.0
Out[21]:
Visualization
Combining prefix caching and speculative decoding keeps latency under 400ms for completions up to about 50 tokens. For longer completions (function bodies), streaming delivery is essential to avoid blocking the developer while still surfacing useful suggestions.