Switching from Direct calls to OpenAI to Langchain

Migrating from Direct API Integration to LangChain

Introduction

When we first integrated AI features into our DevTools, we directly used OpenAI's API. This approach quickly led to code duplication across different tools.

For example:

if (extraContext) {
  parts.push(`extra context: ${extraContext}`);
}

const userMessage = parts.join("\n");

To streamline our process, we began adding a framework that included input schema, output schema, and function calling. Initially, we enforced this with custom code:

args: [{
  name: 'input',
  type: 'json',
  description: 'Input used',
  required: true,
  limitChars: 512
}]

Introducing LangChain

In our search for a better solution, we discovered LangChain, an open-source framework designed to address these challenges. LangChain allows you to define AI chains in a declarative and composable manner.

For example, defining a prompt template:

const jqPrompt = ChatPromptTemplate.fromMessages([
    ["system", "You are a JQ expert. With given inputs and outputs and extra context, you can call a function with a single JQ expression to satisfy the input-output pair."],
    ["user", "Help me write JQ for this input: `${input}`; expected output `${output}`"],
    ["user", "{extraContext}"],
]);

You can define your schema using the 'zod' package:

const jqInputSchema = z.object({
    extraContext: z.string().describe("Prompt to generate Graphviz/Dot markup"),
});

To enforce structured output from OpenAI, we use 'tools/functions', for example:

function buildOneFunctionParams(description, zodSchema) {
    return {
        functions: [{
          name: genericFunctionName,
          description: description,
          parameters: zodToJsonSchema(zodSchema),
        }],
        function_call: { name: genericFunctionName },
    };
}

const jqFunctionSchema = z.object({
    expression: z.string().describe("JQ expression"),
});
const jqLlmParams = buildOneFunctionParams("Accept JQ Expression", jqFunctionSchema);

Then we build chains:

const chatModel = new ChatOpenAI({
    apiKey: process.env['OPENAI_API_KEY'],
    model: 'gpt-4'
});

const regexChain = RunnableSequence.from([
    jqPrompt,
    chatModel.bind(jqLlmParams),
    genericJsonOutputFunctionsParser
]);

The chain returns JSON that matches the schema of a function call, for example:

{
    "expression": ".[] | select(.name == 'John')"
}

Pros and Cons of Using LangChain

Cons

1. Another framework to learn.
2. New OpenAI features are not available immediately.

Pros

1. Less boilerplate code.
2. Potentially less direct dependency on the OpenAI API, allowing for easier swapping with alternatives like LLAMA or other APIs.

Try it out

1. JQ playground with AI
2. MermaidJS playground
3. Graphviz playground
4. Regex playground