드디어 등장!? Firestore의 전문 검색(벡터 검색)

functions-framework==3.*
google-cloud-firestore
google-cloud-aiplatform

import functions_framework
import os
# firestore
from google.cloud import firestore
from google.cloud.firestore_v1.vector import Vector
# Vertex AI
import vertexai
from vertexai.language_models import TextEmbeddingModel

# 프로젝트명(환경 변수에서 취득)
MY_PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT") 

# 전달된 문자열에서 벡터 값을 계산
def text_embedding(text: str) -> list:

    # location은 각자의 로케이션을 설정
    vertexai.init(project=MY_PROJECT_ID, location="asia-northeast1") 

    # 현재 벡터 계산을 위한 최신 AI가 'textembedding-gecko@003'이므로 이를 사용
    model = TextEmbeddingModel.from_pretrained("textembedding-gecko@003")
    embeddings = model.get_embeddings([text])
    for embedding in embeddings:
        vector = embedding.values

    return Vector(vector)

# 메인 처리
# (함수명은 적당히)
@functions_framework.http
def hello_http(request):

    # 요청에서 기사 요약(description)을 취득하고 있습니다
    request_json = request.get_json(silent=True)
    request_args = request.args

    if request_json and 'description' in request_json:
        description = request_json['description']
    elif request_args and 'description' in request_args:
        description = request_args['description']
    else:
        description = 'World'

    # Firestore 클라이언트 초기화
    firestore_client = firestore.Client(project=MY_PROJECT_ID)
    # 컬렉션 참조(컬렉션명은 임의(컬렉션이 없는 경우 사전에 생성해주세요))
    collection = firestore_client.collection("article_collection")

    # embedding을 계산
    embedding_vector = text_embedding(description)

    # Firestore에 추가할 문서 준비
    doc = {
        "description": description,
        "embedding_field": embedding_vector
    }
    # 문서 추가
    collection.add(doc)

    return 'OK!'

curl -m 70 -X POST https://asia-northeast1-python-tool-001.cloudfunctions.net/vector_chenge \
  -H "Authorization: bearer $(gcloud auth print-identity-token)" \
  -H "Content-Type: application/json" \
  -d '{ "description": "<임의의 문자열>"}'

gcloud alpha firestore indexes composite create \
  --collection-group=article_collection \
  --query-scope=COLLECTION \
  --field-config field-path=embedding_field,vector-config='{"dimension":"768", "flat": "{}"}' \
  --database=<데이터베이스 ID>

import functions_framework
import os
# firestore
from google.cloud.firestore import firestore
from google.cloud.firestore_v1.vector import Vector
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
# Vertex AI
import vertexai
from vertexai.language_models import TextEmbeddingModel

# 프로젝트명(환경 변수에서 취득)
MY_PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT") 

# 전달된 문자열에서 벡터 값을 계산
def text_embedding(text: str) -> list:

    # location은 각자의 로케이션을 설정
    vertexai.init(project=MY_PROJECT_ID, location="asia-northeast1") 

    # 현재 벡터 계산을 위한 최신 AI가 'textembedding-gecko@003'이므로 이를 사용
    model = TextEmbeddingModel.from_pretrained("textembedding-gecko@003")
    embeddings = model.get_embeddings([text])
    for embedding in embeddings:
        vector = embedding.values

    return Vector(vector)

# 메인 처리
# (함수명은 적당히)
@functions_framework.http
def hello_http(request):

    # 요청에서 기사 요약(description)을 취득하고 있습니다
    request_json = request.get_json(silent=True)
    request_args = request.args

    if request_json and 'target' in request_json:
        target = request_json['target']
    elif request_args and 'target' in request_args:
        target = request_args['target']
    else:
        target = 'World'

    # Firestore 클라이언트 초기화
    firestore_client = firestore.Client(project=MY_PROJECT_ID)
    # 컬렉션 참조
    collection = firestore_client.collection("article_collection")

    # embedding을 계산
    embedding_vector = text_embedding(target)

    # 벡터 검색 실행
    docs = collection.find_nearest(
        vector_field="embedding_field",
        query_vector=embedding_vector,
        distance_measure=DistanceMeasure.COSINE,
        limit=3
    ).get()

    # 표 형식(여기서는 문자열 형식)으로의 출력용
    output = "Description \n"
    output += "-" * 50 + "\n"
    
    # 벡터 검색으로 취득한 문서의 내용을 출력
    for doc in docs:
        doc_data = doc.to_dict()
        description = doc_data.get("description", "No description")
        # 문서 내용을 문자열에 추가
        output += f"{description[:100]} \n"
    
    return output

curl -m 70 -X POST https://asia-northeast1-python-tool-001.cloudfunctions.net/vector_search 
-H "Authorization: bearer $(gcloud auth print-identity-token)" 
-H "Content-Type: application/json" 
-d '{
  "target": "About Riverpod"
}'

Description 
--------------------------------------------------
What is Riverpod, Flutter's most major state management introduction! StatefulWidget, which we introduced before, is one of the functions that performs state management…
What is Flutter? An overview of Flutter What is Flutter? Flutter" is "useful for developing mobile apps…
What is MVVM, which we often hear about in app development? MVVM (Model-View-ViewModel) is the combination of an app's logic and UI (user interface)…