nexus/tools/kafka_tool.py

"""
Kafka Tool - Produce / consume Kafka messages qua kafka-python.
Author: Hieu Louis (2026)
Operations: produce, consume, list_topics, describe_topic.
"""
from __future__ import annotations

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

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


# Produce = write (cần confirmation) // write ops
WRITE_OPS = {"produce"}


class KafkaTool(Tool):
    """Produce/consume Kafka messages: produce, consume, list_topics, describe_topic."""

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

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

    @property
    def description(self) -> str:
        return (
            "Produce/consume Kafka messages qua kafka-python: produce (sync/async), "
            "consume (poll batch), list_topics, describe_topic."
        )

    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "bootstrap_servers": {
                    "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}],
                    "description": "Kafka brokers, vd: localhost:9092",
                },
                "operation": {
                    "type": "string",
                    "enum": ["produce", "consume", "list_topics", "describe_topic"],
                    "description": "Kafka operation",
                },
                "topic": {"type": "string", "description": "Topic name"},
                "messages": {
                    "oneOf": [{"type": "string"}, {"type": "array", "items": {}}],
                    "description": "Message(s) cần produce (string hoặc list)",
                },
                "key": {"type": "string", "description": "Message key (produce)"},
                "group_id": {"type": "string", "description": "Consumer group (consume)"},
                "max_messages": {"type": "integer", "description": "Giới hạn số message consume (default 100)"},
                "timeout_ms": {"type": "integer", "description": "Poll timeout ms (default 5000)"},
                "sasl_username": {"type": "string", "description": "SASL username (optional)"},
                "sasl_password": {"type": "string", "description": "SASL password (optional)"},
            },
            "required": ["bootstrap_servers", "operation"],
        }

    def validate_args(self, args: Dict[str, Any]) -> Optional[str]:
        if not args.get("bootstrap_servers"):
            return "Missing required arg: bootstrap_servers"
        op = args.get("operation")
        if not op:
            return "Missing required arg: operation"
        if op in {"produce", "consume", "describe_topic"} and not args.get("topic"):
            return f"Operation '{op}' requires 'topic' arg"
        if op == "produce" and args.get("messages") is None:
            return "Operation 'produce' requires 'messages' arg"
        return None

    def execute(self, args: Dict[str, Any], context: ToolContext) -> ToolResult:
        bs = args["bootstrap_servers"]
        op: str = args["operation"]
        sasl_u = args.get("sasl_username")
        sasl_p = args.get("sasl_password")

        # Lazy import kafka-python // lazy import
        try:
            from kafka import KafkaProducer, KafkaConsumer, KafkaAdminClient  # type: ignore
            from kafka.admin import NewTopic  # type: ignore  # noqa: F401
        except ImportError as e:
            return ToolResult(
                success=False,
                error=f"kafka-python not installed: {e}. Cài: pip install kafka-python",
                return_code=127,
            )

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

        common_kwargs: Dict[str, Any] = {
            "bootstrap_servers": bs if isinstance(bs, list) else [bs],
            "request_timeout_ms": (context.timeout * 1000) if context.timeout else 30000,
        }
        if sasl_u and sasl_p:
            common_kwargs.update({
                "security_protocol": "SASL_PLAINTEXT",
                "sasl_mechanism": "PLAIN",
                "sasl_plain_username": sasl_u,
                "sasl_plain_password": sasl_p,
            })

        try:
            if op == "list_topics":
                admin = KafkaAdminClient(**common_kwargs)
                topics = sorted(admin.list_topics())
                admin.close()
                payload = json.dumps(topics, ensure_ascii=False, indent=2)
                return ToolResult(success=True, output=payload, metadata={"count": len(topics)})

            if op == "describe_topic":
                admin = KafkaAdminClient(**common_kwargs)
                # Lấy partitions // get partitions via consumer
                consumer = KafkaConsumer(args["topic"], **{k: v for k, v in common_kwargs.items() if k != "request_timeout_ms"}, request_timeout_ms=common_kwargs["request_timeout_ms"])
                parts = sorted(consumer.partitions_for_topic(args["topic"]) or [])
                consumer.close()
                admin.close()
                info = {"topic": args["topic"], "partitions": list(parts), "partition_count": len(parts)}
                return ToolResult(success=True, output=json.dumps(info, ensure_ascii=False, indent=2), metadata=info)

            if op == "produce":
                producer = KafkaProducer(
                    bootstrap_servers=common_kwargs["bootstrap_servers"],
                    value_serializer=lambda v: v.encode("utf-8") if isinstance(v, str) else json.dumps(v).encode("utf-8"),
                    key_serializer=lambda k: k.encode("utf-8") if isinstance(k, str) else k,
                    request_timeout_ms=common_kwargs["request_timeout_ms"],
                )
                msgs = args["messages"]
                if not isinstance(msgs, list):
                    msgs = [msgs]
                futures = []
                for m in msgs:
                    futures.append(producer.send(args["topic"], key=args.get("key"), value=m))
                # flush chờ gửi xong // wait for all
                producer.flush()
                for f in futures:
                    f.get(timeout=context.timeout or 30)  # raises nếu lỗi
                producer.close()
                return ToolResult(
                    success=True,
                    output=f"Produced {len(msgs)} message(s) to {args['topic']}",
                    metadata={"topic": args["topic"], "count": len(msgs)},
                )

            if op == "consume":
                consumer = KafkaConsumer(
                    args["topic"],
                    group_id=args.get("group_id"),
                    bootstrap_servers=common_kwargs["bootstrap_servers"],
                    auto_offset_reset="earliest",
                    enable_auto_commit=False,
                    consumer_timeout_ms=int(args.get("timeout_ms", 5000)),
                    value_deserializer=lambda v: v.decode("utf-8", errors="replace"),
                )
                max_msgs = int(args.get("max_messages") or 100)
                records: List[Dict[str, Any]] = []
                for msg in consumer:
                    records.append({
                        "topic": msg.topic,
                        "partition": msg.partition,
                        "offset": msg.offset,
                        "key": msg.key.decode("utf-8", errors="replace") if msg.key else None,
                        "value": msg.value,
                    })
                    if len(records) >= max_msgs:
                        break
                consumer.close()
                payload = json.dumps(records, ensure_ascii=False, indent=2)
                return ToolResult(success=True, output=payload, metadata={"consumed": len(records), "topic": args["topic"]})

            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)