huihuiSquare/digital-avatar-app/backend/main.py

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

import os
import logging

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.interval import IntervalTrigger
import redis as redis_lib

from database import init_db, SessionLocal
from models import Avatar, Authorization, Organization, TokenAccount, TokenPlan
from fastapi.staticfiles import StaticFiles
import routers.avatars
import routers.tokens
import routers.authorizations
import routers.organizations
import routers.knowledge
import routers.huihui_auth
import routers.chat
import routers.takeover
from responses import ok

logger = logging.getLogger(__name__)

takeover_scheduler = None
takeover_db = None

app = FastAPI(title="会会数字分身 API", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(routers.avatars.router, prefix="/api")
app.include_router(routers.tokens.router, prefix="/api")
app.include_router(routers.authorizations.router, prefix="/api")
app.include_router(routers.organizations.router, prefix="/api")
app.include_router(routers.knowledge.router, prefix="/api")
app.include_router(routers.huihui_auth.router, prefix="/api")
app.include_router(routers.chat.router, prefix="/api")
app.include_router(routers.takeover.router, prefix="/api")

UPLOAD_DIR = routers.knowledge.UPLOAD_DIR
os.makedirs(UPLOAD_DIR, exist_ok=True)
app.mount("/api/files", StaticFiles(directory=UPLOAD_DIR), name="knowledge-files")


@app.get("/api/health")
def health():
    return ok({"status": "ok"})


def seed():
    db = SessionLocal()
    try:
        if db.query(TokenAccount).first() is None:
            db.add(TokenAccount(balance=1250))

        if db.query(TokenPlan).count() == 0:
            plans = [
                TokenPlan(id="1", name="新手体验", amount=1000, price=9.9, desc="新手体验"),
                TokenPlan(id="2", name="热门套餐", amount=5000, price=39.9, badge="热门"),
                TokenPlan(id="3", name="超值套餐", amount=12000, price=89.9, badge="超值"),
                TokenPlan(id="4", name="企业推荐", amount=30000, price=199, badge="企业推荐", desc="适合高频使用"),
            ]
            db.add_all(plans)

        if db.query(Avatar).count() == 0:
            avatar = Avatar(
                name="我的数字分身",
                display_name="会会助手",
                description="我是您的AI数字分身，可以帮您管理日程、回复消息、处理任务。",
                emoji="🤖",
                status="active",
                token_balance=0,
                config={
                    "replyStyle": "professional",
                    "creativity": 50,
                    "rigor": 50,
                    "humor": 30,
                    "responseLength": "medium",
                    "systemPrompt": "",
                },
            )
            db.add(avatar)
            db.commit()
            db.refresh(avatar)

            if db.query(Authorization).count() == 0:
                auths = [
                    Authorization(avatar_id=avatar.id, target_type="application", target_name="微信小程序", permissions=["read", "reply"], status="active"),
                    Authorization(avatar_id=avatar.id, target_type="user", target_name="张三", permissions=["read"], status="active"),
                    Authorization(avatar_id=avatar.id, target_type="organization", target_name="产品团队", permissions=["read", "edit"], status="inactive"),
                ]
                db.add_all(auths)

        if db.query(Organization).count() == 0:
            orgs = [
                Organization(name="会会增长团队", description="负责会会产品的增长与运营", emoji="🚀", org_type="team", member_count=12),
                Organization(name="AI 实验室", description="探索前沿 AI 能力", emoji="💡", org_type="company", member_count=8),
            ]
            db.add_all(orgs)

        db.commit()
    finally:
        db.close()


@app.on_event("startup")
def on_startup():
    global takeover_scheduler, takeover_db

    init_db()
    seed()

    # Release stale resources when startup is invoked again by a reload/test.
    stop_takeover_scheduler()

    # --- Takeover scheduler ---
    try:
        # Initialize Redis (optional)
        redis_client = None
        redis_url = os.getenv("REDIS_URL", "")
        if redis_url:
            try:
                redis_client = redis_lib.from_url(redis_url)
                redis_client.ping()
            except Exception as e:
                logger.warning(f"Redis connection failed, delayed takeover will degrade to immediate: {e}")

        # Initialize Box IM client
        from services.boxim_client import BoxIMClient
        boxim_config = {
            "HUIHUI_IM_BASE_URL": os.getenv("HUIHUI_IM_BASE_URL", "http://192.168.1.200:60040"),
            "HUIHUI_APP_ID": os.getenv("HUIHUI_APP_ID", ""),
            "HUIHUI_ACCESS_ID": os.getenv("HUIHUI_ACCESS_ID", ""),
            "HUIHUI_ACCESS_SECRET": os.getenv("HUIHUI_ACCESS_SECRET", ""),
        }
        boxim_client = BoxIMClient(boxim_config)

        # Initialize takeover service
        from services.takeover_service import TakeoverService
        takeover_db = SessionLocal()
        takeover_service = TakeoverService(takeover_db, boxim_client, redis_client)

        # AsyncIOScheduler awaits the service coroutine instead of dropping it.
        takeover_scheduler = AsyncIOScheduler()
        takeover_scheduler.add_job(
            takeover_service.poll_and_process_messages,
            trigger=IntervalTrigger(seconds=10),
            id="takeover_message_poll",
        )
        takeover_scheduler.start()
        logger.info("Takeover message polling scheduler started (interval=10s)")
    except Exception as e:
        stop_takeover_scheduler()
        logger.warning(f"Failed to initialize takeover scheduler, app will continue without it: {e}")


def stop_takeover_scheduler():
    global takeover_scheduler, takeover_db

    if takeover_scheduler is not None:
        try:
            if takeover_scheduler.running:
                takeover_scheduler.shutdown(wait=False)
        except Exception as e:
            logger.warning(f"Failed to stop takeover scheduler cleanly: {e}")
        finally:
            takeover_scheduler = None

    if takeover_db is not None:
        try:
            takeover_db.close()
        finally:
            takeover_db = None


@app.on_event("shutdown")
def on_shutdown():
    stop_takeover_scheduler()