ai-tax-agent/libs/config/factories.py

"""Client factories for various services."""

from typing import Any

import boto3  # type: ignore
import hvac
import redis.asyncio as redis
from aiokafka import AIOKafkaConsumer, AIOKafkaProducer  # type: ignore
from minio import Minio
from neo4j import GraphDatabase
from qdrant_client import QdrantClient

from .settings import BaseAppSettings


class VaultClientFactory:  # pylint: disable=too-few-public-methods
    """Factory for creating Vault clients"""

    @staticmethod
    def create_client(settings: BaseAppSettings) -> hvac.Client:
        """Create authenticated Vault client"""
        client = hvac.Client(url=settings.vault_addr)

        if settings.vault_token:
            # Development mode with token
            client.token = settings.vault_token
        elif settings.vault_role_id and settings.vault_secret_id:
            # Production mode with AppRole
            try:
                auth_response = client.auth.approle.login(
                    role_id=settings.vault_role_id, secret_id=settings.vault_secret_id
                )
                client.token = auth_response["auth"]["client_token"]
            except Exception as e:
                raise ValueError("Failed to authenticate with Vault") from e
        else:
            raise ValueError(
                "Either vault_token or vault_role_id/vault_secret_id must be provided"
            )

        if not client.is_authenticated():
            raise ValueError("Failed to authenticate with Vault")

        return client


class MinIOClientFactory:  # pylint: disable=too-few-public-methods
    """Factory for creating MinIO clients"""

    @staticmethod
    def create_client(settings: BaseAppSettings) -> Minio:
        """Create MinIO client"""
        return Minio(
            endpoint=settings.minio_endpoint,
            access_key=settings.minio_access_key,
            secret_key=settings.minio_secret_key,
            secure=settings.minio_secure,
        )


class QdrantClientFactory:  # pylint: disable=too-few-public-methods
    """Factory for creating Qdrant clients"""

    @staticmethod
    def create_client(settings: BaseAppSettings) -> QdrantClient:
        """Create Qdrant client"""
        return QdrantClient(url=settings.qdrant_url, api_key=settings.qdrant_api_key)


class Neo4jDriverFactory:  # pylint: disable=too-few-public-methods
    """Factory for creating Neo4j drivers"""

    @staticmethod
    def create_driver(settings: BaseAppSettings) -> Any:
        """Create Neo4j driver"""
        return GraphDatabase.driver(
            settings.neo4j_uri, auth=(settings.neo4j_user, settings.neo4j_password)
        )


class RedisClientFactory:  # pylint: disable=too-few-public-methods
    """Factory for creating Redis clients"""

    @staticmethod
    async def create_client(settings: BaseAppSettings) -> "redis.Redis[str]":
        """Create Redis client"""
        return redis.from_url(
            settings.redis_url, encoding="utf-8", decode_responses=True
        )


class EventBusFactory:
    """Factory for creating event bus clients"""

    @staticmethod
    def create_kafka_producer(settings: BaseAppSettings) -> AIOKafkaProducer:
        """Create Kafka producer"""
        return AIOKafkaProducer(
            bootstrap_servers=settings.kafka_bootstrap_servers,
            value_serializer=lambda v: v.encode("utf-8") if isinstance(v, str) else v,
        )

    @staticmethod
    def create_kafka_consumer(
        settings: BaseAppSettings, topics: list[str]
    ) -> AIOKafkaConsumer:
        """Create Kafka consumer"""
        return AIOKafkaConsumer(
            *topics,
            bootstrap_servers=settings.kafka_bootstrap_servers,
            value_deserializer=lambda m: m.decode("utf-8") if m else None,
        )

    @staticmethod
    def create_sqs_client(settings: BaseAppSettings) -> Any:
        """Create SQS client"""
        return boto3.client("sqs", region_name=settings.aws_region)

    @staticmethod
    def create_sns_client(settings: BaseAppSettings) -> Any:
        """Create SNS client"""
        return boto3.client("sns", region_name=settings.aws_region)