import argparse
from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain_ollama.llms import OllamaLLM

VECTOR_STORE_PATH = "./vectorstore"
PROMPT_TEMPLATE = """
Answer the question based only on the following context:

{context}

---

Answer the question based on the above context: {question}
"""


def chat(message):
    embeddings = OllamaEmbeddings(model="llama3")

    db = Chroma(persist_directory=VECTOR_STORE_PATH, embedding_function=embeddings)
    
    results = db.similarity_search_with_relevance_scores(message, k=3)
    if len(results) == 0 or results[0][1] < 0 or results[0][1] > 1:
        print(f"Unable to find matching results.")  
        return
    
    breakpoint()
    
    context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
    prompt = prompt_template.format(context=context_text, question=message)

    model = OllamaLLM(model="llama3")
    response_text = model.invoke(prompt)
    sources = [doc.metadata.get("source", None) for doc, _score in results]
    formatted_response = f"{response_text}\n\nSources: {', '.join(filter(None, sources))}"
    print(formatted_response)
    return formatted_response
    #print(formatted_response)

chat("Hola")