service-forge 0.1.0__py3-none-any.whl

service_forge/api/http_api.py

@@ -0,0 +1,138 @@
+from fastapi import FastAPI
+import uvicorn
+from typing import Optional
+from fastapi import APIRouter
+from loguru import logger
+from urllib.parse import urlparse
+from fastapi import HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from service_forge.api.websocket_api import router as websocket_router
+
+
+def is_trusted_origin(origin_host: str, host: str, trusted_root: str = "ring.shiweinan.com") -> bool:
+    """
+    Check if the origin host is trusted based on domain matching.
+    
+    Args:
+        origin_host: The hostname from the origin header
+        host: The hostname from the host header
+        trusted_root: The trusted root domain (can be customized)
+        
+    Returns:
+        bool: True if the origin is trusted, False otherwise
+    """
+    # Convert to lowercase to avoid case sensitivity issues
+    origin_host = origin_host.lower()
+    host = host.lower()
+
+    # Allow same domain, or subdomains under the same trusted root
+    return (
+        origin_host == host or
+        origin_host.endswith("." + trusted_root) or
+        host.endswith("." + trusted_root)
+    )
+
+
+def create_app(
+    app: Optional[FastAPI] = None,
+    routers: Optional[list[APIRouter]] = None,
+    cors_origins: Optional[list[str]] = None,
+    enable_auth_middleware: bool = True,
+    trusted_domain: str = "ring.shiweinan.com"
+) -> FastAPI:
+    """
+    Create or configure a FastAPI app with common middleware and configuration.
+    
+    Args:
+        app: Optional existing FastAPI instance. If None, creates a new one.
+        routers: List of APIRouter instances to include
+        cors_origins: List of allowed CORS origins. Defaults to ["*"]
+        enable_auth_middleware: Whether to enable authentication middleware
+        trusted_domain: Trusted domain for origin validation
+        
+    Returns:
+        FastAPI: Configured FastAPI application instance
+    """
+    if app is None:
+        app = FastAPI()
+    
+    # Configure CORS middleware
+    if cors_origins is None:
+        cors_origins = ["*"]
+        
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=cors_origins,
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    # Include routers if provided
+    if routers:
+        for router in routers:
+            app.include_router(router)
+    
+    # Always include WebSocket router
+    app.include_router(websocket_router)
+    
+    # Add authentication middleware if enabled
+    if enable_auth_middleware:
+        @app.middleware("http")
+        async def auth_middleware(request: Request, call_next):
+            """
+            Authentication middleware for API routes.
+            
+            Validates user authentication for /api routes with origin-based
+            trust verification and X-User-ID header validation.
+            """
+            if request.url.path.startswith("/api"):
+                origin = request.headers.get("origin") or request.headers.get("referer")
+                scheme = request.url.scheme
+                host = request.headers.get("host", "")
+                is_same_origin = False
+                
+                logger.debug(f"origin {origin}, host:{host}")
+                
+                if origin and host:
+                    try:
+                        parsed_origin = urlparse(origin)
+                        parsed_host = urlparse(f"{scheme}://{host}")
+                        is_same_origin = (
+                            parsed_origin.hostname == parsed_host.hostname
+                            and parsed_origin.port == parsed_host.port
+                            and is_trusted_origin(parsed_origin.hostname, parsed_host.hostname, trusted_domain)
+                        )
+                    except Exception:
+                        pass  # If parsing fails, continue with default behavior
+                if not is_same_origin:
+                    headers = request.headers
+                    user_id = headers.get("X-User-ID")
+                    if not user_id:
+                        raise HTTPException(status_code=401, detail="Unauthorized")
+                    
+                    request.state.user_id = user_id
+                else:
+                    # Same-origin requests can skip auth, but still set default user_id
+                    request.state.user_id = "0"  # Can be None or default value as needed
+
+            return await call_next(request)
+    
+    return app
+
+async def start_fastapi_server(host: str, port: int):
+    try:
+        config = uvicorn.Config(
+            fastapi_app,
+            host=host,
+            port=int(port),
+            log_level="info",
+            access_log=True
+        )
+        server = uvicorn.Server(config)
+        await server.serve()
+    except Exception as e:
+        logger.error(f"Server error: {e}")
+        raise
+
+fastapi_app = create_app(enable_auth_middleware=False)

service_forge/api/kafka_api.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+from typing import Callable, Any
+from aiokafka import AIOKafkaConsumer, AIOKafkaProducer, ConsumerRecord
+import asyncio
+import json
+import inspect
+from pydantic import BaseModel
+from loguru import logger
+
+class KafkaApp:
+    def __init__(self, bootstrap_servers: str = None):
+        self.bootstrap_servers = bootstrap_servers
+        self._handlers: dict[str, Callable] = {}
+        self._consumer_tasks: dict[str, asyncio.Task] = {}
+        self._producer: AIOKafkaProducer = None
+        self._lock = asyncio.Lock()
+        self._running = False
+
+    def kafka_input(self, topic: str, data_type: type, group_id: str):
+        def decorator(func: Callable):
+            self._handlers[topic] = (func, data_type, group_id)
+            logger.info(f"Registered Kafka input handler for topic '{topic}', data_type: {data_type}")
+
+            if self._running:
+                asyncio.create_task(self._start_consumer(topic, func, data_type, group_id))
+            return func
+        return decorator
+
+    def set_bootstrap_servers(self, bootstrap_servers: str) -> None:
+        self.bootstrap_servers = bootstrap_servers
+
+    async def start(self):
+        if not self.bootstrap_servers:
+            raise ValueError("bootstrap_servers 未设置")
+
+        logger.info(f"🚀 KafkaApp started with servers: {self.bootstrap_servers}")
+
+        await self._start_producer()
+
+        async with self._lock:
+            for topic, (handler, data_type, group_id) in self._handlers.items():
+                if topic not in self._consumer_tasks:
+                    self._consumer_tasks[topic] = asyncio.create_task(self._start_consumer(topic, handler, data_type, group_id))
+
+        self._running = True
+        while self._running:
+            await asyncio.sleep(1)
+
+    async def _start_consumer(self, topic: str, handler: Callable, data_type: type, group_id: str):
+        consumer = AIOKafkaConsumer(
+            topic,
+            bootstrap_servers=self.bootstrap_servers,
+            group_id=group_id,
+            enable_auto_commit=True,
+            auto_offset_reset="latest",
+        )
+        await consumer.start()
+        logger.info(f"✅ Started consumer for topic: {topic}")
+
+        try:
+            async for msg in consumer:
+                await self._dispatch_message(handler, msg, data_type)
+        except asyncio.CancelledError:
+            logger.warning(f"🛑 Consumer for {topic} cancelled")
+        finally:
+            await consumer.stop()
+
+    async def _dispatch_message(self, handler: Callable, msg: ConsumerRecord, data_type: type):
+        try:
+            data = data_type()
+            data.ParseFromString(msg.value)
+        except Exception as e:
+            print("Error:", e)
+            data = data_type(**json.loads(msg.value.decode("utf-8")))
+        result = handler(data)
+        if inspect.iscoroutine(result):
+            await result
+
+    async def _start_producer(self):
+        if self._producer is None:
+            self._producer = AIOKafkaProducer(
+                bootstrap_servers=self.bootstrap_servers,
+                value_serializer=lambda v: v.SerializeToString(),
+            )
+            await self._producer.start()
+            logger.info("✅ Kafka producer started")
+
+    async def _stop_producer(self):
+        if self._producer is not None:
+            await self._producer.stop()
+            self._producer = None
+            logger.info("✅ Kafka producer stopped")
+
+    async def send_message(self, topic: str, data_type: type, data: Any) -> None:
+        if not self._running:
+            raise RuntimeError("KafkaApp is not running. Call start() first.")
+        
+        if self._producer is None:
+            raise RuntimeError("Kafka producer is not initialized.")
+        
+        try:
+            await self._producer.send_and_wait(topic, data)
+            logger.info(f"✅ 已发送消息到 topic '{topic}', type: {data_type}")
+            
+        except Exception as e:
+            logger.error(f"❌ 发送消息到 topic '{topic}' 失败: {e}")
+            raise
+
+    async def stop(self):
+        logger.info("Stopping KafkaApp ...")
+        self._running = False
+        
+        for t in list(self._consumer_tasks.values()):
+            t.cancel()
+        await asyncio.sleep(0.1)
+        self._consumer_tasks.clear()
+        
+        await self._stop_producer()
+        
+        logger.info("✅ KafkaApp stopped")
+
+kafka_app = KafkaApp()
+
+async def start_kafka_server(bootstrap_servers: str):
+    kafka_app.set_bootstrap_servers(bootstrap_servers)
+    await kafka_app.start()

service_forge/api/task_manager.py

@@ -0,0 +1,141 @@
+
+from __future__ import annotations
+import asyncio
+import uuid
+import datetime
+from typing import Dict, List, Set, Any, Optional
+
+class TaskManager:
+    """任务管理器，用于跟踪任务状态和队列信息"""
+
+    def __init__(self):
+        # 存储所有任务信息: {task_id: task_info}
+        self.tasks: Dict[uuid.UUID, Dict[str, Any]] = {}
+        # 任务队列，按添加顺序排列
+        self.task_queue: List[uuid.UUID] = []
+        # 正在执行的任务ID集合
+        self.running_tasks: Set[uuid.UUID] = set()
+        # 已完成的任务ID集合
+        self.completed_tasks: Set[uuid.UUID] = set()
+        # 客户端与任务的映射: {client_id: set(task_id)}
+        self.client_tasks: Dict[str, Set[uuid.UUID]] = {}
+
+    def add_task(self, task_id: uuid.UUID, client_id: str, workflow_name: str, steps: int) -> Dict[str, Any]:
+        """添加新任务到队列"""
+        current_time = asyncio.get_event_loop().time()
+        task_info = {
+            "task_id": task_id,
+            "client_id": client_id,
+            "workflow_name": workflow_name,
+            "steps": steps,
+            "current_step": 0,  # 当前步骤，从0开始
+            "status": "pending",  # pending, running, completed, failed
+            "created_at": datetime.datetime.fromtimestamp(current_time).isoformat(),
+            "queue_position": len(self.task_queue) + 1
+        }
+        self.tasks[task_id] = task_info
+        self.task_queue.append(task_id)
+
+        # 更新客户端任务映射
+        if client_id not in self.client_tasks:
+            self.client_tasks[client_id] = set()
+        self.client_tasks[client_id].add(task_id)
+
+        return task_info
+
+    def start_task(self, task_id: uuid.UUID) -> bool:
+        """标记任务开始执行"""
+        if task_id not in self.tasks:
+            return False
+
+        current_time = asyncio.get_event_loop().time()
+        self.tasks[task_id]["status"] = "running"
+        self.tasks[task_id]["started_at"] = datetime.datetime.fromtimestamp(current_time).isoformat()
+        self.tasks[task_id]["current_step"] = 1  # 开始执行第一步
+        self.running_tasks.add(task_id)
+
+        # 从队列中移除
+        if task_id in self.task_queue:
+            self.task_queue.remove(task_id)
+
+        # 更新队列中所有任务的位置
+        for i, q_task_id in enumerate(self.task_queue):
+            self.tasks[q_task_id]["queue_position"] = i + 1
+
+        return True
+
+    def complete_task(self, task_id: uuid.UUID) -> bool:
+        """标记任务完成"""
+        if task_id not in self.tasks:
+            return False
+
+        current_time = asyncio.get_event_loop().time()
+        self.tasks[task_id]["status"] = "completed"
+        self.tasks[task_id]["completed_at"] = datetime.datetime.fromtimestamp(current_time).isoformat()
+        self.tasks[task_id]["current_step"] = self.tasks[task_id]["steps"]  # 完成所有步骤
+        self.running_tasks.discard(task_id)
+        self.completed_tasks.add(task_id)
+
+        return True
+
+    def fail_task(self, task_id: uuid.UUID, error: str = None) -> bool:
+        """标记任务失败"""
+        if task_id not in self.tasks:
+            return False
+
+        current_time = asyncio.get_event_loop().time()
+        self.tasks[task_id]["status"] = "failed"
+        self.tasks[task_id]["failed_at"] = datetime.datetime.fromtimestamp(current_time).isoformat()
+        if error:
+            self.tasks[task_id]["error"] = error
+        self.running_tasks.discard(task_id)
+
+        return True
+
+    def get_client_tasks(self, client_id: str) -> List[Dict[str, Any]]:
+        """获取客户端的所有任务"""
+        if client_id not in self.client_tasks:
+            return []
+
+        return [
+            self.tasks[task_id] 
+            for task_id in self.client_tasks[client_id] 
+            if task_id in self.tasks
+        ]
+
+    def get_queue_position(self, task_id: uuid.UUID) -> int:
+        """获取任务在队列中的位置，从1开始计数，如果不在队列中返回-1"""
+        if task_id not in self.tasks:
+            return -1
+
+        return self.tasks[task_id].get("queue_position", -1)
+
+    def get_global_queue_info(self) -> Dict[str, int]:
+        """获取全局队列信息"""
+        return {
+            "total": len(self.running_tasks) + len(self.task_queue),
+            "waiting": len(self.task_queue),
+            "running": len(self.running_tasks),
+        }
+
+    def get_task_info(self, task_id: uuid.UUID) -> Optional[Dict[str, Any]]:
+        """获取特定任务的详细信息"""
+        task_info = self.tasks.get(task_id)
+        if task_info:
+            # 创建任务信息的副本，并将UUID转换为字符串
+            task_copy = task_info.copy()
+            task_copy["task_id"] = str(task_id)
+            return task_copy
+        return None
+        
+    def update_current_step(self, task_id: uuid.UUID, step: int) -> bool:
+        """更新当前任务的步骤"""
+        if task_id not in self.tasks:
+            return False
+            
+        # 确保步骤在有效范围内
+        if step < 0 or step > self.tasks[task_id]["steps"]:
+            return False
+            
+        self.tasks[task_id]["current_step"] = step
+        return True

service_forge/api/websocket_api.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+import uuid
+import asyncio
+import json
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Request, Query
+from fastapi.responses import JSONResponse
+from loguru import logger
+from typing import Dict, Any, Optional
+from .websocket_manager import websocket_manager
+
+router = APIRouter(prefix="/ws", tags=["websocket"])
+
+@router.websocket("/connect")
+async def websocket_endpoint(websocket: WebSocket, client_id: Optional[str] = Query(None)):
+    """WebSocket连接端点，支持指定客户端ID"""
+    client_id = await websocket_manager.connect(websocket, client_id)
+    try:
+        while True:
+            # 接收客户端消息
+            data = await websocket.receive_text()
+            try:
+                message = json.loads(data)
+                await handle_client_message(client_id, message)
+            except json.JSONDecodeError:
+                logger.error(f"从客户端 {client_id} 收到无效JSON消息: {data}")
+                await websocket_manager.send_personal_message(
+                    json.dumps({"error": "Invalid JSON format"}),
+                    client_id
+                )
+    except WebSocketDisconnect:
+        websocket_manager.disconnect(client_id)
+    except Exception as e:
+        logger.error(f"WebSocket连接处理异常: {e}")
+        websocket_manager.disconnect(client_id)
+
+async def handle_client_message(client_id: str, message: Dict[str, Any]):
+    """处理来自客户端的消息"""
+    message_type = message.get("type")
+
+    if message_type == "subscribe":
+        # 客户端订阅任务
+        task_id_str = message.get("task_id")
+        if not task_id_str:
+            await websocket_manager.send_personal_message(
+                json.dumps({"error": "Missing task_id in subscribe message"}),
+                client_id
+            )
+            return
+
+        try:
+            task_id = uuid.UUID(task_id_str)
+            success = await websocket_manager.subscribe_to_task(client_id, task_id)
+            response = {"success": success}
+            await websocket_manager.send_personal_message(json.dumps(response), client_id)
+        except ValueError:
+            await websocket_manager.send_personal_message(
+                json.dumps({"error": "Invalid task_id format"}),
+                client_id
+            )
+
+    elif message_type == "unsubscribe":
+        # 客户端取消订阅任务
+        task_id_str = message.get("task_id")
+        if not task_id_str:
+            await websocket_manager.send_personal_message(
+                json.dumps({"error": "Missing task_id in unsubscribe message"}),
+                client_id
+            )
+            return
+
+        try:
+            task_id = uuid.UUID(task_id_str)
+            success = await websocket_manager.unsubscribe_from_task(client_id, task_id)
+            response = {"success": success}
+            await websocket_manager.send_personal_message(json.dumps(response), client_id)
+        except ValueError:
+            await websocket_manager.send_personal_message(
+                json.dumps({"error": "Invalid task_id format"}),
+                client_id
+            )
+
+    else:
+        await websocket_manager.send_personal_message(
+            json.dumps({"error": f"Unknown message type: {message_type}"}),
+            client_id
+        )