flamepy 0.3.0__py3-none-any.whl

flamepy/service.py

@@ -0,0 +1,225 @@
+"""
+Copyright 2025 The Flame Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import asyncio
+import os
+import grpc
+from abc import ABC, abstractmethod
+from typing import Optional, Dict, Any, Union
+from dataclasses import dataclass
+import logging
+
+from .types import Shim, FlameError, FlameErrorCode
+from .shim_pb2_grpc import GrpcShimServicer, add_GrpcShimServicer_to_server
+from .types_pb2 import Result, EmptyRequest
+from .shim_pb2 import TaskOutput as TaskOutputProto
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class ApplicationContext:
+    """Context for an application."""
+    name: str
+    shim: Shim
+    url: Optional[str] = None
+    command: Optional[str] = None
+
+
+@dataclass
+class SessionContext:
+    """Context for a session."""
+    session_id: str
+    application: ApplicationContext
+    common_data: Optional[bytes] = None
+
+
+@dataclass
+class TaskContext:
+    """Context for a task."""
+    task_id: str
+    session_id: str
+    input: Optional[bytes] = None
+
+
+@dataclass
+class TaskOutput:
+    """Output from a task."""
+    data: Optional[bytes] = None
+
+
+class FlameService:
+    """Base class for implementing Flame services."""
+    
+    @abstractmethod
+    async def on_session_enter(self, context: SessionContext):
+        """
+        Called when entering a session.
+        
+        Args:
+            context: Session context information
+            
+        Returns:
+            True if successful, False otherwise
+        """
+        pass
+    
+    @abstractmethod
+    async def on_task_invoke(self, context: TaskContext) -> TaskOutput:
+        """
+        Called when a task is invoked.
+        
+        Args:
+            context: Task context information
+            
+        Returns:
+            Task output
+        """
+        pass
+    
+    @abstractmethod
+    async def on_session_leave(self):
+        """
+        Called when leaving a session.
+        
+        Returns:
+            True if successful, False otherwise
+        """
+        pass
+
+
+class GrpcShimServicer(GrpcShimServicer):
+    """gRPC servicer implementation for GrpcShim service."""
+    
+    def __init__(self, service: FlameService):
+        self._service = service
+    
+    async def OnSessionEnter(self, request, context):
+        """Handle OnSessionEnter RPC call."""
+        try:
+            # Convert protobuf request to SessionContext
+            app_context = ApplicationContext(
+                name=request.application.name,
+                shim=Shim(request.application.shim),
+                url=request.application.url,
+                command=request.application.command
+            )
+            
+            session_context = SessionContext(
+                session_id=request.session_id,
+                application=app_context,
+                common_data=request.common_data
+            )
+            
+            # Call the service implementation
+            await self._service.on_session_enter(session_context)
+            
+            # Return result
+            return Result(
+                return_code=0,
+            )
+            
+        except Exception as e:
+            return Result(
+                return_code=1,
+                message=f"Session enter error: {str(e)}"
+            )
+    
+    async def OnTaskInvoke(self, request, context):
+        """Handle OnTaskInvoke RPC call."""
+        try:
+            # Convert protobuf request to TaskContext
+            task_context = TaskContext(
+                task_id=request.task_id,
+                session_id=request.session_id,
+                input=request.input
+            )
+            
+            # Call the service implementation
+            output = await self._service.on_task_invoke(task_context)
+            
+            # Return task output
+            return TaskOutputProto(data=output.data)
+            
+        except Exception as e:
+            return TaskOutputProto(data=None)
+    
+    async def OnSessionLeave(self, request, context):
+        """Handle OnSessionLeave RPC call."""
+        try:
+            # Call the service implementation
+            await self._service.on_session_leave()
+            
+            # Return result
+            return Result(
+                return_code=0,
+            )
+            
+        except Exception as e:
+            return Result(
+                return_code=1,
+                message=f"Session leave error: {str(e)}"
+            )
+
+class GrpcShimServer:
+    """Server for gRPC shim services."""
+    
+    def __init__(self, service: FlameService):
+        self._service = service
+        self._server = None
+    
+    async def start(self):
+        """Start the gRPC server."""
+        try:
+            # Create gRPC server
+            self._server = grpc.aio.server()
+            
+            # Add servicer to server
+            shim_servicer = GrpcShimServicer(self._service)
+            add_GrpcShimServicer_to_server(shim_servicer, self._server)
+            
+            # Listen on Unix socket
+            socket_path = f"/tmp/flame/shim/fsi.sock"
+            self._server.add_insecure_port(f"unix://{socket_path}")
+            
+            # Start server
+            await self._server.start()
+            
+            print(f"Flame Python service started on Unix socket: {socket_path}")
+
+            # Keep server running
+            await self._server.wait_for_termination()
+            
+        except Exception as e:
+            raise FlameError(
+                FlameErrorCode.INTERNAL,
+                f"Failed to start gRPC server: {str(e)}"
+            )
+    
+    async def stop(self):
+        """Stop the gRPC server."""
+        if self._server:
+            await self._server.stop(grace=5)
+            print("gRPC shim server stopped")
+
+
+async def run(service: FlameService):
+    """
+    Run a gRPC shim server.
+    
+    Args:
+        service: The shim service implementation
+    """
+
+    server = GrpcShimServer(service)
+    await server.start()
+

flamepy/shim_pb2.py

@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: shim.proto
+# Protobuf Python Version: 6.31.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    6,
+    31,
+    1,
+    '',
+    'shim.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+import flamepy.types_pb2 as types__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\nshim.proto\x12\x05\x66lame\x1a\x0btypes.proto\"}\n\x12\x41pplicationContext\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x19\n\x04shim\x18\x02 \x01(\x0e\x32\x0b.flame.Shim\x12\x12\n\x05image\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x14\n\x07\x63ommand\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\x08\n\x06_imageB\n\n\x08_command\"~\n\x0eSessionContext\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12.\n\x0b\x61pplication\x18\x02 \x01(\x0b\x32\x19.flame.ApplicationContext\x12\x18\n\x0b\x63ommon_data\x18\x03 \x01(\x0cH\x00\x88\x01\x01\x42\x0e\n\x0c_common_data\"P\n\x0bTaskContext\x12\x0f\n\x07task_id\x18\x01 \x01(\t\x12\x12\n\nsession_id\x18\x02 \x01(\t\x12\x12\n\x05input\x18\x04 \x01(\x0cH\x00\x88\x01\x01\x42\x08\n\x06_input\"(\n\nTaskOutput\x12\x11\n\x04\x64\x61ta\x18\x01 \x01(\x0cH\x00\x88\x01\x01\x42\x07\n\x05_data2\xb5\x01\n\x08GrpcShim\x12\x38\n\x0eOnSessionEnter\x12\x15.flame.SessionContext\x1a\r.flame.Result\"\x00\x12\x37\n\x0cOnTaskInvoke\x12\x12.flame.TaskContext\x1a\x11.flame.TaskOutput\"\x00\x12\x36\n\x0eOnSessionLeave\x12\x13.flame.EmptyRequest\x1a\r.flame.Result\"\x00\x42&Z$github.com/flame-sh/flame/sdk/go/rpcb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'shim_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'Z$github.com/flame-sh/flame/sdk/go/rpc'
+  _globals['_APPLICATIONCONTEXT']._serialized_start=34
+  _globals['_APPLICATIONCONTEXT']._serialized_end=159
+  _globals['_SESSIONCONTEXT']._serialized_start=161
+  _globals['_SESSIONCONTEXT']._serialized_end=287
+  _globals['_TASKCONTEXT']._serialized_start=289
+  _globals['_TASKCONTEXT']._serialized_end=369
+  _globals['_TASKOUTPUT']._serialized_start=371
+  _globals['_TASKOUTPUT']._serialized_end=411
+  _globals['_GRPCSHIM']._serialized_start=414
+  _globals['_GRPCSHIM']._serialized_end=595
+# @@protoc_insertion_point(module_scope)

flamepy/shim_pb2_grpc.py

@@ -0,0 +1,184 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+import flamepy.shim_pb2 as shim__pb2
+import flamepy.types_pb2 as types__pb2
+
+GRPC_GENERATED_VERSION = '1.74.0'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + f' but the generated code in shim_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )
+
+
+class GrpcShimStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.OnSessionEnter = channel.unary_unary(
+                '/flame.GrpcShim/OnSessionEnter',
+                request_serializer=shim__pb2.SessionContext.SerializeToString,
+                response_deserializer=types__pb2.Result.FromString,
+                _registered_method=True)
+        self.OnTaskInvoke = channel.unary_unary(
+                '/flame.GrpcShim/OnTaskInvoke',
+                request_serializer=shim__pb2.TaskContext.SerializeToString,
+                response_deserializer=shim__pb2.TaskOutput.FromString,
+                _registered_method=True)
+        self.OnSessionLeave = channel.unary_unary(
+                '/flame.GrpcShim/OnSessionLeave',
+                request_serializer=types__pb2.EmptyRequest.SerializeToString,
+                response_deserializer=types__pb2.Result.FromString,
+                _registered_method=True)
+
+
+class GrpcShimServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def OnSessionEnter(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def OnTaskInvoke(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def OnSessionLeave(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_GrpcShimServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'OnSessionEnter': grpc.unary_unary_rpc_method_handler(
+                    servicer.OnSessionEnter,
+                    request_deserializer=shim__pb2.SessionContext.FromString,
+                    response_serializer=types__pb2.Result.SerializeToString,
+            ),
+            'OnTaskInvoke': grpc.unary_unary_rpc_method_handler(
+                    servicer.OnTaskInvoke,
+                    request_deserializer=shim__pb2.TaskContext.FromString,
+                    response_serializer=shim__pb2.TaskOutput.SerializeToString,
+            ),
+            'OnSessionLeave': grpc.unary_unary_rpc_method_handler(
+                    servicer.OnSessionLeave,
+                    request_deserializer=types__pb2.EmptyRequest.FromString,
+                    response_serializer=types__pb2.Result.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'flame.GrpcShim', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('flame.GrpcShim', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class GrpcShim(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def OnSessionEnter(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/flame.GrpcShim/OnSessionEnter',
+            shim__pb2.SessionContext.SerializeToString,
+            types__pb2.Result.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def OnTaskInvoke(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/flame.GrpcShim/OnTaskInvoke',
+            shim__pb2.TaskContext.SerializeToString,
+            shim__pb2.TaskOutput.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def OnSessionLeave(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/flame.GrpcShim/OnSessionLeave',
+            types__pb2.EmptyRequest.SerializeToString,
+            types__pb2.Result.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

flamepy/test_flame.py

@@ -0,0 +1,195 @@
+#!/usr/bin/env python3
+"""
+Tests for the Flame Python SDK.
+"""
+
+import pytest
+import asyncio
+from unittest.mock import Mock, AsyncMock
+from datetime import datetime
+
+from flamepy import (
+    Connection, SessionAttributes, ApplicationAttributes,
+    Session, Task, Application, Shim, SessionState, TaskState,
+    ApplicationState, FlameError, FlameErrorCode, TaskInformer
+)
+
+
+class TestTypes:
+    """Test type definitions."""
+    
+    def test_session_attributes(self):
+        """Test SessionAttributes."""
+        attrs = SessionAttributes(
+            application="test-app",
+            slots=2,
+            common_data=b"test data"
+        )
+        assert attrs.application == "test-app"
+        assert attrs.slots == 2
+        assert attrs.common_data == b"test data"
+    
+    def test_application_attributes(self):
+        """Test ApplicationAttributes."""
+        attrs = ApplicationAttributes(
+            name="test-app",
+            shim=Shim.SHELL,
+            command="python",
+            arguments=["script.py"],
+            environments=["PATH=/usr/bin"],
+            working_directory="/tmp"
+        )
+        assert attrs.name == "test-app"
+        assert attrs.shim == Shim.SHELL
+        assert attrs.command == "python"
+        assert attrs.arguments == ["script.py"]
+        assert attrs.environments == ["PATH=/usr/bin"]
+        assert attrs.working_directory == "/tmp"
+    
+    def test_session(self):
+        """Test Session."""
+        session = Session(
+            id="session-1",
+            application="test-app",
+            slots=2,
+            state=SessionState.OPEN,
+            creation_time=datetime.now(),
+            pending=1,
+            running=1,
+            succeed=0,
+            failed=0
+        )
+        assert session.id == "session-1"
+        assert session.application == "test-app"
+        assert session.slots == 2
+        assert session.state == SessionState.OPEN
+    
+    def test_task(self):
+        """Test Task."""
+        task = Task(
+            id="task-1",
+            session_id="session-1",
+            state=TaskState.PENDING,
+            creation_time=datetime.now(),
+            input=b"test input"
+        )
+        assert task.id == "task-1"
+        assert task.session_id == "session-1"
+        assert task.state == TaskState.PENDING
+        assert task.input == b"test input"
+        assert not task.is_completed()
+        
+        # Test completed task
+        completed_task = Task(
+            id="task-2",
+            session_id="session-1",
+            state=TaskState.SUCCEED,
+            creation_time=datetime.now()
+        )
+        assert completed_task.is_completed()
+    
+    def test_application(self):
+        """Test Application."""
+        app = Application(
+            id="app-1",
+            name="test-app",
+            shim=Shim.SHELL,
+            state=ApplicationState.ENABLED,
+            creation_time=datetime.now(),
+            command="python",
+            arguments=["script.py"]
+        )
+        assert app.id == "app-1"
+        assert app.name == "test-app"
+        assert app.shim == Shim.SHELL
+        assert app.state == ApplicationState.ENABLED
+
+
+class TestFlameError:
+    """Test FlameError."""
+    
+    def test_flame_error(self):
+        """Test FlameError creation and string representation."""
+        error = FlameError(FlameErrorCode.INVALID_CONFIG, "Test error")
+        assert error.code == FlameErrorCode.INVALID_CONFIG
+        assert error.message == "Test error"
+        assert "Test error" in str(error)
+        assert "INVALID_CONFIG" in str(error)
+
+
+class TestTaskInformer:
+    """Test TaskInformer."""
+    
+    def test_task_informer_default(self):
+        """Test default TaskInformer implementation."""
+        informer = TaskInformer()
+        task = Task(
+            id="task-1",
+            session_id="session-1",
+            state=TaskState.PENDING,
+            creation_time=datetime.now()
+        )
+        error = FlameError(FlameErrorCode.INTERNAL, "Test error")
+        
+        # Should not raise any exceptions
+        informer.on_update(task)
+        informer.on_error(error)
+
+
+class TestConnection:
+    """Test Connection class."""
+    
+    @pytest.mark.asyncio
+    async def test_connection_connect_empty_addr(self):
+        """Test connection with empty address."""
+        with pytest.raises(FlameError) as exc_info:
+            await Connection.connect("")
+        assert exc_info.value.code == FlameErrorCode.INVALID_CONFIG
+        assert "address cannot be empty" in exc_info.value.message
+    
+    @pytest.mark.asyncio
+    async def test_connection_connect_invalid_addr(self):
+        """Test connection with invalid address."""
+        with pytest.raises(FlameError) as exc_info:
+            await Connection.connect("invalid://address")
+        assert exc_info.value.code == FlameErrorCode.INVALID_CONFIG
+
+
+class TestEnums:
+    """Test enum values."""
+    
+    def test_session_state_values(self):
+        """Test SessionState enum values."""
+        assert SessionState.OPEN == 0
+        assert SessionState.CLOSED == 1
+    
+    def test_task_state_values(self):
+        """Test TaskState enum values."""
+        assert TaskState.PENDING == 0
+        assert TaskState.RUNNING == 1
+        assert TaskState.SUCCEED == 2
+        assert TaskState.FAILED == 3
+    
+    def test_application_state_values(self):
+        """Test ApplicationState enum values."""
+        assert ApplicationState.ENABLED == 0
+        assert ApplicationState.DISABLED == 1
+    
+    def test_shim_values(self):
+        """Test Shim enum values."""
+        assert Shim.LOG == 0
+        assert Shim.STDIO == 1
+        assert Shim.WASM == 2
+        assert Shim.SHELL == 3
+        assert Shim.GRPC == 4
+    
+    def test_flame_error_code_values(self):
+        """Test FlameErrorCode enum values."""
+        assert FlameErrorCode.INVALID_CONFIG == 0
+        assert FlameErrorCode.INVALID_STATE == 1
+        assert FlameErrorCode.INVALID_ARGUMENT == 2
+        assert FlameErrorCode.INTERNAL == 3
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])