jflam · January 16, 2023 02:55 · Jan 16, 2023
diff --git a/app.py b/app.py
@@ -0,0 +1,69 @@
+# To run you'll need some secrets:
+# 1. SERPAPI_API_KEY secret in env var - get from https://serpapi.com/
+# 2. OPENAI_API_KEY secret in env var - get from https://openai.com
+
+import streamlit as st
+import json, os
+from langchain.prompts import PromptTemplate
+from langchain.llms import OpenAI
+from serpapi import GoogleSearch
+
+llm = OpenAI(temperature=0.9, max_tokens=1000)
+prompt = PromptTemplate(
+    input_variables=["answer", "number"],
+    template="""
+Below is an answer. It needs citations. I want you to find {number} ideas or
+facts in the answer that could benefit from citations and formulate Google
+queries to retrieve documents that could be used in citations.
+
+Only return your answer in a JSON structure. The format should resemble: 
+{{
+  "citations": [
+    {{
+        "statement": "Statement 1", 
+        "query", "Google query 1" 
+    }}, 
+    {{ 
+        "statement": "Statement 2", 
+        "query", "Google query 2"
+    }},
+   ...
+  ]
+}}
+
+Answer: {answer}
+""")
+
+"""
+## Get citations for an answer from ChatGPT!
+"""
+
+# Formulate the prompt and send to GPT-3
+answer = st.text_area("Answer", 
+    placeholder="Paste answer from ChatGPT here", height=200)
+query = prompt.format(answer=answer, number=3)
+next_queries = llm(query)
+
+# DEBUG only
+# st.write(next_queries)
+
+# For each citation, formulate a Google search query and display the results
+# for that query. There will be a maximum of 5 results per query
+results = json.loads(next_queries)
+for result in results["citations"]:
+    statement = result["statement"]
+    search = GoogleSearch(
+        {
+            "q": result["query"], 
+            "api_key": os.environ["SERPAPI_API_KEY"],
+            "num": 5,
+        }
+    )
+    search_results = search.get_dict()
+    md = f"###### Citations for: '{statement}'\n\n"
+    number = 1
+    for result in search_results["organic_results"]:
+        md += f"\[{number}\] [{result['title']}]({result['link']})\n\n"
+        md += f"*{result['snippet']}*\n\n"
+        number += 1
+    st.markdown(md)