vital-agent-container-sdk 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

vital_agent_container/agent_container_app.py

@@ -0,0 +1,133 @@
+import asyncio
+import gc
+import json
+import sys
+import tracemalloc
+import uvicorn
+import logging
+from logging.handlers import RotatingFileHandler
+from pythonjsonlogger import jsonlogger
+import httpx
+from dotenv import load_dotenv
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from vital_agent_container.tasks.task_manager_async_client import TaskManagerAsyncClient
+from vital_agent_container.utils.aws_utils import AWSUtils
+from vital_agent_container.utils.config_utils import ConfigUtils
+from vital_agent_container.processor.aimp_message_processor import AIMPMessageProcessor
+
+
+logger = logging.getLogger("VitalAgentContainerLogger")
+logger.setLevel(logging.INFO)
+
+formatter = jsonlogger.JsonFormatter('%(asctime)s %(name)s %(levelname)s %(message)s')
+
+stream_handler = logging.StreamHandler(sys.stdout)
+stream_handler.setFormatter(formatter)
+
+log_file = "/var/log/agentcontainer/app.log"
+
+file_handler = RotatingFileHandler(log_file, maxBytes=5 * 1024 * 1024, backupCount=5)
+file_handler.setFormatter(formatter)
+
+logger.addHandler(stream_handler)
+logger.addHandler(file_handler)
+
+service_identifier = AWSUtils.get_task_arn()
+
+
+class AgentContainerApp(FastAPI):
+    def __init__(self, handler, app_home, **kwargs):
+        super().__init__(**kwargs)
+        self.handler = handler
+        self.app_home = app_home
+        load_dotenv()
+        self.config = ConfigUtils.load_config(app_home)
+        self.message_processor = AIMPMessageProcessor()
+        self.add_routes()
+
+    async def process_ws_message(self, client: httpx.AsyncClient, websocket: WebSocket, data: str,
+                                 started_event: asyncio.Event):
+        print(f"process_ws_message: Processing: {data}")
+        await self.message_processor.process_message(self.handler, self.config, client, websocket, data, started_event)
+
+    def add_routes(self):
+        @self.get("/health")
+        async def health_check():
+            logger.info("health check")
+            return {"status": "ok"}
+
+        @self.websocket("/ws")
+        async def websocket_endpoint(websocket: WebSocket):
+            logger.info(f"WebSocket connection from {websocket.client.host}:{websocket.client.port} accepted.")
+            await websocket.accept()
+            client = TaskManagerAsyncClient()
+            client._ws_active = True
+            background_tasks = []
+            try:
+                while True:
+                    data = await websocket.receive_text()
+
+                    logger.info(f"Received message from {websocket.client.host}:{websocket.client.port}: {data}")
+                    message_obj = json.loads(data)
+                    logger.info(f"Received message: {message_obj}")
+                    message_type = message_obj[0].get("type", None)
+                    message_intent = message_obj[0].get("http://vital.ai/ontology/vital-aimp#hasIntent", None)
+                    logger.info(f"message_type: {message_type}")
+                    logger.info(f"message_intent: {message_intent}")
+
+                    if message_intent == "interrupt":
+                        logger.info("Processing interrupted by client.")
+                        client.log_current_tasks()
+                        await client.cancel_current_tasks()
+                        for task in background_tasks:
+                            task.cancel()
+                        try:
+                            await asyncio.gather(*background_tasks, return_exceptions=True)
+                        except Exception as e:
+                            logger.error(f"An error occurred in gather: {e}")
+                        try:
+                            if client.ws_active:
+                                await websocket.close()
+                                client._ws_active = False
+                        except Exception as e:
+                            logger.error(f"An error occurred in interrupt websocket close: {e}")
+                        break
+
+                    if len(background_tasks) > 0:
+                        logger.info("currently processing task, ignoring new request.")
+                        await websocket.send_text("processing task. ignoring message.")
+                    else:
+                        logger.info(f"Processing message: {data}")
+                        started_event = asyncio.Event()
+                        task = asyncio.create_task(self.process_ws_message(client, websocket, data, started_event))
+                        background_tasks.append(task)
+                        await started_event.wait()
+                        logger.info(f"Completed Processing message: {data}")
+                        # break out of infinite loop
+                        break
+            except WebSocketDisconnect:
+                logger.info("WebSocket connection closed by the client.")
+                # still try to close it in finally?
+                # client._ws_active = False
+            except Exception as e:
+                logger.error(f"An error occurred in ws main loop: {e}")
+                try:
+                    await websocket.close()
+                    client._ws_active = False
+                except Exception as e2:
+                    logger.error(f"An error occurred in websocket close: {e2}")
+            finally:
+                try:
+                    await asyncio.gather(*background_tasks)
+                except Exception as e:
+                    logger.error(f"An error occurred in final gather: {e}")
+                try:
+                    if client.ws_active:
+                        await websocket.close()
+                except Exception as e3:
+                    logger.error(f"An error occurred in finally websocket close: {e3}")
+                await client.cancel_current_tasks()
+
+        @self.on_event("shutdown")
+        async def shutdown_event():
+            logger.info("Shutting down...")

vital_agent_container/handler/aimp_message_handler_inf.py

@@ -0,0 +1,14 @@
+import asyncio
+import httpx
+from starlette.websockets import WebSocket
+from abc import ABC, abstractmethod
+
+
+class AIMPMessageHandlerInf(ABC):
+
+    @abstractmethod
+    async def process_message(self, config, client: httpx.AsyncClient, websocket: WebSocket, data: str, started_event: asyncio.Event):
+        pass
+
+
+

vital_agent_container/handler/impl/aimp_echo_message_handler.py

@@ -0,0 +1,23 @@
+import asyncio
+import httpx
+from starlette.websockets import WebSocket
+from vital_agent_container.handler.aimp_message_handler_inf import AIMPMessageHandlerInf
+
+
+class AIMPEchoMessageHandler(AIMPMessageHandlerInf):
+    async def process_message(self, config, client: httpx.AsyncClient, websocket: WebSocket, data: str, started_event: asyncio.Event):
+        try:
+            print(f"Received Message: {data}")
+            await websocket.send_text(data)
+            print(f"Sent Message: {data}")
+            # await websocket.close(1000, "Processing Complete")
+            # print(f"Websocket closed.")
+            started_event.set()
+            print(f"Completed Event.")
+        except asyncio.CancelledError:
+            # log canceling
+            raise
+
+
+
+

vital_agent_container/processor/aimp_message_processor.py

@@ -0,0 +1,18 @@
+import asyncio
+
+from vital_agent_container.handler.aimp_message_handler_inf import AIMPMessageHandlerInf
+
+
+class AIMPMessageProcessor:
+
+    def __init__(self):
+        pass
+
+    async def process_message(self, handler: AIMPMessageHandlerInf, app_config, client, websocket, data, started_event):
+        print(f"Processing: {data}")
+        try:
+            return await handler.process_message(app_config, client, websocket, data, started_event)
+        except asyncio.CancelledError:
+            # log canceling
+            raise
+

vital_agent_container/streaming/custom_streaming_response.py

@@ -0,0 +1,38 @@
+import asyncio
+import logging
+import httpx
+
+
+class CustomStreamingResponse:
+    def __init__(self, response: httpx.Response, end_event: asyncio.Event):
+        self._response = response
+        self._end_event = end_event
+
+    async def aiter_lines(self):
+        try:
+            async for line in self._response.aiter_lines():
+                yield line
+        finally:
+            self.handle_end_of_stream()
+
+    async def aiter_bytes(self):
+        try:
+            async for chunk in self._response.aiter_bytes():
+                yield chunk
+        finally:
+            self.handle_end_of_stream()
+
+    async def aiter_raw(self):
+        try:
+            async for chunk in self._response.aiter_raw():
+                yield chunk
+        finally:
+            self.handle_end_of_stream()
+
+    def handle_end_of_stream(self):
+        self._end_event.set()
+        logger = logging.getLogger("VitalAgentContainerLogger")
+        logger.info("Reached the end of the stream.")
+
+    def __getattr__(self, name):
+        return getattr(self._response, name)

vital_agent_container/tasks/request_task.py

@@ -0,0 +1,21 @@
+import asyncio
+import httpx
+
+
+class RequestTask:
+    def __init__(self):
+        self._request: httpx.Response = None
+        self._event: asyncio.Event = None
+        self._task: asyncio.Task = None
+
+    @property
+    def request(self):
+        return self._request
+
+    @property
+    def task(self):
+        return self._task
+
+    @property
+    def event(self):
+        return self._event

vital_agent_container/tasks/task_manager_async_client.py

@@ -0,0 +1,64 @@
+import asyncio
+import logging
+from typing import Dict
+import httpx
+from vital_agent_container.streaming.custom_streaming_response import CustomStreamingResponse
+from vital_agent_container.tasks.request_task import RequestTask
+
+
+class TaskManagerAsyncClient(httpx.AsyncClient):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._tasks: Dict[RequestTask, str] = {}
+        self._ws_active = None
+
+    @property
+    def ws_active(self):
+        return self._ws_active
+
+    async def send(self, request, *args, **kwargs):
+        is_streaming = kwargs.get("stream", False)
+
+        async def task_coro(rt: RequestTask):
+            response = await super(TaskManagerAsyncClient, self).send(request, *args, **kwargs)
+            if is_streaming:
+                end_event = asyncio.Event()
+                custom_response = CustomStreamingResponse(response, end_event)
+                rt._event = end_event
+                rt._request = custom_response
+                return custom_response
+            else:
+                return response
+
+        request_task = RequestTask()
+        task = asyncio.create_task(task_coro(request_task))
+        request_task._task = task
+        self._tasks[request_task] = request.url.path
+
+        try:
+            return await task
+        finally:
+            # self._tasks.pop(task, None)
+            pass
+
+    def log_current_tasks(self):
+        logger = logging.getLogger("VitalAgentContainerLogger")
+        if self._tasks:
+            logger.info(f"Current tasks: {', '.join(self._tasks.values())}")
+        else:
+            logger.info("No current tasks.")
+
+    async def cancel_current_tasks(self):
+        if self._tasks:
+            for rt in self._tasks.keys():
+                # res = rt.request
+                t = rt.task
+                status = t.cancel()
+                if status:
+                    try:
+                        await t
+                    except asyncio.CancelledError:
+                        pass
+            self._tasks.clear()
+        else:
+            pass

vital_agent_container/utils/aws_utils.py

@@ -0,0 +1,21 @@
+import os
+import requests
+
+
+class AWSUtils:
+
+    @staticmethod
+    def get_task_arn():
+        metadata_uri = os.environ.get('ECS_CONTAINER_METADATA_URI_V4')
+        if not metadata_uri:
+            return "local-instance"
+        task_metadata_url = f"{metadata_uri}/task"
+
+        try:
+            response = requests.get(task_metadata_url)
+            response.raise_for_status()
+            metadata = response.json()
+            task_arn = metadata.get('TaskARN')
+            return task_arn
+        except requests.RequestException as e:
+            return "local-instance"

vital_agent_container/utils/config_utils.py

@@ -0,0 +1,17 @@
+import logging
+import os
+
+import yaml
+
+
+class ConfigUtils:
+
+    @staticmethod
+    def load_config(app_home):
+
+        with open(f"{app_home}/agent_config.yaml", "r") as config_stream:
+            try:
+                return yaml.safe_load(config_stream)
+            except yaml.YAMLError as exc:
+                logger = logging.getLogger("VitalAgentContainerLogger")
+                logger.info("failed to load config file")

vital_agent_container_sdk-0.1.2.dist-info/METADATA

@@ -0,0 +1,35 @@
+Metadata-Version: 2.1
+Name: vital-agent-container-sdk
+Version: 0.1.2
+Summary: Vital Agent Container SDK
+Home-page: https://github.com/vital-ai/vital-agent-container-python
+Author: Marc Hadfield
+Author-email: marc@vital.ai
+License: Apache License 2.0
+Classifier: Programming Language :: Python :: 3.10
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: vital-ai-vitalsigns ==0.1.20
+Requires-Dist: vital-ai-aimp ==0.1.7
+Requires-Dist: httpx ==0.26.0
+Requires-Dist: python-json-logger ==2.0.7
+Requires-Dist: python-dotenv ==1.0.1
+Requires-Dist: uvicorn[standard] ==0.27.0.post1
+Requires-Dist: fastapi ==0.109.2
+Requires-Dist: dataclasses-json ==0.5.7
+Requires-Dist: aiohttp ==3.9.0
+Requires-Dist: aiosignal ==1.2.0
+Requires-Dist: anyio ==4.2.0
+Requires-Dist: async-timeout ==4.0.3
+Requires-Dist: starlette ==0.36.3
+Requires-Dist: marshmallow ==3.19.0
+Requires-Dist: pyyaml ==6.0.1
+Requires-Dist: requests ==2.31.0
+Requires-Dist: Pillow ==10.2.0
+Requires-Dist: websockets ==12.0
+Provides-Extra: dev
+
+# vital-agent-container-python

vital_agent_container_sdk-0.1.2.dist-info/RECORD

@@ -0,0 +1,22 @@
+vital_agent_container/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vital_agent_container/agent_container_app.py,sha256=5lbge8OC8tEkjJ_TsZ9-zkRPPKA_Jr8gwwIFxmo7z6U,5879
+vital_agent_container/config/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vital_agent_container/handler/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vital_agent_container/handler/aimp_message_handler_inf.py,sha256=H9Bey5fAAh13LscsmwlCkK6i-rMKe5iN_vwk279p57s,315
+vital_agent_container/handler/impl/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vital_agent_container/handler/impl/aimp_echo_message_handler.py,sha256=QQl3MbAWONno408wXOXzlubhFRvp9Vjwk5iCiF12Sek,764
+vital_agent_container/processor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vital_agent_container/processor/aimp_message_processor.py,sha256=4aaMEKxpzpXuH-kjLyqGY8t5PkqFif7Dr5ZaFHKEr2o,532
+vital_agent_container/streaming/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vital_agent_container/streaming/custom_streaming_response.py,sha256=KDuhl4Z5opb0pxCymkjFvyBPKvhf64TZatLY-qK0l_w,1059
+vital_agent_container/tasks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vital_agent_container/tasks/request_task.py,sha256=l-cGxSuI0hKk_ZjI5nkMCWtwFYs4WDu2DXVtcWGIXSY,391
+vital_agent_container/tasks/task_manager_async_client.py,sha256=hOkhGYgtdrpOkZygfYhjdENmfJViL7W8HTA2junHrkc,2081
+vital_agent_container/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vital_agent_container/utils/aws_utils.py,sha256=HuwdOs58VU9UDnytx2QlKjJNoPE9k2KIQbKJwgi-9zk,580
+vital_agent_container/utils/config_utils.py,sha256=0iqBXct1O0nMTTBkYSdqIp0WIuVM24fbSpw5uoZQi8k,426
+vital_agent_container_sdk-0.1.2.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+vital_agent_container_sdk-0.1.2.dist-info/METADATA,sha256=UqRcZ3v3Q_jkm_2luMEthBtiSGKzeKVDR265f9uSgrY,1187
+vital_agent_container_sdk-0.1.2.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+vital_agent_container_sdk-0.1.2.dist-info/top_level.txt,sha256=acZrRgjFIj0mWxdLPWONKwMgxnFReQmyjttb8Z76fsc,22
+vital_agent_container_sdk-0.1.2.dist-info/RECORD,,

{vital_agent_container_sdk-0.1.0.dist-info → vital_agent_container_sdk-0.1.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: setuptools (72.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

vital_agent_container_sdk-0.1.0.dist-info/METADATA

@@ -1,19 +0,0 @@
-Metadata-Version: 2.1
-Name: vital-agent-container-sdk
-Version: 0.1.0
-Summary: Vital Agent Container SDK
-Home-page: https://github.com/vital-ai/vital-agent-container-python
-Author: Marc Hadfield
-Author-email: marc@vital.ai
-License: Apache License 2.0
-Classifier: Programming Language :: Python :: 3.10
-Classifier: License :: OSI Approved :: Apache Software License
-Classifier: Operating System :: OS Independent
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: vital-ai-vitalsigns >=0.1.19
-Requires-Dist: vital-ai-aimp >=0.1.6
-Provides-Extra: dev
-
-# vital-agent-container-python

vital_agent_container_sdk-0.1.0.dist-info/RECORD

@@ -1,6 +0,0 @@
-vital_agent_container/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vital_agent_container_sdk-0.1.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-vital_agent_container_sdk-0.1.0.dist-info/METADATA,sha256=hXtZ54QkKEWfPiVd8eBYfx9F2bGF0WpddEw_EPvC1tc,630
-vital_agent_container_sdk-0.1.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-vital_agent_container_sdk-0.1.0.dist-info/top_level.txt,sha256=acZrRgjFIj0mWxdLPWONKwMgxnFReQmyjttb8Z76fsc,22
-vital_agent_container_sdk-0.1.0.dist-info/RECORD,,