import os
from langchain.chat_models import ChatOpenAI
from langchain.schema import (
    SystemMessage,
    HumanMessage,
    AIMessage
)
from datasets import load_dataset
from pinecone import Pinecone


dataset = load_dataset(
    "jamescalam/llama-2-arxiv-papers-chunked",
    split="train"
)

os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY") 

chat = ChatOpenAI(
    openai_api_key = os.environ["OPENAI_API_KEY"],
    model='gpt-3.5-turbo'

)


messages = [
    SystemMessage(content="You are a helpful assistant."),
    HumanMessage(content="Hi AI, how are you today?"),
    AIMessage(content="I'm great thank you. How can I help you?"),
    HumanMessage(content="I'd like to understand string theory.")
]

res = chat(messages)

# add latest AI response to messages
messages.append(res)

# connect to pinecone