技术解读 | 泽拓昆仑Klustron向量数据管理

kl@kunlun1:~$ source data/kl/env.shkl@kunlun1:~$ psql -h 192.168.0.152 -p 47001  postgrespostgres=#  CREATE EXTENSION vector WITH SCHEMA public ;

postgres=# CREATE TABLE items (
id bigserial PRIMARY KEY, 
item text, 
embedding vector(2)
);

postgres=# INSERT INTO items (item, embedding)VALUES  ('苹果', '[1, 1]'),  ('香蕉', '[1.2, 0.8]'),  ('猫', '[6, 0.4]');

postgres=# select * from items ;

postgres=# SELECT item, 1 - (embedding <=> '[6,0]') AS cosine_similarityFROM items ORDER BY cosine_similarity DESC;

root@kunlun1:/home/kl# pip3 install openai==0.28 psycopg2 tiktoken requests beautifulsoup4 numpy

kl@kunlun1:~$ source data/kl/env.sh
kl@kunlun1:~$ psql -h 192.168.0.152 -p 47001  postgres
postgres=#  CREATE TABLE rds_pg_help_docs (
id bigserial PRIMARY KEY, 
title text,      -- 文档标题
description text,     -- 描述
doc_chunk text,     -- 文档分块
token_size int,     -- 文档分块字数
embedding vector(1024));  -- 文本嵌入信息

knowledge_chunk_storage.py 
---------------------------------------------
import openai
import psycopg2
import tiktoken
import requests
from bs4 import BeautifulSoup


EMBEDDING_MODEL = "text-embedding-ada-002"
tokenizer = tiktoken.get_encoding("cl100k_base")


# 连接Klustron数据库
conn = psycopg2.connect(database="postgres",
                        host="192.168.0.152",
                        user="abc",
                        password="abc",
                        port="47001")
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/148038.html'


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 rds_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')

kl@kunlun1:~$ export all_proxy="socks5h://127.0.0.1:1080"   #本地代理配置示例
kl@kunlun1:~$ python3 knowledge_chunk_storage.py

kl@kunlun1:~$ psql -h 192.168.0.152 -p 47001  postgrespostgres=# SELECT id,token_size,SUBSTRING(title FROM -1 FOR 30), SUBSTRING(description FROM -1 FOR 30),  SUBSTRING(doc_chunk FROM -1 FOR 30),SUBSTRING(embedding FROM -1 FOR 50)  from rds_pg_help_docs;

chatbot.py
------------------------------------------
import openai
import psycopg2
from psycopg2.extras import DictCursor


GPT_MODEL = "gpt-3.5-turbo"
EMBEDDING_MODEL = "text-embedding-ada-002"
GPT_COMPLETIONS_MODEL = "text-davinci-003"
MAX_TOKENS = 1024


# OpenAI的API Key
openai.api_key = '<Secret API Key>'


prompt = '如何创建一个RDS PostgreSQL实例'


prompt_response = openai.Embedding.create(
    model=EMBEDDING_MODEL,
    input=prompt,
)
prompt_embedding = prompt_response['data'][0]['embedding']


# 连接Klustron数据库
conn = psycopg2.connect(database="postgres",
                        host="192.168.0.152",
                        user="abc",
                        password="abc",
                        port="47001") 
conn.autocommit = True




def answer(prompt_doc, prompt):
    improved_prompt = f"""
    按下面提供的文档和步骤来回答接下来的问题：
    (1) 首先，分析文档中的内容，看是否与问题相关
    (2) 其次，只能用文档中的内容进行回复,越详细越好，并且以markdown格式输出
    (3) 最后，如果问题与RDS PostgreSQL不相关，请回复"我对RDS 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 rds_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)

kl@kunlun1:~$ python3 chatbot.py