src/main.py

import argparse
import json
from typing import Dict, Any, Optional

from src.agents.legal_agent import LegalAgent

def parse_arguments():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="Agentic Defensor: Legal RAG System")
    parser.add_argument("query", type=str, help="The legal query to process")
    parser.add_argument("--top-k", type=int, default=None, help="Number of chunks to retrieve")
    parser.add_argument("--model", type=str, default=None, help="OpenAI model to use")
    parser.add_argument("--output", type=str, default=None, help="Output file path for saving the response")
    parser.add_argument("--verbose", action="store_true", help="Print detailed information")
    
    return parser.parse_args()

def process_query(query: str, top_k: Optional[int] = None, model: Optional[str] = None) -> Dict[str, Any]:
    """
    Process a query using the legal agent.
    
    Args:
        query: The query text
        top_k: Number of chunks to retrieve (optional)
        model: OpenAI model to use (optional)
        
    Returns:
        Dictionary containing the result
    """
    # Initialize the agent with the specified model if provided
    agent = LegalAgent(model=model) if model else LegalAgent()
    
    # Process the query
    return agent.answer_query(query, top_k)

def save_result(result: Dict[str, Any], output_path: str) -> None:
    """
    Save the result to a JSON file.
    
    Args:
        result: The result dictionary
        output_path: Path to save the JSON file
    """
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    print(f"Result saved to {output_path}")

def main():
    """Main entry point."""
    args = parse_arguments()
    
    # Process the query
    print(f"Processing query: {args.query}")
    result = process_query(args.query, args.top_k, args.model)
    
    # Print the answer
    print("\n--- Answer ---")
    print(result["answer"])
    
    # Print additional information if verbose
    if args.verbose:
        print("\n--- Query Information ---")
        print(f"Model used: {result['model_used']}")
        print(f"Retrieved chunks: {len(result['retrieved_chunks'])}")
    
    # Save the result if output path is provided
    if args.output:
        save_result(result, args.output)

if __name__ == "__main__":
    main()