沉浸式学习PostgreSQL|PolarDB 9: AI大模型+向量数据库, 提升AI通用机器人在专业领域的精准度, 完美诠释柏拉图提出的“知识是回忆而不是知觉”

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);  
-- 或者使用hnsw算法  
CREATE INDEX ON polardb_pg_help_docs USING hnsw (embedding vector_cosine_ops) WITH (m = 12, ef_construction=40;

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)

docker run -d -it --cap-add=SYS_PTRACE --cap-add SYS_ADMIN --privileged=true --name pg registry.cn-hangzhou.aliyuncs.com/digoal/opensource_database:pg14_with_exts

docker exec -ti pg bash

psql

# apt install -y python3-pip    
# pip3 install openai    
root@689ed216de12:/tmp# psql    
psql (14.8 (Debian 14.8-1.pgdg110+1))    
Type "help" for help.    
postgres=# create extension plpython3u ;    
CREATE EXTENSION    
postgres=# create extension vector ;    
CREATE EXTENSION

create table tbl_faq (    
  id serial8 primary key,    
  f text,  -- 问题    
  q text,  -- 标准答案    
  v vector(1536)  -- 文本格式: faq (textcat('title: '||f, ' --- '||q)) 文本向量    
);

create or replace function get_v (faq text) returns vector as $$    
  import openai    
  import os    
  text = faq    
  openai.api_key = os.getenv("OPENAI_API_KEY")    
  response = openai.Embedding.create(    
      model = "text-embedding-ada-002",    
      input = text.replace("\n", " ")    
    )    
  embedding = response['data'][0]['embedding']    
  return embedding    
$$ language plpython3u;

insert into tbl_faq(f,q,v) select f,q,get_v(textcat('title: '||f, ' --- '||q)) from 外部表;

create index on tbl_faq using ivfflat (v vector_cosine_ops);    
analyze tbl_faq;

select get_v('用户输入的问题1');

create or replace function get_faq(    
  v vector(1536),   -- 用户抛出问题向量    
  th float,  -- 相似度阈值    
  cnt int    -- 返回多少条    
)    
returns table (    
  id int8,   -- 辅助表ID    
  faq text,   -- 辅助表问题+答案    
  similarity float   -- 相似度    
)    
as $$    
  select     
    tbl_faq.id,     
    textcat('title: '||tbl_faq.f, ' --- '||tbl_faq.q) as faq,    
    1 - (tbl_faq.v <=> v) as similarity    
  from tbl_faq    
  where 1 - (tbl_faq.v <=> v) > th    
  order by similarity desc    
  limit cnt;    
$$ language sql strict stable;

select t.id, t.faq, t.similarity    
from get_faq(    
  (select get_v('用户输入的问题1')),    
  0.8,   -- 相似度阈值    
  1      -- 返回最相似的1条    
) as t;

create or replace function ask_openai(    
  user_input text,  -- 用户输入问题    
  faq text   -- get_faq()得到的参考问题和答案    
)    
returns text as    
$$    
  import openai    
  import os    
  openai.api_key = os.getenv("OPENAI_API_KEY")    
  search_string = user_input    
  docs_text = faq    
  messages = [{"role": "system", "content": "You concisely answer questions based on text provided to you."}]    
  prompt = """Answer the user's prompt or question:     
  {search_string}    
  by summarising the following text:    
  {docs_text}    
  Keep your answer direct and concise. Provide code snippets where applicable.    
  The question is about a Greenplum/PostgreSQL/PolarDB database. You can enrich the answer with other     
  Greenplum or PostgreSQ-relevant details if applicable.""".format(search_string=search_string, docs_text=docs_text)    
  messages.append({"role": "user", "content": prompt})    
  response = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=messages)    
  return response.choices[0]["message"]["content"]    
$$ language plpython3u;

create or replace function good_ai_assistant(    
  user_input text  -- 只需要用户输入    
)    
returns table (    
  content text    
)    
language sql stable    
as $$    
  select ask_openai(user_input,    
    (select t.faq from get_faq( (select get_v('用户输入的问题1')), 0.8, 1 ) as t)    
  );    
$$;

select content from good_ai_assistant('用户输入的问题');