| 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 |
| |
|
|
| chat("Hola") |
