service-forge 0.1.18__py3-none-any.whl

service_forge/db/database.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+import redis
+import pymongo
+import psycopg2
+from typing import AsyncGenerator
+from loguru import logger
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
+from service_forge.service_config import ServiceConfig
+
+class PostgresDatabase:
+    def __init__(
+        self,
+        name: str,
+        postgres_user: str,
+        postgres_password: str,
+        postgres_host: str,
+        postgres_port: int,
+        postgres_db: str,
+    ) -> None:
+        self.name = name
+        self.postgres_user = postgres_user
+        self.postgres_password = postgres_password
+        self.postgres_host = postgres_host
+        self.postgres_port = postgres_port
+        self.postgres_db = postgres_db
+        self.engine = None
+        self.session_factory = None
+        self.test_connection()
+
+    @property
+    def database_url(self) -> str:
+        return f"postgresql+asyncpg://{self.postgres_user}:{self.postgres_password}@{self.postgres_host}:{self.postgres_port}/{self.postgres_db}"
+
+    @property
+    def database_base_url(self) -> str:
+        return f"postgresql+asyncpg://{self.postgres_user}:{self.postgres_password}@{self.postgres_host}:{self.postgres_port}/postgres"
+
+    async def init(self) -> None:
+        if self.engine is None:
+            self.engine = await self.create_engine()
+            self.session_factory = async_sessionmaker(bind=self.engine, class_=AsyncSession, expire_on_commit=False)
+
+    async def close(self) -> None:
+        if self.engine:
+            await self.engine.dispose()
+            self.engine = None
+            self.session_factory = None
+            logger.info("Database connection closed")
+
+    async def create_engine(self) -> AsyncEngine:
+        if not all([self.postgres_user, self.postgres_host, self.postgres_port, self.postgres_db]):
+            raise ValueError("Missing required database configuration. Please check your .env file or configuration.")
+        logger.info(f"Creating database engine: {self.database_url}")
+        return create_async_engine(self.database_url)
+
+    async def get_async_session(self) -> AsyncGenerator[AsyncSession, None]:
+        if self.session_factory is None:
+            await self.init()
+        
+        if self.session_factory is None:
+            raise RuntimeError("Session factory is not initialized")
+        
+        async with self.session_factory() as session:
+            try:
+                yield session
+            except Exception:
+                await session.rollback()
+                raise
+            finally:
+                await session.close()
+            yield session
+
+    async def get_session_factory(self) -> async_sessionmaker[AsyncSession]:
+        if self.engine is None:
+            await self.init()
+
+        if self.session_factory is None:
+            raise RuntimeError("Session factory is not initialized")
+
+        return self.session_factory
+
+    def test_connection(self) -> bool:
+        try:
+            conn = psycopg2.connect(
+                host=self.postgres_host,
+                port=self.postgres_port,
+                user=self.postgres_user,
+                password=self.postgres_password,
+                database=self.postgres_db,
+                connect_timeout=5
+            )
+            conn.close()
+            logger.info(f"PostgreSQL connection test successful for database '{self.name}'")
+            return True
+        except Exception as e:
+            logger.warning(f"PostgreSQL connection test failed for database '{self.name}': {e}")
+            return False
+
+class MongoDatabase:
+    def __init__(
+        self,
+        name: str,
+        mongo_host: str,
+        mongo_port: int,
+        mongo_user: str,
+        mongo_password: str,
+        mongo_db: str,
+    ) -> None:
+        self.name = name
+        self.mongo_host = mongo_host
+        self.mongo_port = mongo_port
+        self.mongo_user = mongo_user
+        self.mongo_password = mongo_password
+        self.mongo_db = mongo_db or ""
+        self.client = pymongo.MongoClient(self.database_url)
+        self.test_connection()
+
+    @property
+    def database_url(self) -> str:
+        return f"mongodb://{self.mongo_user}:{self.mongo_password}@{self.mongo_host}:{self.mongo_port}/{self.mongo_db}"
+
+    def test_connection(self) -> bool:
+        try:
+            self.client.admin.command('ping')
+            logger.info(f"MongoDB connection test successful for database '{self.name}'")
+            return True
+        except Exception as e:
+            logger.error(f"MongoDB connection test failed for database '{self.name}': {e}")
+            return False
+
+class RedisDatabase:
+    def __init__(
+        self,
+        name: str,
+        redis_host: str,
+        redis_port: int,
+        redis_password: str,
+    ) -> None:
+        self.name = name
+        self.redis_host = redis_host
+        self.redis_port = redis_port
+        self.redis_password = redis_password
+        self.client = redis.Redis(host=redis_host, port=redis_port, password=redis_password)
+        self.test_connection()
+
+    def test_connection(self) -> bool:
+        try:
+            self.client.ping()
+            logger.info(f"Redis connection test successful for database '{self.name}'")
+            return True
+        except Exception as e:
+            logger.error(f"Redis connection test failed for database '{self.name}': {e}")
+            return False
+
+
+class DatabaseManager:
+    def __init__(
+        self,
+        postgres_databases: list[PostgresDatabase],
+        mongo_databases: list[MongoDatabase],
+        redis_databases: list[RedisDatabase],
+    ) -> None:
+        self.postgres_databases = postgres_databases
+        self.mongo_databases = mongo_databases
+        self.redis_databases = redis_databases
+
+    def get_database(self, name: str) ->  PostgresDatabase | MongoDatabase | RedisDatabase | None:
+        for database in self.postgres_databases:
+            if database.name == name:
+                return database
+        return None
+
+    def get_default_postgres_database(self) -> PostgresDatabase | None:
+        if len(self.postgres_databases) > 0:
+            return self.postgres_databases[0]
+        return None
+
+    def get_default_mongo_database(self) -> MongoDatabase | None:
+        if len(self.mongo_databases) > 0:
+            return self.mongo_databases[0]
+        return None
+
+    def get_default_redis_database(self) -> RedisDatabase | None:
+        if len(self.redis_databases) > 0:
+            return self.redis_databases[0]
+        return None
+
+    @staticmethod
+    def from_config(config_path: str = None, config: ServiceConfig = None) -> DatabaseManager:
+        if config is None:
+            config = ServiceConfig.from_yaml_file(config_path)
+
+        postgres_databases = []
+        mongo_databases = []
+        redis_databases = []
+
+        databases_config = config.databases
+
+        if databases_config is not None:
+            for database_config in databases_config:
+                if all([database_config.postgres_host is None, database_config.mongo_host is None, database_config.redis_host is None]):
+                    raise ValueError(f"Database '{database_config.name}' is missing required configuration. Please check your service.yaml file.")
+
+                if (database_config.postgres_host is not None) + (database_config.mongo_host is not None) + (database_config.redis_host is not None) > 1:
+                    raise ValueError(f"Database '{database_config['name']}' has multiple host configurations. Please check your service.yaml file.")
+
+                if database_config.postgres_host is not None:
+                    postgres_databases.append(PostgresDatabase(
+                        name=database_config.name,
+                        postgres_user=database_config.postgres_user,
+                        postgres_password=database_config.postgres_password,
+                        postgres_host=database_config.postgres_host,
+                        postgres_port=database_config.postgres_port,
+                        postgres_db=database_config.postgres_db,
+                    ))
+                elif database_config.mongo_host is not None:
+                    mongo_databases.append(MongoDatabase(
+                        name=database_config.name,
+                        mongo_host=database_config.mongo_host,
+                        mongo_port=database_config.mongo_port,
+                        mongo_user=database_config.mongo_user,
+                        mongo_password=database_config.mongo_password,
+                        mongo_db=database_config.mongo_db,
+                    ))
+                elif database_config.redis_host is not None:
+                    redis_databases.append(RedisDatabase(
+                        name=database_config.name,
+                        redis_host=database_config.redis_host,
+                        redis_port=database_config.redis_port,
+                        redis_password=database_config.redis_password,
+                    ))
+
+        return DatabaseManager(postgres_databases=postgres_databases, mongo_databases=mongo_databases, redis_databases=redis_databases)
+
+def create_database_manager(config_path: str = None, config = None) -> DatabaseManager:
+    return DatabaseManager.from_config(config_path, config)

service_forge/db/migrations/feedback_migration.py

@@ -0,0 +1,154 @@
+"""
+数据库迁移脚本 - 创建 feedback 表
+
+使用方法:
+1. 确保你的 service.yaml 中配置了数据库连接
+2. 运行此脚本: python -m src.service_forge.db.migrations.feedback_migration
+
+此脚本会:
+- 检查 feedback 表是否存在
+- 如果不存在,创建 feedback 表及索引
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+# 添加项目根目录到 Python 路径
+project_root = Path(__file__).parent.parent.parent.parent
+sys.path.insert(0, str(project_root))
+
+from loguru import logger
+from sqlalchemy import text
+from src.service_forge.db.database import DatabaseManager
+from src.service_forge.db.models.feedback import Base, FeedbackBase
+
+
+async def create_database_if_not_exists(db):
+    """如果数据库不存在，则创建数据库"""
+    from sqlalchemy import create_engine
+    from sqlalchemy.exc import OperationalError
+    import asyncpg
+
+    try:
+        # 先尝试连接目标数据库
+        await db.init()
+        logger.info(f"✓ 数据库 '{db.postgres_db}' 已存在")
+        await db.close()
+        return True
+    except Exception as e:
+        logger.info(f"数据库 '{db.postgres_db}' 不存在，准备创建...")
+
+        try:
+            # 连接到 postgres 默认数据库
+            conn = await asyncpg.connect(
+                host=db.postgres_host,
+                port=db.postgres_port,
+                user=db.postgres_user,
+                password=db.postgres_password,
+                database='postgres'
+            )
+
+            # 创建数据库
+            await conn.execute(f'CREATE DATABASE {db.postgres_db}')
+            await conn.close()
+
+            logger.info(f"✓ 数据库 '{db.postgres_db}' 创建成功!")
+            return True
+
+        except Exception as create_error:
+            logger.error(f"创建数据库失败: {create_error}")
+            return False
+
+
+async def create_feedback_table(database_manager: DatabaseManager):
+    """创建 feedback 表"""
+    db = database_manager.get_default_postgres_database()
+
+    if db is None:
+        logger.error("未找到默认 PostgreSQL 数据库配置")
+        return False
+
+    try:
+        # 1. 先确保数据库存在
+        db_created = await create_database_if_not_exists(db)
+        if not db_created:
+            return False
+
+        # 2. 连接到数据库
+        await db.init()
+        engine = db.engine
+
+        logger.info(f"连接到数据库: {db.database_url}")
+
+        # 3. 检查表是否存在
+        async with engine.begin() as conn:
+            result = await conn.execute(
+                text(
+                    "SELECT EXISTS (SELECT FROM information_schema.tables "
+                    "WHERE table_schema = 'public' AND table_name = 'feedback');"
+                )
+            )
+            table_exists = result.scalar()
+
+            if table_exists:
+                logger.info("✓ feedback 表已存在,跳过创建")
+                return True
+
+            logger.info("创建 feedback 表...")
+
+            # 4. 创建表
+            await conn.run_sync(Base.metadata.create_all)
+
+            logger.info("✓ feedback 表创建成功!")
+            logger.info("表结构:")
+            logger.info("  - feedback_id: UUID (主键)")
+            logger.info("  - task_id: VARCHAR(255) (索引)")
+            logger.info("  - workflow_name: VARCHAR(255) (索引)")
+            logger.info("  - rating: INTEGER (可选, 1-5)")
+            logger.info("  - comment: TEXT (可选)")
+            logger.info("  - metadata: JSONB (可选)")
+            logger.info("  - created_at: TIMESTAMP (索引)")
+            logger.info("  - updated_at: TIMESTAMP (可选)")
+
+        await db.close()
+        return True
+
+    except Exception as e:
+        logger.error(f"创建 feedback 表失败: {e}")
+        return False
+
+
+async def main():
+    """主函数"""
+    logger.info("=== Feedback 表迁移脚本 ===")
+
+    # 提示用户提供配置文件路径
+    if len(sys.argv) > 1:
+        config_path = sys.argv[1]
+    else:
+        logger.info("请提供 service 配置文件路径:")
+        logger.info("  python -m src.service_forge.db.migrations.feedback_migration <config_path>")
+        logger.info("例如:")
+        logger.info("  python -m src.service_forge.db.migrations.feedback_migration configs/service/my_service.yaml")
+        return
+
+    logger.info(f"读取配置文件: {config_path}")
+
+    try:
+        database_manager = DatabaseManager.from_config(config_path=config_path)
+        success = await create_feedback_table(database_manager)
+
+        if success:
+            logger.info("✓ 迁移完成!")
+        else:
+            logger.error("✗ 迁移失败")
+            sys.exit(1)
+
+    except Exception as e:
+        logger.error(f"迁移过程出错: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

service_forge/db/models/feedback.py

@@ -0,0 +1,33 @@
+from datetime import datetime
+from sqlalchemy import Column, String, Integer, DateTime, Text
+from sqlalchemy.dialects.postgresql import UUID, JSONB
+from sqlalchemy.ext.declarative import declarative_base
+import uuid
+
+Base = declarative_base()
+
+
+class FeedbackBase(Base):
+    """反馈数据表模型"""
+    __tablename__ = "feedback"
+
+    feedback_id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    task_id = Column(String(255), nullable=False, index=True)
+    workflow_name = Column(String(255), nullable=False, index=True)
+    rating = Column(Integer, nullable=True)
+    comment = Column(Text, nullable=True)
+    extra_metadata = Column("metadata", JSONB, nullable=True, default={})
+    created_at = Column(DateTime, nullable=False, default=datetime.now, index=True)
+    updated_at = Column(DateTime, nullable=True, onupdate=datetime.now)
+
+    def to_dict(self):
+        """转换为字典格式"""
+        return {
+            "feedback_id": str(self.feedback_id),
+            "task_id": self.task_id,
+            "workflow_name": self.workflow_name,
+            "rating": self.rating,
+            "comment": self.comment,
+            "metadata": self.extra_metadata or {},
+            "created_at": self.created_at,
+        }

service_forge/llm/__init__.py

@@ -0,0 +1,67 @@
+import os
+from .llm import LLM
+from enum import Enum
+from typing import Iterator
+
+_llm_dicts = {}
+
+class Model(Enum):
+    GPT_4_1_NANO = "gpt-4.1-nano"
+    QWEN_TURBO_LATEST = "qwen-turbo-latest"
+    QWEN_PLUS_LATEST = "qwen-plus-latest"
+    QWEN_MAX_LATEST = "qwen-max-latest"
+    DOUBO_SEED_1_6_250615 = "doubao-seed-1-6-250615"
+    DOUBO_SEED_1_6_THINKING_250615 = "doubao-seed-1-6-thinking-250615"
+    DOUBO_SEED_1_6_FLASH_250615 = "doubao-seed-1-6-flash-250615"
+    DEEPSEEK_V3_250324 = "deepseek-v3-250324"
+    AZURE_GPT_4O_MINI = "azure-gpt-4o-mini"
+    GEMINI = "gemini-2.5-flash"
+
+    def provider(self) -> str:
+        if self.value.startswith("gpt"):
+            return "openai"
+        elif self.value.startswith("qwen"):
+            return "dashscope"
+        elif self.value.startswith("doubao"):
+            return "doubao"
+        elif self.value.startswith("deepseek"):
+            return "deepseek"
+        elif self.value.startswith("azure"):
+            return "azure"
+        elif self.value.startswith("gemini"):
+            return "gemini"
+        raise ValueError(f"Invalid model: {self.value}")
+
+def get_model(model: str) -> Model:
+    if model in Model.__members__:
+        return Model[model]
+
+    model = model.upper().replace("-", "_")
+    if model in Model.__members__:
+        return Model[model]
+
+    raise ValueError(f"Invalid model: {model}")
+            
+def get_llm(provider: str) -> LLM:
+    if provider not in _llm_dicts:
+        if provider == "openai":
+            _llm_dicts[provider] = LLM(os.environ.get("OPENAI_API_KEY", ""), os.environ.get("OPENAI_BASE_URL", ""), int(os.environ.get("OPENAI_TIMEOUT", 2000)))
+        elif provider == "doubao":
+            _llm_dicts[provider] = LLM(os.environ.get("DOUBAO_API_KEY", ""), os.environ.get("DOUBAO_BASE_URL", ""), int(os.environ.get("DOUBAO_TIMEOUT", 2000)))
+        elif provider == "dashscope":
+            _llm_dicts[provider] = LLM(os.environ.get("DASHSCOPE_API_KEY", ""), os.environ.get("DASHSCOPE_BASE_URL", ""), int(os.environ.get("DASHSCOPE_TIMEOUT", 2000)))
+        elif provider == "deepseek":
+            _llm_dicts[provider] = LLM(os.environ.get("DEEPSEEK_API_KEY", ""), os.environ.get("DEEPSEEK_BASE_URL", ""), int(os.environ.get("DEEPSEEK_TIMEOUT", 2000)))
+        elif provider == "azure":
+            _llm_dicts[provider] = LLM(os.environ.get("AZURE_API_KEY", ""), os.environ.get("AZURE_BASE_URL", ""), int(os.environ.get("AZURE_TIMEOUT", 2000)), os.environ.get("AZURE_API_VERSION", ""))
+        elif provider == "gemini":
+            _llm_dicts[provider] = LLM(os.environ.get("GEMINI_API_KEY", ""), os.environ.get("GEMINI_BASE_URL", ""), int(os.environ.get("GEMINI_TIMEOUT", 2000)))
+        else:
+            raise ValueError(f"Invalid provider: {provider}")
+    return _llm_dicts[provider]
+
+def chat(input: str, system_prompt: str, model: Model, temperature: float) -> str:
+    return get_llm(model.provider()).chat(input, system_prompt, model.value, temperature)
+
+def chat_stream(input: str, system_prompt: str, model: Model, temperature: float) -> Iterator[str]:
+    return get_llm(model.provider()).chat_stream(input, system_prompt, model.value, temperature)

service_forge/llm/llm.py

@@ -0,0 +1,56 @@
+import random
+from openai import OpenAI
+from openai import AzureOpenAI
+from typing import Iterator
+
+class LLM():
+    def __init__(self, api_key: str, base_url: str, timeout: int, api_version: str | None = None):
+        if api_version is not None:
+            self.client = AzureOpenAI(
+                api_key=api_key,
+                azure_endpoint=base_url,
+                timeout=timeout,
+                api_version=api_version,
+            )
+        else:
+            self.client = OpenAI(
+                api_key=api_key,
+                base_url=base_url,
+                timeout=timeout,
+            )
+
+    def chat(self, input: str, system_prompt: str, model: str, temperature: float) -> str:
+        if model.startswith("azure"):
+            model = model.replace("azure-", "")
+
+        response = self.client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": input},
+            ],
+            temperature=temperature,
+        )
+
+        if response.choices[0].message.content is None:
+            return "Error"
+        else:
+            return response.choices[0].message.content
+
+    def chat_stream(self, input: str, system_prompt: str, model: str, temperature: float) -> Iterator[str]:
+        if model.startswith("azure"):
+            model = model.replace("azure-", "")
+
+        stream = self.client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": input},
+            ],
+            temperature=temperature,
+            stream=True,
+        )
+
+        for chunk in stream:
+            if chunk.choices[0].delta.content is not None:
+                yield chunk.choices[0].delta.content

service_forge/model/feedback.py

@@ -0,0 +1,30 @@
+from pydantic import BaseModel, Field
+from typing import Optional, Any
+from datetime import datetime
+from uuid import UUID
+
+
+class FeedbackCreate(BaseModel):
+    """创建反馈的请求模型"""
+    task_id: str = Field(..., description="工作流任务ID")
+    workflow_name: str = Field(..., description="工作流名称")
+    rating: Optional[int] = Field(None, ge=1, le=5, description="评分 (1-5)")
+    comment: Optional[str] = Field(None, description="用户评论")
+    metadata: Optional[dict[str, Any]] = Field(default_factory=dict, description="额外的元数据")
+
+
+class FeedbackResponse(BaseModel):
+    """反馈响应模型"""
+    feedback_id: str = Field(..., description="反馈ID")
+    task_id: str = Field(..., description="工作流任务ID")
+    workflow_name: str = Field(..., description="工作流名称")
+    rating: Optional[int] = Field(None, description="评分")
+    comment: Optional[str] = Field(None, description="用户评论")
+    metadata: dict[str, Any] = Field(default_factory=dict, description="元数据")
+    created_at: datetime = Field(..., description="创建时间")
+
+
+class FeedbackListResponse(BaseModel):
+    """反馈列表响应模型"""
+    total: int = Field(..., description="总数")
+    feedbacks: list[FeedbackResponse] = Field(..., description="反馈列表")

service_forge/model/websocket.py

@@ -0,0 +1,13 @@
+
+from pydantic import BaseModel
+
+class WebSocketMessage(BaseModel):
+    """WebSocket消息模型"""
+    message: str = None
+    client_id: str
+
+class WebSocketResponse(BaseModel):
+    """WebSocket响应模型"""
+    status: str
+    message: str
+    data: dict = None

service_forge/proto/foo_input.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+class FooInput(BaseModel):
+    user_id: int
+    data: str