kangda-robot-backend/ruoyi-fastapi-backend/sub_applications/kb_es/create_index.py

81 lines
2.3 KiB
Python

import os
from typing import Any, Dict
import httpx
def _bool_env(name: str, default: bool) -> bool:
val = os.getenv(name)
if val is None:
return default
return val.strip().lower() in {"1", "true", "yes", "y", "on"}
def _get_settings() -> Dict[str, Any]:
es_url = os.getenv("KB_ES_URL") or os.getenv("KB_OS_URL")
if not es_url:
raise SystemExit("Missing KB_ES_URL (or KB_OS_URL)")
return {
"es_url": es_url.rstrip("/"),
"index": os.getenv("KB_ES_INDEX", "kb_chunks_v1"),
"username": os.getenv("KB_ES_USERNAME", ""),
"password": os.getenv("KB_ES_PASSWORD", ""),
"verify_ssl": _bool_env("KB_ES_VERIFY_SSL", True),
}
def _mapping() -> Dict[str, Any]:
return {
"settings": {
"analysis": {
"analyzer": {
"kb_default": {
"type": "standard",
}
}
}
},
"mappings": {
"dynamic": "true",
"properties": {
"doc_id": {"type": "keyword"},
"chunk_id": {"type": "keyword"},
"title": {"type": "text", "analyzer": "kb_default"},
"content": {"type": "text", "analyzer": "kb_default"},
"source": {"type": "keyword"},
"tags": {"type": "keyword"},
"category": {"type": "keyword"},
"updated_at": {"type": "date", "ignore_malformed": True},
"version": {"type": "keyword"},
},
},
}
async def main() -> None:
cfg = _get_settings()
url = f"{cfg['es_url']}/{cfg['index']}"
auth = None
if cfg["username"] and cfg["password"]:
auth = (cfg["username"], cfg["password"])
async with httpx.AsyncClient(timeout=10.0, verify=cfg["verify_ssl"]) as client:
head = await client.head(url, auth=auth)
if head.status_code == 200:
print(f"Index already exists: {cfg['index']}")
return
resp = await client.put(url, json=_mapping(), auth=auth)
if resp.status_code >= 300:
raise SystemExit(f"Create index failed HTTP {resp.status_code}: {resp.text}")
print(f"Created index: {cfg['index']}")
if __name__ == "__main__":
import asyncio
asyncio.run(main())