# 1) Cài đặt
# pip install llama-index llama-index-embeddings-openai llama-index-llms-openai
import os
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
# 2) Cấu hình model + embeddings của bạn
os.environ["OPENAI_API_KEY"] = "YOUR_KEY" # hoặc sử dụng bất kỳ nhà cung cấp LLM/embedding được hỗ trợ nào
llm = OpenAI(model="gpt-4o-mini")
embed_model = OpenAIEmbedding(model="text-embedding-3-small")
# 3) Tải tài liệu (ví dụ: ./data/*.pdf, .md, .txt)
docs = SimpleDirectoryReader("./data").load_data
# 4) Xây dựng một chỉ mục
index = VectorStoreIndex.from_documents(docs, embed_model=embed_model)
# 5) Tạo một query engine và đặt câu hỏi
query_engine = index.as_query_engine(llm=llm)
response = query_engine.query("What are the key security practices mentioned in the docs?")
print(response)

from llama_index.core import SimpleDirectoryReader
from llama_index.readers.web import SimpleWebPageReader
file_docs = SimpleDirectoryReader("./policies").load_data
web_docs = SimpleWebPageReader(html_to_text=True).load_data
all_docs = file_docs + web_docs

from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import Document
parser = SentenceSplitter(chunk_size=800, chunk_overlap=100)
nodes = []
for d in all_docs:
nodes.extend(parser.get_nodes_from_documents([Document(text=d.text, metadata=d.metadata)]))

from llama_index.core import VectorStoreIndex, SummaryIndex
from llama_index.core.retrievers import BM25Retriever
from llama_index.core.query_engine import RetrieverQueryEngine
# Vector index từ các node đã được phân tích cú pháp trước
v_index = VectorStoreIndex(nodes)
# BM25 keyword retriever
bm25_retriever = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=6)
# Hybrid: hợp nhất các ứng viên, sau đó rerank
from llama_index.core.retrievers import RouterRetriever
from llama_index.retrievers.merge import MergerRetriever
v_retriever = v_index.as_retriever(similarity_top_k=6)
hybrid = MergerRetriever(retrievers=[v_retriever, bm25_retriever], top_k=8)
query_engine = RetrieverQueryEngine.from_args(retriever=hybrid)

from llama_index.postprocessor.flag_embedding_reranker import FlagEmbeddingReranker
from llama_index.core.query_engine import RetrieverQueryEngine
reranker = FlagEmbeddingReranker(top_n=5, model="BAAI/bge-reranker-base")
query_engine = v_index.as_query_engine(
similarity_top_k=12,
node_postprocessors=[reranker]
)

from llama_index.core.response_synthesizers import get_response_synthesizer
from llama_index.core import ServiceContext
synth = get_response_synthesizer(response_mode="tree_summarize")
query_engine = v_index.as_query_engine(response_synthesizer=synth)
ans = query_engine.query("Summarize the onboarding steps and cite sources.")
print(ans)

from llama_index.core.prompts import PromptTemplate
qa_tmpl = PromptTemplate(
"""
Bạn là một trợ lý ngắn gọn, ưu tiên bằng chứng. Chỉ sử dụng ngữ cảnh được cung cấp.
Nếu không chắc chắn, hãy nói rằng bạn không biết. Trả về JSON với các khóa: answer, sources.
Question: {query_str}
Context: {context_str}
"""
)
query_engine = v_index.as_query_engine(text_qa_template=qa_tmpl)

from llama_index.core.agent import ReActAgent
from llama_index.tools.sql import SQLQueryEngineTool
from sqlalchemy import create_engine
engine = create_engine("sqlite:///analytics.db")
sql_tool = SQLQueryEngineTool.from_engine(engine)
agent = ReActAgent.from_tools([sql_tool], llm=llm, verbose=True)
agent.chat("What was monthly churn in Q2 2025? If needed, query the DB.")

from llama_index.core.evaluation import FaithfulnessEvaluator, RelevancyEvaluator
faith = FaithfulnessEvaluator(llm=llm)
rel = RelevancyEvaluator(llm=llm)
pred = query_engine.query("List SOC 2 control families in our policy.")
print("faithful?", faith.evaluate_response(pred))
print("relevant?", rel.evaluate_response(pred))

# Ví dụ: Chroma
# pip install chromadb llama-index-vector-stores-chroma
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext
import chromadb
chroma_client = chromadb.PersistentClient(path="./chroma_store")
collection = chroma_client.get_or_create_collection("company_knowledge")
vector_store = ChromaVectorStore(chroma_collection=collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(all_docs, storage_context=storage_context)

# Ví dụ: lọc dựa trên metadata tại thời điểm truy vấn
retriever = index.as_retriever(similarity_top_k=8)
retriever.metadata_filters = {"department": ["legal", "security"], "published": [True]}

# FastAPI wrapper tối thiểu
# pip install fastapi uvicorn
from fastapi import FastAPI
app = FastAPI
qe = index.as_query_engine(llm=llm)
@app.post("/query")
async def query(payload: dict):
q = payload.get("q", "")
resp = qe.query(q)
return {"answer": str(resp), "sources": [s.node.metadata for s in resp.source_nodes]}