Elasticsearch Goクライアントを使用してElasticsearchでベクトル検索を実行する

func GetElasticsearchClient() (*elasticsearch.TypedClient, error) {
	var cloudID = os.Getenv("ELASTIC_CLOUD_ID")
	var apiKey = os.Getenv("ELASTIC_API_KEY")

	var es, err = elasticsearch.NewTypedClient(elasticsearch.Config{
		CloudID: cloudID,
		APIKey:  apiKey,
		Logger:  &elastictransport.ColorLogger{os.Stdout, true, true},
	})

	if err != nil {
		return nil, fmt.Errorf("unable to connect: %w", err)
	}

	return es, nil
}

docker run -it --rm --network host \
    docker.elastic.co/eland/eland \
    eland_import_hub_model \
      --cloud-id $ELASTIC_CLOUD_ID \
      --es-api-key $ELASTIC_API_KEY \
      --hub-model-id sentence-transformers/msmarco-MiniLM-L-12-v3 \
      --task-type text_embedding

PUT _ingest/pipeline/search-rodents-vector-embedding-pipeline
{
  "processors": [
    {
      "inference": {
        "model_id": "sentence-transformers__msmarco-minilm-l-12-v3",
        "target_field": "text_embedding",
        "field_map": {
          "body_content": "text_field"
        }
      }
    }
  ]
}

PUT vector-search-rodents
{
  "mappings": {
    "properties": {
      "text_embedding.predicted_value": {
        "type": "dense_vector",
        "dims": 384,
        "index": true,
        "similarity": "cosine"
      },
      "text": {
        "type": "text"
      }
    }
  }
}

POST _reindex
{
  "source": {
    "index": "search-rodents"
  },
  "dest": {
    "index": "vector-search-rodents",
    "pipeline": "search-rodents-vector-embedding-pipeline"
  }
}

func VectorSearch(client *elasticsearch.TypedClient, term string) ([]Rodent, error) {
  var k = 10
	var numCandidates = 10

	res, err := client.Search().
		Index("vector-search-rodents").
		Knn(types.KnnSearch{
      # Field in document containing vector
			Field:         "text_embedding.predicted_value",
      # Number of neighbors to return
			K:             &k,
      # Number of candidates to evaluate in comparison
			NumCandidates: &numCandidates,
      # Generate query vector using the same model used in the inference processor
			QueryVectorBuilder: &types.QueryVectorBuilder{
				TextEmbedding: &types.TextEmbedding{
					ModelId:   "sentence-transformers__msmarco-minilm-l-12-v3",
					ModelText: term,
				},
			}}).Do(context.Background())

	if err != nil {
		return nil, fmt.Errorf("error in rodents vector search: %w", err)
	}

	return getRodents(res.Hits.Hits)
}

[
  {ID:64f74ecd4acb3df024d91112 Title:Gopher - Wikipedia Url:https://en.wikipedia.org/wiki/Gopher} 
  {ID:64f74ed34acb3d71aed91fcd Title:Squirrel - Wikipedia Url:https://en.wikipedia.org/wiki/Squirrel} 
  //Other results omitted
]

// HuggingFace text embedding helper
func GetTextEmbeddingForQuery(term string) []float32 {
    // HTTP endpoint
    model := "sentence-transformers/msmarco-minilm-l-12-v3"
    posturl := fmt.Sprintf("https://api-inference.huggingface.co/pipeline/feature-extraction/%s", model)

    // JSON body
    body := []byte(fmt.Sprintf(`{
        "inputs": "%s",
        "options": {"wait_for_model":True}
    }`, term))

    // Create a HTTP post request
    r, err := http.NewRequest("POST", posturl, bytes.NewBuffer(body))

    if err != nil {
        log.Fatal(err)
        return nil
    }

    token := os.Getenv("HUGGING_FACE_TOKEN")
    r.Header.Add("Authorization", fmt.Sprintf("Bearer %s", token))

    client := &http.Client{}
    res, err := client.Do(r)
    if err != nil {
        panic(err)
    }

    defer res.Body.Close()

    var post []float32
    derr := json.NewDecoder(res.Body).Decode(&post)

    if derr != nil {
        log.Fatal(derr)
        return nil
    }

    return post
}

func VectorSearchWithGeneratedQueryVector(client *elasticsearch.TypedClient, term string) ([]Rodent, error) {
	vector, err := GetTextEmbeddingForQuery(term)
	if err != nil {
		return nil, err
	}

	if vector == nil {
		return nil, fmt.Errorf("unable to generate vector: %w", err)
	}

  var k = 10
	var numCandidates = 10

	res, err := client.Search().
		Index("vector-search-rodents").
		Knn(types.KnnSearch{
      # Field in document containing vector
			Field:         "text_embedding.predicted_value",
      # Number of neighbors to return
			K:             &k,
      # Number of candidates to evaluate in comparison
			NumCandidates: &numCandidates,
      # Query vector returned from Hugging Face inference API
			QueryVector:   vector,
		}).
		Do(context.Background())

	if err != nil {
		return nil, err
	}

	return getRodents(res.Hits.Hits)
}