# import basics
import os
from dotenv import load_dotenv

# import pinecone
from pinecone import Pinecone, ServerlessSpec

# import langchain
from langchain_pinecone import PineconeVectorStore
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_core.documents import Document

load_dotenv()

# initialize pinecone database
pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))

# set the pinecone index

index_name = "sample-index"
index = pc.Index(index_name)

# initialize embeddings model + vector store

embeddings = GoogleGenerativeAIEmbeddings(model="models/gemini-embedding-001")
vector_store = PineconeVectorStore(index=index, embedding=embeddings)

# retrieval
'''

###### add docs to db ##############################
results = vector_store.similarity_search_with_score(
    "what did you have for breakfast?",
    #k=2,
    filter={"source": "tweet"},
)

print("RESULTS:")

for res in results:
    print(f"* {res[0].page_content} [{res[0].metadata}] -- {res[1]}")

'''

retriever = vector_store.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={"k": 5, "score_threshold": 0.6},
)
results = retriever.invoke("what did you have for breakfast?")

print("RESULTS:")

for res in results:
    print(f"* {res.page_content} [{res.metadata}]")

#'''