amq-adapter-python/cleverswarm/clever_swarm_adapter.py

import logging
import pathlib
import re
from typing import Tuple, AnyStr, Any
from aiohttp import ClientSession, ClientResponse
from starlette.datastructures import UploadFile, Headers
import jwt

from datetime import datetime, timedelta

import core.deps
from amqp.adapter.amq_response_factory import AMQResponseFactory
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, call_cleverthis_get_api, \
    call_cleverthis_post_put_api, parse_request_body
from amqp.service.multipart_parser import CleverMultiPartParser
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQMessage, AMQResponse

import rest.v0.endpoints.benchmark as benchmark
import rest.v0.endpoints.jobs as jobs
import rest.v0.endpoints.unstructured as unstructured
from amqp.service.amq_service import AMQService
from core.configs import settings
from schemas.user_schema import UserSchema


BENCHMARK = '/benchmark'
JOBS = '/jobs/'
JOB_DETAILS = '/jobs/details/'
UNSTRUCTURED = '/unstructured'

GET_BENCHMARK_KG = r".*/benchmark/(?P<job_id>[^\?]+)\?index=(?P<index>\d+)"
GET_JOB_STATUS = r".*/jobs/(?P<job_id>.*)"
GET_KG = r".*/unstructured/(?P<job_id>.*)"
PUT_BENCHMARK_KG = r".*/benchmark/(?P<job_id>.*)"
DUMMY = r"(?P<job_id>.*)"


#
# Following methods are wrappers to actual CleverSwarm endpoints, so I can refer to the CleverSwarm method and
# pass it as input parameter using general Callable[[Any, UserSchema], str] type
#
async def cleverswarm_create_benchmark_job(args: Any, authenticated_user: UserSchema) -> str:
    return str(await benchmark.create_benchmark_job(**args, authenticated_user=authenticated_user))


async def cleverswarm_create_extract_with_ontology(args: Any, authenticated_user: UserSchema) -> str:
    return str(await unstructured.create_extract_with_ontology(**args, authenticated_user=authenticated_user))


async def cleverswarm_create_extract_with_wildcards(args: Any, authenticated_user: UserSchema) -> str:
    return str(await unstructured.create_extract_with_wildcards(**args, authenticated_user=authenticated_user))


async def cleverswarm_update_benchmark_job(args: Any, authenticated_user: UserSchema) -> str:
    return str(await benchmark.update_benchmark_job(**args, authenticated_user=authenticated_user))


async def cleverswarm_retry_job(args: Any, authenticated_user: UserSchema) -> str:
    return str(await jobs.retry_job(args[0], authenticated_user=authenticated_user))


async def cleverswarm_delete_job(args: Any, authenticated_user: UserSchema) -> str:
    return str(await jobs.delete_job(args[0], authenticated_user=authenticated_user))


async def cleverswarm_get_benchmark_kg(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
    return str(await benchmark.get_benchmark_kg(args[0], args[1], authenticated_user=authenticated_user))


async def cleverswarm_get_job_status(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
    return str(await jobs.get_job_status(args[0], authenticated_user=authenticated_user))

async def cleverswarm_get_job_details(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
    return str(await jobs.get_job_details(args[0], authenticated_user=authenticated_user))


async def cleverswarm_jobs_get_requests(args: Tuple[AnyStr | Any, ...],
                                        authenticated_user: UserSchema) -> str:
    return str(await jobs.get_requests(authenticated_user=authenticated_user))


async def cleverswarm_get_kg(args: Tuple[AnyStr | Any, ...], authenticated_user: UserSchema) -> str:
    return str(await unstructured.get_kg(args[0], authenticated_user=authenticated_user))


class CleverSwarmAmqpAdapter(CleverThisServiceAdapter):
    def __init__(self, amq_configuration: AMQConfiguration, session: ClientSession = None):
        super().__init__(amq_configuration, session)
        amq_service: AMQService = AMQService(amq_configuration,self)
        amq_service.rabbit_mq_consumer.channel.start_consuming()

    async def get_current_user(self, token: str) -> UserSchema:
        if token == 'DEBUG_NO_AUTH':
            # Payload with username and expiration
            payload = {
                "sub": "0",  # Standard claim for subject (user)
                "username": "user",  # Custom claim (optional)
                "iat": datetime.utcnow(),  # Issued at time
                "exp": datetime.utcnow() + timedelta(days=365),  # Expires in 1 year
            }
            # Generate the JWT token
            token = jwt.encode(payload, settings.JWT_SECRET, algorithm=settings.ALGORITHM)

        return await core.deps.get_current_user(token)
    # ==================================================================================================
    # =====================  C l e v e r S w a r m   A P I   m e t h o d s  ============================
    # ==================================================================================================

    async def get(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
        """
        Handles the GET requests for the CleverSwarm API.
        :param message: message from the AMQP that contains the GET request
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """
        if BENCHMARK in message.path:
            return await call_cleverthis_get_api(
                message,
                GET_BENCHMARK_KG,
                cleverswarm_get_benchmark_kg,
                authenticated_user=auth_user
            )

        if JOBS in message.path:
            if JOB_DETAILS in message.path:
                return await call_cleverthis_get_api(
                    message,
                    DUMMY,
                    cleverswarm_get_job_status,
                    authenticated_user=auth_user
                )
            return await call_cleverthis_get_api(
                message,
                GET_JOB_STATUS,
                cleverswarm_get_job_status,
                authenticated_user=auth_user
            )

        if message.path.endswith('/jobs'):
            return await call_cleverthis_get_api(
                message,
                DUMMY,
                cleverswarm_jobs_get_requests,
                authenticated_user=auth_user
            )

        if UNSTRUCTURED in message.path:
            return await call_cleverthis_get_api(
                message,
                GET_KG,
                cleverswarm_get_kg,
                authenticated_user=auth_user
            )

        # unmatched path - try to invoke the service directly on this path
        url = f"http://{self.service_host}:{self.service_port}{message.path}"
        headers = {**message.headers, **message.trace_info}
        logging.info(f" [*] SEND HEADERS out: {', '.join(f'{k}={v}' for k, v in headers.items())}, REQ={url}")
        if self.session:
            _http_resp: ClientResponse = await self.session.get(url, headers=headers)
            return AMQResponseFactory.of(message.id, _http_resp.status, _http_resp.content_type,
                                         await _http_resp.read(),
                                         message.trace_info)
        return AMQResponseFactory.of(message.id, 404, "text/plain",
                                                str(f"Destination location [{url}] does not exists").encode('utf-8'),
                                                message.trace_info)

    async def put(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
        """
        Handles the PUT requests for the CleverSwarm API.
        :param message: message from the AMQP that contains the PUT request
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """
        parser: CleverMultiPartParser = CleverMultiPartParser(message)
        if BENCHMARK in message.path:
            match = re.search(PUT_BENCHMARK_KG, message.path)
            if match:
                parser.form_data['job_id'] = match.group('job_id')
            return await call_cleverthis_post_put_api(
                message,
                parser.form_data,
                cleverswarm_update_benchmark_job,
                201,
                authenticated_user=auth_user
            )

        if JOBS in message.path:
            return await call_cleverthis_get_api(
                message,
                GET_JOB_STATUS,
                cleverswarm_retry_job,
                authenticated_user=auth_user
            )
        url = f"http://{self.service_host}:{self.service_port}{message.path}"
        return await self.handle_possible_form(
            self.session.put(
                url,
                data=message.body,
                headers={**message.headers, **message.trace_info},
            ),
            auth_user
        ) if self.session else AMQResponseFactory.of(message.id, 404, "text/plain",
                                                str(f"Destination location [{url}] does not exists").encode('utf-8'),
                                                message.trace_info)

    async def post(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
        """
        Handles the POST requests for the CleverSwarm API.
        :param message: message from the AMQP that contains the POST request, including body and headers
        :param auth_user: user authenticated by the API Gateway, as UserSchem
        :return: AMQResponse to be sent back to the API Gateway via AMQP
        """
        _args = parse_request_body(message)
        args = {}
        if isinstance(_args, dict):
            for key, file in _args['files'].items():
                if len(file) > 0:
                    filepath = message.extra_data[key]
                    ufile = UploadFile(
                        file=pathlib.Path(filepath).open("rb"),
                        size=file[0]['size'],
                        filename=file[0]['filename'],
                        headers=Headers({"Content-Type": file[0]['mediaType']})
                    )
                    args[key] = ufile
            for key, obj in _args['form'].items():
                if key not in args:
                    args[key] = obj[0] if isinstance(obj, list) else obj

        if BENCHMARK in message.path:
            return await call_cleverthis_post_put_api(
                message,
                args,
                cleverswarm_create_benchmark_job,
                201,
                authenticated_user=auth_user
            )
        if UNSTRUCTURED in message.path:
            if 'with_ontology' in message.path:
                return await call_cleverthis_post_put_api(
                    message,
                    args,
                    cleverswarm_create_extract_with_ontology,
                    200,
                    authenticated_user=auth_user
                )
            if 'with_wildcards' in message.path:
                return await call_cleverthis_post_put_api(
                    message,
                    args,
                    cleverswarm_create_extract_with_wildcards,
                    200,
                    authenticated_user=auth_user
                )
        url = f"http://{self.service_host}:{self.service_port}{message.path}"
        return await self.handle_possible_form(
            self.session.post(
                url,
                data=message.body,
                headers={**message.headers, **message.trace_info},
            ),
            auth_user
        ) if self.session else AMQResponseFactory.of(message.id, 404, "text/plain",
                                                str(f"Destination location [{url}] does not exists").encode('utf-8'),
                                                message.trace_info)

    async def handle_possible_form(self,
                                   message: AMQMessage,
                                   request_coroutine,
                                   auth_user: UserSchema) -> AMQResponse:
        return await request_coroutine()

    async def head(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
        raise NotImplementedError

    async def delete(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
        if JOBS in message.path:
            return await call_cleverthis_get_api(
                message,
                GET_JOB_STATUS,
                cleverswarm_retry_job,
                authenticated_user=auth_user
            )

        raise NotImplementedError