main.py

#!/usr/bin/env python3
import os
import sys
import json
import argparse
from pathlib import Path
from dotenv import load_dotenv

load_dotenv()

class AIProcessor:
    """Base processor class for AI text operations."""
    
    def __init__(self, provider='openai'):
        self.provider = provider
        self.api_key = self._get_api_key()
        
    def _get_api_key(self):
        if self.provider == 'openai':
            return os.getenv('OPENAI_API_KEY')
        elif self.provider == 'anthropic':
            return os.getenv('ANTHROPIC_API_KEY')
        return None
    
    def process(self, text, command, **kwargs):
        raise NotImplementedError("Subclasses must implement process()")


class OpenAIProcessor(AIProcessor):
    def process(self, text, command, **kwargs):
        try:
            import openai
            client = openai.OpenAI(api_key=self.api_key)
            
            prompts = {
                'summarize': f'Summarize the following text concisely:\n\n{text}',
                'translate': f'Translate the following text to {kwargs.get("target", "English")}:\n\n{text}',
                'sentiment': f'Analyze the sentiment of the following text. Return only: positive, negative, or neutral.\n\n{text}',
                'keywords': f'Extract main keywords from the following text. Return as comma-separated list:\n\n{text}',
                'fix': f'Fix grammar and improve readability of the following text:\n\n{text}',
            }
            
            if command not in prompts:
                raise ValueError(f"Unknown command: {command}")
            
            response = client.chat.completions.create(
                model=kwargs.get('model', 'gpt-3.5-turbo'),
                messages=[{'role': 'user', 'content': prompts[command]}],
                max_tokens=kwargs.get('max_tokens', 500),
            )
            
            return response.choices[0].message.content.strip()
            
        except ImportError:
            print("Error: openai package not installed. Run: pip install openai")
            sys.exit(1)


class AnthropicProcessor(AIProcessor):
    def process(self, text, command, **kwargs):
        try:
            import anthropic
            client = anthropic.Anthropic(api_key=self.api_key)
            
            prompts = {
                'summarize': f'Summarize the following text concisely:\n\n{text}',
                'translate': f'Translate the following text to {kwargs.get("target", "English")}:\n\n{text}',
                'sentiment': f'Analyze the sentiment of the following text. Return only: positive, negative, or neutral.\n\n{text}',
                'keywords': f'Extract main keywords from the following text. Return as comma-separated list:\n\n{text}',
                'fix': f'Fix grammar and improve readability of the following text:\n\n{text}',
            }
            
            if command not in prompts:
                raise ValueError(f"Unknown command: {command}")
            
            response = client.messages.create(
                model=kwargs.get('model', 'claude-3-sonnet-20240229'),
                max_tokens=kwargs.get('max_tokens', 500),
                messages=[{'role': 'user', 'content': prompts[command]}]
            )
            
            return response.content[0].text.strip()
            
        except ImportError:
            print("Error: anthropic package not installed. Run: pip install anthropic")
            sys.exit(1)


class LocalProcessor(AIProcessor):
    """Processor for local models using Ollama or similar."""
    
    def process(self, text, command, **kwargs):
        try:
            import requests
            
            prompts = {
                'summarize': f'Summarize the following text concisely:\n\n{text}',
                'translate': f'Translate the following text to {kwargs.get("target", "English")}:\n\n{text}',
                'sentiment': f'Analyze the sentiment of the following text. Return only: positive, negative, or neutral.\n\n{text}',
                'keywords': f'Extract main keywords from the following text. Return as comma-separated list:\n\n{text}',
                'fix': f'Fix grammar and improve readability of the following text:\n\n{text}',
            }
            
            if command not in prompts:
                raise ValueError(f"Unknown command: {command}")
            
            response = requests.post(
                'http://localhost:11434/api/generate',
                json={
                    'model': kwargs.get('model', 'llama2'),
                    'prompt': prompts[command],
                    'stream': False
                },
                timeout=kwargs.get('timeout', 60)
            )
            
            if response.status_code == 200:
                return response.json().get('response', '').strip()
            else:
                raise Exception(f"API request failed: {response.status_code}")
                
        except ImportError:
            print("Error: requests package not installed. Run: pip install requests")
            sys.exit(1)
        except requests.exceptions.ConnectionError:
            print("Error: Cannot connect to local model server. Make sure Ollama is running.")
            sys.exit(1)


def get_processor(provider):
    if provider == 'openai':
        return OpenAIProcessor(provider)
    elif provider == 'anthropic':
        return AnthropicProcessor(provider)
    elif provider == 'local':
        return LocalProcessor(provider)
    else:
        raise ValueError(f"Unknown provider: {provider}")


def read_text_from_file(file_path):
    path = Path(file_path)
    if not path.exists():
        raise FileNotFoundError(f"File not found: {file_path}")
    return path.read_text(encoding='utf-8')


def batch_process(file_list, command, provider, **kwargs):
    """Process multiple files."""
    processor = get_processor(provider)
    results = {}
    
    for file_path in file_list:
        try:
            text = read_text_from_file(file_path)
            result = processor.process(text, command, **kwargs)
            results[file_path] = result
            print(f"✓ {file_path}")
        except Exception as e:
            results[file_path] = f"Error: {e}"
            print(f"✗ {file_path}: {e}")
    
    return results


def main():
    parser = argparse.ArgumentParser(description='AI Text Processor')
    subparsers = parser.add_subparsers(dest='command', required=True,
                                       help='Processing command')
    
    # Common arguments
    parser.add_argument('--provider', '-p', default='openai',
                       choices=['openai', 'anthropic', 'local'],
                       help='AI provider')
    parser.add_argument('--model', '-m',
                       help='Model name (e.g., gpt-4, claude-3-opus)')
    
    # Summarize command
    sum_parser = subparsers.add_parser('summarize', help='Summarize text')
    sum_parser.add_argument('--file', '-f', help='Input file')
    sum_parser.add_argument('--text', '-t', help='Input text directly')
    sum_parser.add_argument('--max-tokens', type=int, default=500)
    
    # Translate command
    trans_parser = subparsers.add_parser('translate', help='Translate text')
    trans_parser.add_argument('--file', '-f', help='Input file')
    trans_parser.add_argument('--text', '-t', help='Input text directly')
    trans_parser.add_argument('--target', required=True,
                            help='Target language (e.g., zh, en, es)')
    trans_parser.add_argument('--max-tokens', type=int, default=500)
    
    # Sentiment command
    sent_parser = subparsers.add_parser('sentiment', help='Analyze sentiment')
    sent_parser.add_argument('--file', '-f', help='Input file')
    sent_parser.add_argument('--text', '-t', help='Input text directly')
    
    # Keywords command
    key_parser = subparsers.add_parser('keywords', help='Extract keywords')
    key_parser.add_argument('--file', '-f', help='Input file')
    key_parser.add_argument('--text', '-t', help='Input text directly')
    
    # Fix command
    fix_parser = subparsers.add_parser('fix', help='Fix grammar')
    fix_parser.add_argument('--file', '-f', help='Input file')
    fix_parser.add_argument('--text', '-t', help='Input text directly')
    fix_parser.add_argument('--max-tokens', type=int, default=500)
    
    # Batch command
    batch_parser = subparsers.add_parser('batch', help='Process multiple files')
    batch_parser.add_argument('--files', '-f', nargs='+', required=True,
                             help='List of files to process')
    batch_parser.add_argument('--output', '-o', help='Output JSON file')
    batch_parser.add_argument('--max-tokens', type=int, default=500)
    
    args = parser.parse_args()
    
    # Get processor
    processor = get_processor(args.provider)
    
    # Process input
    if args.command == 'batch':
        results = batch_process(
            args.files, 'summarize', args.provider,
            model=args.model, max_tokens=args.max_tokens
        )
        
        if args.output:
            with open(args.output, 'w', encoding='utf-8') as f:
                json.dump(results, f, indent=2, ensure_ascii=False)
            print(f"\nResults saved to {args.output}")
        else:
            print("\n" + json.dumps(results, indent=2, ensure_ascii=False))
    else:
        # Get text from file or command line
        if args.file:
            text = read_text_from_file(args.file)
        elif args.text:
            text = args.text
        else:
            print("Error: Please provide --file or --text")
            sys.exit(1)
        
        # Process text
        kwargs = {
            'model': args.model,
            'max_tokens': getattr(args, 'max_tokens', 500),
            'target': getattr(args, 'target', None),
        }
        
        result = processor.process(text, args.command, **kwargs)
        print(f"\n{result}")


if __name__ == '__main__':
    main()