odoo-ai/agent_service/main.py

from __future__ import annotations
import asyncio
import logging
import sys
from contextlib import asynccontextmanager

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

from .config import get_settings
from . import app_state
from .routers import dispatch, approval, registry, sweep, health

logger = logging.getLogger(__name__)


async def _init_db(settings) -> asyncpg.Pool:
    pool = await asyncpg.create_pool(
        host=settings.postgres_host,
        port=settings.postgres_port,
        database=settings.postgres_db,
        user=settings.postgres_user,
        password=settings.postgres_password,
        min_size=settings.postgres_min_connections,
        max_size=settings.postgres_max_connections,
        max_inactive_connection_lifetime=300,
    )
    logger.info('DB pool created (min=%d max=%d)', settings.postgres_min_connections, settings.postgres_max_connections)
    return pool


async def _db_health_loop(pool: asyncpg.Pool) -> None:
    while True:
        await asyncio.sleep(60)
        try:
            async with pool.acquire(timeout=5) as conn:
                await conn.fetchval('SELECT 1')
        except Exception as exc:
            logger.warning('DB health check failed: %s', exc)


@asynccontextmanager
async def lifespan(app: FastAPI):
    settings = get_settings()
    _configure_logging(settings)

    # 1. Database
    try:
        pool = await _init_db(settings)
        app_state.set_db_pool(pool)
        asyncio.create_task(_db_health_loop(pool))
    except Exception as exc:
        logger.error('Failed to connect to database: %s', exc)
        pool = None

    # 2. LLM Router
    try:
        from .llm.ollama_backend import OllamaBackend
        from .llm.llm_config_store import LLMConfigStore
        from .llm.llm_router import LLMRouter

        ollama = OllamaBackend(
            base_url=settings.ollama_url,
            model=settings.ollama_model,
            timeout=settings.ollama_timeout,
        )
        config_store = LLMConfigStore(pool) if pool else None
        claude = None
        if settings.llm_privacy_mode != 'local' and settings.anthropic_api_key:
            from .llm.claude_backend import ClaudeBackend
            claude = ClaudeBackend(api_key=settings.anthropic_api_key, model=settings.claude_model)

        llm_router = LLMRouter(
            ollama=ollama,
            claude=claude,
            config_store=config_store,
            privacy_mode=settings.llm_privacy_mode,
            env_overrides={
                name: settings.agent_backend_override(name)
                for name in [
                    'finance_agent', 'accounting_agent', 'crm_agent', 'sales_agent',
                    'project_agent', 'elearning_agent', 'expenses_agent', 'employees_agent',
                ]
                if settings.agent_backend_override(name)
            },
        )
        app_state.set_llm_router(llm_router)
        logger.info('LLM router ready (mode=%s)', settings.llm_privacy_mode)
    except Exception as exc:
        logger.error('Failed to init LLM router: %s', exc)
        llm_router = None

    # 3. Odoo client
    try:
        from .tools.odoo_client import OdooClient
        odoo = OdooClient(
            url=settings.odoo_url,
            db=settings.odoo_db,
            api_key=settings.odoo_api_key,
        )
        logger.info('Odoo client initialised (%s)', settings.odoo_url)
    except Exception as exc:
        logger.error('Failed to init Odoo client: %s', exc)
        odoo = None

    # 4. Agent registry
    try:
        from .agents.registry import AgentRegistry
        agent_registry = AgentRegistry(odoo=odoo, pool=pool)
        app_state.set_agent_registry(agent_registry)
    except Exception as exc:
        logger.error('Failed to init agent registry: %s', exc)
        agent_registry = None

    # 5. Memory manager
    try:
        from .memory.memory_manager import MemoryManager
        memory_mgr = MemoryManager(pool=pool, llm=llm_router) if pool else None
    except Exception as exc:
        logger.error('Failed to init memory manager: %s', exc)
        memory_mgr = None

    # 6. Peer bus + specialist agents
    try:
        from .agents.peer_bus import PeerBus
        peer_bus = PeerBus()
        _register_specialist_agents(peer_bus, odoo, llm_router)
    except Exception as exc:
        logger.error('Failed to init peer bus / specialist agents: %s', exc)
        peer_bus = None

    # 7. Master agent
    try:
        from .agents.master_agent import MasterAgent
        master = MasterAgent(
            odoo=odoo,
            llm=llm_router,
            memory=memory_mgr,
            peer_bus=peer_bus,
            registry=agent_registry,
        )
        app_state.set_master_agent(master)
        logger.info('MasterAgent ready')
    except Exception as exc:
        logger.error('Failed to init MasterAgent: %s', exc)

    # 8. Sweep coordinator (lazy import — defined in Step 16)
    try:
        from .agents.sweep_coordinator import SweepCoordinator
        sweep_coord = SweepCoordinator(peer_bus=peer_bus)
        app_state.set_sweep_coordinator(sweep_coord)
    except ImportError:
        pass  # not yet implemented
    except Exception as exc:
        logger.warning('Sweep coordinator not available: %s', exc)

    logger.info('ActiveBlue AI agent service started on port %d', settings.agent_service_port)
    yield

    # Shutdown
    if pool:
        await pool.close()
    logger.info('Agent service shut down')


def _register_specialist_agents(peer_bus, odoo, llm_router) -> None:
    try:
        from .agents.finance_agent import FinanceAgent
        peer_bus.register('finance_agent', FinanceAgent(odoo=odoo, llm=llm_router, peer_bus=peer_bus))
    except Exception as exc:
        logger.warning('Could not register finance_agent: %s', exc)

    specialist_map = {
        'accounting_agent': 'AccountingAgent',
        'crm_agent': 'CrmAgent',
        'sales_agent': 'SalesAgent',
        'project_agent': 'ProjectAgent',
        'elearning_agent': 'ElearningAgent',
        'expenses_agent': 'ExpensesAgent',
        'employees_agent': 'EmployeesAgent',
    }
    for agent_name, class_name in specialist_map.items():
        module_name = agent_name.replace('_agent', '_agent')
        try:
            import importlib
            mod = importlib.import_module(f'.agents.{agent_name}', package='agent_service')
            cls = getattr(mod, class_name)
            peer_bus.register(agent_name, cls(odoo=odoo, llm=llm_router, peer_bus=peer_bus))
        except ImportError:
            logger.debug('%s module not yet implemented, skipping', agent_name)
        except Exception as exc:
            logger.warning('Could not register %s: %s', agent_name, exc)


def _configure_logging(settings) -> None:
    level = getattr(logging, settings.log_level.upper(), logging.INFO)
    if settings.log_format == 'json':
        try:
            import json_log_formatter
            formatter = json_log_formatter.JSONFormatter()
            handler = logging.StreamHandler(sys.stdout)
            handler.setFormatter(formatter)
            logging.root.handlers = [handler]
        except ImportError:
            logging.basicConfig(level=level, stream=sys.stdout)
    else:
        logging.basicConfig(
            level=level,
            stream=sys.stdout,
            format='%(asctime)s %(name)s %(levelname)s %(message)s',
        )
    logging.root.setLevel(level)


def create_app() -> FastAPI:
    app = FastAPI(
        title='ActiveBlue AI Agent Service',
        version='0.1.0',
        docs_url='/docs',
        redoc_url=None,
        lifespan=lifespan,
    )
    app.add_middleware(
        CORSMiddleware,
        allow_origins=['*'],
        allow_methods=['GET', 'POST', 'DELETE'],
        allow_headers=['*'],
    )
    app.include_router(dispatch.router)
    app.include_router(approval.router)
    app.include_router(registry.router)
    app.include_router(sweep.router)
    app.include_router(health.router)
    return app


app = create_app()