agentkernel 0.3.1__tar.gz → 0.3.2__tar.gz

{agentkernel-0.3.1 → agentkernel-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: agentkernel
-Version: 0.3.1
+Version: 0.3.2
 Summary: Agent Kernel - Unified AI Agents Runtime
 Author: Yaala Labs
 Author-email: Yaala Labs <agentkernel@yaalalabs.com>

{agentkernel-0.3.1 → agentkernel-0.3.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "agentkernel"
-version = "0.3.1"
+version = "0.3.2"
 description = "Agent Kernel - Unified AI Agents Runtime"
 readme = "README.md"
 requires-python = ">=3.12,<3.14"

{agentkernel-0.3.1 → agentkernel-0.3.2}/src/agentkernel/deployment/aws/__init__.py

@@ -11,5 +11,6 @@ try:
 except importlib.metadata.PackageNotFoundError:
     __version__ = "0.1.0"
 
+from .core.sqs_handler import SQSHandler
 from .serverless import APIGatewayAuthorizer, Lambda, ResponseHandler, ServerlessAgentRunner
 from .serverless.core import LambdaSQSConsumer

{agentkernel-0.3.1 → agentkernel-0.3.2}/src/agentkernel/deployment/aws/core/sqs_handler.py

@@ -7,9 +7,22 @@ from typing import Any, Dict, Mapping, Optional
 import boto3
 from pydantic import BaseModel, ConfigDict
 
+from ....core.config import AKConfig
+
 
 class SQSHandler:
-    """Shared helper for building and sending SQS messages."""
+    """Shared helper for building and sending SQS messages.
+
+    When used in a Non-Agent Kernel lambda/environment, the following environment variables
+    must be exported:
+    - AK_EXECUTION__QUEUES__INPUT__URL
+    - AK_EXECUTION__QUEUES__OUTPUT__URL
+    """
+
+    _sqs_client = None
+    _config = None
+    _input_queue_url = None
+    _output_queue_url = None
 
     class AttributeDataType(str, Enum):
         STRING = "String"
@@ -33,10 +46,38 @@ class SQSHandler:
         value: Any
         datatype: "SQSHandler.AttributeDataType"
 
-    _sqs_client = None
+    @classmethod
+    def _get_config(cls):
+        """Return a cached AKConfig instance.
+
+        :return: A lazily created AKConfig instance.
+        """
+        if cls._config is None:
+            cls._config = AKConfig.get()
+        return cls._config
+
+    @classmethod
+    def get_input_queue_url(cls):
+        """Return the cached input queue URL from config.
+
+        :return: The input queue URL string.
+        """
+        if cls._input_queue_url is None:
+            cls._input_queue_url = cls._get_config().execution.queues.input.url
+        return cls._input_queue_url
+
+    @classmethod
+    def get_output_queue_url(cls):
+        """Return the cached output queue URL from config.
+
+        :return: The output queue URL string.
+        """
+        if cls._output_queue_url is None:
+            cls._output_queue_url = cls._get_config().execution.queues.output.url
+        return cls._output_queue_url
 
     @classmethod
-    def _get_sqs_client(cls):
+    def get_sqs_client(cls):
         """Return a cached boto3 SQS client.
 
         :return: A lazily created boto3 SQS client instance.
@@ -196,7 +237,7 @@ class SQSHandler:
             message_attributes=message_attributes,
             **extra_kwargs,
         )
-        return cls._get_sqs_client().send_message(QueueUrl=queue_url, **message_kwargs)
+        return cls.get_sqs_client().send_message(QueueUrl=queue_url, **message_kwargs)
 
     @classmethod
     def send_prepared_message(cls, queue_url: str, message_kwargs: Mapping[str, Any]):
@@ -209,7 +250,93 @@ class SQSHandler:
         :param message_kwargs: A mapping of boto3 send_message keyword arguments.
         :return: The boto3 send_message response.
         """
-        return cls._get_sqs_client().send_message(QueueUrl=queue_url, **dict(message_kwargs))
+        return cls.get_sqs_client().send_message(QueueUrl=queue_url, **dict(message_kwargs))
+
+    @classmethod
+    def send_message_to_input_queue(
+        cls,
+        message_group_id: Optional[str] = None,
+        message_deduplication_id: Optional[str] = None,
+        message_body: Optional[Any] = None,
+        request_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        **extra_kwargs: Any,
+    ):
+        """Send a message to the input queue with standard custom attributes.
+
+        This method handles the common pattern of sending messages to the input queue
+        with request_id and user_id as custom message attributes.
+
+        :param message_group_id: The FIFO message group id, if required.
+        :param message_deduplication_id: The FIFO deduplication id, if required.
+        :param message_body: The payload to send.
+        :param request_id: Optional request ID custom attribute.
+        :param user_id: Optional user ID custom attribute.
+        :param extra_kwargs: Additional send_message keyword arguments to include.
+        :return: The boto3 send_message response.
+        """
+        queue_url = cls.get_input_queue_url()
+        if not queue_url:
+            raise ValueError("Input queue URL is not configured in AKConfig")
+
+        # Build custom message attributes
+        message_attributes = []
+        if request_id is not None:
+            message_attributes.append(cls.CustomAttribute(name="request_id", value=request_id, datatype=cls.AttributeDataType.STRING))
+        if user_id is not None:
+            message_attributes.append(cls.CustomAttribute(name="user_id", value=user_id, datatype=cls.AttributeDataType.STRING))
+
+        return cls.send_message(
+            queue_url=queue_url,
+            message_body=message_body,
+            message_group_id=message_group_id,
+            message_deduplication_id=message_deduplication_id,
+            message_attributes=message_attributes if message_attributes else None,
+            **extra_kwargs,
+        )
+
+    @classmethod
+    def send_message_to_output_queue(
+        cls,
+        message_group_id: Optional[str] = None,
+        message_deduplication_id: Optional[str] = None,
+        message_body: Optional[Any] = None,
+        request_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        **extra_kwargs: Any,
+    ):
+        """Send a message to the output queue with standard custom attributes.
+
+        This method handles the common pattern of sending messages to the output queue
+        with request_id and user_id as custom message attributes.
+
+        :param message_group_id: The FIFO message group id, if required.
+        :param message_deduplication_id: The FIFO deduplication id, if required.
+        :param message_body: The payload to send.
+        :param request_id: Optional request ID custom attribute.
+        :param user_id: Optional user ID custom attribute.
+        :param extra_kwargs: Additional send_message keyword arguments to include.
+        :return: The boto3 send_message response.
+        """
+        queue_url = cls.get_output_queue_url()
+        if not queue_url:
+            raise ValueError("Output queue URL is not configured in AKConfig")
+
+        # Build custom message attributes
+        message_attributes = []
+        if request_id is not None:
+            message_attributes.append(cls.CustomAttribute(name="request_id", value=request_id, datatype=cls.AttributeDataType.STRING))
+        if user_id is not None:
+            message_attributes.append(cls.CustomAttribute(name="user_id", value=user_id, datatype=cls.AttributeDataType.STRING))
+
+        return cls.send_message(
+            queue_url=queue_url,
+            message_body=message_body,
+            message_group_id=message_group_id,
+            message_deduplication_id=message_deduplication_id,
+            message_attributes=message_attributes if message_attributes else None,
+            **extra_kwargs,
+        )
 
 
 # Tell Pydantic to resolve the string annotation for CustomAttribute.datatype after the class is fully defined, which allows us to reference AttributeDataType before it's defined in the class body.

{agentkernel-0.3.1 → agentkernel-0.3.2}/src/agentkernel/deployment/aws/serverless/akagentrunner.py

@@ -20,34 +20,29 @@ class ServerlessAgentRunner(LambdaSQSConsumer):
         return cls._chat_service
 
     @classmethod
-    def _get_output_queue_url(cls) -> str:
-        return AKConfig.get().execution.queues.output.url
-
-    @classmethod
-    def _construct_queue_input_message(cls, raw_queue_message: dict, queue_input_message_body: dict) -> dict:
+    def _get_record_attributes(cls, raw_queue_message: dict) -> dict:
         """
-        Construct the SQS message payload for the response queue.
+        Extract attributes from the raw SQS message record.
         :param raw_queue_message: Original SQS message (``dict``) received by the Lambda function.
-        :param queue_input_message_body: JSON-serializable body (``dict``) to be sent to the response queue.
-        :return: Dictionary (``dict``) of parameters for ``boto3`` ``send_message`` (excluding ``QueueUrl``).
+        :return: Dictionary (``dict``) containing extracted attributes.
         """
         attributes = SQSHandler.get_message_system_attributes(raw_queue_message)
         message_attributes = SQSHandler.get_message_custom_attributes(raw_queue_message)
         request_id = message_attributes.get("request_id")
         user_id = message_attributes.get("user_id")
+
         if not request_id:
             raise ValueError("request_id is required")
-        message_attributes = [SQSHandler.CustomAttribute(name="request_id", value=request_id, datatype=SQSHandler.AttributeDataType.STRING)]
-        if user_id is not None:
-            message_attributes.append(SQSHandler.CustomAttribute(name="user_id", value=user_id, datatype=SQSHandler.AttributeDataType.STRING))
-        queue_input_message = SQSHandler.build_send_message_kwargs(
-            message_body=queue_input_message_body,
-            message_group_id=attributes["MessageGroupId"],
-            message_deduplication_id=attributes.get("MessageDeduplicationId"),
-            message_attributes=message_attributes,
-        )
-        cls._log.info(f"Constructed queue input message: {queue_input_message}")
-        return queue_input_message
+
+        record_attributes = {
+            "message_group_id": attributes["MessageGroupId"],
+            "message_deduplication_id": attributes.get("MessageDeduplicationId"),
+            "request_id": request_id,
+            "user_id": user_id,
+        }
+
+        cls._log.info(f"Extracted record attributes: {record_attributes}")
+        return record_attributes
 
     @classmethod
     def _construct_error_message_body(cls, error_msg: str) -> dict:
@@ -59,13 +54,20 @@ class ServerlessAgentRunner(LambdaSQSConsumer):
         return {"error": error_msg}
 
     @classmethod
-    def _send_to_output_queue(cls, queue_input_message: dict) -> None:
+    def _send_to_output_queue(cls, message_body: dict, record_attributes: dict) -> None:
         """
-        Send a prepared message to the configured response SQS queue.
-        :param queue_input_message: Message payload (``dict``) as constructed by
+        Send a prepared message to the configured response SQS queue using send_message_to_output_queue.
+        :param message_body: Message body (``dict``) to be sent to the response queue.
+        :param record_attributes: Extracted attributes (``dict``) from the record.
         :return: None.
         """
-        SQSHandler.send_prepared_message(queue_url=cls._get_output_queue_url(), message_kwargs=queue_input_message)
+        SQSHandler.send_message_to_output_queue(
+            message_group_id=record_attributes["message_group_id"],
+            message_deduplication_id=record_attributes["message_deduplication_id"],
+            message_body=message_body,
+            request_id=record_attributes["request_id"],
+            user_id=record_attributes["user_id"],
+        )
 
     @classmethod
     def _parse_body(cls, record: dict) -> dict:
@@ -90,15 +92,15 @@ class ServerlessAgentRunner(LambdaSQSConsumer):
         if session_id != message_group_id:
             cls._log.info(f"Session ID mismatch: message body session_id {session_id} does not match MessageGroupId {message_group_id}")
             error_message_body = cls._construct_error_message_body(error_msg="Session ID mismatch")
-            queue_input_message = cls._construct_queue_input_message(raw_queue_message=record, queue_input_message_body=error_message_body)
-            cls._send_to_output_queue(queue_input_message=queue_input_message)
-            cls._log.info(f"Sent Session ID Mismatch message to Output Queue: '{cls._get_output_queue_url()}'")
+            record_attributes = cls._get_record_attributes(raw_queue_message=record)
+            cls._send_to_output_queue(message_body=error_message_body, record_attributes=record_attributes)
+            cls._log.info(f"Sent Session ID Mismatch message to Output Queue: '{SQSHandler.get_output_queue_url()}'")
             return
         _, agent_response = cls._get_chat_service().process_chat_request(body=body)
         cls._log.info(f"Chat service response: '{agent_response}'")
-        queue_input_message = cls._construct_queue_input_message(raw_queue_message=record, queue_input_message_body=agent_response)
-        cls._send_to_output_queue(queue_input_message=queue_input_message)
-        cls._log.info(f"Sent Response message to Output Queue: '{cls._get_output_queue_url()}'")
+        record_attributes = cls._get_record_attributes(raw_queue_message=record)
+        cls._send_to_output_queue(message_body=agent_response, record_attributes=record_attributes)
+        cls._log.info(f"Sent Response message to Output Queue: '{SQSHandler.get_output_queue_url()}'")
 
     @classmethod
     def on_permanent_failure(cls, record: dict) -> None:
@@ -110,10 +112,10 @@ class ServerlessAgentRunner(LambdaSQSConsumer):
         cls._log.info(f"Permanent failure: {record}: Retried message {cls.max_receive_count} times. Sending error message to Output Queue`")
         try:
             error_message_body = cls._construct_error_message_body(error_msg=f"Failed to process message. Retried {cls.max_receive_count} times")
-            queue_input_message = cls._construct_queue_input_message(raw_queue_message=record, queue_input_message_body=error_message_body)
-            cls._send_to_output_queue(queue_input_message=queue_input_message)
-            cls._log.info(f"Sent Permanent Failure message to Output Queue: '{cls._get_output_queue_url()}'")
+            record_attributes = cls._get_record_attributes(raw_queue_message=record)
+            cls._send_to_output_queue(message_body=error_message_body, record_attributes=record_attributes)
+            cls._log.info(f"Sent Permanent Failure message to Output Queue: '{SQSHandler.get_output_queue_url()}'")
         except Exception as e:
             # Message comes to this function only if the message has reached its maximum no of retries
             # Catching the error here so that this message will not be returned as batchItemFailures for another retry.
-            cls._log.info(f"Failed sending permanent failure message to Output Queue '{cls._get_output_queue_url()}' due to error: '{str(e)}'")
+            cls._log.info(f"Failed sending permanent failure message to Output Queue '{SQSHandler.get_output_queue_url()}' due to error: '{str(e)}'")

{agentkernel-0.3.1 → agentkernel-0.3.2}/src/agentkernel/deployment/aws/serverless/core/default_endpoints.py

@@ -30,10 +30,6 @@ class DefaultEndpointsHandler:
     def _get_execution_mode(cls):
         return cls._get_config().execution.mode
 
-    @classmethod
-    def _get_input_queue_url(cls):
-        return cls._get_config().execution.queues.input.url
-
     @classmethod
     def _get_default_user_polling_method(cls):
         return "GET" if cls._get_execution_mode() == ExecutionMode.REST_ASYNC else None
@@ -69,7 +65,7 @@ class DefaultEndpointsHandler:
         :return: Dictionary mapping paths → HTTP methods → handler functions
         """
 
-        input_queue_url = cls._get_input_queue_url()
+        input_queue_url = SQSHandler.get_input_queue_url()
         exec_mode = cls._get_execution_mode()
 
         if not input_queue_url:
@@ -158,16 +154,12 @@ class DefaultEndpointsHandler:
         if not session_id:
             raise ValueError("session_id is required")
 
-        message_attributes = [SQSHandler.CustomAttribute(name="request_id", value=payload.request_id, datatype=SQSHandler.AttributeDataType.STRING)]
-        if payload.user_id is not None:
-            message_attributes.append(SQSHandler.CustomAttribute(name="user_id", value=payload.user_id, datatype=SQSHandler.AttributeDataType.STRING))
-
-        response = SQSHandler.send_message(
-            queue_url=cls._get_input_queue_url(),
+        response = SQSHandler.send_message_to_input_queue(
             message_body=request_body,
             message_group_id=session_id,
             message_deduplication_id=payload.request_id,
-            message_attributes=message_attributes,
+            request_id=payload.request_id,
+            user_id=payload.user_id,
         )
         return response