from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

prompt = ChatPromptTemplate.from_template("用一句话解释{topic}")
chain = prompt | ChatOpenAI(model="gpt-4o-mini") | StrOutputParser()

print(chain.invoke({"topic": "RAG"}))
# 输出: RAG 是把外部知识检索和大模型生成结合的技术,让回答更准确。

import os
os.environ["OPENAI_API_KEY"] = "sk-xxx"  # 或读环境变量

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# 1) 准备"知识库"(实际可换成 PDF 加载)
docs = [
    Document(page_content="RAG 是检索增强生成,先检索后回答。"),
    Document(page_content="Advanced RAG 加入查询改写、Rerank、压缩等优化。"),
    Document(page_content="LangChain 是大模型应用编排框架,支持 LCEL 链式调用。"),
]

# 2) 建向量库
embeddings = OpenAIEmbeddings()
vectorstore = FAISS.from_documents(docs, embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 2})

# 3) 用 LCEL 串起"检索 → 提示 → 生成"三步
prompt = ChatPromptTemplate.from_template(
    "你是RAG讲师。基于以下资料回答问题,分点列出。\n"
    "资料:{context}\n问题:{question}"
)
def format_docs(docs):
    return "\n".join(d.page_content for d in docs)

chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt | ChatOpenAI(model="gpt-4o-mini") | StrOutputParser()
)

# 4) 试跑
print(chain.invoke("什么是RAG?"))
# 输出示例:
# - RAG 是检索增强生成技术。
# - LangChain 是大模型应用编排框架,支持 LCEL 链式调用。

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor

compressor = LLMChainExtractor.from_llm(llm)
smart_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=retriever,
)
# 以后用 smart_retriever 代替 retriever 即可

from langchain.memory import ConversationBufferWindowMemory
from langchain.chains import ConversationChain

memory = ConversationBufferWindowMemory(k=3, return_messages=True)
chat = ConversationChain(llm=llm, memory=memory)

print(chat.run("我叫小李"))           # 记住名字
print(chat.run("我最关心RAG优化"))    # 记住偏好
print(chat.run("我刚才说我关心什么?"))  # 答:RAG 优化

prompt = ChatPromptTemplate.from_template("""
你是严谨的企业知识助手。
【铁律】
1. 只能基于"以下资料"回答,不要用任何外部知识;
2. 每条结论后必须标注 [来源: {source}];
3. 资料里没有的信息,直接说"依据不足"。

【问题】{question}
【资料】{context}
""")

import os
os.environ["OPENAI_API_KEY"] = "sk-xxx"

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.settings import Settings
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

# 1) 配模型(全局生效)
Settings.llm = OpenAI(model="gpt-4o-mini", temperature=0)
Settings.embed_model = OpenAIEmbedding()

# 2) 一行加载目录里的所有文档
docs = SimpleDirectoryReader("./data", recursive=True).load_data()

# 3) 一行建索引
index = VectorStoreIndex.from_documents(docs)

# 4) 一行变问答引擎
query_engine = index.as_query_engine(similarity_top_k=3)

# 5) 直接问!
print(query_engine.query("我们公司去年的差旅报销标准是什么?"))

import os
os.environ["OPENAI_API_KEY"] = "sk-xxx"

from llama_index.core import (
    VectorStoreIndex, KeywordTableIndex, SimpleDirectoryReader,
    Settings, StorageContext
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

# ============ 1) 全局配置 ============
Settings.llm = OpenAI(model="gpt-4o-mini", temperature=0)
Settings.embed_model = OpenAIEmbedding()
Settings.chunk_size = 512
Settings.chunk_overlap = 50

# ============ 2) 加载文档(支持 PDF/Word/HTML/MD) ============
docs = SimpleDirectoryReader(
    input_dir="./data",
    recursive=True,
    required_exts=[".pdf", ".docx", ".md", ".txt"]  # 只看这些
).load_data()
print(f"📚 加载了 {len(docs)} 份文档")

# ============ 3) 自定义分块(按"句号"切) ============
parser = SentenceSplitter(chunk_size=512, chunk_overlap=50)
nodes = parser.get_nodes_from_documents(docs)
print(f"✂️ 切成了 {len(nodes)} 个 Node")

# ============ 4) 给 Node 贴"标签"(元数据) ============
for i, node in enumerate(nodes):
    node.metadata["chunk_id"] = i
    node.metadata["source"] = node.metadata.get("file_name", "unknown")

# ============ 5) 建"双索引":语义 + 关键词 ============
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)

# 5.1 语义索引(按"意思"找)
vector_index = VectorStoreIndex(
    nodes, storage_context=storage_context, show_progress=True
)

# 5.2 关键词索引(按"术语"找)
keyword_index = KeywordTableIndex(
    nodes, storage_context=storage_context, show_progress=True
)

# ============ 6) 配置 QueryEngine(返回 Top5 相似片段) ============
vector_engine = vector_index.as_query_engine(
    similarity_top_k=5,
    response_mode="tree_summarize",  # 长答案汇总
)

# ============ 7) 试跑 ============
question = "公司去年在新能源板块投入了多少研发资金?"
response = vector_engine.query(question)
print(f"\n❓ 问题: {question}")
print(f"✅ 答案: {response}\n")
print("📎 参考资料:")
for i, src in enumerate(response.source_nodes, 1):
    print(f"  [{i}] {src.metadata.get('source', 'unknown')} - {src.text[:80]}...")

# ====== 第 1 步:LlamaIndex 负责"做索引" ======
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

docs = SimpleDirectoryReader("./kb").load_data()
index = VectorStoreIndex.from_documents(docs)

# 把 LlamaIndex 的检索器导出
li_retriever = index.as_retriever(similarity_top_k=6)

# ====== 第 2 步:套个壳,变 LangChain 能用的 Retriever ======
from langchain_core.documents import Document as LCDoc
class LIRetrieverAsLC:
    def __init__(self, li_retriever):
        self.li_retriever = li_retriever
    def get_relevant_documents(self, query):
        results = self.li_retriever.retrieve(query)
        return [LCDoc(page_content=r.get_text(), metadata=r.metadata)
                for r in results]

lc_retriever = LIRetrieverAsLC(li_retriever)

# ====== 第 3 步:LangChain 负责"做编排" ======
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains import RetrievalQA

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是企业知识助手。仅基于提供的上下文回答,列出引用出处。"),
    ("user", "问题:{query}\n上下文:{context}")
])

qa = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=lc_retriever,  # 关键:这里塞的是 LlamaIndex 的检索器
    chain_type="stuff",
)

print(qa.run("解释Advanced RAG的关键组成,并给出处"))