from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_community.vectorstores import FAISS
from langchain_core.output_parsers import StrOutputParser
from langchain_ollama import OllamaLLM, OllamaEmbeddings

llm = OllamaLLM(model="llama3.2", base_url="http://studio1.dorceus.local:11434" )

## Load existing Vector DB FAISS, that we created in the section "Create a
# Vector Store (Local Vector Store)"
# embedding_model = AzureOpenAI(deployment_name="dp-text-embedding-ada-002",
# model_name="text-embedding-ada-002")
embedding_model = OllamaEmbeddings(model="mxbai-embed-large", base_url="http://studio1.dorceus.local:11434")

# Adding New Data
new_texts = ["harrison worked at google", "harrison likes spicy food"]
vectorstore = FAISS.from_texts(
    new_texts,
    embedding_model
)
vectorstore.add_texts(new_texts)
vectorstore.save_local("./faiss_store")

vectorstore = FAISS.load_local(
    "./faiss_store/",
    embedding_model, 
    allow_dangerous_deserialization=True
)

retriever = vectorstore.as_retriever()
template = """Answer the question based only on the
following context: {context}
Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)
model = llm
chain = (
    RunnableParallel( {"context": retriever, "question": RunnablePassthrough()} )
    | prompt
    | model
    | StrOutputParser()
)
result = chain.invoke("where did harrison work?")
print(result)