search.py

from datetime import datetime, timedelta
from whoosh import fields, index
from whoosh.qparser import QueryParser
from whoosh.qparser.dateparse import DateParserPlugin

import os
import openai

# token
openai.api_key = "sk-<snip>"

# names
USER = "kordless"
BOT = "joyous-capybara"

# db schema
schema = fields.Schema(
	text=fields.TEXT(stored=True),
	author=fields.TEXT(stored=True),
	updated=fields.DATETIME(stored=True)
)

# indexer
index = index.create_in("indexes", schema)
writer = index.writer()

# prompts
user_prompt = "Hello %s. I am %s, a user." % (BOT, USER)
bot_prompt = "Hello %s. I am %s, a conversational agent." % (USER, BOT)

# output
print(">", user_prompt)
print(bot_prompt)

# add a few documents
writer.add_document(
	text=bot_prompt,
	author=BOT,
	updated=datetime.now()
)
writer.add_document(
	text=user_prompt, 
	author=USER,
	updated=datetime.now()
)
writer.commit(optimize=True)

# create the parser
parser = QueryParser("text", schema=index.schema)
parser.add_plugin(DateParserPlugin(free=True))

while True:
	line = input("> ")
	
	# check for exit or log entry
	if line == "exit":
		break
	else:
		writer = index.writer()
		writer.add_document(
			text=line,
			author=USER,
			updated=datetime.now()
		)
		writer.commit(optimize=True)

		# timestamps for searching
		timestring="%Y%m%d%H%M%S"
		now = datetime.now()
		then = now - timedelta(minutes=4)
		start = then.strftime(timestring)
		end = now.strftime(timestring)

		# query
		query = parser.parse(u"updated:[%s to %s]" % (start, end))

		# build prompt
		prompt = ""
		with index.searcher() as s:
			results = s.search(query, limit=20)
			for result in reversed(results):
				if result['author'] == USER:
					prompt = prompt + "%s\n" % (result['text'])
				else:
					prompt = prompt + "%s\n\n" % (result['text'])

		response = openai.Completion.create(
			model="text-davinci-002",
			prompt=prompt,
			temperature=0.7,
			max_tokens=256,
			top_p=1,
			frequency_penalty=0,
			presence_penalty=0
		)
		line = response.get('choices')[0].get('text').strip("\n")
		
		writer = index.writer()
		writer.add_document(
			text=line,
			author=BOT,
			updated=datetime.now()
		)
		writer.commit(optimize=True)
		print(line)

# timestamps for searching
timestring="%Y%m%d%H%M%S"
now = datetime.now()
then = now - timedelta(minutes=20)
start = then.strftime(timestring)
end = now.strftime(timestring)

# query
query = parser.parse(u"updated:[%s to %s]" % (start, end))

with index.searcher() as s:
	prompt = ""
	results = s.search(query, limit=1000)
	for result in reversed(results):
		if result['author'] == USER:
			prompt = prompt + "%s\n" % (result['text'])
		else:
			prompt = prompt + "%s\n\n" % (result['text'])

print(prompt)