Add support for file streaming upload, update to async pika lib

2025-03-26 13:54:30 +00:00
parent e8f4abdfaf
commit a2ddce1a2e
22 changed files with 1369 additions and 395 deletions
@@ -0,0 +1,287 @@
+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