nexus/tools/elasticsearch_tool.py

"""
Elasticsearch Tool - Quản lý Elasticsearch qua elasticsearch-py.
Author: Hieu Louis (2026)
Operations: search, index, create, update, delete, bulk.
"""
from __future__ import annotations

import json
from datetime import datetime, date
from typing import Any, Dict, List, Optional

from .base import Tool, ToolResult, ToolContext, ToolCategory, ToolSafety


# Write/delete ops (cần confirmation) // write ops
WRITE_OPS = {"index", "create", "update", "delete", "bulk"}


def _json_default(o: Any) -> Any:
    if isinstance(o, (datetime, date)):
        return o.isoformat()
    if isinstance(o, bytes):
        try:
            return o.decode("utf-8")
        except Exception:
            return o.hex()
    return str(o)


class ElasticsearchTool(Tool):
    """Quản lý Elasticsearch: search, index, create, update, delete, bulk."""

    category = ToolCategory.DATABASE
    safety = ToolSafety.DANGEROUS
    requires_confirmation = True

    @property
    def name(self) -> str:
        return "elasticsearch"

    @property
    def description(self) -> str:
        return (
            "Quản lý Elasticsearch qua elasticsearch-py: search, index, create, "
            "update, delete, bulk. Hỗ trợ auth header và cloud_id."
        )

    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "hosts": {
                    "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}],
                    "description": "ES host(s), vd: http://localhost:9200",
                },
                "api_key": {"type": "string", "description": "API key (optional)"},
                "operation": {
                    "type": "string",
                    "enum": ["search", "index", "create", "update", "delete", "bulk"],
                    "description": "ES operation",
                },
                "index": {"type": "string", "description": "Tên ES index"},
                "id": {"type": "string", "description": "Document ID (create/update/delete)"},
                "body": {"type": "object", "description": "Document body hoặc query DSL"},
                "query": {"type": "object", "description": "Query DSL (search)"},
                "doc": {"type": "object", "description": "Partial doc (update)"},
                "actions": {"type": "array", "items": {"type": "object"}, "description": "Bulk actions (vd: [{index:{_id:1}}, {doc:..}])"},
                "size": {"type": "integer", "description": "Số kết quả trả về (default 50)"},
            },
            "required": ["hosts", "operation"],
        }

    def validate_args(self, args: Dict[str, Any]) -> Optional[str]:
        if not args.get("hosts"):
            return "Missing required arg: hosts"
        op = args.get("operation")
        if not op:
            return "Missing required arg: operation"
        if op in {"search", "index", "create", "update", "delete"} and not args.get("index"):
            return f"Operation '{op}' requires 'index' arg"
        if op in {"create", "update", "delete"} and not args.get("id"):
            return f"Operation '{op}' requires 'id' arg"
        if op == "search" and not (args.get("query") or args.get("body")):
            return "Operation 'search' requires 'query' or 'body' arg"
        if op == "bulk" and not args.get("actions"):
            return "Operation 'bulk' requires 'actions' arg"
        return None

    def execute(self, args: Dict[str, Any], context: ToolContext) -> ToolResult:
        hosts = args["hosts"]
        api_key = args.get("api_key")
        op: str = args["operation"]

        # Lazy import elasticsearch // lazy import
        try:
            from elasticsearch import Elasticsearch  # type: ignore
        except ImportError as e:
            return ToolResult(
                success=False,
                error=f"elasticsearch not installed: {e}. Cài: pip install elasticsearch",
                return_code=127,
            )

        # Dry-run cho write ops // dry-run
        if context.dry_run and op in WRITE_OPS:
            return ToolResult(
                success=True,
                output=f"[dry-run] Would run {op} on index={args.get('index')!r}",
                metadata={"dry_run": True, "operation": op, "index": args.get("index")},
            )

        try:
            client = Elasticsearch(hosts, api_key=api_key, request_timeout=context.timeout)

            if op == "search":
                body = args.get("body") or {"query": args["query"]}
                body.setdefault("size", int(args.get("size") or 50))
                resp = client.search(index=args["index"], body=body)
                hits = resp.get("hits", {}).get("hits", [])
                payload = json.dumps(resp, default=_json_default, ensure_ascii=False, indent=2)
                return ToolResult(success=True, output=payload, metadata={"hits": len(hits), "took": resp.get("took")})

            if op == "index":
                resp = client.index(index=args["index"], id=args.get("id"), document=args["body"])
                return ToolResult(
                    success=resp.get("result") in {"created", "updated"},
                    output=json.dumps(resp, default=_json_default, ensure_ascii=False, indent=2),
                    metadata={"result": resp.get("result"), "id": resp.get("_id")},
                )

            if op == "create":
                resp = client.create(index=args["index"], id=args["id"], document=args["body"])
                return ToolResult(
                    success=resp.get("result") == "created",
                    output=json.dumps(resp, default=_json_default, ensure_ascii=False, indent=2),
                    metadata={"result": resp.get("result"), "id": resp.get("_id")},
                )

            if op == "update":
                resp = client.update(index=args["index"], id=args["id"], doc=args.get("doc") or args.get("body"))
                return ToolResult(
                    success=resp.get("result") == "updated",
                    output=json.dumps(resp, default=_json_default, ensure_ascii=False, indent=2),
                    metadata={"result": resp.get("result")},
                )

            if op == "delete":
                resp = client.delete(index=args["index"], id=args["id"])
                return ToolResult(
                    success=resp.get("result") == "deleted",
                    output=json.dumps(resp, default=_json_default, ensure_ascii=False, indent=2),
                    metadata={"result": resp.get("result")},
                )

            if op == "bulk":
                # actions là list các cặp action_header/doc // actions list
                body_lines: List[str] = []
                for item in args["actions"]:
                    body_lines.append(json.dumps(item, default=_json_default))
                body = "\n".join(body_lines) + "\n"
                resp = client.bulk(index=args.get("index"), body=body)
                errors = bool(resp.get("errors"))
                return ToolResult(
                    success=not errors,
                    output=json.dumps(resp, default=_json_default, ensure_ascii=False, indent=2),
                    metadata={"errors": errors, "took": resp.get("took")},
                )

            return ToolResult(success=False, error=f"Unknown operation: {op}", return_code=1)
        except Exception as e:
            return ToolResult(success=False, error=str(e), return_code=1)