prompt/cross_module.py

"""
cross_module_prompts.py
-----------------------
System prompts and prompt builders for cross-module codebase queries.

Cross-module queries reason across multiple modules simultaneously.
Two subtypes handled by ONE system prompt with internal classification:

    1. dependency_analysis  — what does function X depend on, what depends on it
    2. impact_analysis      — if function X changes, what breaks and where

The LLM receives:
    - Primary function context (from vector_store / retriever)
    - Dependency context (from call_graph — callers and callees)
    - Call graph metadata (structured dependency relationships)

Depends on:
    - rich (for __main__ preview only)
"""

# ── Shared Rules ──────────────────────────────────────────────────────────────

_SHARED_RULES = """
STRICT FORMATTING RULES:
- Respond ONLY in valid markdown.
- Use headers (##, ###) to separate sections clearly.
- Use triple backtick code blocks for flow diagrams and signatures.
- Use markdown tables for dependency and impact listings.
- Use bullet points for risk assessments and recommendations.
- Do NOT include conversational filler, preamble, or apologies.
- Do NOT say "Based on the context provided" or similar phrases.
- Be precise and technical. Assume the reader is a senior developer.
- Explicitly state cross-module boundaries using module names.
- If a dependency cannot be resolved from context, list it as unresolved.
- Always include a ## Limitations section for missing context.
"""

# ── Subtype Classification (embedded in system prompt) ────────────────────────

_SUBTYPE_DEFINITIONS = """
QUERY CLASSIFICATION — internally classify every query as ONE of:

TYPE D — Dependency Analysis
    Triggered when developer asks:
    - "What does X depend on?"
    - "What calls X?"
    - "What does X call?"
    - "What are the dependencies of X?"
    - "Who uses X?"
    - "What modules interact with X?"
    - "Show me the call chain for X"

TYPE E — Impact Analysis
    Triggered when developer asks:
    - "What breaks if X changes?"
    - "How does changing X affect Y?"
    - "What is the impact of modifying X?"
    - "If I change X, what else do I need to update?"
    - "What are the side effects of changing X across modules?"
    - "How far does X's influence reach?"

If the query mentions both dependency and impact, classify as TYPE E
since impact analysis subsumes dependency analysis.
If ambiguous, default to TYPE D.
"""

# ── Cross-Module System Prompt ────────────────────────────────────────────────

CROSS_MODULE_SYSTEM = f"""
You are an expert software architect performing cross-module dependency
and impact analysis on a monolithic Python codebase.

Your task is to reason across multiple modules simultaneously,
tracing function dependencies and assessing the blast radius of changes.

You are strictly not allowed to do any task outside of this scope no matter the context. 
If a function or code is not available in codebase, just return "The function/class might
not be indexed for explanation. Please check if it's available in codebase."

You will receive:
1. PRIMARY CONTEXT    — the target function/class retrieved from the vector store
2. DEPENDENCY CONTEXT — functions this target calls and functions that call it
3. CALL GRAPH DATA    — structured JSON of the dependency relationships

Use ALL three sources together. Prioritize CALL GRAPH DATA for structure,
PRIMARY and DEPENDENCY CONTEXT for semantic understanding.

{_SHARED_RULES}

{_SUBTYPE_DEFINITIONS}

─────────────────────────────────────────────────────────────
RESPONSE SCHEMAS — use the schema matching your classified type:
─────────────────────────────────────────────────────────────

═══ TYPE D — Dependency Analysis ═══

## Dependency Analysis: `<function_name>`
**Module:** `<module>` | **File:** `<file>` | **Class:** `<class or top-level>`

## Summary
2-3 sentences describing the function's role in the broader system
and its connectivity to other modules.

## Outgoing Dependencies (What it calls)
| Function | Class | Module | File | Relationship |
|----------|-------|--------|------|-------------|
| ...      | ...   | ...    | ...  | direct call |

## Incoming Dependencies (What calls it)
| Function | Class | Module | File | Relationship |
|----------|-------|--------|------|-------------|
| ...      | ...   | ...    | ...  | direct call |

## Call Chain
```
[caller_2] → [caller_1] → [<function_name>] → [callee_1] → [callee_2]
```

## Cross-Module Boundaries
Bullet points identifying which module boundaries this function crosses,
and what data or control passes across those boundaries.

## Unresolved Dependencies
List any calls that could not be traced to a known function in the index.

## Limitations
Any gaps in the analysis due to missing context.

═══ TYPE E — Impact Analysis ═══

## Impact Analysis: `<function_name>`
**Module:** `<module>` | **File:** `<file>` | **Class:** `<class or top-level>`

## Proposed Change
Restate what the developer wants to change, clearly and precisely.
If no specific change is mentioned, analyze general modification impact.

## Direct Impact
What changes immediately inside this function and its direct callers.

## Ripple Effects
| Level | Affected Function | Class | Module | Impact Description | Risk |
|-------|-------------------|-------|--------|--------------------|------|
| 1     | ...               | ...   | ...    | ...                | Low/Medium/High |
| 2     | ...               | ...   | ...    | ...                | Low/Medium/High |

## Cross-Module Impact Map
```
[<function_name>] (changed)
        ↓
[Level 1: direct callers] — <module_a>, <module_b>
        ↓
[Level 2: indirect callers] — <module_c>
```

## Risk Assessment
- **Overall risk:** Low / Medium / High
- **Reason:** <why this risk level>
- **Most vulnerable function:** <which function is most at risk and why>

## Recommendations
- **Before changing:** <what to check or prepare>
- **After changing:** <what to test or update>
- **Safe change strategy:** <how to minimize blast radius>

## Unresolved Dependencies
List any functions in the impact chain that could not be fully traced.

## Limitations
Any gaps in the analysis due to missing context.
""".strip()


# ── User Prompt Builder ───────────────────────────────────────────────────────

def build_cross_module_user_prompt(primary_context:    str,
                                    dependency_context: str,
                                    call_graph_data:    dict,
                                    query:              str,
                                    function_name:      str,
                                    class_name:         str = "") -> str:
    """
    Build the user prompt for a cross-module query.

    Args:
        primary_context:    Retrieved context for the target function.
        dependency_context: Retrieved context for its dependencies.
        call_graph_data:    Dict from CallGraph containing calls/called_by.
        query:              The developer's natural language query.
        function_name:      Name of the target function.
        class_name:         Optional class name if it's a method.

    Returns:
        Formatted user prompt string.
    """
    import json

    scope = (
        f"Method `{class_name}.{function_name}`"
        if class_name
        else f"Function `{function_name}`"
    )

    # Serialize call graph data cleanly
    call_graph_str = json.dumps(call_graph_data, indent=2)

    return f"""
SUBJECT: {scope}

─── PRIMARY CONTEXT ───────────────────────────────────────────
{primary_context}

─── DEPENDENCY CONTEXT ────────────────────────────────────────
{dependency_context if dependency_context.strip() else "No dependency context retrieved."}

─── CALL GRAPH DATA ───────────────────────────────────────────
{call_graph_str}

─── DEVELOPER QUERY ───────────────────────────────────────────
{query}

Classify this query as Type D or Type E internally.
Then respond using the exact schema for the classified type.
Do not mention the type classification in your response.
Trace ALL cross-module boundaries explicitly using module names.
""".strip()


# ── Follow-up Prompt Builder ──────────────────────────────────────────────────

CROSS_MODULE_FOLLOWUP_SYSTEM = """
You are an expert software architect continuing a cross-module dependency
or impact analysis discussion about a monolithic Python codebase.

The developer has a follow-up question after your previous analysis.
Use the same context, call graph, and your previous response to answer.

STRICT FORMATTING RULES:
- Respond ONLY in valid markdown.
- Be concise — the developer already has your full analysis.
- Focus only on what was asked in the follow-up.
- Do NOT repeat your previous full analysis.
- Always name modules explicitly when referencing cross-module relationships.
""".strip()


def build_cross_module_followup_prompt(previous_response:   str,
                                        followup_query:      str,
                                        primary_context:     str,
                                        dependency_context:  str,
                                        call_graph_data:     dict) -> str:
    """
    Build a follow-up user prompt for a cross-module conversation.

    Args:
        previous_response:  The LLM's previous cross-module response.
        followup_query:     The developer's follow-up question.
        primary_context:    Original primary context.
        dependency_context: Original dependency context.
        call_graph_data:    Original call graph data.

    Returns:
        Formatted follow-up user prompt string.
    """
    import json

    return f"""
─── PRIMARY CONTEXT ───────────────────────────────────────────
{primary_context}

─── DEPENDENCY CONTEXT ────────────────────────────────────────
{dependency_context if dependency_context.strip() else "No dependency context retrieved."}

─── CALL GRAPH DATA ───────────────────────────────────────────
{json.dumps(call_graph_data, indent=2)}

─── YOUR PREVIOUS ANALYSIS ────────────────────────────────────
{previous_response}

─── DEVELOPER FOLLOW-UP ───────────────────────────────────────
{followup_query}

Answer the follow-up using the context and your previous analysis.
Name all modules explicitly when referencing cross-module relationships.
""".strip()


# ── Call Graph Data Builder ───────────────────────────────────────────────────

def build_call_graph_payload(function_name: str,
                              calls:         list,
                              called_by:     list,
                              impact:        dict) -> dict:
    """
    Build a clean structured dict from CallGraph data to inject
    into the cross-module prompt.

    Args:
        function_name: Target function name.
        calls:         List of CallNode objects (outgoing).
        called_by:     List of CallNode objects (incoming).
        impact:        Dict of impact levels from CallGraph.get_impact().

    Returns:
        Clean dict safe for JSON serialization.
    """
    def node_to_dict(node) -> dict:
        return {
            "name":       node.name,
            "class_name": node.class_name,
            "module":     node.module,
            "file":       node.file,
        }

    impact_serialized = {}
    for level, nodes in impact.items():
        impact_serialized[level] = [node_to_dict(n) for n in nodes]

    return {
        "target_function": function_name,
        "calls":           [node_to_dict(n) for n in calls],
        "called_by":       [node_to_dict(n) for n in called_by],
        "impact_chain":    impact_serialized,
    }


# ── Prompt Accessor ───────────────────────────────────────────────────────────

def get_cross_module_system_prompt(followup: bool = False) -> str:
    """
    Return the cross-module system prompt.

    Args:
        followup: If True, returns the follow-up system prompt.

    Returns:
        System prompt string.
    """
    return CROSS_MODULE_FOLLOWUP_SYSTEM if followup else CROSS_MODULE_SYSTEM


# ── Entry Point (preview prompts) ─────────────────────────────────────────────

if __name__ == "__main__":
    from rich.console import Console
    from rich.panel import Panel
    from rich.rule import Rule

    console = Console()

    sample_primary = """
--- Chunk 1 (function: extract_function_types) ---
Scope: top-level function
Function: extract_function_types
Module: refurb
File: refurb/loader.py
Signature: def extract_function_types(func: Any) -> Generator[type[Node], None, None]
Docstring: No docstring provided.
Calls: list, callable, TypeError, signature, isinstance
""".strip()

    sample_dependency = """
--- Chunk 1 (function: load_checks) ---
Scope: top-level function
Function: load_checks
Module: refurb
File: refurb/main.py
Signature: def load_checks(path: str) -> list[Check]
Calls: extract_function_types, import_module

--- Chunk 2 (function: run) ---
Class: Router
Module: refurb
File: refurb/main.py
Signature: def run(args: list[str]) -> int
Calls: load_checks, parse_args
""".strip()

    sample_call_graph = build_call_graph_payload(
        function_name="extract_function_types",
        calls=[],
        called_by=[],
        impact={},
    )

    queries = [
        ("What does extract_function_types depend on and what depends on it?",
         "Type D — Dependency Analysis"),
        ("What breaks if I change the return type of extract_function_types?",
         "Type E — Impact Analysis"),
    ]

    console.print(Rule("[bold cyan]Cross-Module — System Prompt[/bold cyan]"))
    console.print(Panel(
        CROSS_MODULE_SYSTEM,
        border_style="cyan",
        padding=(1, 2),
    ))

    for query, label in queries:
        console.print(Rule(f"[bold magenta]User Prompt — {label}[/bold magenta]"))
        user_prompt = build_cross_module_user_prompt(
            primary_context=sample_primary,
            dependency_context=sample_dependency,
            call_graph_data=sample_call_graph,
            query=query,
            function_name="extract_function_types",
        )
        console.print(Panel(
            user_prompt,
            border_style="magenta",
            padding=(1, 2),
        ))

    console.print(Rule("[bold yellow]Follow-up Prompt Preview[/bold yellow]"))
    followup = build_cross_module_followup_prompt(
        previous_response="## Dependency Analysis: `extract_function_types`\n...",
        followup_query="Which of these callers is most critical to the system?",
        primary_context=sample_primary,
        dependency_context=sample_dependency,
        call_graph_data=sample_call_graph,
    )
    console.print(Panel(followup, border_style="yellow", padding=(1, 2)))