60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
"""
|
||
Minimal Retrieval‑Augmented Generation example using LangChain, Chroma and Ollama.
|
||
|
||
The script demonstrates:
|
||
1. Splitting a sample text into chunks.
|
||
2. Creating embeddings with Ollama.
|
||
3. Storing chunks in a local Chroma vector store.
|
||
4. Querying the store and generating a response with an Ollama LLM.
|
||
|
||
Prerequisites:
|
||
- Ollama must be installed and a model (e.g. tinyllama) available.
|
||
- Python dependencies from requirements.txt must be installed.
|
||
"""
|
||
|
||
from pathlib import Path
|
||
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain_chroma import Chroma
|
||
from langchain_ollama import OllamaEmbeddings, ChatOllama
|
||
|
||
# 1. Sample text
|
||
SAMPLE_TEXT = (
|
||
"LangChain is a framework for developing applications powered by language models. "
|
||
"It provides abstractions for prompt construction, chaining, and memory. "
|
||
"Chroma is a fast, lightweight vector database that can be used as a backend "
|
||
"for retrieval‑augmented generation. Ollama offers locally hosted LLMs that "
|
||
"can be used for embeddings and generation.")
|
||
|
||
# 2. Split text into chunks
|
||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=20)
|
||
documents = text_splitter.split_text(SAMPLE_TEXT)
|
||
|
||
# 3. Create embeddings
|
||
embeddings = OllamaEmbeddings(model="tinyllama")
|
||
|
||
# 4. Create a Chroma store in a temporary directory
|
||
persist_dir = Path("./chroma_db")
|
||
persist_dir.mkdir(exist_ok=True)
|
||
|
||
vector_store = Chroma.from_texts(
|
||
texts=documents,
|
||
embedding=embeddings,
|
||
persist_directory=str(persist_dir),
|
||
)
|
||
vector_store.persist()
|
||
|
||
# 5. Query the store
|
||
query = "What is LangChain used for?"
|
||
results = vector_store.similarity_search(query, k=1)
|
||
retrieved_text = results[0].page_content
|
||
|
||
# 6. Generate a response using Ollama LLM
|
||
llm = ChatOllama(model="tinyllama")
|
||
prompt = f"Answer the following question using the provided context:\n\nContext: {retrieved_text}\n\nQuestion: {query}\nAnswer:"
|
||
response = llm.invoke(prompt)
|
||
print("\n--- Generated Response ---\n")
|
||
print(response)
|
||
|
||
if __name__ == "__main__":
|
||
pass |