From eea57a264fc99ee48ebe06348defce47bc71251e Mon Sep 17 00:00:00 2001
From: Nicole White <nicole@autoblocks.ai>
Date: Mon, 23 Oct 2023 16:12:59 -0400
Subject: [PATCH] Improve langchain examples (#16)

Making these examples slightly more verbose / complicated so they can
better demonstrate AB
---
 JavaScript/langchain/src/index.js | 47 ++++++++++++++++++++++++++-----
 Python/langchain/main.py          | 32 +++++++++++++++++++--
 2 files changed, 70 insertions(+), 9 deletions(-)

diff --git a/JavaScript/langchain/src/index.js b/JavaScript/langchain/src/index.js
index 313b83a8..1f385b27 100644
--- a/JavaScript/langchain/src/index.js
+++ b/JavaScript/langchain/src/index.js
@@ -1,16 +1,49 @@
 import { AutoblocksCallbackHandler } from '@autoblocks/client/langchain';
 
-import { LLMChain } from 'langchain/chains';
-import { OpenAI } from 'langchain/llms/openai';
-import { PromptTemplate } from 'langchain/prompts';
+import { SimpleSequentialChain, LLMChain } from "langchain/chains";
+import { OpenAI } from "langchain/llms/openai";
+import { PromptTemplate } from "langchain/prompts";
 
 const main = async () => {
-  const llm = new OpenAI({ temperature: 0 });
-  const prompt = PromptTemplate.fromTemplate('2 + {number} =');
-  const chain = new LLMChain({ prompt, llm });
+  // This is an LLMChain to write a synopsis given a title of a play.
+  const synopsisLLM = new OpenAI({ temperature: 0.7 });
+  const titleTemplate = `You are a playwright. Given the title of a play, it is your job to write a synopsis for that title.
   
+Title: {title}
+Playwright: This is a synopsis for the above play:`;
+  const synopsisTemplate = new PromptTemplate({
+    template: titleTemplate,
+    inputVariables: ["title"],
+  });
+  const synopsisChain = new LLMChain({ llm: synopsisLLM, prompt: synopsisTemplate });
+
+  // This is an LLMChain to write a review of a play given a synopsis.
+  const reviewLLM = new OpenAI({ temperature: 0.7 });
+  const reviewTemplate = `You are a play critic from the New York Times. Given the synopsis of a play, it is your job to write a review for that play.
+  
+Play Synopsis:
+{synopsis}
+Review from a New York Times play critic of the above play:`;
+  const reviewPromptTemplate = new PromptTemplate({
+    template: reviewTemplate,
+    inputVariables: ["synopsis"],
+  });
+  const reviewChain = new LLMChain({
+    llm: reviewLLM,
+    prompt: reviewPromptTemplate,
+  });
+
+  // This is the overall chain where we run these two chains in sequence.
+  const overallChain = new SimpleSequentialChain({
+    chains: [synopsisChain, reviewChain],
+  });
+
+  // Run the chain
   const handler = new AutoblocksCallbackHandler();
-  await chain.call({ number: 2 }, { callbacks: [handler] });
+  const review = await overallChain.run("Tragedy at sunset on the beach", { callbacks: [handler] });
+
+  console.log('Review:');
+  console.log(review);
 }
 
 main();
diff --git a/Python/langchain/main.py b/Python/langchain/main.py
index bc130d9b..1c9c763d 100644
--- a/Python/langchain/main.py
+++ b/Python/langchain/main.py
@@ -1,11 +1,39 @@
 import dotenv
 from autoblocks.vendor.langchain import AutoblocksCallbackHandler
+from langchain.chains import LLMChain
+from langchain.chains import SimpleSequentialChain
 from langchain.llms import OpenAI
+from langchain.prompts import PromptTemplate
 
 dotenv.load_dotenv("../../.env")
 
 
 if __name__ == "__main__":
-    llm = OpenAI()
+    # This is an LLMChain to write a synopsis given a title of a play.
+    synopsis_llm = OpenAI(temperature=0.7)
+    synopsis_template = """You are a playwright. Given the title of a play, it is your job to write a synopsis for that title.
+
+Title: {title}
+Playwright: This is a synopsis for the above play:"""
+    synopsis_prompt_template = PromptTemplate(input_variables=["title"], template=synopsis_template)
+    synopsis_chain = LLMChain(llm=synopsis_llm, prompt=synopsis_prompt_template)
+
+    # This is an LLMChain to write a review of a play given a synopsis.
+    review_llm = OpenAI(temperature=0.7)
+    review_template = """You are a play critic from the New York Times. Given the synopsis of a play, it is your job to write a review for that play.
+
+Play Synopsis:
+{synopsis}
+Review from a New York Times play critic of the above play:"""
+    review_prompt_template = PromptTemplate(input_variables=["synopsis"], template=review_template)
+    review_chain = LLMChain(llm=review_llm, prompt=review_prompt_template)
+
+    # This is the overall chain where we run these two chains in sequence.
+    overall_chain = SimpleSequentialChain(chains=[synopsis_chain, review_chain])
+
+    # Run the chain
     handler = AutoblocksCallbackHandler()
-    llm.predict("hello, world!", callbacks=[handler])
+    review = overall_chain.run("Tragedy at sunset on the beach", callbacks=[handler])
+
+    print("Review:")
+    print(review)