File size: 1,366 Bytes
9718baa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | 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") |