# .env
PINECONE_API_KEY = pcsk_**************************

from pinecone import Pinecone, ServerlessSpec
from dotenv import load_dotenv
import os
 
 
load_dotenv() # 환경변수 불러오기
 
pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))
 
index_name = "quickstart"
if index_name not in [index_info["name"] for index_info in pc.list_indexes()]:
    pc.create_index(
        name=index_name,
        dimension=4096, # Upstage embedding api의 임베딩 차원, 사용하는 임베딩 API의 차원으로 설정
        metric="dotproduct", # euclidean, cosine
        spec=ServerlessSpec(
            cloud="aws",
            region="us-east-1"
        ) 
    )
    print(f"{index_name} has been successfully created")
 
else:
    print(f"{index_name} is already exists.")
 
print(pc.list_indexes())

# .env
UPSTAGE_API_KEY = up_***************************

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyMuPDFLoader
import glob
 
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=150)
 
split_docs = []
files = sorted(glob.glob("data/*.pdf"))
 
for file in files:
    loader = PyMuPDFLoader(file)
    split_docs.extend(loader.load_and_split(text_splitter))
 
# 문서 개수 확인
print(len(split_docs))
 
def preprocess_docs(docs, keys = ['source', 'page'], min_len=10):
    result_docs = []
    for doc in docs:
        doc.metadata = {key: doc.metadata[key] for key in keys}
        if len(doc.page_content) > min_len:
            result_docs.append(doc)
    return result_docs
 
split_docs = preprocess_docs(split_docs)

from langchain_pinecone import PineconeVectorStore
from langchain_upstage import ChatUpstage, UpstageEmbeddings
 
embeddings_passage = UpstageEmbeddings(model="embedding-passage") #4096
embeddings_query = UpstageEmbeddings(model="embedding-query") #4096
 
docsearch = PineconeVectorStore.from_existing_index(index_name=index_name, embedding=embeddings_passage)
docsearch.add_documents(split_docs)

docsearch = PineconeVectorStore.from_existing_index(index_name=index_name, embedding=embeddings_query) # query 임베딩으로 변경
query = "범죄의 성립요건 세가지는?"
docs = docsearch.similarity_search(query=query)
print(docs)

from langchain_upstage import ChatUpstage
llm = ChatUpstage(api_key=os.environ.get("UPSTAGE_API_KEY"), temperature=0)
 
response_llm = llm.invoke(query)
print(response_llm.content)
 
prompt = f"Question: {query}\n Reference: {docs[0][0].page_content}\n Answer:"
response_llm_with_RAG = llm.invoke(prompt)
print(response_llm_with_RAG.content)

# import library
from pinecone import Pinecone
from langchain_upstage import ChatUpstage, UpstageEmbeddings
from dotenv import load_dotenv
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
import pandas as pd
import os
 
load_dotenv()
 
index_name = "quickstart"
client = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"), source_tag="langchain")
llm_upstage = ChatUpstage(api_key=os.environ.get("UPSTAGE_API_KEY"), temperature=0)
embeddings_query = UpstageEmbeddings(model="embedding-query") #4096
 
retriever_prompt_template = """
\n\nHuman: Here is the context information, inside <context></context> XML tags.
 
<context>
{context}
</context>
 
Given the context information and not prior knowledge.
generate only questions based on the below query.
 
You are a Professor. Your task is to setup \
{num_questions_per_chunk} questions for an upcoming \
quiz/examination. The questions should be diverse in nature \
across the document. The questions should not contain options, start with "-"
Restrict the questions to the context information provided.
 
\n\nAssistant:"""
 
PROMPT_RETRIEVER = PromptTemplate(
    template=retriever_prompt_template, input_variables=["context", "num_questions_per_chunk"]
)
 
generation_prompt_template = """
Here is the context, inside <context></context> XML tags.
 
<context>
{context}
</context>
 
Only using the context as above, answer the following question with the rules as below:
    - Don't insert XML tag such as <context> and </context> when answering.
    - Write as much as you can
    - Be courteous and polite
    - Only answer the question if you can find the answer in the context with certainty.
    - Skip the preamble
    - Use three sentences maximum and keep the answer concise.
    - If the answer is not in the context, just say "Could not find answer in given contexts."
 
Question:
{question}
 
Answer:"""
 
PROMPT_GENERATION = PromptTemplate(
    template=generation_prompt_template, input_variables=["context", "question"]
)
 
 
def GTGenerator(index, llm_retriever, llm_generation, prompt_retriever, prompt_generation, batch_size=50, max_batch=2, num_questions_per_chunk=1):
 
    llm_chain_retriever = LLMChain(llm=llm_retriever, prompt=prompt_retriever)
    llm_chain_generation = LLMChain(llm=llm_generation, prompt=prompt_generation)
    gt = [] # [question, 정답 id, 정답 text]
 
    all_ids = list(index.list(limit=batch_size))
 
    if max_batch < len(all_ids):
        all_ids = all_ids[:max_batch]    
 
    # ID를 기반으로 데이터를 하나씩 가져오기
    for fetched_ids in all_ids:
        fetched_docs = index.fetch(ids=fetched_ids)
        fetched_docs = fetched_docs.vectors
 
        for doc_id in fetched_ids:
            doc_text =  fetched_docs[doc_id]["metadata"]["text"]
 
            questions = llm_chain_retriever.predict(context=doc_text, num_questions_per_chunk=str(num_questions_per_chunk))
 
            questions = questions.split("\n\n-")
            if len(questions) <= num_questions_per_chunk + 1:
 
                if len(questions) == num_questions_per_chunk:
                    questions = list(map(lambda x:x.strip(), questions))
                else:
                    questions = list(map(lambda x:x.strip(), questions[1:]))
                for q in questions:
                    answer = llm_chain_generation.predict(question=q, context=doc_text)
                    answer = answer.strip()
                    gt.append([q, answer, doc_id, doc_text])
            else:
                print ("err")
                print (questions)
 
    return gt
 
index = client.Index(index_name)
gt = GTGenerator(
    index = index,
    llm_retriever=llm_upstage,
    llm_generation=llm_upstage,
    prompt_retriever=PROMPT_RETRIEVER,
    prompt_generation=PROMPT_GENERATION,
    batch_size=50,
    max_batch=3,
    num_questions_per_chunk=1
)
 
eval_dataset_retriever = pd.DataFrame(gt, columns=["question", "answer", "doc_id", "doc"])
eval_dataset_retriever.to_csv("eval_dataset.csv", index=False)