mkdir local-rag-assistant
cd local-rag-assistant
pip install langchain langchain-community chromadb streamlit
pip install pypdf python-docx tiktoken sentence-transformers

pip install -r requirements.txt

# 下载嵌入模型（用于文档向量化）
ollama pull nomic-embed-text

# 下载对话模型
ollama pull llama3

# 验证模型已下载
ollama list

"""
本地 AI 知识库助手
基于 Ollama + ChromaDB + Streamlit
"""
import streamlit as st
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.chat_models import ChatOllama
from langchain.chains import RetrievalQA
import tempfile
import os

# 配置页面
st.set_page_config(page_title="本地知识库助手", page_icon="🤖")
st.title("🤖 本地 AI 知识库助手")

# 初始化向量数据库
@st.cache_resource
def get_vectorstore():
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    vectorstore = Chroma(
        persist_directory="./chroma_db",
        embedding_function=embeddings
    )
    return vectorstore

# 初始化聊天模型
@st.cache_resource
def get_llm():
    return ChatOllama(model="llama3", temperature=0.7)

# 上传文档
def process_document(uploaded_file):
    # 保存上传文件到临时目录
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
        tmp_file.write(uploaded_file.getvalue())
        tmp_path = tmp_file.name
    
    # 加载文档
    if uploaded_file.name.endswith(".pdf"):
        loader = PyPDFLoader(tmp_path)
    else:
        from langchain_community.document_loaders import TextLoader
        loader = TextLoader(tmp_path)
    
    documents = loader.load()
    
    # 文本分块
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50
    )
    chunks = text_splitter.split_documents(documents)
    
    # 向量化并存入数据库
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory="./chroma_db"
    )
    
    os.unlink(tmp_path)
    return len(chunks)

# 问答
def answer_question(question, vectorstore, llm):
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever
    )
    return qa_chain.run(question)

# 侧边栏 - 上传文档
with st.sidebar:
    st.header("📁 上传文档")
    uploaded_file = st.file_uploader(
        "选择 PDF 或 TXT 文件",
        type=["pdf", "txt"],
        help="支持 PDF 和纯文本文件"
    )
    
    if uploaded_file and st.button("处理文档", type="primary"):
        with st.spinner("正在处理文档..."):
            chunk_count = process_document(uploaded_file)
            st.success(f"✅ 文档处理完成！共 {chunk_count} 个文本块")
    
    st.divider()
    st.markdown("### 使用说明")
    st.markdown("""
    1. 上传 PDF 或 TXT 文档
    2. 点击「处理文档」按钮
    3. 在下方输入问题并提问
    """)

# 主界面 - 问答
if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

question = st.chat_input("请输入你的问题...")

if question:
    with st.chat_message("user"):
        st.markdown(question)
    st.session_state.messages.append({"role": "user", "content": question})
    
    with st.spinner("思考中..."):
        try:
            vectorstore = get_vectorstore()
            llm = get_llm()
            answer = answer_question(question, vectorstore, llm)
            
            with st.chat_message("assistant"):
                st.markdown(answer)
            st.session_state.messages.append({"role": "assistant", "content": answer})
        except Exception as e:
            st.error(f"出错了: {str(e)}")

langchain>=0.1.0
langchain-community>=0.0.10
chromadb>=0.4.0
streamlit>=1.28.0
pypdf>=3.0.0
python-docx>=1.0.0
tiktoken>=0.5.0
sentence-transformers>=2.2.0

# 在另一个终端启动 Ollama（如果未运行）
ollama serve

# 启动 Streamlit 应用
streamlit run app.py --server.port 8501