main.py

from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain, SequentialChain
from dotenv import load_dotenv

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--task", help="Task to be performed", type=str, default="add numbers 0 to 10")
parser.add_argument("--language", help="Language to be used", type=str, default="python")
args = parser.parse_args()


load_dotenv()

# API key is stored in .env file
llm= OpenAI()

code_prompt = PromptTemplate(
  template="Write a very short {language} function that will {task}.",
  input_variables=[
    "language", "task"
  ]
)

test_prompt = PromptTemplate(
  input_variables=[
    "language", "code"
  ],
  template="Write a unit test for the following {language} code:\n{code}"
)

code_chain = LLMChain(
  llm=llm,
  prompt=code_prompt,
  output_key="code"
)

test_chain = LLMChain(
  llm=llm,
  prompt=test_prompt,
  output_key="test"
)

chain = SequentialChain(
  chains=[
    code_chain,
    test_chain
  ],
  input_variables=["language", "task"],
  output_variables=["code", "test"]
)

result = chain({
  "language": args.language,
  "task": args.task
})

print(result)