prompt = f"""You are improving the '{section_name}' section for an investment memo about {company_name}.

INVESTMENT TYPE: {investment_type.upper()}
MEMO MODE: {memo_mode.upper()}

# Build disambiguation block if we have identifying info
disambiguation_context = ""
if company_description or company_url:
    disambiguation_context = f"""
CRITICAL - ENTITY DISAMBIGUATION:
There may be multiple companies named "{company_name}". You MUST research the CORRECT company:
...
DISAMBIGUATION RULES:
1. ONLY use sources that reference THIS specific company
2. If you find data for a DIFFERENT company with the same name, DISCARD IT
...
"""

prompt = f"""You are improving the '{section_name}' section for an investment memo about {company_name}.

INVESTMENT TYPE: {investment_type.upper()}
MEMO MODE: {memo_mode.upper()}
{disambiguation_context}

def search(
    self,
    query: str,
    max_results: int = 10,
    sources: Optional[List[str]] = None,
    disambiguation_context: Optional[str] = None  # NEW
) -> List[Dict[str, str]]:

system_prompt = "You are a research assistant providing detailed, cited information..."

if disambiguation_context:
    system_prompt += f"""

CRITICAL DISAMBIGUATION:
{disambiguation_context}

ONLY return information about the CORRECT company. If search results contain
data about a DIFFERENT company with a similar name, DISCARD that data and
state that you could not find information for the correct entity."""

disambiguation_context = None
if company_description or company_url:
    disambiguation_parts = [f"Target company: {company_name}"]
    if company_description:
        disambiguation_parts.append(f"Description: {company_description}")
    if company_url:
        parsed = urlparse(company_url)
        domain = parsed.netloc if parsed.netloc else company_url
        disambiguation_parts.append(f"Website: {company_url}")
        disambiguation_parts.append(f"Domain: {domain}")
    if research_notes:
        disambiguation_parts.append(f"Notes: {research_notes}")
    disambiguation_context = "\n".join(disambiguation_parts)

if isinstance(search_provider, PerplexityProvider):
    results = search_provider.search(
        query,
        max_results=5 if idx > 1 else max_results,
        sources=research_sources,
        disambiguation_context=disambiguation_context  # NEW
    )

def build_section_research_query(
    section_def: Any,
    company_name: str,
    company_description: str,
    general_research: Dict[str, Any],
    memo_mode: str,
    deck_draft_content: str = "",
    company_url: str = "",        # NEW
    research_notes: str = ""      # NEW
) -> str:

disambiguation_block = ""
if company_url or research_notes:
    disambiguation_block = f"""
CRITICAL - ENTITY DISAMBIGUATION:
There may be multiple companies named "{company_name}". You MUST research the CORRECT company:
- Company website: {company_url or 'See description'}
- Description: {company_description}
"""
    if research_notes:
        disambiguation_block += f"- Research notes: {research_notes}\n"

    disambiguation_block += """
DISAMBIGUATION RULES:
1. ONLY use sources that reference THIS specific company
2. If you find funding/revenue data for a DIFFERENT company with the same name, DISCARD IT
3. Cross-reference company website to verify you have the correct entity
4. If unsure, state "Data not verified for this entity" rather than include wrong data
"""

company_url = state.get("company_url", "")
research_notes = state.get("research_notes", "")

query = build_section_research_query(
    section_def=section_def,
    company_name=company_name,
    company_description=company_description,
    general_research=general_research,
    memo_mode=memo_mode,
    deck_draft_content=deck_draft_content,
    company_url=company_url,           # NEW
    research_notes=research_notes      # NEW
)

python -m cli.improve_section "Sava" "Funding & Terms"

{
  "description": "Specific description with unique identifiers",
  "url": "https://official-website.com",
  "notes": "Research X company only, NOT Y company (describe confusion risk)"
}

{
  "description": "Sava is a financial and legal technology platform specialized in quickly setting up and managing Trusts",
  "url": "https://www.savahq.com",
  "notes": "Ensure research is about SavaHQ (savahq.com), NOT other companies named Sava."
}

# Test section improvement with disambiguation
python -m cli.improve_section "Sava" "Funding & Terms"

# Re-run full memo generation (future test)
python -m src.main "Sava"

# Check for wrong-entity contamination
grep -i "biosensor\|glucose\|CGM\|Imperial College\|Balderton" output/Sava-v0.0.x/2-sections/*.md