"""
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