docker-compose.yml

services:
  # MySQL数据库
  mysql:
    image: docker1.aeko.cn/library/mysql:8
    container_name: threatrag-mysql
    restart: unless-stopped
    environment:
      MYSQL_DATABASE: knowledge_db
      MYSQL_USER: mysql
      MYSQL_PASSWORD: 12345678
      MYSQL_ROOT_PASSWORD: 12345678
    ports:
      - "3309:3306"
    volumes:
      - ./data/mysql:/var/lib/mysql
      - ./scripts/init-mysql.sql:/docker-entrypoint-initdb.d/init-mysql.sql
    networks:
      - threatrag-network
    healthcheck:
      test: [ "CMD-SHELL", "mysqladmin ping -h 127.0.0.1 -u root -p$MYSQL_ROOT_PASSWORD || exit 1" ]
      interval: 10s
      timeout: 5s
      retries: 5

  # Redis缓存数据库
  redis:
    image: docker1.aeko.cn/library/redis:7-alpine
    container_name: threatrag-redis
    restart: unless-stopped
    ports:
      - "6379:6379"
    volumes:
      - ./data/redis:/data
    networks:
      - threatrag-network
    healthcheck:
      test: [ "CMD", "redis-cli", "ping" ]
      interval: 10s
      timeout: 5s
      retries: 3

  # Neo4j图数据库
  neo4j:
    image: docker1.aeko.cn/neo4j:5.15-community
    container_name: threatrag-neo4j
    restart: unless-stopped
    environment:
      NEO4J_AUTH: neo4j/12345678
      NEO4J_PLUGINS: '["graph-data-science"]'
      NEO4J_dbms_security_procedures_unrestricted: gds.*
      NEO4J_dbms_security_procedures_allowlist: gds.*
    ports:
      - "7475:7474" # HTTP
      - "7688:7687" # Bolt
    volumes:
      - ./data/neo4j/data:/data
      - ./data/neo4j/logs:/logs
    networks:
      - threatrag-network
    healthcheck:
      test: [ "CMD-SHELL", "echo 'RETURN 1;' | cypher-shell -u neo4j -p 12345678" ]
      interval: 30s
      timeout: 10s
      retries: 5

  # Etcd (Milvus依赖)
  etcd:
    container_name: milvus-etcd
    image: quay.io/coreos/etcd:v3.5.5
    restart: unless-stopped
    environment:
      - ETCD_AUTO_COMPACTION_MODE=revision
      - ETCD_AUTO_COMPACTION_RETENTION=1000
      - ETCD_QUOTA_BACKEND_BYTES=4294967296
      - ETCD_SNAPSHOT_COUNT=50000
    volumes:
      - ./data/etcd:/etcd
    command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd
    networks:
      - threatrag-network

  # MinIO (Milvus依赖)
  minio:
    container_name: milvus-minio
    image: docker1.aeko.cn/minio/minio:RELEASE.2023-03-20T20-16-18Z
    restart: unless-stopped
    environment:
      MINIO_ACCESS_KEY: minioadmin
      MINIO_SECRET_KEY: minioadmin
    ports:
      - "9001:9001"
      - "9000:9000"
    volumes:
      - ./data/minio:/minio_data
    command: minio server /minio_data --console-address ":9001"
    networks:
      - threatrag-network
    healthcheck:
      test: [ "CMD", "curl", "-f", "http://localhost:9000/minio/health/live" ]
      interval: 30s
      timeout: 20s
      retries: 3

  # Milvus向量数据库
  milvus-standalone:
    container_name: milvus-standalone
    image: docker1.aeko.cn/milvusdb/milvus:v2.3.4
    restart: unless-stopped
    command: [ "milvus", "run", "standalone" ]
    environment:
      ETCD_ENDPOINTS: etcd:2379
      MINIO_ADDRESS: minio:9000
    volumes:
      - ./data/milvus:/var/lib/milvus
    ports:
      - "19530:19530"
      - "9091:9091"
    depends_on:
      - "etcd"
      - "minio"
    networks:
      - threatrag-network
    healthcheck:
      test: [ "CMD", "curl", "-f", "http://localhost:9091/healthz" ]
      interval: 30s
      start_period: 90s
      timeout: 20s
      retries: 3

  # Ollama 本地模型服务（GPU版本）
  ollama:
    image: docker1.aeko.cn/ollama/ollama:latest
    container_name: threatrag-ollama
    restart: unless-stopped
    ports:
      - "11434:11434"
    volumes:
      - ./data/ollama:/root/.ollama
    networks:
      - threatrag-network
    # GPU 支持已启用
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all  # 使用所有可用GPU
              capabilities: [gpu]
    healthcheck:
      test: ["CMD-SHELL", "ollama list || exit 1"]
      interval: 30s
      timeout: 10s
      retries: 5
      start_period: 60s  # 给 Ollama 更多启动时间

  # ThreatRAG API服务
  threatrag:
    build:
      context: .
      dockerfile: Dockerfile
    container_name: threatrag-api
    restart: unless-stopped
    ports:
      - "8006:8000"
    env_file:
      - .env
    environment:
      - PYTHONPATH=/app
      - FASTAPI_ENV=${FASTAPI_ENV:-production}
      - NEO4J_URL=bolt://neo4j:7687
      - NEO4J_USERNAME=${NEO4J_USERNAME:-neo4j}
      - NEO4J_PASSWORD=${NEO4J_PASSWORD:-12345678}
      - MILVUS_HOST=milvus-standalone
      - MILVUS_PORT=19530
      - MYSQL_HOST=mysql
      - MYSQL_PORT=3306
      - MYSQL_DB=knowledge_db
      - MYSQL_USER=mysql
      - MYSQL_PASSWORD=12345678
      - REDIS_URL=redis://redis:6379
      - BASE_MODEL=${BASE_MODEL:-deepseek-ai/DeepSeek-V3}
      - OPENAI_API_KEY=${OPENAI_API_KEY}
      - ZHIPUAI_API_KEY=${ZHIPUAI_API_KEY}
      - DEEPSEEK_API_KEY=${DEEPSEEK_API_KEY}
      - SILICONFLOW_API_KEY=${SILICONFLOW_API_KEY}
      - SILICONFLOW_API_BASE=${SILICONFLOW_API_BASE:-https://api.siliconflow.cn/v1}
      - OLLAMA_API_BASE=${OLLAMA_API_BASE:-http://ollama:11434}
      - SESSION_EXPIRE_TIME=${SESSION_EXPIRE_TIME:-3600}
      - MAX_HISTORY_LENGTH=${MAX_HISTORY_LENGTH:-50}
      - MAX_CONCURRENT_CHATS=${MAX_CONCURRENT_CHATS:-20}
    volumes:
      - ./data:/app/data
      - ./models:/app/models
      - ./config.yaml:/app/config.yaml
      - ./scripts:/app/scripts
    depends_on:
      mysql:
        condition: service_healthy
      redis:
        condition: service_healthy
      neo4j:
        condition: service_healthy
      milvus-standalone:
        condition: service_healthy
      ollama:
        condition: service_healthy
    networks:
      - threatrag-network
    healthcheck:
      test: [ "CMD", "curl", "-f", "http://localhost:8000/health" ]
      interval: 300s
      start_period: 600s
      timeout: 10s
      retries: 3
    logging:
      driver: "json-file"
      options:
        max-size: "10m"
        max-file: "3"

networks:
  threatrag-network:
    driver: bridge

# 所有数据现在挂载到本地 ./data 目录
# volumes 部分已移除，改为使用本地目录挂载