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