Citations needed

app.py

# 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)