from langchain import OpenAI, LLMChain
from langchain.callbacks import StdOutCallbackHandler
from langchain.chat_models import ChatOpenAI

from src.agents.chat_chain import ChatChain
from src.agents.graphdb_traversal_chain import GraphDBTraversalChain, mem_query_template, mem_system_message
from src.memory.triple_modal_memory import TripleModalMemory

import os
from dotenv import load_dotenv


# Set up the cache
import langchain
from langchain.cache import SQLiteCache
langchain.llm_cache = SQLiteCache(database_path=".langchain.db")

# initialize the memory
load_dotenv()
uri = os.getenv("NEO4J_URI")
user = os.getenv("NEO4J_USER")
password = os.getenv("NEO4J_PASSWORD")

mem = TripleModalMemory(uri, user, password)

# Create memory from docks or load from file if it exists
ingested = os.path.exists('../data/triple_modal_memory.faiss')
if not ingested:
    knowledge_path = r'C:\Users\colli\Documents\AIPapers'
    mem.ingest_docs(knowledge_path)
    mem.save()
    print("Memory initialized and saved.")

else:
    mem.load()
    print("Memory loaded.")

handler = StdOutCallbackHandler()

llm = ChatOpenAI(
    model_name="gpt-4", #"gpt-3.5-turbo"
    temperature=0,
    verbose=True
)
chain = ChatChain(llm=llm, prompt=mem_query_template, callbacks=[handler], system_message=mem_system_message)
knowledge_base_query_agent = GraphDBTraversalChain(llm_chain=chain, graph_vector_store=mem.vector_store)

# Example Research questions:
# What are different methods of providing language models with additional context to better answer questions?
# How can semantic search be used in conjunction with large language models in order to better answer questions?
# What are some techniques for achieving better general intelligence in language models?

def main_loop():
    try:
        while True:
            question = input("Enter a question: ")
            print(knowledge_base_query_agent.run(question))


    except KeyboardInterrupt:
        print("Shutdown: Saving...")
        mem.save()
        print("Shutdown: Complete")

    else:
        print("Completed all tasks.")

if __name__ == '__main__':
    main_loop()
langchain.callbacks.StdOutCallbackHandler langchain.chat_models.ChatOpenAI langchain.cache.SQLiteCache
[((495, 537), 'langchain.cache.SQLiteCache', 'SQLiteCache', ([], {'database_path': '""".langchain.db"""'}), "(database_path='.langchain.db')\n", (506, 537), False, 'from langchain.cache import SQLiteCache\n'), ((563, 576), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (574, 576), False, 'from dotenv import load_dotenv\n'), ((583, 605), 'os.getenv', 'os.getenv', (['"""NEO4J_URI"""'], {}), "('NEO4J_URI')\n", (592, 605), False, 'import os\n'), ((613, 636), 'os.getenv', 'os.getenv', (['"""NEO4J_USER"""'], {}), "('NEO4J_USER')\n", (622, 636), False, 'import os\n'), ((648, 675), 'os.getenv', 'os.getenv', (['"""NEO4J_PASSWORD"""'], {}), "('NEO4J_PASSWORD')\n", (657, 675), False, 'import os\n'), ((683, 721), 'src.memory.triple_modal_memory.TripleModalMemory', 'TripleModalMemory', (['uri', 'user', 'password'], {}), '(uri, user, password)\n', (700, 721), False, 'from src.memory.triple_modal_memory import TripleModalMemory\n'), ((792, 843), 'os.path.exists', 'os.path.exists', (['"""../data/triple_modal_memory.faiss"""'], {}), "('../data/triple_modal_memory.faiss')\n", (806, 843), False, 'import os\n'), ((1074, 1097), 'langchain.callbacks.StdOutCallbackHandler', 'StdOutCallbackHandler', ([], {}), '()\n', (1095, 1097), False, 'from langchain.callbacks import StdOutCallbackHandler\n'), ((1105, 1164), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-4"""', 'temperature': '(0)', 'verbose': '(True)'}), "(model_name='gpt-4', temperature=0, verbose=True)\n", (1115, 1164), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1204, 1309), 'src.agents.chat_chain.ChatChain', 'ChatChain', ([], {'llm': 'llm', 'prompt': 'mem_query_template', 'callbacks': '[handler]', 'system_message': 'mem_system_message'}), '(llm=llm, prompt=mem_query_template, callbacks=[handler],\n    system_message=mem_system_message)\n', (1213, 1309), False, 'from src.agents.chat_chain import ChatChain\n'), ((1335, 1410), 'src.agents.graphdb_traversal_chain.GraphDBTraversalChain', 'GraphDBTraversalChain', ([], {'llm_chain': 'chain', 'graph_vector_store': 'mem.vector_store'}), '(llm_chain=chain, graph_vector_store=mem.vector_store)\n', (1356, 1410), False, 'from src.agents.graphdb_traversal_chain import GraphDBTraversalChain, mem_query_template, mem_system_message\n')]