“今天天氣真好，我們出去放風(fēng)箏吧”
“今天天氣真好，我們出去散散步吧”
“這么大的雨，我們還是在家呆著吧”

.
├── package.json                              // ts依賴包
├── docs
│   ├── PostgreSQL15.mdx                      // 知識庫文檔
├── script
│   ├── generate-embeddings.ts                // 構(gòu)建知識庫
│   ├── queryGPT.py                           // 問答腳本

create extension if not exists vector;

-- 記錄文檔信息
create table docs (
  id bigserial primary key,
  -- 父文檔ID
  parent_doc bigint references docs,
  -- 文檔路徑
  path text not null unique,
  -- 文檔校驗值
  checksum text
);
-- 記錄chunk信息
create table doc_chunks (
  id bigserial primary key,
  doc_id bigint not null references docs on delete cascade, -- 文檔ID
  content text, -- chunk內(nèi)容
  token_count int, -- chunk中的token數(shù)量
  embedding vector(1536), -- chunk轉(zhuǎn)化成的embedding向量
  slug text, -- 為標題生成唯一標志
  heading text -- 標題
);

curl -fsSL https://get.pnpm.io/install.sh | sh -

sudo apt-get update
sudo apt-get install -y ca-certificates curl gnupg
sudo mkdir -p /etc/apt/keyrings
curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key | sudo gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg
NODE_MAJOR=16
echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_$NODE_MAJOR.x nodistro main" | sudo tee /etc/apt/sources.list.d/nodesource.list
sudo apt-get update
sudo apt-get install nodejs -y

pnpm run setup
pnpm install tsx

#這里需要將user、passwd、127.0.0.1、5432 替換為實際數(shù)據(jù)庫用戶、密碼、數(shù)據(jù)庫地址、端口
const postgresql_url = 'pg://user:passwd@127.0.0.1:5432/database';
const openai_key = '-------------';
const SOURCE_DIR = path.join(__dirname, 'document path');

pnpm tsx script/generate-embeddings.ts

create or replace function match_chunks(chunck_embedding vector(1536), threshold float, count int, min_length int)
returns table (id bigint, content text, similarity float)
language plpgsql
as $$
begin
  return query
  select
    doc_chunks.id,
    doc_chunks.content,
    (doc_chunks.embedding <#> chunck_embedding) * -1 as similarity
  from doc_chunks

  -- chunk內(nèi)容大于設(shè)定的長度
  where length(doc_chunks.content) >= min_length

  -- The dot product is negative because of a Postgres limitation, so we negate it
  and (doc_chunks.embedding <#> chunck_embedding) * -1 > threshold
  order by doc_chunks.embedding <#> chunck_embedding
  
  limit count;
end;
$$;

import os, psycopg2, openai
def query_handler(query = None):
    if query is None or query == "":
        print('請輸入有效問題')
        return
    query = query.strip().replace('\n', ' ')
    embedding = None
    try:
        # 使用 GPT 將提問轉(zhuǎn)化為 embedding 向量
        response = openai.Embedding.create(
            engine="text-embedding-ada-002",  # 固定為text-embedding-ada-002
            input=[query],
        )
        embedding = response.data[0].embedding
    except Exception as ex:
        print(ex)
        return
    content = ""
    con = None
    try:
        # 處理 postgres 配置，連接數(shù)據(jù)庫
        # host:127.0.0.1,port:5432,user:test,password:test,database:test
        params = postgresql_url.split(',')
        database, user, password, host, port = "test", "test", "test", "127.0.0.1", "5432"
        for param in params:
            pair = param.split(':')
            if len(pair) != 2:
                print('POSTGRESQL_URL error: ' + postgresql_url)
                return
            k, v = pair[0].strip(), pair[1].strip()
            if k == 'database':
                database = v
            elif k == 'user':
                user = v
            elif k == 'password':
                password = v
            elif k == 'host':
                host = v
            elif k == 'port':
                port = v
        # connect postgres
        con = psycopg2.connect(database=database, user=user, password=password, host=host, port=port)
        cur = con.cursor()
        # 從數(shù)據(jù)庫查詢?nèi)舾蓷l最接近提問的 chunk
        sql = "select match_chunks('[" + ','.join([str(x) for x in embedding]) + "]', 0.78, 5, 50)"
        cur.execute(sql)
        rows = cur.fetchall()
        for row in rows:
            row = row[0][1:-2].split(',')[-2][1:-2].strip()
            content = content + row + "\n---\n"
    except Exception as ex:
        print(ex)
        return
    finally:
        if con is not None:
            con.close()
    try:
        # 組織提問和 chunk 內(nèi)容，發(fā)送給 GPT
        prompt = '''Pretend you are GPT-4 model , Act an database expert.
        I will introduce a database scenario for which you will provide advice and related sql commands.
        Please only provide advice related to this scenario. Based on the specific scenario from the documentation,
        answer the question only using that information. Please note that if there are any updates to the database
        syntax or usage rules, the latest content shall prevail. If you are uncertain or the answer is not explicitly
        written in the documentation, please respond with "I'm sorry, I cannot assist with this.\n\n''' + "Context sections:\n" + \
        content.strip().replace('\n', ' ') + "\n\nQuestion:"""" + query.replace('\n', ' ') + """"\n\nAnswer:"
        print('\n正在處理，請稍后。。。\n')
        response = openai.ChatCompletion.create(
            engine="gpt_openapi",  # 固定為gpt_openapi
            messages=[
                {"role": "user", "content": prompt}
            ],
            model="gpt-35-turbo",
            temperature=0,
        )
        print('回答:')
        print(response['choices'][0]['message']['content'])

    except Exception as ex:
        print(ex)
        return
os.environ['OPENAI_KEY'] = '-----------------------'
os.environ['POSTGRESQL_URL'] = 'host:127.0.0.1,port:5432,user:test,password:test,database:test'
openai_key = os.getenv('OPENAI_KEY')
postgresql_url = os.getenv('POSTGRESQL_URL')
# openai config
openai.api_type = "azure"
openai.api_base = "https://example-endpoint.openai.azure.com"
openai.api_version = "2023-XX"
openai.api_key = openai_key
def main():
    if openai_key is None or postgresql_url is None:
        print('Missing environment variable OPENAI_KEY, POSTGRESQL_URL(host:127.XX.XX.XX,port:5432,user:XX,password:XX,database:XX)')
        return
    print('我是您的PostgreSQL AI助手，請輸入您想查詢的問題，例如：\n1、如何創(chuàng)建table？\n2、給我解釋一下select語句？\n3、如何創(chuàng)建一個存儲過程？')
    while True:
        query = input("\n輸入您的問題:")
        query_handler(query)
if __name__ == "__main__":
    main()

os.environ['OPENAI_KEY'] = '-----------------------'
os.environ['POSTGRESQL_URL'] = 'host:127.0.0.1,port:5432,user:test,password:test,database:test'

openai.api_type = "azure"
openai.api_base = "https://example-endpoint.openai.azure.com"
openai.api_version = "2023-XX"

prompt = '''Pretend you are GPT-4 model , Act an database expert.
        I will introduce a database scenario for which you will provide advice and related sql commands.
        Please only provide advice related to this scenario. Based on the specific scenario from the documentation,
        answer the question only using that information. Please note that if there are any updates to the database
        syntax or usage rules, the latest content shall prevail. If you are uncertain or the answer is not explicitly
        written in the documentation, please respond with "I'm sorry, I cannot assist with this.\n\n''' + "Context sections:\n" + \
        content.strip().replace('\n', ' ') + "\n\nQuestion:"""" + query.replace('\n', ' ') + """"\n\nAnswer:"

pip install psycopg2-binary
pip install openai
pip install 'openai[datalib]'