Blame for knowledge_extractor.py

"""
Knowledge Base Extractor for Astrology Aspects
Extracts professional aspect interpretations from markdown knowledge base
"""

import re
from pathlib import Path
from typing import Dict, List, Optional

# Knowledge base directory (relative to this file)
KNOWLEDGE_BASE_DIR = Path(__file__).parent / "knowledge-base" / "aspects"

# Aspect type mapping
ASPECT_SYMBOLS = {
    "conjunction": "☌",
    "sextile": "⚹",
    "square": "□",
    "trine": "△",
    "opposition": "☍"
}

# Planet symbols
PLANET_SYMBOLS = {
    "Sun": "☉",
    "Moon": "☽",
    "Mercury": "☿",
    "Venus": "♀",
    "Mars": "♂",
    "Jupiter": "♃",
    "Saturn": "♄",
    "Uranus": "♅",
    "Neptune": "♆",
    "Pluto": "♇"
}


def normalize_aspect_name(planet1: str, planet2: str, aspect: str) -> str:
    """
    Normalize aspect name to match knowledge base format.

    Args:
        planet1: First planet name (e.g., "Sun")
        planet2: Second planet name (e.g., "Moon")
        aspect: Aspect type (e.g., "trine")

    Returns:
        Normalized aspect name (e.g., "Sun Trine Moon")
    """
    # Capitalize aspect type
    aspect_capitalized = aspect.capitalize()

    # Format: "Planet1 Aspect Planet2"
    return f"{planet1} {aspect_capitalized} {planet2}"


def get_aspect_filename(planet: str) -> str:
    """
    Get the knowledge base filename for a planet.

    Args:
        planet: Planet name (e.g., "Sun", "Moon")

    Returns:
        Filename (e.g., "sun-aspects.md")
    """
    return f"{planet.lower()}-aspects.md"


def extract_aspect_interpretation(planet1: str, planet2: str, aspect: str) -> Optional[str]:
    """
    Extract aspect interpretation from knowledge base.

    Args:
        planet1: First planet name
        planet2: Second planet name
        aspect: Aspect type

    Returns:
        Full aspect interpretation text, or None if not found
    """
    # Determine which file to read (aspects are stored by the first planet)
    # For aspects like Sun-Moon, we read sun-aspects.md
    # For aspects like Moon-Venus, we read moon-aspects.md

    knowledge_file = KNOWLEDGE_BASE_DIR / get_aspect_filename(planet1)

    if not knowledge_file.exists():
        print(f"Knowledge file not found: {knowledge_file}")
        return None

    # Read the knowledge base file
    try:
        with open(knowledge_file, 'r', encoding='utf-8') as f:
            content = f.read()
    except Exception as e:
        print(f"Error reading {knowledge_file}: {e}")
        return None

    # Create the aspect header to search for
    # Format in knowledge base: "## Sun Trine Moon ☉△☽"
    aspect_name = normalize_aspect_name(planet1, planet2, aspect)

    # Get symbols
    aspect_symbol = ASPECT_SYMBOLS.get(aspect.lower(), "")
    planet1_symbol = PLANET_SYMBOLS.get(planet1, "")
    planet2_symbol = PLANET_SYMBOLS.get(planet2, "")

    # Try to find the aspect section with symbols
    header_pattern = f"## {aspect_name} {planet1_symbol}{aspect_symbol}{planet2_symbol}"

    # Alternative: try without symbols
    if header_pattern not in content:
        header_pattern = f"## {aspect_name}"

    # Find the section
    if header_pattern not in content:
        print(f"Aspect not found in knowledge base: {aspect_name}")
        return None

    # Extract the section (from this header to the next ## or end of file)
    # Split by ## headers
    sections = re.split(r'\n## ', content)

    target_section = None
    for section in sections:
        if section.startswith(aspect_name):
            target_section = "## " + section
            break

    if target_section:
        # Remove the next section if it exists
        next_header = re.search(r'\n## ', target_section[4:])  # Skip first "## "
        if next_header:
            target_section = target_section[:next_header.start() + 4]

        return target_section.strip()

    return None


def extract_knowledge_for_aspects(aspects: List[Dict]) -> str:
    """
    Extract knowledge base context for multiple aspects.

    Args:
        aspects: List of aspect dictionaries from API
                 Example: [{"planet1": "Sun", "planet2": "Moon", "aspect": "trine", "orb": 2.5}, ...]

    Returns:
        Combined knowledge context string ready for AI prompt
    """
    knowledge_sections = []

    for aspect_data in aspects:
        planet1 = aspect_data.get("planet1")
        planet2 = aspect_data.get("planet2")
        aspect_type = aspect_data.get("aspect")
        orb = aspect_data.get("orb", 0)

        # Extract interpretation
        interpretation = extract_aspect_interpretation(planet1, planet2, aspect_type)

        if interpretation:
            # Add orb information
            aspect_with_orb = f"{interpretation}\n\n**Orb in your chart**: {orb:.1f}°"
            knowledge_sections.append(aspect_with_orb)
        else:
            # If not found, add a placeholder
            aspect_name = normalize_aspect_name(planet1, planet2, aspect_type)
            knowledge_sections.append(
                f"## {aspect_name}\n\n"
                f"*Knowledge base entry not yet available for this aspect.*\n"
                f"Orb: {orb:.1f}°"
            )

    # Combine all sections
    if knowledge_sections:
        header = "# Professional Aspect Interpretations\n\n"
        header += "The following interpretations are based on classical and modern Western astrology texts.\n\n"
        header += "---\n\n"

        return header + "\n\n---\n\n".join(knowledge_sections)

    return "No aspect interpretations available."


def get_aspect_summary(aspects: List[Dict], limit: int = 5) -> str:
    """
    Get a summary of the most important aspects (by smallest orb).

    Args:
        aspects: List of aspect dictionaries
        limit: Maximum number of aspects to include

    Returns:
        Summary string with top aspects
    """
    # Sort by orb (smallest = strongest)
    sorted_aspects = sorted(aspects, key=lambda x: abs(x.get("orb", 99)))

    # Take top N
    top_aspects = sorted_aspects[:limit]

    summary_lines = [f"**Top {len(top_aspects)} Strongest Aspects** (by orb):\n"]

    for i, aspect in enumerate(top_aspects, 1):
        planet1 = aspect.get("planet1")
        planet2 = aspect.get("planet2")
        aspect_type = aspect.get("aspect", "").capitalize()
        orb = aspect.get("orb", 0)

        summary_lines.append(
            f"{i}. **{planet1} {aspect_type} {planet2}** (orb: {orb:.1f}°)"
        )

    return "\n".join(summary_lines)


# Example usage and testing
if __name__ == "__main__":
    # Test with sample aspect data
    test_aspects = [
        {"planet1": "Sun", "planet2": "Moon", "aspect": "trine", "orb": 2.5},
        {"planet1": "Venus", "planet2": "Mars", "aspect": "square", "orb": 3.1},
        {"planet1": "Moon", "planet2": "Mercury", "aspect": "conjunction", "orb": 1.2}
    ]

    print("=" * 60)
    print("TESTING KNOWLEDGE EXTRACTOR")
    print("=" * 60)

    # Test single extraction
    print("\n1. Testing single aspect extraction:")
    print("-" * 60)
    result = extract_aspect_interpretation("Sun", "Moon", "trine")
    if result:
        print(f"Found! Length: {len(result)} characters")
        print(result[:200] + "...")
    else:
        print("Not found")

    # Test multiple aspects
    print("\n\n2. Testing multiple aspects extraction:")
    print("-" * 60)
    knowledge_context = extract_knowledge_for_aspects(test_aspects)
    print(f"Total context length: {len(knowledge_context)} characters")
    print("\nFirst 500 characters:")
    print(knowledge_context[:500] + "...")

    # Test summary
    print("\n\n3. Testing aspect summary:")
    print("-" * 60)
    summary = get_aspect_summary(test_aspects)
    print(summary)