import chainlit as cl
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

llm = llm = ChatOpenAI(
    api_key="ollama",
    model='llama3.2',
    base_url="http://localhost:11434/v1",
    temperature=0
)

# Load the persisted Chroma database
persist_directory = 'mydb'
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings)

# Create a retriever from the vector store
retriever = vectordb.as_retriever()

# Set up the QA chain
qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type='stuff', retriever=retriever)

# Define the Chainlit app
@cl.on_message
def main(message):
    response = qa_chain.run(message.content)
    cl.Message(content=response).send()