vital-agent-container-sdk 0.1.0__tar.gz → 0.1.1__tar.gz

vital_agent_container_sdk-0.1.1/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.1
+Name: vital-agent-container-sdk
+Version: 0.1.1
+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
+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.0 → vital_agent_container_sdk-0.1.1}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='vital-agent-container-sdk',
-    version='0.1.0',
+    version='0.1.1',
     author='Marc Hadfield',
     author_email='marc@vital.ai',
     description='Vital Agent Container SDK',
@@ -21,8 +21,24 @@ setup(
     },
     license='Apache License 2.0',
     install_requires=[
-        'vital-ai-vitalsigns>=0.1.19',
-        'vital-ai-aimp>=0.1.6',
+        'vital-ai-vitalsigns==0.1.19',
+        'vital-ai-aimp==0.1.6',
+        'httpx==0.26.0',
+        'python-json-logger==2.0.7',
+        'python-dotenv==1.0.1',
+        'uvicorn[standard]==0.27.0.post1',
+        'fastapi==0.109.2',
+        'dataclasses-json==0.5.7',
+        'aiohttp==3.9.0',
+        'aiosignal==1.2.0',
+        'anyio==4.2.0',
+        'async-timeout==4.0.3',
+        'starlette==0.36.3',
+        'marshmallow==3.19.0',
+        'pyyaml==6.0.1',
+        'requests==2.31.0',
+        'Pillow==10.2.0',
+        'websockets==12.0'
     ],
     extras_require={
         'dev': [

vital_agent_container_sdk-0.1.1/test/test_app.py

@@ -0,0 +1,24 @@
+import os
+import uvicorn
+from vital_agent_container.agent_container_app import AgentContainerApp
+from vital_agent_container.handler.impl.aimp_echo_message_handler import AIMPEchoMessageHandler
+
+
+def main():
+    print('Test Agent Container App')
+
+    current_file_directory = os.path.dirname(os.path.abspath(__file__))
+    parent_directory = os.path.dirname(current_file_directory)
+
+    app_home = parent_directory
+
+    handler = AIMPEchoMessageHandler()
+
+    agent_container_app = AgentContainerApp(handler, app_home)
+
+    uvicorn.run(host="0.0.0.0", port=6006, app=agent_container_app)
+
+
+if __name__ == "__main__":
+
+    main()

vital_agent_container_sdk-0.1.1/vital_agent_container/agent_container_app.py

@@ -0,0 +1,127 @@
+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.add_routes()
+
+    async def process_ws_message(self, client: httpx.AsyncClient, websocket: WebSocket, data: str,
+                                 started_event: asyncio.Event):
+        await AIMPMessageProcessor.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)
+                    message_type = message_obj[0].get("type")
+                    message_intent = message_obj[0].get("http://vital.ai/ontology/vital-aimp#hasIntent")
+                    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(self.handler, client, websocket, data, started_event))
+                        background_tasks.append(task)
+                        await started_event.wait()
+            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_sdk-0.1.1/vital_agent_container/handler/aimp_message_handler_inf.py

@@ -0,0 +1,5 @@
+
+
+class AIMPMessageHandlerInf:
+    pass
+

vital_agent_container_sdk-0.1.1/vital_agent_container/handler/impl/aimp_echo_message_handler.py

@@ -0,0 +1,7 @@
+from vital_agent_container.handler.aimp_message_handler_inf import AIMPMessageHandlerInf
+
+
+class AIMPEchoMessageHandler(AIMPMessageHandlerInf):
+    pass
+
+

vital_agent_container_sdk-0.1.1/vital_agent_container/processor/aimp_message_processor.py

@@ -0,0 +1,9 @@
+
+
+class AIMPMessageProcessor:
+
+    def __init__(self):
+        pass
+
+    async def process_message(self, handler, app_config, client, websocket, data, started_event):
+        pass

vital_agent_container_sdk-0.1.1/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_sdk-0.1.1/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_sdk-0.1.1/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_sdk-0.1.1/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_sdk-0.1.1/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.1/vital_agent_container_sdk.egg-info/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.1
+Name: vital-agent-container-sdk
+Version: 0.1.1
+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
+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.1/vital_agent_container_sdk.egg-info/SOURCES.txt

@@ -0,0 +1,26 @@
+LICENSE
+README.md
+setup.py
+test/test_app.py
+vital_agent_container/__init__.py
+vital_agent_container/agent_container_app.py
+vital_agent_container/config/__init__.py
+vital_agent_container/handler/__init__.py
+vital_agent_container/handler/aimp_message_handler_inf.py
+vital_agent_container/handler/impl/__init__.py
+vital_agent_container/handler/impl/aimp_echo_message_handler.py
+vital_agent_container/processor/__init__.py
+vital_agent_container/processor/aimp_message_processor.py
+vital_agent_container/streaming/__init__.py
+vital_agent_container/streaming/custom_streaming_response.py
+vital_agent_container/tasks/__init__.py
+vital_agent_container/tasks/request_task.py
+vital_agent_container/tasks/task_manager_async_client.py
+vital_agent_container/utils/__init__.py
+vital_agent_container/utils/aws_utils.py
+vital_agent_container/utils/config_utils.py
+vital_agent_container_sdk.egg-info/PKG-INFO
+vital_agent_container_sdk.egg-info/SOURCES.txt
+vital_agent_container_sdk.egg-info/dependency_links.txt
+vital_agent_container_sdk.egg-info/requires.txt
+vital_agent_container_sdk.egg-info/top_level.txt

vital_agent_container_sdk-0.1.1/vital_agent_container_sdk.egg-info/requires.txt

@@ -0,0 +1,20 @@
+vital-ai-vitalsigns==0.1.19
+vital-ai-aimp==0.1.6
+httpx==0.26.0
+python-json-logger==2.0.7
+python-dotenv==1.0.1
+uvicorn[standard]==0.27.0.post1
+fastapi==0.109.2
+dataclasses-json==0.5.7
+aiohttp==3.9.0
+aiosignal==1.2.0
+anyio==4.2.0
+async-timeout==4.0.3
+starlette==0.36.3
+marshmallow==3.19.0
+pyyaml==6.0.1
+requests==2.31.0
+Pillow==10.2.0
+websockets==12.0
+
+[dev]

vital_agent_container_sdk-0.1.0/PKG-INFO

@@ -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/vital_agent_container_sdk.egg-info/PKG-INFO

@@ -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/vital_agent_container_sdk.egg-info/SOURCES.txt

@@ -1,9 +0,0 @@
-LICENSE
-README.md
-setup.py
-vital_agent_container/__init__.py
-vital_agent_container_sdk.egg-info/PKG-INFO
-vital_agent_container_sdk.egg-info/SOURCES.txt
-vital_agent_container_sdk.egg-info/dependency_links.txt
-vital_agent_container_sdk.egg-info/requires.txt
-vital_agent_container_sdk.egg-info/top_level.txt

vital_agent_container_sdk-0.1.0/vital_agent_container_sdk.egg-info/requires.txt

@@ -1,4 +0,0 @@
-vital-ai-vitalsigns>=0.1.19
-vital-ai-aimp>=0.1.6
-
-[dev]