service-forge 0.1.18__py3-none-any.whl

service_forge/workflow/triggers/a2a_api_trigger.py

@@ -0,0 +1,257 @@
+from __future__ import annotations
+from loguru import logger
+from service_forge.workflow.trigger import Trigger
+from typing import AsyncIterator, Any
+from service_forge.workflow.port import Port
+from google.protobuf.message import Message
+from google.protobuf.json_format import MessageToJson
+from fastapi import FastAPI
+from a2a.types import (
+    AgentCapabilities,
+    AgentCard,
+    AgentSkill,
+)
+from a2a.server.apps import A2AStarletteApplication
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.tasks import InMemoryTaskStore
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.events import EventQueue
+from a2a.utils import new_agent_text_message
+from a2a.utils.constants import DEFAULT_RPC_URL, EXTENDED_AGENT_CARD_PATH, AGENT_CARD_WELL_KNOWN_PATH
+
+import json
+import uuid
+import asyncio
+from service_forge.workflow.workflow_event import WorkflowResult
+
+class A2AAgentExecutor(AgentExecutor):
+    def __init__(self, trigger: A2AAPITrigger):
+        self.trigger = trigger
+
+    @staticmethod
+    def serialize_result(result: Any) -> str:
+        if isinstance(result, Message):
+            return MessageToJson(
+                result,
+                preserving_proto_field_name=True
+            )
+        return json.dumps(result)
+
+    async def send_event(self, event_queue: EventQueue, item: WorkflowResult) -> None:
+        if item.is_error:
+            result = {
+                'event': 'error',
+                'detail': str(item.result)
+            }
+            await event_queue.enqueue_event(new_agent_text_message(json.dumps(result)))
+
+        if item.is_end:
+            result = {
+                'event': 'end',
+                'detail': self.serialize_result(item.result)
+            }
+            await event_queue.enqueue_event(new_agent_text_message(json.dumps(result)))
+
+        result = {
+            'event': 'data',
+            'data': self.serialize_result(item.result)
+        }
+        await event_queue.enqueue_event(new_agent_text_message(json.dumps(result)))
+
+    async def execute(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ) -> None:
+        task_id = uuid.uuid4()
+        self.trigger.result_queues[task_id] = asyncio.Queue()
+
+        self.trigger.trigger_queue.put_nowait({
+            'id': task_id,
+            'context': context,
+        })
+
+        # TODO: support stream output
+        if False:
+            self.trigger.stream_queues[task_id] = asyncio.Queue()
+            while True:
+                item = await self.trigger.stream_queues[task_id].get()
+                await self.send_event(event_queue, item)
+
+                if item.is_error or item.is_end:
+                    break
+
+            if task_id in self.trigger.stream_queues:
+                del self.trigger.stream_queues[task_id]
+        else:
+            result = await self.trigger.result_queues[task_id].get()
+            await self.send_event(event_queue, result)
+
+        if task_id in self.trigger.result_queues:
+            del self.trigger.result_queues[task_id]
+
+    async def cancel(
+        self, context: RequestContext, event_queue: EventQueue
+    ) -> None:
+        raise Exception('cancel not supported')
+
+
+class A2AAPITrigger(Trigger):
+    DEFAULT_INPUT_PORTS = [
+        Port("app", FastAPI),
+        Port("path", str),
+        Port("skill_id", str, is_extended=True),
+        Port("skill_name", str, is_extended=True),
+        Port("skill_description", str, is_extended=True),
+        Port("skill_tags", list[str], is_extended=True),
+        Port("skill_examples", list[str], is_extended=True),
+        Port("agent_name", str),
+        Port("agent_url", str),
+        Port("agent_description", str),
+        Port("agent_version", str),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("trigger", bool),
+        Port("context", RequestContext),
+    ]
+
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.events = {}
+        self.is_setup_handler = False
+        self.agent_card: AgentCard | None = None
+
+    @staticmethod
+    def serialize_result(result: Any):
+        if isinstance(result, Message):
+            return MessageToJson(
+                result,
+                preserving_proto_field_name=True
+            )
+        return result
+
+    def _setup_handler(
+        self,
+        app: FastAPI,
+        path: str,
+        skill_id: list[tuple[int, str]],
+        skill_name: list[tuple[int, str]],
+        skill_description: list[tuple[int, str]],
+        skill_tags: list[tuple[int, list[str]]],
+        skill_examples: list[tuple[int, list[str]]],
+        agent_name: str,
+        agent_url: str,
+        agent_description: str,
+        agent_version: str,
+    ) -> None:
+    
+        skills_config = []
+        for i in range(len(skill_id)):
+            skills_config.append({
+                'id': '',
+                'name': '',
+                'description': '',
+                'tags': [],
+                'examples': [],
+            })
+
+        for i in range(len(skill_id)):
+            skills_config[skill_id[i][0]]['id'] = skill_id[i][1]
+            skills_config[skill_name[i][0]]['name'] = skill_name[i][1]
+            skills_config[skill_description[i][0]]['description'] = skill_description[i][1]
+            skills_config[skill_tags[i][0]]['tags'] = skill_tags[i][1]
+            skills_config[skill_examples[i][0]]['examples'] = skill_examples[i][1]
+    
+        skills = []
+        for config in skills_config:
+            skills.append(AgentSkill(
+                id=config['id'],
+                name=config['name'],
+                description=config['description'],
+                tags=config['tags'],
+                examples=config['examples'],
+            ))
+
+        agent_card = AgentCard(
+            name=agent_name,
+            description=agent_description,
+            url=agent_url,
+            version=agent_version,
+            default_input_modes=['text'],
+            default_output_modes=['text'],
+            capabilities=AgentCapabilities(streaming=True),
+            skills=skills,
+            supports_authenticated_extended_card=False,
+        )
+        
+        self.agent_card = agent_card
+
+        request_handler = DefaultRequestHandler(
+            agent_executor=A2AAgentExecutor(self),
+            task_store=InMemoryTaskStore(),
+        )
+
+        try:
+            server = A2AStarletteApplication(
+                agent_card=agent_card,
+                http_handler=request_handler,
+            )
+            
+            server.add_routes_to_app(
+                app,
+                agent_card_url="/a2a" + path + AGENT_CARD_WELL_KNOWN_PATH,
+                rpc_url="/a2a" + path + DEFAULT_RPC_URL,
+                extended_agent_card_url="/a2a" + path + EXTENDED_AGENT_CARD_PATH,
+            )
+
+        except Exception as e:
+            logger.error(f"Error adding A2A routes: {e}")
+            raise
+
+    async def _run(
+        self,
+        app: FastAPI,
+        path: str,
+        skill_id: list[tuple[int, str]],
+        skill_name: list[tuple[int, str]],
+        skill_description: list[tuple[int, str]],
+        skill_tags: list[tuple[int, list[str]]],
+        skill_examples: list[tuple[int, list[str]]],
+        agent_name: str,
+        agent_url: str,
+        agent_description: str,
+        agent_version: str,
+    ) -> AsyncIterator[bool]:
+        if len(skill_id) != len(skill_name) or len(skill_id) != len(skill_description) or len(skill_id) != len(skill_tags) or len(skill_id) != len(skill_examples):
+            raise ValueError("skill_id, skill_name, skill_description, skill_tags, skill_examples must have the same length")
+
+        if not self.is_setup_handler:
+            self._setup_handler(
+                app,
+                path,
+                skill_id,
+                skill_name,
+                skill_description,
+                skill_tags,
+                skill_examples,
+                agent_name,
+                agent_url,
+                agent_description,
+                agent_version,
+            )
+            self.is_setup_handler = True
+
+        logger.info(f"A2A Trigger {self.name} is running")
+
+        while True:
+            try:
+                trigger = await self.trigger_queue.get()
+                self.prepare_output_edges('context', trigger['context'])
+                yield self.trigger(trigger['id'])
+            except Exception as e:
+                logger.error(f"Error in A2AAPITrigger._run: {e}")
+                continue
+
+    async def _stop(self) -> AsyncIterator[bool]:
+        pass

service_forge/workflow/triggers/fast_api_trigger.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+import uuid
+import asyncio
+import json
+from loguru import logger
+from service_forge.workflow.trigger import Trigger
+from typing import AsyncIterator, Any
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+from service_forge.workflow.port import Port
+from service_forge.utils.default_type_converter import type_converter
+from service_forge.api.routers.websocket.websocket_manager import websocket_manager
+from fastapi import HTTPException
+from google.protobuf.message import Message
+from google.protobuf.json_format import MessageToJson
+
+class FastAPITrigger(Trigger):
+    DEFAULT_INPUT_PORTS = [
+        Port("app", FastAPI),
+        Port("path", str),
+        Port("method", str),
+        Port("data_type", type),
+        Port("is_stream", bool),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("trigger", bool),
+        Port("user_id", int),
+        Port("data", Any),
+    ]
+
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.events = {}
+        self.is_setup_route = False
+        self.app = None
+        self.route_path = None
+        self.route_method = None
+
+    @staticmethod
+    def serialize_result(result: Any):
+        if isinstance(result, Message):
+            return MessageToJson(
+                result,
+                preserving_proto_field_name=True
+            )
+        return result
+
+    async def handle_request(
+        self,
+        request: Request,
+        data_type: type,
+        extract_data_fn: callable[[Request], dict],
+        is_stream: bool,
+    ):
+        task_id = uuid.uuid4()
+        self.result_queues[task_id] = asyncio.Queue()
+
+        body_data = await extract_data_fn(request)
+        converted_data = data_type(**body_data)
+
+        client_id = (
+            body_data.get("client_id")
+            or request.query_params.get("client_id")
+            or request.headers.get("X-Client-ID")
+        )
+        if client_id:
+            workflow_name = getattr(self.workflow, "name", "Unknown")
+            steps = len(self.workflow.nodes) if hasattr(self.workflow, "nodes") else 1
+            websocket_manager.create_task_with_client(task_id, client_id, workflow_name, steps)
+
+        self.trigger_queue.put_nowait({
+            "id": task_id,
+            "user_id": getattr(request.state, "user_id", None),
+            "data": converted_data,
+        })
+
+        if is_stream:
+            self.stream_queues[task_id] = asyncio.Queue()
+
+            async def generate_sse():
+                try:
+                    while True:
+                        item = await self.stream_queues[task_id].get()
+
+                        if item.is_error:
+                            yield f"event: error\ndata: {json.dumps({'detail': str(item.result)})}\n\n"
+                            break
+                        
+                        if item.is_end:
+                            # TODO: send the result?
+                            break
+
+                        # TODO: modify
+                        serialized = self.serialize_result(item.result)
+                        if isinstance(serialized, str):
+                            data = serialized
+                        else:
+                            data = json.dumps(serialized)
+                        
+                        yield f"data: {data}\n\n"
+                    
+                except Exception as e:
+                    yield f"event: error\ndata: {json.dumps({'detail': str(e)})}\n\n"
+                finally:
+                    if task_id in self.stream_queues:
+                        del self.stream_queues[task_id]
+                    if task_id in self.result_queues:
+                        del self.result_queues[task_id]
+            
+            return StreamingResponse(
+                generate_sse(),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "X-Accel-Buffering": "no",
+                }
+            )
+        else:
+            result = await self.result_queues[task_id].get()
+            del self.result_queues[task_id]
+
+            if result.is_error:
+                if isinstance(result.result, HTTPException):
+                    raise result.result
+                else:
+                    raise HTTPException(status_code=500, detail=str(result.result))
+
+            return self.serialize_result(result.result)
+
+    def _setup_route(self, app: FastAPI, path: str, method: str, data_type: type, is_stream: bool) -> None:
+        async def get_data(request: Request) -> dict:
+            return dict(request.query_params)
+
+        async def body_data(request: Request) -> dict:
+            raw = await request.body()
+            if not raw:
+                return {}
+            return json.loads(raw.decode("utf-8"))
+
+        extractor = get_data if method == "GET" else body_data
+
+        async def handler(request: Request):
+            return await self.handle_request(request, data_type, extractor, is_stream)
+
+        # Save route information for cleanup
+        self.app = app
+        self.route_path = path
+        self.route_method = method.upper()
+
+        if method == "GET":
+            app.get(path)(handler)
+        elif method == "POST":
+            app.post(path)(handler)
+        elif method == "PUT":
+            app.put(path)(handler)
+        elif method == "DELETE":
+            app.delete(path)(handler)
+        else:
+            raise ValueError(f"Invalid method {method}")
+
+    async def _run(self, app: FastAPI, path: str, method: str, data_type: type, is_stream: bool = False) -> AsyncIterator[bool]:
+        if not self.is_setup_route:
+            self._setup_route(app, path, method, data_type, is_stream)
+            self.is_setup_route = True
+
+        while True:
+            try:
+                trigger = await self.trigger_queue.get()
+                self.prepare_output_edges(self.get_output_port_by_name('user_id'), trigger['user_id'])
+                self.prepare_output_edges(self.get_output_port_by_name('data'), trigger['data'])
+                yield self.trigger(trigger['id'])
+            except Exception as e:
+                logger.error(f"Error in FastAPITrigger._run: {e}")
+                continue
+
+    async def _stop(self) -> AsyncIterator[bool]:
+        if self.is_setup_route:
+            # Remove the route from the app
+            if self.app and self.route_path and self.route_method:
+                # Find and remove matching route
+                routes_to_remove = []
+                for route in self.app.routes:
+                    if hasattr(route, "path") and hasattr(route, "methods"):
+                        if route.path == self.route_path and self.route_method in route.methods:
+                            routes_to_remove.append(route)
+                
+                # Remove found routes
+                for route in routes_to_remove:
+                    try:
+                        self.app.routes.remove(route)
+                        logger.info(f"Removed route {self.route_method} {self.route_path} from FastAPI app")
+                    except ValueError:
+                        logger.warning(f"Route {self.route_method} {self.route_path} not found in app.routes")
+            
+            # Reset route information
+            self.app = None
+            self.route_path = None
+            self.route_method = None
+            self.is_setup_route = False

service_forge/workflow/triggers/kafka_api_trigger.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+import uuid
+from typing import Any
+from service_forge.workflow.trigger import Trigger
+from typing import AsyncIterator
+from service_forge.workflow.port import Port
+from service_forge.api.kafka_api import KafkaApp
+
+class KafkaAPITrigger(Trigger):
+    DEFAULT_INPUT_PORTS = [
+        Port("app", KafkaApp),
+        Port("topic", str),
+        Port("data_type", type),
+        Port("group_id", str),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("trigger", bool),
+        Port("data", Any),
+    ]
+    
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.events = {}
+        self.is_setup_kafka_input = False
+
+    def _setup_kafka_input(self, app: KafkaApp, topic: str, data_type: type, group_id: str) -> None:
+        @app.kafka_input(topic, data_type, group_id)
+        async def handle_message(data):
+            task_id = uuid.uuid4()
+            self.trigger_queue.put_nowait({
+                "id": task_id,
+                "data": data,
+            })
+
+    async def _run(self, app: KafkaApp, topic: str, data_type: type, group_id: str) -> AsyncIterator[bool]:
+        if not self.is_setup_kafka_input:
+            self._setup_kafka_input(app, topic, data_type, group_id)
+            self.is_setup_kafka_input = True
+
+        while True:
+            trigger = await self.trigger_queue.get()
+            self.prepare_output_edges(self.get_output_port_by_name('data'), trigger['data'])
+            yield self.trigger(trigger['id'])
+
+    async def _stop(self) -> AsyncIterator[bool]:
+        pass

service_forge/workflow/triggers/once_trigger.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+from service_forge.workflow.node import Node
+from service_forge.workflow.port import Port
+from service_forge.workflow.trigger import Trigger
+from typing import AsyncIterator
+import uuid
+
+class OnceTrigger(Trigger):
+    DEFAULT_INPUT_PORTS = [
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("trigger", bool),
+    ]
+
+    def __init__(self, name: str):
+        super().__init__(name)
+
+    async def _run(self) -> AsyncIterator[bool]:
+        yield self.trigger(uuid.uuid4())
+
+    async def _stop(self) -> AsyncIterator[bool]:
+        pass

service_forge/workflow/triggers/period_trigger.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+import asyncio
+from service_forge.workflow.node import Node
+from service_forge.workflow.port import Port
+from service_forge.workflow.trigger import Trigger
+from typing import AsyncIterator
+import uuid
+
+class PeriodTrigger(Trigger):
+    DEFAULT_INPUT_PORTS = [
+        Port("TRIGGER", bool),
+        Port("period", float),
+        Port("times", int),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("trigger", bool),
+    ]
+
+    def __init__(self, name: str):
+        super().__init__(name)
+
+    async def _run(self, times: int, period: float) -> AsyncIterator[bool]:
+        for _ in range(times):
+            await asyncio.sleep(period)
+            yield uuid.uuid4()
+
+    async def _stop(self) -> AsyncIterator[bool]:
+        pass