75 lines
2.6 KiB
Python
75 lines
2.6 KiB
Python
import sys
|
|
import json
|
|
from model_downloader import download_model
|
|
from ai_agent import AIAgent
|
|
import argparse
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='AI 웹 정보 수집 시스템')
|
|
parser.add_argument('--topics', nargs='+', help='수집할 주제 목록', default=None)
|
|
parser.add_argument('--config', default='./config.json', help='설정 파일 경로')
|
|
parser.add_argument('--auto-topics', action='store_true', help='AI가 스스로 주제를 선정하여 조사')
|
|
parser.add_argument('--save-path', help='데이터 저장 경로 (예: /content/drive/MyDrive/MyFolder)', default=None)
|
|
|
|
args = parser.parse_args()
|
|
|
|
print("AI 웹 정보 수집 시스템 시작")
|
|
|
|
# 1. 설정 파일 로드
|
|
with open(args.config, 'r') as f:
|
|
config = json.load(f)
|
|
|
|
# 2. 모델 다운로드 (필요한 경우)
|
|
print("모델 확인 중...")
|
|
model, tokenizer = download_model(args.config)
|
|
|
|
# 모델 로딩은 AIAgent에서 수행하므로, 다운로드만 성공해도 계속 진행
|
|
# 로컬 경로에 스냅샷 파일 유무로 성공 여부 확인
|
|
local_model_path = config.get('model_local_path', './models/model')
|
|
has_files = False
|
|
try:
|
|
has_files = any(True for _ in __import__('os').scandir(local_model_path))
|
|
except Exception:
|
|
has_files = False
|
|
|
|
if (model is None and not has_files):
|
|
print("모델 다운로드 실패. 프로그램을 종료합니다.")
|
|
sys.exit(1)
|
|
|
|
if args.save_path:
|
|
config['data_storage']['drive_mount_path'] = args.save_path
|
|
# 수정된 config 저장
|
|
with open(args.config, 'w') as f:
|
|
json.dump(config, f, indent=2)
|
|
print(f"저장 경로 설정됨: {args.save_path}")
|
|
|
|
# 3. AI 에이전트 초기화
|
|
print("AI 에이전트 초기화 중...")
|
|
agent = AIAgent(args.config)
|
|
|
|
# 3. 주제 결정
|
|
if args.auto_topics or args.topics is None:
|
|
print("AI가 스스로 주제를 선정합니다...")
|
|
topics = agent.generate_topics(num_topics=3)
|
|
print(f"선정된 주제: {topics}")
|
|
else:
|
|
topics = args.topics
|
|
|
|
# 4. 정보 수집 실행
|
|
print(f"다음 주제들에 대해 정보를 수집합니다: {topics}")
|
|
results = agent.collect_information(topics)
|
|
|
|
# 5. 결과 출력
|
|
print("\n=== 수집 결과 ===")
|
|
for result in results:
|
|
print(f"주제: {result['topic']}")
|
|
print(f"응답: {result['response']}")
|
|
print("-" * 50)
|
|
|
|
# 6. 정리
|
|
agent.close()
|
|
print("프로그램 완료")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|