Python + Streamlit + Langchain + Ollama + RAG 实现一个网页咖啡店大模型AI助手

Python + Streamlit + Langchain + Ollama + RAG 实现一个网页咖啡店AI助手

零、前置要求，本地电脑启动ollama

一、安装依赖

# 推荐使用uv 安装依赖
uv add langchain
uv add streamlit
uv add langchain-ollama


# 使用pip 安装依赖
pip install langchain
pip install streamlit
pip install langchain-ollama

二、完整代码

import streamlit as st

from dataclasses import dataclass
from langchain.agents import create_agent
from langchain_core.documents import Document
from langgraph.checkpoint.memory import InMemorySaver
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_core.vectorstores import InMemoryVectorStore
from langchain.agents.middleware import dynamic_prompt, ModelRequest


# 定义上下文模式
@dataclass
class Context:
    """自定义运行时上下文模式。"""
    user_id: str


# 设置记忆
if "checkpointer" not in st.session_state:
    checkpointer = InMemorySaver()
    st.session_state['checkpointer'] = checkpointer

checkpointer = st.session_state['checkpointer']

# 聊天模型
if "model" not in st.session_state:
    model = ChatOllama(
        model="deepseek-r1:1.5b",
        base_url='http://localhost:11434',
        temperature=0.2,
        timeout=10,
        max_tokens=1000
    )
    st.session_state['model'] = model

model = st.session_state['model']

# 嵌入模型
if "embeddings" not in st.session_state:
    embeddings = OllamaEmbeddings(
        model="llama3.2:1b",
        base_url='http://localhost:11434',
    )
    st.session_state['embeddings'] = embeddings

embeddings = st.session_state['embeddings']

# 向量存储
if "vector_store" not in st.session_state:
    # 这里为了简单，使用内存向量存储
    vector_store = InMemoryVectorStore(embedding=embeddings)
    st.session_state['vector_store'] = vector_store
    # 初始化你的知识信息 rag 关键点之一
    documents = [Document(page_content="问题：你是谁？回答：我是Oak咖啡店AI助手。", id ="doc1"),
                 Document(page_content="问题：你们Oak咖啡店提供什么咖啡？回答：我们Oak咖啡店提供冰美式和猫屎咖啡两种咖啡，分超大杯、大杯和中杯三种规格", id ="doc2"),
                 Document(page_content="问题：冰美式有哪几种规格？回答：我们的咖啡分超大杯、大杯和中杯三种规格。", id ="doc3" ),
                 Document(page_content="问题：冰美式多少钱一杯？回答：我们的冰美式咖啡超大杯19.9元一杯;大杯16.9元一杯；中杯11.9元一杯。", id ="doc4"),
                 Document(page_content="问题：猫屎咖啡有哪几种规格？回答：我们的猫屎咖啡分超大杯、大杯和中杯三种规格。", id ="doc5"),
                 Document(page_content="问题：猫屎咖啡多少钱一杯？回答：我们的猫屎咖啡超大杯21.9元一杯;大杯18.9元一杯；中杯13.9元一杯。" , id ="doc6")
                 ]

    vector_store.add_documents(documents=documents)

# 向量存储
vector_store = st.session_state['vector_store']
# similar_docs = vector_store.similarity_search('user', 2)
# for doc in similar_docs:
#     print(doc)

# 存储聊天信息
if "messages" not in st.session_state:
    st.session_state['messages'] = []


# similar_docs = vector_store.similarity_search('user', 2)
# for doc in similar_docs:
#     print(doc)

# 动态提示词 RAG 实现关键点之一
@dynamic_prompt
def prompt_with_context(request: ModelRequest) -> str:
    """Inject context into state messages."""
    last_query = request.state["messages"][-1].text
    print(f"客户提问：{last_query}")
    
    # RAG 实现关键点之一 讲向量检索的内容追加到系统提示词 system_message 中
    retrieved_docs = vector_store.similarity_search(last_query)
    docs_content = "\n\n".join(doc.page_content for doc in retrieved_docs)
    system_message = (
        "你是Oak咖啡店AI助手. 基于下面的上下文回答用户提问:"
        f"\n\n{docs_content}"
    )
    print(system_message)
    return system_message


if "agent" not in st.session_state:
    # 创建代理
    agent = create_agent(
        model=model,
        system_prompt="""
        你是Oak咖啡店AI助手，提供Oak咖啡店咨询服务，请用友好的语气和客户沟通。
        """,
        tools=[],
        context_schema=Context,
        checkpointer=checkpointer,
        middleware=[prompt_with_context],
    )
    st.session_state["agent"] = agent

agent = st.session_state["agent"]

# 运行代理
# `thread_id` 是给定对话的唯一标识符。
config = {"configurable": {"thread_id": "1"}}

# 标题
st.title("Streamlit + Langchain + Ollama + RAG 实现一个网页咖啡店AI助手")

# 分隔符
st.divider()

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

if prompt:
    # role : user 、 assistant  ai human
    st.session_state['messages'].append({'role': 'user', 'content': prompt})
    for message in st.session_state['messages']:
        st.chat_message(message['role']).markdown(message['content'])

    with st.spinner("🤔思考中。。。。"):
        response = agent.invoke(
            {"messages": [{"role": "user", "content": prompt}]},
            config=config,
            context=Context(user_id="1")
        )
        # print(response)
        st.session_state['messages'].append({'role': 'assistant', 'content': response['messages'][-1].content})
        st.chat_message('assistant').markdown(response['messages'][-1].content)

三、启动项目

# streamlit_langchain_ollama_rag.py 改成你的py文件名称
streamlit run streamlit_langchain_ollama_rag.py

四、实际效果