import re
import json
import logging
from django.core.management.base import BaseCommand
from bazi.models import Person
from django.utils import timezone
logger = logging.getLogger(__name__)
class Command(BaseCommand):
help = 'Fixes existing AI analysis records by extracting thinking content and properly formatting JSON'
def add_arguments(self, parser):
parser.add_argument(
'--dry-run',
action='store_true',
dest='dry_run',
help='Show what would be updated without actually making changes',
)
parser.add_argument(
'--person-id',
type=int,
help='Specify a specific Person ID to fix (optional)',
)
parser.add_argument(
'--verbose',
action='store_true',
dest='verbose',
help='Show detailed information about changes',
)
parser.add_argument(
'--preview',
action='store_true',
dest='preview',
help='Show detailed preview of the content before and after changes',
)
parser.add_argument(
'--limit',
type=int,
default=10,
help='Limit the number of records to process (default: 10, use 0 for all)',
)
def handle(self, *args, **options):
dry_run = options['dry_run']
person_id = options['person_id']
verbose = options['verbose']
preview = options['preview']
limit = options['limit']
if dry_run:
self.stdout.write(self.style.WARNING('Running in dry-run mode - no changes will be made'))
# Get the queryset of Person records with AI analysis
queryset = Person.objects.exclude(ai_analysis__isnull=True).exclude(ai_analysis={})
if person_id:
queryset = queryset.filter(id=person_id)
# Apply limit if specified
if limit > 0 and not person_id:
queryset = queryset[:limit]
self.stdout.write(f'Limiting to {limit} records')
self.stdout.write(f'Found {queryset.count()} Person records with AI analysis')
fixed_count = 0
error_count = 0
no_change_count = 0
for person in queryset:
try:
# Check if the record needs fixing (has thinking tags or triple backticks)
ai_analysis = person.ai_analysis
if not ai_analysis or not isinstance(ai_analysis, dict):
self.stdout.write(self.style.WARNING(f'Person {person.id}: Invalid AI analysis format, skipping'))
continue
bazi_analysis = ai_analysis.get('bazi_analysis', '')
raw_analysis = None
# If bazi_analysis is a string, it might need fixing
if isinstance(bazi_analysis, str):
raw_analysis = bazi_analysis
# If it's already a dict, convert it back to a string to process potential triple backticks/think tags
elif isinstance(bazi_analysis, dict):
raw_analysis = json.dumps(bazi_analysis)
else:
self.stdout.write(self.style.WARNING(f'Person {person.id}: Unexpected bazi_analysis type: {type(bazi_analysis)}'))
continue
# Save original for preview
original_analysis = raw_analysis
# Initialize analysis and thinking content
analysis = raw_analysis.strip()
thinking_content = None
needs_update = False
# Extract thinking content if present
thinking_match = re.search(r'(.*?)', analysis, re.DOTALL)
if thinking_match:
thinking_content = thinking_match.group(1).strip()
# Remove the thinking section from the analysis
analysis = re.sub(r'.*?', '', analysis, flags=re.DOTALL).strip()
needs_update = True
# Extract content from triple backticks if present
backticks_match = re.search(r'```(?:json)?(.*?)```', analysis, re.DOTALL)
if backticks_match:
# Get the content inside backticks
backticks_content = backticks_match.group(1).strip()
analysis = backticks_content
needs_update = True
# If the analysis is now a string and it looks like JSON, try to parse it
if isinstance(analysis, str):
try:
json_data = json.loads(analysis)
# If successful parsing, use the parsed JSON object
analysis = json_data
needs_update = True
except json.JSONDecodeError:
# If not valid JSON, keep as string
if verbose:
self.stdout.write(self.style.WARNING(f'Person {person.id}: Could not parse as JSON'))
# If no changes are needed, skip this record
if not needs_update:
no_change_count += 1
if verbose:
self.stdout.write(f'Person {person.id}: No changes needed')
continue
# Create the updated ai_analysis dictionary
updated_ai_analysis = ai_analysis.copy()
updated_ai_analysis['bazi_analysis'] = analysis
# Add thinking content if it was found
if thinking_content:
updated_ai_analysis['think'] = thinking_content
# Show preview if requested
if preview:
self.stdout.write(self.style.NOTICE(f'\nPerson {person.id} Preview:'))
self.stdout.write(self.style.NOTICE('=' * 80))
self.stdout.write(self.style.WARNING('BEFORE:'))
# Truncate very long content for readability
before_preview = original_analysis
if len(before_preview) > 500:
before_preview = before_preview[:250] + '...' + before_preview[-250:]
self.stdout.write(before_preview)
self.stdout.write('\n' + self.style.SUCCESS('AFTER:'))
# Show the thinking content if extracted
if thinking_content:
thinking_preview = thinking_content
if len(thinking_preview) > 200:
thinking_preview = thinking_preview[:100] + '...' + thinking_preview[-100:]
self.stdout.write(self.style.NOTICE('Thinking Content:'))
self.stdout.write(thinking_preview)
self.stdout.write('')
# Show the cleaned analysis content
self.stdout.write(self.style.NOTICE('Analysis Content:'))
if isinstance(analysis, dict):
# Pretty print JSON for readability
after_preview = json.dumps(analysis, ensure_ascii=False, indent=2)
else:
after_preview = str(analysis)
# Truncate very long content for readability
if len(after_preview) > 500:
after_preview = after_preview[:250] + '...' + after_preview[-250:]
self.stdout.write(after_preview)
self.stdout.write(self.style.NOTICE('=' * 80))
# Show the changes if in verbose mode
elif verbose:
self.stdout.write(f'Person {person.id}:')
if thinking_content:
self.stdout.write(f' - Extracted thinking content ({len(thinking_content)} chars)')
self.stdout.write(f' - Updated bazi_analysis format')
# Apply the changes unless in dry-run mode
if not dry_run:
person.ai_analysis = updated_ai_analysis
person.analysis_timestamp = timezone.now()
person.save(update_fields=['ai_analysis', 'analysis_timestamp'])
fixed_count += 1
except Exception as e:
error_count += 1
self.stdout.write(self.style.ERROR(f'Error processing Person {person.id}: {str(e)}'))
self.stdout.write(self.style.SUCCESS(f'\nSummary:'))
self.stdout.write(f' - Records processed: {queryset.count()}')
self.stdout.write(f' - Records updated: {fixed_count}')
self.stdout.write(f' - Records unchanged: {no_change_count}')
self.stdout.write(f' - Errors: {error_count}')
if dry_run:
self.stdout.write(self.style.WARNING('This was a dry run - no changes were made'))