ai_pipeline/src/pipeline/core/es_client.py

46 lines
1.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import httpx
from pipeline.config import config
class AsyncES:
def __init__(self):
self.base = f"{config['es_host']}:{config['es_port']}"
self.auth = (config["es_user"], config['es_password'])
async def create_index(self, index: str):
"""
创建 index知识库
"""
async with httpx.AsyncClient() as client:
resp = await client.put(
f"{self.base}/{index}",
auth=self.auth,
json={
"settings": {"analysis": {"analyzer": {"default": {"type": "standard"}}}},
"mappings": {
"properties": {
"text": {"type": "text"},
"kb": {"type": "keyword"},
}
},
},
)
return resp.json()
async def add_doc(self, index: str, text: str, kb: str):
"""
写入文档(普通文本)
"""
async with httpx.AsyncClient() as client:
resp = await client.post(f"{self.base}/{index}/_doc", auth=self.auth, json={"text": text, "kb": kb})
return resp.json()
async def bm25_search(self, index: str, query: str, top_k: int = 10):
"""
BM25 搜索
"""
async with httpx.AsyncClient() as client:
resp = await client.post(f"{self.base}/{index}/_search", auth=self.auth, json={"size": top_k, "query": {"match": {"text": query}}})
return resp.json()
es_client = AsyncES