46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
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
|