Redis Search と OpenAI Embedding で日本国憲法に無理な質問してみる

OpenAI のサンプルにある [chatbot-kickstarter](https://github.com/openai/openai-cookbook/tree/main/apps/chatbot-kickstarter) では、Redis と Embedding を使ってカスタムデータに対するチャット応答を実現しています。この説明がめちゃめちゃ分かりにくいのでまとめました。手頃なテキストデータの例として日本国憲法を使いました。


## 利用の準備

[redis-stack](https://hub.docker.com/r/redis/redis-stack) を参考に Docker で redis-stack を起動します。

```
docker run -d --name redis-stack -p 6379:6379 -p 8001:8001 redis/redis-stack:latest
```

実は redis-stack の Docker には Redis Insight という Web ベースのクライアントが付属しているので、`redis-cli` を使う必要がありません。クエリが複雑になると `redis-cli` では見づらいし、brew などで redis をインストールした人はポートが衝突するので `brew uninstall redis` してしまった方が良いかも知れません。

http://localhost:8001/ で Redis Insight を起動し [Workbench](http://localhost:8001/redis-stack/workbench) に移動すると Redis にコマンドを入力できます。右側の緑の三角を押すと実行されます。Mac の場合 Command + Return が実行ショートカットキーになっています。

または、CLI が好みの人は [Browser](http://localhost:8001/redis-stack/browser) 画面の左下にある `>_CLI` ボタンで CLI を開く事もできます。

## Redis の基本

Redis の使い方の基本は以前 [Redis](https://qiita.com/propella/items/ee2fff343a834c87abae) に書きましたが、忘れたので再確認します。`>` の行がユーザの入力です。

```shell
> ping // 導通確認
"PONG"
> set mykey somevalue // 文字列の保存
"OK"
> get mykey // 文字列の取得
"somevalue"

> mset a 10 b 20 c 30 // 複数文字列の保存
"OK"
> mget a b c // 複数文字列の取得。リストとして返ります。
1) "10"
2) "20"
3) "30"
> INCR a // 値の文字列を数字として解釈できる時は INCR や DECR などの演算を行える。演算はアトミック。
(integer) 11
```

奇妙ですが、Redis はキーにリストやハッシュなどのデータ構造を保存する事ができます。この後で使うのでハッシュの保存と取得の例を挙げます。

```shell
> hset bike:1 model Deimos brand Ergonom price 4972 // ハッシュに値を保存するには、`hset key field value field value...` のように指定します。
(integer) 3
> hget bike:1 model // ハッシュの特定の値を取り出すには、`hget key field` とします。
"Deimos"
> hgetall bike:1 // ハッシュの全部の値を hgetall で取り出せます。リストとして返ります。
1) "model"
2) "Deimos"
3) "brand"
4) "Ergonom"
5) "price"
6) "4972"
> get bike:1 // ちなみに、ハッシュの値を get で取り出すとエラーが出ます。
"WRONGTYPE Operation against a key holding the wrong kind of value"
```

今まで保存したデータを全部消すには `FLUSHALL` を実行します。

## Redis Embedding

ここから先は CLI ではキツイので Python でやります。必要なパッケージをインストールしたり `export OPENAI_API_KEY=sk-ほにゃらら` で OpenAPI の API Key を設定する必要があります。以下必要な import 文。

```python
import openai
import numpy as np
import re
import requests
import xml.etree.ElementTree as ET

import redis
from redis.commands.search.query import Query
from redis.commands.search.field import TextField, VectorField
from redis.commands.search.indexDefinition import IndexDefinition, IndexType
```

### 日本国憲法をダウンロードして各条文ごとに Embedding を作成する

日本国憲法は e-gov.go.jp から XML でダウンロードできるので、適当に加工して条文のリストを作ります。

```python
response = requests.get("https://elaws.e-gov.go.jp/api/1/lawdata/321CONSTITUTION")
response.raise_for_status()
root = ET.fromstring(response.text)
preamble = root.find(".//Preamble")
preamble_text = "前文 " + re.sub(r'\s+', ' ',"".join(preamble.itertext()).strip()) 

articles = root.findall(".//Article")
article_texts = [re.sub(r'\s+', ' ',"".join(article.itertext()).strip()) for article in articles]

constitution = [preamble_text, *article_texts]
constitution[-3:]
```

```
['第百一条 この憲法施行の際、参議院がまだ成立してゐないときは、その成立するまでの間、衆議院は、国会としての権限を行ふ。',
 '第百二条 この憲法による第一期の参議院議員のうち、その半数の者の任期は、これを三年とする。 その議員は、法律の定めるところにより、これを定める。',
 '第百三条 この憲法施行の際現に在職する国務大臣、衆議院議員及び裁判官並びにその他の公務員で、その地位に相応する地位がこの憲法で認められてゐる者は、法律で特別の定をした場合を除いては、この憲法施行のため、当然にはその地位を失ふことはない。 但し、この憲法によつて、後任者が選挙又は任命されたときは、当然その地位を失ふ。']
```

これを、一つずつ OpenAI の embedding API に与えると Embedding という 1536 次元ベクトルが返ります。

```python
embeddings = []
for text in constitution:
    embedding = openai.Embedding.create(input = [text], model="text-embedding-ada-002")['data'][0]['embedding']
    embeddings.append(embedding)

embeddings[0][:3] # 一応中を確認
```

```
[-0.00279120821505785, -0.004937539808452129, -0.004875657614320517]
```

### Redis にデータと Embedding を保存

Redis に接続し導通確認を行います。

```python
r = redis.Redis() # default settings | localhost:6379 password=""
r.ping()
```

```
True
```

次に index というのを作成します。普通 Redis では key を元に検索しますが、Redis Stack では別に index を使っても検索できます。

```python
r.ft("idx:constitution").create_index( # idx:constitution という名前のインデックスを作成します。
    fields = [
        TextField("text"), # 条文本体を保存する text という名前の field を作成します。
        VectorField("embedding", # embedding を保存する embedding という名前の field を作成します。
        "HNSW", {
            "TYPE": "FLOAT32",
            "DIM": 1536, # OpenAI text-embedding-ada-002 の場合の次元数 1536 です。
            "DISTANCE_METRIC": "COSINE" # 類似度の計算に OpenAI おすすめの COSINE を使います。
        }),
    ],
    definition = IndexDefinition(
        prefix="constitution", # Redis に constitution から始まるデータを保存するとこの index が作成されます。
        index_type=IndexType.HASH, # ここではデータ型として hash を使います。json も使えます。
    )
)
```

Redis の HSET コマンドを使って、`constitution` から始まる key に条文と Embedding を保存します。

```python
for idx, text in enumerate(constitution):
    key = f"constitution:{idx}"
    vector = embeddings[idx]
    vector_bytes = np.array(vector, dtype='float32').tobytes()
    
    r.hset(key, mapping={
        "text": text,
        "embedding": vector_bytes
    })
```

### あいまい検索を実行する。

いよいよ検索を実行します。クエリの文法は結構難しいので、[Query | Redis](https://redis.io/docs/interact/search-and-query/query/) を参照してください。

```python
def query(text, top_k=2):
    
    # 検索文字列もデータ同様 Embedding 化します。
    embedding = openai.Embedding.create(input = [text], model="text-embedding-ada-002")['data'][0]['embedding']
    embedded_query = np.array(embedding, dtype='float32').tobytes()
    
    # クエリを作成します。検索文字に近い条文をマッチした順に top_k 個返します。
    q = Query(f'*=>[KNN {top_k} @embedding $vec_param AS vector_score]').sort_by('vector_score').paging(0,top_k).return_fields('vector_score','text').dialect(2) 
    params_dict = {"vec_param": embedded_query}

    # インデックスに対してクエリを発行します。
    results = r.ft("idx:constitution").search(q, query_params = params_dict)
    return results

def print_query(text, top_k=2):
    result = query(text, top_k)
    for doc in result.docs:
        print(f"score: {doc['vector_score']}")
        print(f"text: {doc['text']}")


print_query("5000兆円欲しい", 5)
```

```
score: 0.188872158527
text: 第八十五条 国費を支出し、又は国が債務を負担するには、国会の議決に基くことを必要とする。
score: 0.193886220455
text: 第四十九条 両議院の議員は、法律の定めるところにより、国庫から相当額の歳費を受ける。
score: 0.199605345726
text: 第八十三条 国の財政を処理する権限は、国会の議決に基いて、これを行使しなければならない。
score: 0.19975990057
text: 第八十八条 すべて皇室財産は、国に属する。 すべて皇室の費用は、予算に計上して国会の議決を経なければならない。
score: 0.200573921204
text: 第八十七条 予見し難い予算の不足に充てるため、国会の議決に基いて予備費を設け、内閣の責任でこれを支出することができる。 すべて予備費の支出については、内閣は、事後に国会の承諾を得なければならない。
```

「5000兆円欲しい」と質問すると、関連する憲法の条文を答えてくれました。結構あってる気がする。調子に乗って色々質問してみます。

```python
print_query("君達はどう生きるか", 5)

score: 0.203023433685
text: 第二十二条 何人も、公共の福祉に反しない限り、居住、移転及び職業選択の自由を有する。 何人も、外国に移住し、又は国籍を離脱する自由を侵されない。
score: 0.204712867737
text: 第二十五条 すべて国民は、健康で文化的な最低限度の生活を営む権利を有する。 国は、すべての生活部面について、社会福祉、社会保障及び公衆衛生の向上及び増進に努めなければならない。
score: 0.211046934128
text: 第十三条 すべて国民は、個人として尊重される。 生命、自由及び幸福追求に対する国民の権利については、公共の福祉に反しない限り、立法その他の国政の上で、最大の尊重を必要とする。
score: 0.214408934116
text: 第二条 皇位は、世襲のものであつて、国会の議決した皇室典範の定めるところにより、これを継承する。
score: 0.216077923775
text: 第二十三条 学問の自由は、これを保障する。
```

```python
print_query("SDGs", 5)

score: 0.240178346634
text: 第二十五条 すべて国民は、健康で文化的な最低限度の生活を営む権利を有する。 国は、すべての生活部面について、社会福祉、社会保障及び公衆衛生の向上及び増進に努めなければならない。
score: 0.250034213066
text: 第八十五条 国費を支出し、又は国が債務を負担するには、国会の議決に基くことを必要とする。
score: 0.256796717644
text: 第九条 日本国民は、正義と秩序を基調とする国際平和を誠実に希求し、国権の発動たる戦争と、武力による威嚇又は武力の行使は、国際紛争を解決する手段としては、永久にこれを放棄する。 前項の目的を達するため、陸海空軍その他の戦力は、これを保持しない。 国の交戦権は、これを認めない。
score: 0.257199466228
text: 第十三条 すべて国民は、個人として尊重される。 生命、自由及び幸福追求に対する国民の権利については、公共の福祉に反しない限り、立法その他の国政の上で、最大の尊重を必要とする。
score: 0.260500133038
text: 第九十二条 地方公共団体の組織及び運営に関する事項は、地方自治の本旨に基いて、法律でこれを定める。
```

と言う事で、日本国憲法は現在でも様々な問いに答えてくれる事が分かりました。

## 参考

* [Get started with Redis | Redis](https://redis.io/docs/getting-started/)
* [Redis hashes | Redis](https://redis.io/docs/data-types/hashes/)
* [Vector similarity | Redis](https://redis.io/docs/interact/search-and-query/search/vectors/)
* [FT.SEARCH | Redis](https://redis.io/commands/ft.search/)
* [Query | Redis](https://redis.io/docs/interact/search-and-query/query/)

## おまけ: 法令 API について

法令 API というのを使うと、API を使って日本の各法令をダウンロードできます。以下に curl を使った例を挙げます。どう言うわけか User-Agent がデフォルトのままだと `The requested URL was rejected. ` と怒られてしまうため、`-A ""` で消しています。

    curl -A "" https://elaws.e-gov.go.jp/api/1/lawdata/321CONSTITUTION

法令を指定するには法令IDを使います。日本国憲法の場合 `321CONSTITUTION` です。下記「法令IDについて」に法令IDの仕様がありますが、法令を e-gov.go.jp で検索して出て来るページの lawid が法令IDです。例えば民法 https://elaws.e-gov.go.jp/document?lawid=129AC0000000089 の場合は `129AC0000000089` が法令IDです。

* [法令API | e-Gov法令検索](https://elaws.e-gov.go.jp/apitop/)
* [e-Gov 法令 API(Version 1)仕様書](https://elaws.e-gov.go.jp/file/houreiapi_shiyosyo.pdf)
* [法令IDについて](https://elaws.e-gov.go.jp/file/LawIdNamingConvention.pdf)