テックブログ

import base64
import json
import mimetypes
import os
import tempfile

import functions_framework
import urllib3
from langchain.retrievers import GoogleVertexAISearchRetriever
from langchain.callbacks import get_openai_callback
from langchain.chains.summarize import load_summarize_chain
from langchain.chat_models import ChatOpenAI, ChatVertexAI
from langchain.document_loaders import PyPDFLoader, GCSFileLoader
from langchain.prompts import PromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.chains import LLMChain

PROJECT_ID = os.environ.get("PROJECT_ID")
DATA_STORE_ID = os.environ.get("DATA_STORE_ID")
BUCKET_ID = os.environ.get("BUCKET_ID")

http = urllib3.PoolManager()
SLACK_BOT_TOKEN = os.environ.get("SLACK_BOT_TOKEN")
WEBHOOK_URL = os.environ.get("WEBHOOK_URL")

def search(query):
    retriever = GoogleVertexAISearchRetriever(
        project_id=PROJECT_ID,
        location_id=global,
        data_store_id=DATA_STORE_ID,
        max_documents=1,
        engine_data_type=0,
        max_extractive_answer_count=1,
        get_extractive_answers=False,
        query_expansion_condition=2,
    )

    result = retriever.get_relevant_documents(query)
    if query == "":
        slack_message = {"text": "質問文を書いてください。"}
        encoded_msg = json.dumps(slack_message, default=str).encode("utf-8")
        http.request("POST", WEBHOOK_URL, body=encoded_msg)
    elif len(result) == 0:
        slack_message = {"text": "検索でエラーが発生しました。少ししてからやり直してください。"}
        encoded_msg = json.dumps(slack_message, default=str).encode("utf-8")
        http.request("POST", WEBHOOK_URL, body=encoded_msg)

    url = result[0].metadata["source"]
    document = result[0].metadata["source"][result[0].metadata["source"].rfind("/") + 1:]
    # print(document)
    return url, document


def load_pdf(file_path):
    return PyPDFLoader(file_path)


def load_text(message):
    url, document = search(message)

   
    loader = GCSFileLoader(project_name=PROJECT_ID, bucket=BUCKET_ID, blob=document, loader_func=load_pdf)
    text = loader.load()

    return text, url


def summarize_palm2(message, docs, url):
    prompt_template = """\nUse the following text to answer the question. If there is no information relevant to your question, please say "No information".:\n# Question\n""" + message + """\n# text\n{text}"""

    PROMPT = PromptTemplate(template=prompt_template, input_variables=["text"])
    llm = ChatVertexAI(model="chat-bison@001", temperature=0)

    llm_chain = LLMChain(llm=llm, prompt=PROMPT)
    stuff_chain = StuffDocumentsChain(
        llm_chain=llm_chain,
        document_variable_name="text",
    )

    return stuff_chain.run(docs)

def summarize_gpt(message, docs, url, model):
    prompt_template = """
    You are a professional editor. Please extract sentences related to the following questions. However, the following constraints must be followed:

    # Constraints
    - Do not omit any important keywords.
    - Extract without compromising the original intent of the text.
    - Avoid using fictional expressions or terms.
    - Do not modify any numbers within the text.
    - Ensure that the extracted sentences contain detailed information.
    # Question\n""" + message + """\n# Input Text\n{text}"""

    # ただし、以下のテキストに以下の質問に関連する情報がない場合は「出力フォーマット1」に従い、以下のテキストに以下の質問に関連する情報がある場合は「出力フォーマット2」に従ってください。
    combine_prompt = """
    あなたはプロの編集者です。以下のテキストを使って、以下の質問に以下の出力のフォーマットに従って「回答」と「参考」を答えてください。:
    
    # 質問\n""" + message + """\n# テキスト\n{text}""" + """
    
    ### 出力フォーマット
    回答：（質問に対する回答を示す）
    参考：（質問に関連する情報がない場合は無し、質問に関連する情報がある場合は右のurlをそのまま出力してください）""" + url

    PROMPT = PromptTemplate(template=prompt_template, input_variables=["text"])
    COMBINE_PROMPT = PromptTemplate(template=combine_prompt, input_variables=["text"])

    llm = ChatOpenAI(model=model, temperature=0)

    chain = load_summarize_chain(
        llm, chain_type="map_reduce", verbose=True, map_prompt=PROMPT, combine_prompt=COMBINE_PROMPT
    )
    response = chain(
        {
            "input_documents": docs,
            "token_max": 7000,
        },
        return_only_outputs=True,
    )
    
    return response["output_text"]


# Triggered from a message on a Cloud Pub/Sub topic.
@functions_framework.cloud_event
def send_slack_message(cloud_event):
    # Decode the Pub/Sub message data
    pubsub_message = base64.b64decode(cloud_event.data["message"]["data"]).decode("utf-8")
    pubsub_message_data = json.loads(pubsub_message)
    message = pubsub_message_data["data"]["message"]
    message = message.strip("<>")
    command = pubsub_message_data["data"]["command"]
    doc, url = load_text(message)
    
    if command == "/palm2":
        response = summarize_palm2(message, doc, url)
        slack_message = {"text": "モデル：PaLM2\n質問：" + message + "\n回答：" + response + "\n参考：" + url} 
    elif command == "/gpt-3":
        response = summarize_gpt(message, doc, url, model="gpt-3.5-turbo")
        slack_message = {"text": "モデル：GPT-3.5\n質問：" + message + "\n" + response}
    elif command == "/gpt-4":
        response = summarize_gpt(message, doc, url, model="gpt-4")
        slack_message = {"text": "モデル：GPT-4\n質問：" + message + "\n" + response}
    else:
        slack_message = {"text": "コマンドがありません。"}

    try:
        encoded_msg = json.dumps(slack_message, default=str).encode("utf-8")
        http.request("POST", WEBHOOK_URL, body=encoded_msg)
        return "Message sent to Slack"
    except Exception as e:
        print(e)
        encoded_msg = json.dumps(e, default=str).encode("utf-8")
        http.request("POST", WEBHOOK_URL, body=encoded_msg)
        return "Message sent to Slack"

functions-framework==3.*
azure-core==1.26.4
langchain==0.0.312
openai==0.27.7
pypdf==3.13.0
slack-bolt==1.18.0
slack-sdk==3.21.3
requests==2.29.0
urllib3==1.26.15
tiktoken==0.4.0
google-cloud-discoveryengine==0.11.2
google-cloud-storage==2.10.0
google-cloud-aiplatform>=1.33.0