Code Generation with Codex: Building Smart Python Code with AI

Why Use Codex for Code Generation?

OpenAI Codex, a descendant of GPT-3, excels at translating natural language into code, trained on billions of lines from public repositories like GitHub—see What Is Generative AI and Why Use It?. Prompt it with “Write a vector search function,” and it delivers working Python code, complete with logic and syntax. It’s fast, turning ideas into scripts in seconds, versatile, supporting languages beyond Python (e.g., JavaScript, Go), and accessible, with a free tier ($5 credit) or low cost (~$0.002/1000 tokens).

Why choose it? It boosts productivity, automating repetitive coding tasks. It’s educational, showing how prompts shape output, and scalable, with embeddings storable in Weaviate for search or analysis. Validating with unit tests and testing complex prompts ensures reliability. Let’s set it up naturally.

Step 1: Use OpenAI Codex API for Python Code

Start by setting up the OpenAI Codex API to generate Python code from a natural language prompt.

Coding with Codex

Ensure Python 3.8+ and pip are installed—see Setting Up Your Creative Environment—and install libraries:

Copy codepip install openai python-dotenv

Get an API key from platform.openai.com—free tier offers $5 credit—create a key (e.g., “Codex2025”), and add it to .env in a folder like “CodexCode”:

Copy codemkdir CodexCode
cd CodexCode

Add to .env:

Copy codeOPENAI_API_KEY=sk-abc123xyz

Create codex_gen.py:

Copy codeimport openai
from dotenv import load_dotenv
import os

# Load API key
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")

# Generate code with Codex
prompt = "Write a Python function to perform a vector search."
response = openai.Completion.create(
    model="code-davinci-002",
    prompt=prompt,
    max_tokens=150
)

# Display generated code
code = response.choices[0].text.strip()
print("Generated Code:")
print(code)

Run python codex_gen.py, and expect:

Copy codeGenerated Code:
def vector_search(query, vectors):
    """Perform a simple vector search using dot product similarity."""
    best_match = None
    highest_score = float('-inf')
    for vector in vectors:
        score = sum(q * v for q, v in zip(query, vector))
        if score > highest_score:
            highest_score = score
            best_match = vector
    return best_match

How It Works

• .env file: Keeps your OPENAI_API_KEY safe, loading it securely into the script.
• load_dotenv(): Pulls the key from .env, and os.getenv() sets it for openai.api_key to connect to the API.
• prompt: Tells Codex to write a vector search function in Python, guiding its output.
• openai.Completion.create: Calls the API with parameters:
• model="code-davinci-002": Uses Codex’s code-davinci-002, optimized for code generation. Advanced use: Switch to text-davinci-003 for more nuanced code if available.
• prompt=prompt: Feeds the natural language request, shaping the function’s purpose.
• max_tokens=150: Limits output to 150 tokens (~120-150 words), enough for a small function. Advanced use: Set to 500 for a detailed implementation with comments.
• code = response.choices[0].text.strip(): Grabs the generated code, trimming extra spaces for clarity.

This generates your Python function—next, validate it.

Step 2: Prompt “Write a Vector Search Function”

Prompt Codex specifically with “Write a vector search function” and refine the output.

Coding the Prompt

Update codex_gen.py:

Copy codeimport openai
from dotenv import load_dotenv
import os

# Load API key
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")

# Specific prompt
prompt = "Write a Python function to perform a vector search using cosine similarity."
response = openai.Completion.create(
    model="code-davinci-002",
    prompt=prompt,
    max_tokens=200,
    temperature=0.5,
    top_p=0.9,
    frequency_penalty=0.2,
    presence_penalty=0.1,
    n=1,
    stop=None
)

# Display generated code
code = response.choices[0].text.strip()
print("Generated Vector Search Function:")
print(code)

Run python codex_gen.py, and expect:

Copy codeGenerated Vector Search Function:
import math

def vector_search(query, vectors):
    """Perform vector search using cosine similarity."""
    def cosine_similarity(v1, v2):
        dot_product = sum(a * b for a, b in zip(v1, v2))
        norm_v1 = math.sqrt(sum(a * a for a in v1))
        norm_v2 = math.sqrt(sum(b * b for b in v2))
        return dot_product / (norm_v1 * norm_v2) if norm_v1 * norm_v2 != 0 else 0

    best_match = None
    highest_similarity = -1
    for vector in vectors:
        similarity = cosine_similarity(query, vector)
        if similarity > highest_similarity:
            highest_similarity = similarity
            best_match = vector
    return best_match

How It Works

• prompt: Specifies “cosine similarity” for a precise vector search function, guiding Codex’s logic.
• openai.Completion.create: Generates code with all parameters:
• model="code-davinci-002": Uses Codex’s coding model. Advanced use: Test text-davinci-003 for improved logic if accessible.
• prompt=prompt: Drives the function’s purpose with the specific request.
• max_tokens=200: Allows ~150-200 words, fitting a function with helper logic. Advanced use: Increase to 400 for a version with edge cases handled.
• temperature=0.5: Keeps output focused and technical, avoiding wild variations. Advanced use: Set to 0.8 for a less strict, exploratory version.
• top_p=0.9: Selects from the top 90% probable tokens, balancing variety and accuracy. Advanced use: Drop to 0.6 for stricter, formal code phrasing.
• frequency_penalty=0.2: Reduces word repetition (0-2 range), favoring “similarity” over redundant terms. Advanced use: Raise to 1.0 for diverse terms in complex algorithms.
• presence_penalty=0.1: Encourages new ideas (0-2 range), adding slight novelty. Advanced use: Set to 0.5 for broader logic in research code.
• n=1: Produces one response. Advanced use: Set to 3 for multiple drafts to refine.
• stop=None: Runs to max_tokens. Advanced use: Use ["def"] to stop at function ends for multi-function outputs.
• code: Extracts the cleaned-up code, ready for validation.

This refines your vector search function—next, test it.

Step 3: Validate Code with Unit Tests in Python

Validate the generated code with unit tests to ensure it works correctly.

Coding Unit Tests

Update codex_gen.py:

Copy codeimport openai
from dotenv import load_dotenv
import os
import math
import unittest

# Load API key
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")

# Generate code
prompt = "Write a Python function to perform a vector search using cosine similarity."
response = openai.Completion.create(
    model="code-davinci-002",
    prompt=prompt,
    max_tokens=200,
    temperature=0.5,
    top_p=0.9,
    frequency_penalty=0.2,
    presence_penalty=0.1,
    n=1,
    stop=None
)
code = response.choices[0].text.strip()

# Define the generated function (exec for dynamic loading)
exec(code)

# Unit tests
class TestVectorSearch(unittest.TestCase):
    def test_identical_vectors(self):
        query = [1, 0, 0]
        vectors = [[1, 0, 0], [0, 1, 0]]
        result = vector_search(query, vectors)
        self.assertEqual(result, [1, 0, 0], "Should match identical vector")

    def test_different_vectors(self):
        query = [1, 0, 0]
        vectors = [[0, 1, 0], [0, 0, 1]]
        result = vector_search(query, vectors)
        self.assertEqual(result, [0, 1, 0], "Should pick closest vector")

# Run tests
if __name__ == "__main__":
    print("Generated Code:")
    print(code)
    print("\nRunning Unit Tests:")
    unittest.main(argv=[''], exit=False)

Run python codex_gen.py, and expect:

Copy codeGenerated Code:
import math

def vector_search(query, vectors):
    """Perform vector search using cosine similarity."""
    def cosine_similarity(v1, v2):
        dot_product = sum(a * b for a, b in zip(v1, v2))
        norm_v1 = math.sqrt(sum(a * a for a in v1))
        norm_v2 = math.sqrt(sum(b * b for b in v2))
        return dot_product / (norm_v1 * norm_v2) if norm_v1 * norm_v2 != 0 else 0

    best_match = None
    highest_similarity = -1
    for vector in vectors:
        similarity = cosine_similarity(query, vector)
        if similarity > highest_similarity:
            highest_similarity = similarity
            best_match = vector
    return best_match

Running Unit Tests:
..
----------------------------------------------------------------------
Ran 2 tests in 0.001s

OK

How It Works

• exec(code): Loads the generated function dynamically into the script, making it testable.
• TestVectorSearch: Defines a test class with:
• test_identical_vectors: Checks if the function picks an identical vector (cosine similarity = 1).
• test_different_vectors: Tests if it selects the closest vector among non-identical options.
• unittest.main(...): Runs the tests, showing “OK” if both pass, confirming the function works.

This validates your code’s correctness—next, store embeddings.

Step 4: Store Code Embeddings in Weaviate

Convert the code to embeddings and store them in Weaviate, referencing Configure Weaviate Multimodal.

Coding Weaviate Storage

Install Weaviate client and run a local instance (e.g., Docker: docker run -p 8080:8080 semitechnologies/weaviate):

Copy codepip install weaviate-client

Update codex_gen.py:

Copy codeimport openai
import weaviate
from dotenv import load_dotenv
import os
import math
import unittest

# Load API keys
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")

# Generate code
prompt = "Write a Python function to perform a vector search using cosine similarity."
response = openai.Completion.create(
    model="code-davinci-002",
    prompt=prompt,
    max_tokens=200,
    temperature=0.5,
    top_p=0.9,
    frequency_penalty=0.2,
    presence_penalty=0.1,
    n=1,
    stop=None
)
code = response.choices[0].text.strip()
exec(code)

# Connect to Weaviate
client = weaviate.Client("http://localhost:8080")

# Define schema (run once)
schema = {
    "classes": [{
        "class": "CodeSnippet",
        "properties": [
            {"name": "code", "dataType": ["text"]},
            {"name": "prompt", "dataType": ["string"]}
        ],
        "vectorizer": "none"
    }]
}
if not client.schema.contains(schema):
    client.schema.create(schema)

# Generate embedding
embedding_response = openai.Embedding.create(model="text-embedding-ada-002", input=code)
vector = embedding_response["data"][0]["embedding"]

# Store in Weaviate
client.data_object.create(
    data_object={"code": code, "prompt": prompt},
    class_name="CodeSnippet",
    vector=vector
)

# Confirm storage
print("Generated Code Stored in Weaviate:")
print(code[:50] + "...")
print(f"Metadata: { {'prompt': '{prompt}'} }")

Run python codex_gen.py, and expect:

Copy codeGenerated Code Stored in Weaviate:
import math

def vector_search(query, vect...
Metadata: {'prompt': 'Write a Python function to perform a vector search using cosine similarity.'}

How It Works

• client = weaviate.Client(...): Connects to a local Weaviate instance at port 8080, ready to store data.
• schema: Sets up a “CodeSnippet” class with code and prompt properties, defining how data is organized.
• embedding_response: Creates a 1536D vector from the code using text-embedding-ada-002, capturing its meaning.
• client.data_object.create(...): Stores the code, prompt, and vector in Weaviate, linking them for retrieval.
• Print: Shows a snippet and metadata, confirming storage.

This saves your code as vectors—next, test with complex prompts.

Step 5: Test with Complex Algorithm Prompts

Test Codex with complex algorithm prompts to push its limits and validate versatility.

Coding Complex Prompt Test

Update codex_gen.py:

Copy codeimport openai
import weaviate
from dotenv import load_dotenv
import os
import math
import unittest

# Load API keys
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")

# Generate complex code
complex_prompt = "Write a Python function to implement Dijkstra's algorithm for shortest path in a graph."
response = openai.Completion.create(
    model="code-davinci-002",
    prompt=complex_prompt,
    max_tokens=400,
    temperature=0.5,
    top_p=0.9,
    frequency_penalty=0.2,
    presence_penalty=0.1,
    n=1,
    stop=None
)
complex_code = response.choices[0].text.strip()

# Display complex code
print("Complex Algorithm Output (Dijkstra's):")
print(complex_code)

Run python codex_gen.py, and expect:

Copy codeComplex Algorithm Output (Dijkstra's):
def dijkstra(graph, start):
    """Implement Dijkstra's algorithm to find shortest paths from start node."""
    distances = {node: float('infinity') for node in graph}
    distances[start] = 0
    unvisited = set(graph.keys())
    previous = {node: None for node in graph}

    while unvisited:
        current = min(unvisited, key=lambda node: distances[node])
        if distances[current] == float('infinity'):
            break

        unvisited.remove(current)
        for neighbor, weight in graph[current].items():
            if neighbor in unvisited:
                new_distance = distances[current] + weight
                if new_distance < distances[neighbor]:
                    distances[neighbor] = new_distance
                    previous[neighbor] = current

    return distances, previous

# Example usage
graph = {
    'A': {'B': 4, 'C': 2},
    'B': {'A': 4, 'C': 1, 'D': 5},
    'C': {'A': 2, 'B': 1, 'D': 8},
    'D': {'B': 5, 'C': 8}
}
distances, previous = dijkstra(graph, 'A')
print(distances)

How It Works

• complex_prompt: Asks for Dijkstra’s algorithm, a complex graph problem, testing Codex’s reasoning.
• openai.Completion.create: Generates code with parameters:
• model="code-davinci-002": Relies on Codex’s coding strength. Advanced use: Try text-davinci-003 for better handling of complex logic.
• prompt=complex_prompt: Drives the algorithm’s implementation.
• max_tokens=400: Allows ~300-400 words, fitting a full algorithm with example. Advanced use: Set to 1000 for a version with detailed comments and edge cases.
• temperature=0.5: Ensures precise, technical output. Advanced use: Raise to 0.7 for a more interpretive take with alternative approaches.
• top_p=0.9: Balances variety and accuracy. Advanced use: Lower to 0.5 for strict adherence to standard Dijkstra’s.
• frequency_penalty=0.2: Reduces repetition, favoring “distances” over redundant terms. Advanced use: Increase to 0.8 for varied variable names.
• presence_penalty=0.1: Adds slight novelty. Advanced use: Set to 0.4 for exploring unconventional graph methods.
• n=1: Produces one version. Advanced use: Set to 2 for comparing implementations.
• stop=None: Runs to max_tokens. Advanced use: Use ["#"] to stop at comments for modular code.
• complex_code: Extracts the full algorithm, ready to run or store.

This tests Codex’s complexity handling—your pipeline’s versatile!

Next Steps: Expanding Your Codex Skills

Your Codex setup is generating, validating, and storing code smoothly! Try prompts like “Implement a binary search tree” or integrate with Text Embeddings with OpenAI. You’ve mastered code generation, so keep coding and exploring!

FAQ: Common Questions About Code Generation with Codex

1. Can Codex handle other languages?

Yes, it supports JavaScript, Go, and more—tweak the prompt to specify.

2. Why validate with unit tests?

It ensures the code works, catching errors Codex might miss.

3. What if the code fails tests?

Refine the prompt or lower temperature for tighter logic.

4. How does Weaviate help?

It stores code embeddings for search or analysis—see Configure Weaviate Multimodal.

5. Can I use complex prompts directly?

Yes, but start simple—complexity works better with clear instructions.

6. Why test with algorithms?

It pushes Codex’s limits, showing its real-world potential.

Your questions are covered—code with confidence!