main.py

#!/usr/bin/env python3
"""
메인 실행 파일
"""
import json
import logging
import argparse
from pathlib import Path
from typing import List, Dict, Any

from src.data_processor import DataProcessor
from src.chain_maker import ChainMaker


logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
if not logger.handlers:
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)



def load_tools(tools_file: str = None) -> List[Dict[str, Any]]:
    """도구 정의 로드"""
    if tools_file and Path(tools_file).exists():
        with open(tools_file, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    # 기본 도구 세트
    return [
        {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Get the current weather in a given location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The city and state, e.g. San Francisco, CA"
                        }
                    },
                    "required": ["location"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "send_email",
                "description": "Send an email to a given recipient",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "recipient": {
                            "type": "string",
                            "description": "The email address of the recipient"
                        },
                        "subject": {
                            "type": "string",
                            "description": "The subject of the email"
                        },
                        "body": {
                            "type": "string",
                            "description": "The body of the email"
                        }
                    },
                    "required": ["recipient", "subject", "body"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "search_flights",
                "description": "Search for flights between two cities",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "departure": {
                            "type": "string",
                            "description": "The departure city"
                        },
                        "destination": {
                            "type": "string",
                            "description": "The destination city"
                        }
                    },
                    "required": ["departure", "destination"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "book_hotel",
                "description": "Book a hotel for a given location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The location of the hotel"
                        }
                    },
                    "required": ["location"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "find_restaurant",
                "description": "Find a restaurant in a given location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The location of the restaurant"
                        },
                        "cuisine": {
                            "type": "string",
                            "description": "The cuisine of the restaurant"
                        },
                        "rating": {
                            "type": "number",
                            "description": "The rating of the restaurant"
                        }
                    },
                    "required": ["location", "cuisine", "rating"]
                }
            }
        }
    ]


def main():
    parser = argparse.ArgumentParser(description="Data Synthesis Pipeline")
    
    # 입출력 관련
    parser.add_argument("--input", type=str, default="data/questions.jsonl", help="Input JSONL file or directory (default: data/questions.jsonl)")
    parser.add_argument("--output", type=str, help="Output file or directory")
    parser.add_argument("--pattern", type=str, default="*.jsonl", help="File pattern for directory processing")
    
    # 처리 관련
    parser.add_argument("--num-workers", type=int, help="Number of parallel workers")
    parser.add_argument("--max-items", type=int, help="Maximum number of items to process")
    parser.add_argument("--method", type=str, choices=["single", "multiple"], help="Generation method")
    parser.add_argument("--n-responses", type=int, default=4, help="Number of responses for multiple method")
    
    # 도구 관련
    parser.add_argument("--tools", type=str, help="Path to tools JSON file")
    parser.add_argument("--no-tools", action="store_true", help="Disable tool usage")
    
    # 설정 파일
    parser.add_argument("--config", type=str, help="Path to config file")
    
    # 모드
    parser.add_argument("--mode", type=str, choices=["batch", "directory", "single"], 
                       default="batch", help="Processing mode")
    parser.add_argument("--query", type=str, help="Single query to process (for single mode)")
    
    # 사용자 정보 (single mode용)
    parser.add_argument("--company", type=str, help="Company name for single mode")
    parser.add_argument("--departments", type=str, help="Departments for single mode")
    parser.add_argument("--position", type=str, help="Position for single mode")
    parser.add_argument("--grade", type=str, help="Grade for single mode")
    parser.add_argument("--user-name", type=str, help="User name for single mode")
    
    args = parser.parse_args()
    
    # 도구 로드 (기본값은 None - questions.jsonl의 total_tools를 사용)
    tools = None if args.no_tools else (load_tools(args.tools) if args.tools else None)
    
    if args.mode == "single":
        # 단일 쿼리 처리
        if not args.query:
            logger.error("Error: --query is required for single mode")
            return
        
        # 사용자 정보 구성
        user_info = {}
        if args.company:
            user_info["company"] = args.company
        if args.departments:
            user_info["departments"] = args.departments
        if args.position:
            user_info["position"] = args.position
        if args.grade:
            user_info["grade"] = args.grade
        if args.user_name:
            user_info["user_name"] = args.user_name
        if args.employee_id:
            user_info["employee_id"] = args.employee_id

        chain_maker = ChainMaker(args.config, user_info=user_info if user_info else None)
        result = chain_maker.generate_chain(
            query=args.query,
            tools=tools,
            method=args.method,
            n_responses=args.n_responses
        )
        
        # 결과 출력 (토큰 수와 추론 시간 포함)
        logger.info("Result:")
        logger.info(f"Method: {result['method']}")
        logger.info(f"Total turns: {result['total_turns']}")
        logger.info(f"\nOverall Performance Stats:")
        logger.info(f"  - Total tokens: {result['stats'].get('total_tokens', 0):,}")
        logger.info(f"  - Total inference time: {result['stats'].get('total_inference_time', 0):.2f} seconds")
        logger.info(f"  - Total API calls: {result['stats'].get('total_calls', 0)}")
        if result['stats'].get('total_inference_time', 0) > 0:
            logger.info(f"  - Tokens per second: {result['stats'].get('total_tokens', 0) / result['stats'].get('total_inference_time', 1):.1f}")
        
        # 턴별 메타데이터 출력 (agent 타입만 표시)
        if result.get('turn_metadata'):
            agent_turns = [turn for turn in result['turn_metadata'] if turn.get('turn_type') == 'agent']
            if agent_turns:
                logger.info(f"Agent Turn Stats (Final Path - {len(agent_turns)} agent turns):")
                for i, turn in enumerate(agent_turns, 1):
                    logger.info(f"Agent Turn {i}:")
                    if turn.get('query'):
                        logger.info(f"Query: {turn['query'][:50]}...")
                    if turn.get('tokens'):
                        logger.info(f"Tokens: {turn['tokens'].get('total', 0)}")
                        logger.info(f"Inference time: {turn.get('inference_time', 0):.2f}s")
                        if turn.get('tokens_per_sec'):
                            logger.info(f"Tokens/sec: {turn['tokens_per_sec']:.1f}")
                    if turn.get('selected_from') and turn['selected_from'] > 1:
                        logger.info(f"(Selected from {turn['selected_from']} candidates)")
        
        logger.info("\nChat History:")
        for msg in result['chat_history']:
            logger.info(f"\n[{msg['role'].upper()}]")
            logger.info(msg['content'][:500] + "..." if len(msg['content']) > 500 else msg['content'])
        
        # 파일로 저장
        if args.output:
            with open(args.output, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
            logger.info(f"\nResult saved to: {args.output}")
    
    elif args.mode == "batch":
        # 배치 처리 (questions.jsonl이 기본값)
        processor = DataProcessor(args.config)
        stats = processor.process_batch(
            input_file=args.input,
            output_file=args.output,
            tools=tools,  # None이면 각 데이터의 total_tools 사용
            num_workers=args.num_workers,
            max_items=args.max_items
        )
        
        # 배치 처리 통계 출력
        logger.info("\n" + "="*50)
        logger.info("Batch Processing Complete")
        logger.info(f"Total processed: {stats['total_processed']}")
        logger.info(f"Success: {stats['total_success']}")
        logger.info(f"Failed: {stats['total_failed']}")
        if stats.get('end_time') and stats.get('start_time'):
            duration = (stats['end_time'] - stats['start_time']).total_seconds()
            logger.info(f"Total duration: {duration:.2f} seconds")
            logger.info(f"Average time per item: {duration/max(stats['total_processed'], 1):.2f} seconds")
        
    elif args.mode == "directory":
        # 디렉토리 처리
        if not args.input or args.input == "data/questions.jsonl":
            # 기본값이면 data 디렉토리 사용
            input_dir = "data"
        else:
            input_dir = args.input
            
        processor = DataProcessor(args.config)
        processor.process_directory(
            input_dir=input_dir,
            output_dir=args.output,
            pattern=args.pattern,
            tools=tools,  # None이면 각 데이터의 total_tools 사용
            num_workers=args.num_workers
        )


if __name__ == "__main__":
    main()