Building a Basic Code Translation Pipeline

Welcome to the second lesson of our course, Laying the Foundations for Code Translation with Haystack! In our first lesson, you learned how to use Haystack to translate natural language text. Now, we're ready to tackle a more challenging and practical task: building a pipeline that translates code from one programming language to another.

Automated code translation is a powerful tool for developers. Whether you're porting a project, learning a new language, or collaborating with teams using different tech stacks, being able to translate code quickly and accurately can save hours of work. By the end of this lesson, you'll have a working code translation pipeline that you can use and extend for your own projects.

The heart of our pipeline is the prompt we give to the language model. A good prompt makes the model's job clear and helps ensure accurate translations. Let's see how to build a prompt template that's both flexible and effective:

from haystack.components.builders import PromptBuilder

# Create a prompt template for code translation
prompt_template = (
    "Translate the following {{source_lang}} code to {{target_lang}}:\n"
    "{{code}}"
)

# Initialize the PromptBuilder with required variables
prompt_builder = PromptBuilder(
    template=prompt_template,
    required_variables=["source_lang", "target_lang", "code"]
)

Here, we're telling the model exactly what we want: a translation from one programming language to another. By using variables like {{source_lang}}, {{target_lang}}, and {{code}}, we make our prompt reusable for any language pair and any code snippet. This approach keeps our pipeline adaptable and easy to extend.

With our prompt in place, the next step is to set up the components that will power our translation pipeline. We'll use Haystack's modular design to keep things organized and maintainable.

from haystack import Pipeline

def build_code_translation_pipeline(model_name="gpt-4o-mini"):
    # Set up the prompt and generator as before
    prompt_template = (
        "Translate the following {{source_lang}} code to {{target_lang}}:\n"
        "{{code}}"
    )
    prompt_builder = PromptBuilder(
        template=prompt_template,
        required_variables=["source_lang", "target_lang", "code"]
    )
    generator = OpenAIGenerator(model=model_name)

    # Assemble the pipeline and connect components
    pipeline = Pipeline()
    pipeline.add_component("prompt_builder", prompt_builder)
    pipeline.add_component("llm", generator)
    pipeline.connect("prompt_builder.prompt", "llm.prompt")
    return pipeline

This function creates a pipeline that takes in code and language specifications, builds a prompt, and sends it to the language model using the OpenAIGenerator. By wrapping this logic in a function, we make it easy to create pipelines for any language pair.