Add main.py

This commit is contained in:
2026-06-02 06:06:24 +00:00
parent 3ebe04dd70
commit 57cbd87e0a
+60
View File
@@ -0,0 +1,60 @@
"""
Minimal RetrievalAugmented 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 retrievalaugmented 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