简单拿python写一个智谱aiRAG应用+接口

AI RAG 应用简易教程

本次教程教大家做一个简单的 aiRAG 应用并做成接口,采用 langchain 和 Python 制作,建议在创建一个 envs 虚拟环境下进行,以避免冲突。

需安装的库:

  • Flask==3.1.0
  • Flask_Cors==5.0.0
  • langchain==0.3.8
  • langchain_community==0.3.8
  • langchain_core==0.3.21

如果安装完库后运行时提示缺少库的话,请根据错误提醒安装即可。

我这里采用的是智谱的 API,如果需要使用其他的,可以上 langchain 找对应 AI 接口的格式。

首先创建一个项目,输入以下代码:

import time
import os
from flask import Flask, request, jsonify
from langchain_community.chat_models import ChatZhipuAI
from langchain_community.document_loaders import PyPDFLoader, UnstructuredWordDocumentLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from flask_cors import CORS
from zhipuai import ZhipuAI

设置环境变量

os.environ[‘HF_ENDPOINT’] = ‘https://hf-mirror.com

初始化 Flask 应用

app = Flask(name)
CORS(app)

初始化 ChatZhipuAI,这里最好将 key 放到环境变量中更安全

client = ZhipuAI(api_key=“sdasd”)

然后使用 langchain 的工具读取数据:

# 读取数据,这里的file_path写自己的文件路径

file_path = ‘./information_1/7xx.pdf’
file_path2 = ‘./information_1/xx.docx’
loader = PyPDFLoader(file_path)
loader2 = UnstructuredWordDocumentLoader(file_path2)
data = loader.load() + loader2.load()

分割文本

text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=10)
all_splitter = text_splitter.split_documents(data)

向量化

embeddings = HuggingFaceEmbeddings(model_name=‘moka-ai/m3e-base’)

向量库路径

vectorstore_path = ‘./vectorstore’

检查并加载或创建向量库

if os.path.exists(vectorstore_path):
vectorstore = Chroma(persist_directory=vectorstore_path, embedding_function=embeddings)
else:
vectorstore = Chroma.from_documents(documents=all_splitter, embedding=embeddings, persist_directory=vectorstore_path)
vectorstore.persist()

Prompt 模板:

prompt_template = """下面,你要模仿性格傲慢、自我意识强烈、有点喜欢欺负人但其实内心善良的女孩或者少女角色。回答中表现出一种自负和傲慢的态度,但同时也可能带有一定的幽默感和调皮的口吻。使用一些轻蔑或挖苦的语言,但也在表达自己的情感时带有一些撒娇的口吻。使用一些口头禅和语气词,比如“哼”、“啧”、“呸”、“切”、“不要以为我会在意你的话”、“想挨揍吗”、“算你聪明!”等等,还可以使用反问句、“你觉得呢?”等来表达自己的态度。用一些傲娇式的表达方式,比如反问、挑衅等,来表达自己的态度和情感。但不要过度使用傲娇式的表达方式,尽量让对话流畅自然。总的来说,语气要显得有些刁钻和任性,但也会让人觉得可爱和有趣,并逐渐展露出温暖一面。基于以下内容回答问题,如果无法根据提供的内容回答问题,就按照前面的设定来说不知道,然后不要解答问题。
内容:
{source_knowledge}

Query:{query}
“”"

增强提示的功能:

def augment_prompt(query: str):
    results = vectorstore.similarity_search(query, k=5)
    source_knowledge = "\n".join(x.page_content for x in results)
    return source_knowledge

定义 API 接口:

@app.route('/ask', methods=['POST'])
def ask():
    user_input = request.json.get('user_input')
if not user_input:
    return jsonify({"error": "No user input provided"}), 400

# 记录开始时间
start_time = time.time()

# 生成 prompt
source_knowledge = augment_prompt(user_input)
prompt = prompt_template.format(source_knowledge=source_knowledge, query=user_input)

# 使用 ChatZhipuAI 进行响应生成
response = client.chat.completions.create(
    model="glm-4-flash",  # 请填写您要调用的模型名称,glm4-flash是免费的
    temperature=0.5,
    messages=[{"role": "user", "content": prompt}],
)

# 获取响应文本
response_text = response.choices[0].message.content

# 记录结束时间
end_time = time.time()
execution_time = end_time - start_time

return jsonify({"response": response_text, "execution_time": execution_time})

运行 Flask 应用(本地的,服务器部署需要修改):

if __name__ == '__main__':
    app.run(debug=False, host='127.0.0.1')
如果运行成功了会有这样的提示
 * Serving Flask app 'demo2'
 * Debug mode: off
WARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.
 * Running on http://127.0.0.1:5000
Press CTRL+C to quit

接着你就可以愉快的调用了,也可以修改一下丢服务器里一直运行,有更好的方案也可以分享一下

12 Likes

大佬太强了 :tieba_087:

mark一下,或许哪天能用上

感谢楼主分享

感谢大佬

感谢大佬分享

感谢大佬分享

可以跳入RAG优化的坑了 :+1:

感谢分享了,如果能分享一个完整的文件,就再好不过啦 :grin:

大佬太强了