34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
|
||
from django.conf import settings
|
||
from django.core.management.base import BaseCommand, CommandError
|
||
|
||
from review_agent.regulatory_review.services.rag_embedding import get_embedding_provider
|
||
from review_agent.regulatory_review.services.rag_index import build_chroma_index
|
||
from review_agent.regulatory_review.services.rule_loader import load_rule_file
|
||
|
||
|
||
class Command(BaseCommand):
|
||
help = "构建 NMPA 法规材料本地 ChromaDB RAG 索引。"
|
||
|
||
def add_arguments(self, parser):
|
||
parser.add_argument("--provider", default=None, help="覆盖 REGULATORY_RAG_PROVIDER。")
|
||
|
||
def handle(self, *args, **options):
|
||
rule_set = load_rule_file()
|
||
source_dir = Path(settings.BASE_DIR) / rule_set["source_material_dir"]
|
||
if not source_dir.exists():
|
||
raise CommandError(f"法规材料目录不存在：{source_dir}")
|
||
try:
|
||
provider = get_embedding_provider(options["provider"])
|
||
count = build_chroma_index(source_dir=source_dir, embedding_provider=provider)
|
||
except Exception as exc:
|
||
raise CommandError(str(exc)) from exc
|
||
self.stdout.write(
|
||
self.style.SUCCESS(
|
||
f"已构建法规 RAG 索引：collection={settings.REGULATORY_RAG_COLLECTION}, chunks={count}"
|
||
)
|
||
)
|