ai-tagging/agent/tests/test_es.py

from agent.core.es import es
from agent.core.vector import get_embeddings
query = '职业:电动机加工制造,投向:电动机加工制造,用途:电动机加工制造'
resp_bm25 = es.search(
        index='ai-tagging',
        size=10,
        query={
            "multi_match": {
                "query": query,
                "fields": ["tag_path", "tag_remark", "tag_prompt"]
            }
        }
    )
for rank, hit in enumerate(resp_bm25['hits']['hits'], start=1):
    hit["_source"]["tag_vector"] = None
    doc_id = hit['_id']
    tag_nm = hit['_source']['tag_name']
    print(doc_id,tag_nm)
print('-----------------')

query_vector = get_embeddings([query])[0]
resp_vector = es.search(
        index='ai-tagging',
        size=10,
        knn={
            "field": "tag_vector",
            "query_vector": query_vector,
            "k": 10,
            "num_candidates": 100
        }
    )
for rank, hit in enumerate(resp_vector['hits']['hits'], start=1):
    hit["_source"]["tag_vector"] = None
    doc_id = hit['_id']
    tag_nm = hit['_source']['tag_name']
    print(doc_id,tag_nm)