基于PolarDB PostgreSQL版和LLM构建企业专属Chatbot

CREATE DATABASE testdb;

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE polardb_pg_help_docs (
  id bigserial PRIMARY KEY, 
  title text,     -- 文档标题
  description text,     -- 描述
  doc_chunk text,     -- 文档分块
  token_size int,     -- 文档分块字数
  embedding vector(1536));  -- 文本嵌入信息

CREATE INDEX ON polardb_pg_help_docs USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

pip install openai psycopg2 tiktoken requests beautifulsoup4 numpy

import openai
import psycopg2
import tiktoken
import requests
from bs4 import BeautifulSoup
EMBEDDING_MODEL = "text-embedding-ada-002"
tokenizer = tiktoken.get_encoding("cl100k_base")
# 连接PolarDB-PG数据库
conn = psycopg2.connect(database="<数据库名>",
                        host="<PolarDB PostgreSQL版集群连接地址>",
                        user="<用户名>",
                        password="<密码>",
                        port="<数据库端口>")
conn.autocommit = True
# OpenAI的API Key
openai.api_key = '<Secret API Key>'
# 自定义拆分方法（仅为示例）
def get_text_chunks(text, max_chunk_size):
    chunks_ = []
    soup_ = BeautifulSoup(text, 'html.parser')
    content = ''.join(soup_.strings).strip()
    length = len(content)
    start = 0
    while start < length:
        end = start + max_chunk_size
        if end >= length:
            end = length
        chunk_ = content[start:end]
        chunks_.append(chunk_)
        start = end
    return chunks_
# 指定需要拆分的网页
url = 'https://help.aliyun.com/document_detail/602217.html?spm=a2c4g.468881.0.0.5a2c72c2cnmjaL'
response = requests.get(url)
if response.status_code == 200:
    # 获取网页内容
    web_html_data = response.text
    soup = BeautifulSoup(web_html_data, 'html.parser')
    # 获取标题（H1标签）
    title = soup.find('h1').text.strip()
    # 获取描述（class为shortdesc的p标签内容）
    description = soup.find('p', class_='shortdesc').text.strip()
    # 拆分并存储
    chunks = get_text_chunks(web_html_data, 500)
    for chunk in chunks:
        doc_item = {
            'title': title,
            'description': description,
            'doc_chunk': chunk,
            'token_size': len(tokenizer.encode(chunk))
        }
        query_embedding_response = openai.Embedding.create(
            model=EMBEDDING_MODEL,
            input=chunk,
        )
        doc_item['embedding'] = query_embedding_response['data'][0]['embedding']
        cur = conn.cursor()
        insert_query = '''
        INSERT INTO polardb_pg_help_docs 
            (title, description, doc_chunk, token_size, embedding) VALUES (%s, %s, %s, %s, %s);
        '''
        cur.execute(insert_query, (
            doc_item['title'], doc_item['description'], doc_item['doc_chunk'], doc_item['token_size'],
            doc_item['embedding']))
        conn.commit()
else:
    print('Failed to fetch web page')

SELECT * FROM polardb_pg_help_docs;

# 连接PolarDB PostgreSQL版集群数据库
conn = psycopg2.connect(database="<数据库名>",
                        host="<PolarDB PostgreSQL版集群连接地址>",
                        user="<用户名>",
                        password="<密码>",
                        port="<数据库端口>")
conn.autocommit = True
def answer(prompt_doc, prompt):
    improved_prompt = f"""
    按下面提供的文档和步骤来回答接下来的问题：
    (1) 首先，分析文档中的内容，看是否与问题相关
    (2) 其次，只能用文档中的内容进行回复,越详细越好，并且以markdown格式输出
    (3) 最后，如果问题与PolarDB PostgreSQL版不相关，请回复"我对PolarDB PostgreSQL版以外的知识不是很了解"
    文档:
    \"\"\"
    {prompt_doc}
    \"\"\"
    问题: {prompt}
    """
    response = openai.Completion.create(
        model=GPT_COMPLETIONS_MODEL,
        prompt=improved_prompt,
        temperature=0.2,
        max_tokens=MAX_TOKENS
    )
    print(f"{response['choices'][0]['text']}\n")
similarity_threshold = 0.78
max_matched_doc_counts = 8
# 通过pgvector过滤出相似度大于一定阈值的文档块
similarity_search_sql = f'''
SELECT doc_chunk, token_size, 1 - (embedding <=> '{prompt_embedding}') AS similarity 
FROM polardb_pg_help_docs WHERE 1 - (embedding <=> '{prompt_embedding}') > {similarity_threshold} ORDER BY id LIMIT {max_matched_doc_counts};
'''
cur = conn.cursor(cursor_factory=DictCursor)
cur.execute(similarity_search_sql)
matched_docs = cur.fetchall()
total_tokens = 0
prompt_doc = ''
print('Answer: \n')
for matched_doc in matched_docs:
    if total_tokens + matched_doc['token_size'] <= 1000:
        prompt_doc += f"\n---\n{matched_doc['doc_chunk']}"
        total_tokens += matched_doc['token_size']
        continue
    answer(prompt_doc,prompt)
    total_tokens = 0
    prompt_doc = ''
answer(prompt_doc,prompt)