src/report_generator.py

# src/report_generator.py

from typing import Dict, List
import json


SYSTEM_PROMPT = """You are a clinical AI assistant specialized in radiological reporting of kidney stones.
Your role is to generate preliminary AI-assisted reports based on automated detection results, grounded in provided clinical guidelines.

CRITICAL RULES:

1. GROUNDING: Only make clinical claims that are directly supported by the provided clinical guideline sources. Do not introduce medical information from your general knowledge.

2. CITATIONS: When making clinical recommendations or interpretations, reference the source guidelines (e.g., "per EAU Guidelines 2026" or "based on AUA recommendations").

3. CAUTIOUS LANGUAGE: Use appropriate radiological hedging:
   - "Hyperdense focus suggestive of..." instead of "Stone is..."
   - "Management options may include..." instead of "Treatment is..."
   - "Morphology pattern suggests..." instead of "Shape is..."

4. UNCERTAINTY: Clearly state limitations:
   - Size measurements are approximate (pixel-based estimation, assumed spacing 0.7 mm/px)
   - Equivalent diameter is derived from segmentation mask area, not direct caliper measurement
   - Density values are relative pixel intensities (0-255 PNG scale), not true Hounsfield Units
   - Morphology metrics (circularity, solidity) are image-based approximations
   - AI detection is not a substitute for radiologist interpretation

5. STRUCTURE: Follow the exact report format provided.

6. NO FABRICATION: If the provided sources do not cover a specific clinical question, acknowledge this rather than inventing information. In the REFERENCES section, list ONLY the exact source names shown in the [Source N] headers of the CLINICAL KNOWLEDGE CONTEXT. Do not add, modify, or invent any source names.

7. PROFESSIONAL TONE: Use formal clinical language appropriate for a radiology report.

8. COMPLETION: You MUST complete ALL 6 report sections without exception. Never end a response mid-sentence. If guideline context is insufficient for a section, write "Insufficient guideline context available" rather than truncating."""


REPORT_TEMPLATE = """Generate a clinical AI-assisted kidney stone detection report using the structure below.

=== DETECTION RESULTS (Automated Analysis) ===
{detection_json}

=== CLINICAL KNOWLEDGE CONTEXT (Retrieved from Guidelines) ===
{retrieved_context}

=== REPORT STRUCTURE ===

Generate the report with the following sections:

**1. FINDINGS**
- Number of detected hyperdense foci suggestive of urolithiasis
- For each stone: equivalent diameter (mm), estimated area (mm²), image quadrant location
- If multiple stones: total stone area and largest stone size

**2. IMAGE CHARACTERISTICS**
- For each stone: size category, shape category, relative density category
- Detection confidence score
- If morphology available: circularity score, solidity, eccentricity, orientation angle
- Density profile: mean pixel intensity, intensity homogeneity, contrast to surrounding tissue

**3. CLINICAL INTERPRETATION**
- Size-based assessment of spontaneous passage likelihood (reference EAU/AUA size thresholds)
- Shape-based passage prognosis: irregular or low-solidity stones may indicate increased obstruction risk
- Density homogeneity assessment: note relative density pattern (cautiously, as true HU not available)
- Reference specific guideline sources for each clinical claim

**4. MANAGEMENT CONSIDERATIONS**
- Treatment options based on size category, per retrieved guidelines
- Medical expulsive therapy (MET) eligibility if applicable
- Surgical options (ESWL/URS/PCNL) if applicable, noting morphology factors where relevant
- Always cite the source guideline

**5. LIMITATIONS & DISCLAIMERS**
- This is an AI-generated preliminary assessment only
- Size is derived from segmentation mask equivalent diameter (pixel-based, assumed 0.7 mm/px spacing)
- Density values are relative pixel intensities from PNG images, not true Hounsfield Units
- Morphology metrics are image-based approximations, not direct anatomical measurements
- Clinical correlation and radiologist review required
- Not a substitute for professional medical interpretation

**6. REFERENCES**
- List unique sources cited in the report

IMPORTANT: Complete all 6 sections in full before ending your response.

Generate the report now:"""


NO_STONE_PROMPT = """Generate a brief AI-assisted report stating that no kidney stones were detected in the analyzed image.

Include:
1. FINDINGS: State that no hyperdense foci suggestive of urolithiasis were identified
2. LIMITATIONS: AI detection has inherent false negative risk; clinical correlation needed
3. RECOMMENDATION: If clinical suspicion persists despite negative AI findings, further imaging or clinical evaluation is warranted

Keep the report concise but professional."""


def format_retrieved_context(retrieved_docs: List[Dict]) -> str:
    """Retrieved chunk'ları LLM için format'la."""
    if not retrieved_docs:
        return "No specific guideline context available."

    MAX_CONTENT_CHARS = 900
    formatted = []
    for i, doc in enumerate(retrieved_docs, 1):
        content = doc['content']
        if len(content) > MAX_CONTENT_CHARS:
            content = content[:MAX_CONTENT_CHARS] + "..."
        entry = (
            f"[Source {i}] {doc['source']} (Page {doc['page']})\n"
            f"Relevance Score: {doc['similarity_score']:.3f}\n"
            f"Content: {content}\n"
        )
        formatted.append(entry)

    return "\n---\n".join(formatted)


def format_detection_features(features: Dict) -> str:
    """Detection/segmentation features'ı LLM için temiz JSON olarak format'la."""
    clean_features = {
        "stone_detected": features.get("stone_detected"),
        "count": features.get("count"),
        "multiple_stones": features.get("multiple_stones", False),
    }

    if features.get("largest_stone_mm"):
        clean_features["largest_stone_mm"] = features["largest_stone_mm"]

    if features.get("total_stone_area_mm2"):
        clean_features["total_stone_area_mm2"] = features["total_stone_area_mm2"]

    if features.get("stones"):
        clean_features["stones"] = []
        for stone in features["stones"]:
            clean_stone = {
                "stone_id": stone["stone_id"],
                "size": stone["size"],
                "location_quadrant": stone["location"]["quadrant"],
                "shape": stone["shape"],
                "density": stone["density"],
                "detection_confidence": stone["detection_confidence"]
            }
            if stone.get("morphology"):
                clean_stone["morphology"] = stone["morphology"]
            clean_features["stones"].append(clean_stone)

    return json.dumps(clean_features, indent=2)


def build_report_prompt(
    features: Dict,
    retrieved_docs: List[Dict]
) -> tuple[str, str]:
    """
    Complete prompt'u inşa et.
    Returns: (system_prompt, user_prompt) tuple
    """
    if not features.get("stone_detected", False):
        return SYSTEM_PROMPT, NO_STONE_PROMPT

    detection_json = format_detection_features(features)
    retrieved_context = format_retrieved_context(retrieved_docs)

    user_prompt = REPORT_TEMPLATE.format(
        detection_json=detection_json,
        retrieved_context=retrieved_context
    )

    return SYSTEM_PROMPT, user_prompt


if __name__ == "__main__":
    mock_features = {
        "stone_detected": True,
        "count": 1,
        "stones": [{
            "stone_id": 1,
            "size": {
                "equivalent_diameter_mm": 7.2,
                "max_dimension_mm": 7.2,
                "area_mm2": 40.7,
                "size_category": "medium"
            },
            "location": {"quadrant": "upper-left"},
            "shape": "irregular",
            "density": {
                "category": "high",
                "mean_intensity": 210.4,
                "std_intensity": 18.2,
                "homogeneity": "homogeneous",
                "contrast_to_surrounding": 45.3
            },
            "detection_confidence": 0.87,
            "morphology": {
                "area_px": 831.2,
                "area_mm2": 40.7,
                "perimeter_px": 112.4,
                "circularity": 0.826,
                "solidity": 0.91,
                "eccentricity": 0.44,
                "orientation_deg": 32.1
            }
        }],
        "multiple_stones": False
    }

    mock_retrieved = [
        {
            "source": "EAU-Guidelines-2026.pdf",
            "page": 17,
            "similarity_score": 0.61,
            "content": "α-blockers are recommended for 5-10mm distal ureter stones..."
        }
    ]

    system, user = build_report_prompt(mock_features, mock_retrieved)
    print("=== SYSTEM PROMPT ===")
    print(system)
    print("\n=== USER PROMPT ===")
    print(user)