RAGのメイン処理

RAGのメイン処理

    # RAG検索
    query_embed = ollama_embed(prompt)
    results = st.session_state.collection.query(
        query_embeddings=[query_embed],
        n_results=2
    )

query_embed = ollama_embed(prompt)

でユーザのプロンプトをベクトル化

session_state.collection.query

を使うことで
chromaDBのメソッドを使うことで
似ている情報をベクトル同士を比べることで似ている情報を得ることが可能になる

この中で

query_embeddings=[query_embed],

でプロンプトを指定して

n_results=2

で上位何個を該当させるか指定できる
今回は上位２つまで

このときの結果は

    # {
    #     "ids":
    #     "documents": [["doc1", "doc2"]]
    #     "distances": [[XXX, XXX]]
    # }

というように返ってくる
Documents はテキスト内容
Distancesはベクトルの数値
なので必要なのは documentsの値になるので
これがあるかで判定すればいい

なので
Documentsの０番目のインデックスを取得すれば
最初のリストが取得できる

あとは中身を joinで結合すればいいので

    if results["documents"]:
        context_text = "\n".join(results["documents"][0])

というようにして
変数に結果を格納する

次に、RAG部分のプロンプトの作成

        rag_prompt = f"""
        以下は関連ドキュメントの抜粋です。
        {context_text}
        この情報を参考に以下の質問に答えてください。
        {prompt}
        """

こうすることで
質問を入力すれば、事前に参照するプロンプトが既に組み込まれているので
簡単にDB参照の機能が追加された状態で質問が構築される

もし、documentsがない、つまり該当する知識がないのなら

        final_user_prompt = rag_prompt
    else:
        final_user_prompt = prompt

というようにすれば分岐処理になって
RAGなしのプロンプトが実行されるようになる

あとは

    st.session_state.messages.append({"role": "user", "content": final_user_prompt})

で情報を履歴として保持するようにする

あとは
システムプロンプトを入れて過去の履歴を追加する

    if system_prompt.strip():
        messages = [{"role": "system", "content": system_prompt}] + st.session_state.messages
    else:
        messages = st.session_state.messages
    

これで RAGを使ったmessages にプロンプトが格納される

それを ollamaに渡す

    # LLMの返答を表示
    with st.chat_message("assistant"):
        placeholder = st.empty()
        stream_response = ""
        stream = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=temperature,
            stream=True
        )
        for chunk in stream:
            stream_response += chunk.choices[0].delta.content
            placeholder.write(stream_response)

これで streamlitで回答を表示
回答内容を保存もしている