Stanislav Hejny
388 lines
16 KiB
Python
388 lines
16 KiB
Python
import logging
|
|
import re
|
|
from typing import Dict, List, Tuple, Callable, Coroutine, AnyStr, Any
|
|
|
|
import python_multipart
|
|
from aiohttp import ClientSession, ClientResponse
|
|
from fastapi import HTTPException, UploadFile
|
|
from python_multipart.multipart import Field, File
|
|
|
|
from amqp.adapter.amq_response_factory import AMQResponseFactory
|
|
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 core.deps import get_current_user
|
|
from schemas.user_schema import UserSchema
|
|
|
|
import json
|
|
from urllib.parse import parse_qs
|
|
from xml.etree import ElementTree
|
|
import io
|
|
|
|
|
|
class CleverMultiPartParser:
|
|
def __init__(self, message: AMQMessage):
|
|
self.form_data = {}
|
|
# Convert each list of strings to a single string joined by newline, then to bytes
|
|
headers = {key: "\n".join(value).encode("utf-8") for key, value in message.headers.items()}
|
|
python_multipart.parse_form(headers, io.BytesIO(message.body()), self.on_field, self.on_file)
|
|
|
|
def on_field(self, field: Field):
|
|
print(field)
|
|
self.form_data[field.field_name.decode('utf-8')] = field.value
|
|
print(self.form_data)
|
|
|
|
def on_file(self, file: File):
|
|
print(file)
|
|
self.form_data[file.field_name.decode('utf-8')] = \
|
|
UploadFile(file.file_object, size=file.size, filename=file.file_name)
|
|
print(self.form_data)
|
|
|
|
|
|
def parse_request_body(message: AMQMessage):
|
|
"""
|
|
Parses the HTTP POST request body based on the MIME type.
|
|
|
|
Args:
|
|
body (str): The raw request body as a string.
|
|
content_type (str): The Content-Type header of the request.
|
|
|
|
Returns:
|
|
dict: Parsed data as a dictionary.
|
|
"""
|
|
content_type = message.headers.get('Content-Type', '')
|
|
if not content_type:
|
|
raise ValueError("Content-Type header is required")
|
|
|
|
# Parse JSON
|
|
if content_type == "application/json":
|
|
try:
|
|
return json.loads(message.body())
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Invalid JSON: {e}")
|
|
|
|
# Parse URL-encoded form data
|
|
elif content_type == "application/x-www-form-urlencoded":
|
|
try:
|
|
return {k: v[0] if len(v) == 1 else v for k, v in parse_qs(message.body()).items()}
|
|
except Exception as e:
|
|
raise ValueError(f"Invalid URL-encoded form data: {e}")
|
|
|
|
# Parse multipart/form-data
|
|
elif content_type.startswith("multipart/form-data"):
|
|
try:
|
|
parser: CleverMultiPartParser = CleverMultiPartParser(message)
|
|
return parser.form_data
|
|
except Exception as e:
|
|
raise ValueError(f"Invalid multipart/form-data: {e}")
|
|
|
|
# Parse plain text
|
|
elif content_type == "text/plain":
|
|
return {"text": message.body()}
|
|
|
|
# Parse XML
|
|
elif content_type == "application/xml":
|
|
try:
|
|
root = ElementTree.fromstring(message.body())
|
|
return {elem.tag: elem.text for elem in root}
|
|
except ElementTree.ParseError as e:
|
|
raise ValueError(f"Invalid XML: {e}")
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported Content-Type: {content_type}")
|
|
|
|
|
|
BENCHMARK = '/benchmark'
|
|
JOBS = '/jobs/'
|
|
UNSTRUCTURED = '/unstructured'
|
|
|
|
GET_BENCHMARK_KG = r".*/benchmark/(?P<request_id>[^\?]+)\?index=(?P<index>\d+)"
|
|
GET_JOB_STATUS = r".*/jobs/(?P<request_id>.*)"
|
|
GET_KG = r".*/unstructured/(?P<request_id>.*)"
|
|
DUMMY = r"(?P<request_id>.*)"
|
|
|
|
|
|
# ============= POST Helper function =============
|
|
|
|
async def call_cleverswarm_post_put_api(
|
|
message: AMQMessage,
|
|
args: Any,
|
|
api_method: Callable[[Any, UserSchema], Coroutine],
|
|
success_code: int,
|
|
authenticated_user: UserSchema
|
|
) -> AMQResponse:
|
|
try:
|
|
coroutine_call: Coroutine = api_method(args, authenticated_user)
|
|
result: str = await coroutine_call
|
|
return AMQResponseFactory.of(
|
|
message.id, success_code, message.headers.get("Content-Type", "application/json"),
|
|
result.encode('utf-8'), message.trace_info
|
|
)
|
|
except HTTPException as http_error:
|
|
return AMQResponseFactory.of(message.id, http_error.status_code, "text/plain",
|
|
str(http_error.detail).encode('utf-8'), message.trace_info)
|
|
except Exception as error:
|
|
return AMQResponseFactory.of(message.id, 500, "text/plain", str(error).encode('utf-8'), message.trace_info)
|
|
|
|
|
|
#
|
|
# 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_ontology(**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))
|
|
|
|
|
|
# ============= GET Helper function =============
|
|
async def call_cleverswarm_get_api(
|
|
message: AMQMessage,
|
|
pattern: str,
|
|
api_method: Callable[[Tuple[AnyStr | Any, ...], UserSchema], Coroutine],
|
|
authenticated_user: UserSchema
|
|
) -> AMQResponse:
|
|
try:
|
|
match = re.search(pattern, message.path)
|
|
if match:
|
|
crt: Coroutine = api_method(match.groups(), authenticated_user)
|
|
result: str = await crt
|
|
return AMQResponseFactory.of(message.id, 200, "text/plain", result.encode('utf-8'), message.trace_info)
|
|
else:
|
|
_msg = f"Unexpected URL: {message.path}, expected pattern: {pattern}"
|
|
return AMQResponseFactory.of(message.id, 404, "text/plain", _msg.encode('utf-8'), message.trace_info)
|
|
except HTTPException as http_error:
|
|
return AMQResponseFactory.of(message.id, http_error.status_code, "text/plain",
|
|
str(http_error.detail).encode('utf-8'), message.trace_info)
|
|
except Exception as error:
|
|
return AMQResponseFactory.of(message.id, 500, "text/plain", str(error).encode('utf-8'), message.trace_info)
|
|
|
|
|
|
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_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 CleverSwarmAdapter:
|
|
def __init__(self, amq_configuration: AMQConfiguration, session: ClientSession = None):
|
|
self.amq_configuration = amq_configuration
|
|
self.session = session
|
|
self.service_host = self.amq_configuration.dispatch.service_host
|
|
self.service_port = self.amq_configuration.dispatch.service_port
|
|
|
|
async def on_message(self, message) -> AMQResponse:
|
|
DEBUG_NO_AUTH = 1
|
|
_auth: str = message.headers.get('Authorization', '')
|
|
_auth_user: UserSchema | None = None
|
|
if 'Bearer' in _auth:
|
|
try:
|
|
token = _auth.split(' ')[1]
|
|
_auth_user = await get_current_user(token)
|
|
except Exception as error:
|
|
logging.error(f"Error getting user: {error}")
|
|
if _auth_user or DEBUG_NO_AUTH == 1:
|
|
method = message.method.lower()
|
|
if method == 'get':
|
|
return await self.get(message, _auth_user)
|
|
elif method == 'put':
|
|
return await self.put(message, _auth_user)
|
|
elif method == 'post':
|
|
return await self.post(message, _auth_user)
|
|
elif method == 'head':
|
|
return await self.head(message, _auth_user)
|
|
elif method == 'delete':
|
|
return await self.delete(message, _auth_user)
|
|
else:
|
|
raise ValueError(f"Unexpected HTTP method: {message.method}")
|
|
else:
|
|
return AMQResponseFactory.of(message.id, 401, "text/plain", b"Unauthorized", message.trace_info)
|
|
|
|
# ==================================================================================================
|
|
# ===================== 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_cleverswarm_get_api(
|
|
message,
|
|
GET_BENCHMARK_KG,
|
|
cleverswarm_get_benchmark_kg,
|
|
authenticated_user=auth_user
|
|
)
|
|
|
|
if JOBS in message.path:
|
|
return await call_cleverswarm_get_api(
|
|
message,
|
|
GET_JOB_STATUS,
|
|
cleverswarm_get_job_status,
|
|
authenticated_user=auth_user
|
|
)
|
|
|
|
if message.path.endswith('/jobs'):
|
|
return await call_cleverswarm_get_api(
|
|
message,
|
|
DUMMY,
|
|
cleverswarm_jobs_get_requests,
|
|
authenticated_user=auth_user
|
|
)
|
|
|
|
if UNSTRUCTURED in message.path:
|
|
return await call_cleverswarm_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("Destination location 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:
|
|
return await call_cleverswarm_post_put_api(
|
|
message,
|
|
parser.form_data,
|
|
cleverswarm_create_benchmark_job,
|
|
201,
|
|
authenticated_user=auth_user
|
|
)
|
|
|
|
if JOBS in message.path:
|
|
return await call_cleverswarm_get_api(
|
|
message,
|
|
GET_JOB_STATUS,
|
|
cleverswarm_retry_job,
|
|
authenticated_user=auth_user
|
|
)
|
|
|
|
return await self.handle_possible_form(
|
|
self.session.put(
|
|
f"http://{self.service_host}:{self.service_port}{message.path}",
|
|
data=message.body,
|
|
headers={**message.headers, **message.trace_info},
|
|
),
|
|
auth_user
|
|
)
|
|
|
|
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
|
|
"""
|
|
parser: CleverMultiPartParser = CleverMultiPartParser(message)
|
|
if BENCHMARK in message.path:
|
|
return await call_cleverswarm_post_put_api(
|
|
message,
|
|
parser.form_data,
|
|
cleverswarm_create_benchmark_job,
|
|
201,
|
|
authenticated_user=auth_user
|
|
)
|
|
if UNSTRUCTURED in message.path:
|
|
if 'with_ontology' in message.path:
|
|
return await call_cleverswarm_post_put_api(
|
|
message,
|
|
parser.form_data,
|
|
cleverswarm_create_extract_with_ontology,
|
|
200,
|
|
authenticated_user=auth_user
|
|
)
|
|
if 'with_wildcards' in message.path:
|
|
return await call_cleverswarm_post_put_api(
|
|
message,
|
|
parser.form_data,
|
|
cleverswarm_create_extract_with_wildcards,
|
|
200,
|
|
authenticated_user=auth_user
|
|
)
|
|
|
|
return await self.handle_possible_form(
|
|
self.session.post(
|
|
f"http://{self.service_host}:{self.service_port}{message.path}",
|
|
data=message.body,
|
|
headers={**message.headers, **message.trace_info},
|
|
),
|
|
auth_user
|
|
)
|
|
|
|
async def handle_possible_form(self, 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:
|
|
parser: CleverMultiPartParser = CleverMultiPartParser(message)
|
|
if JOBS in message.path:
|
|
return await call_cleverswarm_get_api(
|
|
message,
|
|
GET_JOB_STATUS,
|
|
cleverswarm_retry_job,
|
|
authenticated_user=auth_user
|
|
)
|
|
|
|
raise NotImplementedError
|
|
|
|
def get_content_type(self, message_headers: Dict[str, List[str]]) -> str:
|
|
for header, values in message_headers.items():
|
|
if header.lower() == 'content-type':
|
|
return values[0].split(';')[0]
|
|
return 'application/json'
|