server.py

import asyncio
import json
import requests
from pathlib import Path
from datetime import datetime
from shapely.geometry import shape, Point
from mcp.server import Server
from mcp.types import Tool, TextContent

# ============= CONFIGURATION =============
# Change these settings to adapt for your city
CITY_NAME = "Washington DC"  # Used to add city context to address searches
CITY_KEYWORDS = ['washington', 'dc', 'district of columbia']  # Keywords to detect if city is already in address
USE_QUADRANTS = True  # Set to False for cities without quadrant system
QUADRANTS = ['NW', 'NE', 'SE', 'SW']  # Quadrants to search if USE_QUADRANTS is True

BOUNDARY_DIR = "boundaries"
BOUNDARY_FILE = f"{BOUNDARY_DIR}/city_boundary.geojson"
WARD_FILE = f"{BOUNDARY_DIR}/wards.geojson"
LOG_DIR = "logs"
# =========================================

# Create directories if they don't exist
Path(BOUNDARY_DIR).mkdir(exist_ok=True)
Path(LOG_DIR).mkdir(exist_ok=True)

def log_interaction(tool_name, arguments, response):
    """Log tool calls and responses"""
    timestamp = datetime.now().isoformat()
    log_file = f"{LOG_DIR}/interactions.jsonl"
    
    log_entry = {
        "timestamp": timestamp,
        "tool": tool_name,
        "arguments": arguments,
        "response": response
    }
    
    with open(log_file, "a") as f:
        f.write(json.dumps(log_entry) + "\n")

def load_boundary(filepath):
    """Load GeoJSON boundary file"""
    if not Path(filepath).exists():
        return None
    with open(filepath) as f:
        data = json.load(f)
        if data.get('type') == 'FeatureCollection':
            return shape(data['features'][0]['geometry'])
        return shape(data['geometry'])

def load_districts(filepath):
    """Load GeoJSON district boundaries"""
    if not Path(filepath).exists():
        return None
    with open(filepath) as f:
        data = json.load(f)
        districts = []
        for feature in data['features']:
            districts.append({
                'name': feature['properties'].get('NAME') or feature['properties'].get('WARD'),
                'geometry': shape(feature['geometry'])
            })
        return districts

# Load boundaries on startup
city_boundary = load_boundary(BOUNDARY_FILE)
districts = load_districts(WARD_FILE)

# Create MCP server
app = Server("geocoding-server")

def geocode_with_all_variants(address):
    """
    Geocode an address, searching all quadrants if ambiguous.
    Returns standardized result format with all matches.
    """
    # Check if quadrant is specified (only relevant if USE_QUADRANTS is True)
    has_quadrant = False
    if USE_QUADRANTS:
        has_quadrant = any(quad in address.upper() for quad in QUADRANTS)
    
    # Extract base address without city context
    base_address = address
    for keyword in CITY_KEYWORDS:
        base_address = base_address.replace(f', {keyword}', '').replace(f',{keyword}', '')
        base_address = base_address.replace(f', {keyword.title()}', '').replace(f',{keyword.title()}', '')
    base_address = base_address.strip()
    
    # Determine if we should search multiple variants
    search_all_variants = False
    
    if USE_QUADRANTS and has_quadrant:
        # If quadrant specified, check if it's a simple address that might be ambiguous
        base_no_quad = base_address
        for quad in QUADRANTS:
            if f' {quad}' in base_address.upper():
                base_no_quad = base_address[:base_address.upper().rindex(f' {quad}')].strip()
                break
        
        # Simple heuristic: if address is just number + street name, likely ambiguous
        parts = base_no_quad.split()
        if len(parts) <= 4:
            search_all_variants = True
            base_address = base_no_quad
    
    # If no quadrant OR detected ambiguous address, search all possibilities
    if (not has_quadrant or search_all_variants) and USE_QUADRANTS:
        all_matches = []
        
        # Try without quadrant first
        no_quad_result = census_geocode(f"{base_address}, {CITY_NAME}")
        if no_quad_result.get('success') and not no_quad_result.get('multiple_matches'):
            all_matches.append({
                'lat': no_quad_result['lat'],
                'lng': no_quad_result['lng'],
                'matched_address': no_quad_result['matched_address']
            })
        elif no_quad_result.get('multiple_matches'):
            all_matches.extend(no_quad_result['matches'])
        
        # Try all configured quadrants
        for quadrant in QUADRANTS:
            quad_result = census_geocode(f"{base_address} {quadrant}, {CITY_NAME}")
            if quad_result.get('success') and not quad_result.get('multiple_matches'):
                is_duplicate = any(
                    abs(m['lat'] - quad_result['lat']) < 0.0001 and 
                    abs(m['lng'] - quad_result['lng']) < 0.0001 
                    for m in all_matches
                )
                if not is_duplicate:
                    all_matches.append({
                        'lat': quad_result['lat'],
                        'lng': quad_result['lng'],
                        'matched_address': quad_result['matched_address']
                    })
            elif quad_result.get('multiple_matches'):
                for match in quad_result['matches']:
                    is_duplicate = any(
                        abs(m['lat'] - match['lat']) < 0.0001 and 
                        abs(m['lng'] - match['lng']) < 0.0001 
                        for m in all_matches
                    )
                    if not is_duplicate:
                        all_matches.append(match)
        
        if len(all_matches) > 1:
            return {
                'success': True,
                'multiple_matches': True,
                'matches': all_matches
            }
        elif len(all_matches) == 1:
            return {
                'success': True,
                'multiple_matches': False,
                'lat': all_matches[0]['lat'],
                'lng': all_matches[0]['lng'],
                'matched_address': all_matches[0]['matched_address']
            }
        else:
            return {'success': False, 'error': 'No address match found'}
    else:
        # Either quadrant specified OR city doesn't use quadrants - just geocode as-is
        if not any(keyword in address.lower() for keyword in CITY_KEYWORDS):
            search_address = f"{address}, {CITY_NAME}"
        else:
            search_address = address
        return census_geocode(search_address)

def census_geocode(address):
    """Geocode address using Census Bureau API (free, no key required)"""
    url = "https://geocoding.geo.census.gov/geocoder/locations/onelineaddress"
    params = {
        'address': address,
        'benchmark': 'Public_AR_Current',
        'format': 'json'
    }
    
    try:
        response = requests.get(url, params=params, timeout=10)
        response.raise_for_status()
        data = response.json()
        
        if data.get('result', {}).get('addressMatches'):
            matches = data['result']['addressMatches']
            
            # If multiple matches, return all
            if len(matches) > 1:
                results = []
                for match in matches:
                    coords = match['coordinates']
                    results.append({
                        'lat': coords['y'],
                        'lng': coords['x'],
                        'matched_address': match['matchedAddress']
                    })
                return {
                    'success': True,
                    'multiple_matches': True,
                    'count': len(results),
                    'matches': results,
                    'note': 'Multiple addresses found. Please specify quadrant (NW, NE, SE, SW) for DC addresses.'
                }
            
            # Single match
            match = matches[0]
            coords = match['coordinates']
            return {
                'lat': coords['y'],
                'lng': coords['x'],
                'matched_address': match['matchedAddress'],
                'success': True,
                'multiple_matches': False
            }
        return {'success': False, 'error': 'No address match found'}
    except Exception as e:
        return {'success': False, 'error': str(e)}

@app.list_tools()
async def list_tools() -> list[Tool]:
    tools = [
        Tool(
            name="geocode_address",
            description="Geocode an address to latitude/longitude coordinates. " \
            "IMPORTANT: If an abbreviated address (like 'Penn Ave') fails to geocode, try the full street name (like 'Pennsylvania Avenue'). " \
            "Returns all matching addresses if multiple are found.",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "Address to geocode - use EXACTLY as user provided initially, then try variations if needed"}
                },
                "required": ["address"]
            }
        )
    ]
    
    if city_boundary:
        tools.append(Tool(
            name="check_residency",
            description="" \
            "Verify if an address is within the city boundary for residency validation. " \
            "This tool geocodes the address and checks all matches against the city boundary. " \
            "Returns JSON with is_resident (true/false), matched_address(es), and coordinates.",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "Address to check - use EXACTLY as user provided"}
                },
                "required": ["address"]
            }
        ))
    
    if districts:
        tools.append(Tool(
            name="get_district",
            description="" \
            "Get the district/ward for an address. " \
            "This tool geocodes the address, verifies it's in the city, and identifies which district it's in. " \
            "Returns JSON with is_resident (true/false), district name, matched_address(es), and coordinates.",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "Address to get district - use EXACTLY as user provided"}
                },
                "required": ["address"]
            }
        ))
    
    return tools

@app.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    result = None
    
    if name == "geocode_address":
        address = arguments["address"]
        result = geocode_with_all_variants(address)
        
        # If geocoding failed, provide helpful guidance
        if not result.get('success'):
            result['error'] = f'No match found for "{address}". If this address contains abbreviations (like Penn, Ave, St, etc), try calling geocode_address again with the full street name expanded.'
        
        log_interaction(name, arguments, result)
        return [TextContent(type="text", text=json.dumps(result, indent=2))]
    
    elif name == "check_residency":
        address = arguments["address"]
        
        # Use the geocoding helper
        geo_result = geocode_with_all_variants(address)
        
        if not geo_result['success']:
            result = {
                'is_resident': False,
                'error': geo_result.get('error', 'Geocoding failed')
            }
            log_interaction(name, arguments, result)
            return [TextContent(type="text", text=json.dumps(result, indent=2))]
        
        # Handle both single match and multiple matches
        if geo_result.get('multiple_matches'):
            matches = geo_result['matches']
        else:
            matches = [{
                'lat': geo_result['lat'],
                'lng': geo_result['lng'],
                'matched_address': geo_result['matched_address']
            }]
        
        # Check all matches against city boundary
        matches_in_city = []
        for match in matches:
            point = Point(match['lng'], match['lat'])
            if city_boundary.contains(point):
                matches_in_city.append(match)
        
        if len(matches_in_city) == 0:
            result = {
                'is_resident': False,
                'address': address,
                'note': f'Address not found in {CITY_NAME} or outside city boundary'
            }
        elif len(matches_in_city) == 1:
            result = {
                'is_resident': True,
                'address': address,
                'matched_address': matches_in_city[0]['matched_address'],
                'coordinates': {'lat': matches_in_city[0]['lat'], 'lng': matches_in_city[0]['lng']}
            }
        else:
            result = {
                'is_resident': True,
                'multiple_matches': True,
                'count': len(matches_in_city),
                'address': address,
                'matches': matches_in_city,
                'warning': 'CRITICAL: MULTIPLE ADDRESSES FOUND - DO NOT PICK JUST ONE',
                'note': f'Found {len(matches_in_city)} DIFFERENT addresses in {CITY_NAME}. ALL {len(matches_in_city)} addresses must be presented to the user. User must specify quadrant (NW, NE, SE, or SW).'
            }
        
        log_interaction(name, arguments, result)
        return [TextContent(type="text", text=json.dumps(result, indent=2))]
    
    elif name == "get_district":
        address = arguments["address"]
        
        # Use the geocoding helper
        geo_result = geocode_with_all_variants(address)
        
        if not geo_result['success']:
            result = {
                'is_resident': False,
                'district': None,
                'error': 'Geocoding failed'
            }
            log_interaction(name, arguments, result)
            return [TextContent(type="text", text=json.dumps(result, indent=2))]
        
        # Handle both single match and multiple matches
        if geo_result.get('multiple_matches'):
            matches = geo_result['matches']
        else:
            matches = [{
                'lat': geo_result['lat'],
                'lng': geo_result['lng'],
                'matched_address': geo_result['matched_address']
            }]
        
        # First check residency - only look up districts for addresses in the city
        matches_in_city = []
        for match in matches:
            point = Point(match['lng'], match['lat'])
            if city_boundary.contains(point):
                matches_in_city.append(match)
        
        # If no matches in city, return not resident
        if len(matches_in_city) == 0:
            result = {
                'is_resident': False,
                'district': None,
                'address': address,
                'note': f'Address not found in {CITY_NAME} or outside city boundary'
            }
            log_interaction(name, arguments, result)
            return [TextContent(type="text", text=json.dumps(result, indent=2))]
        
        # Now check districts only for addresses confirmed in city
        district_matches = []
        for match in matches_in_city:
            point = Point(match['lng'], match['lat'])
            for district in districts:
                if district['geometry'].contains(point):
                    district_matches.append({
                        'matched_address': match['matched_address'],
                        'district': district['name'],
                        'coordinates': {'lat': match['lat'], 'lng': match['lng']}
                    })
                    break
        
        if len(district_matches) == 0:
            result = {
                'is_resident': False,
                'district': None,
                'address': address,
                'note': f'Address not found in any {CITY_NAME} district or outside city boundary'
            }
        elif len(district_matches) == 1:
            result = {
                'is_resident': True,
                'district': district_matches[0]['district'],
                'address': address,
                'matched_address': district_matches[0]['matched_address'],
                'coordinates': district_matches[0]['coordinates']
            }
        else:
            result = {
                'is_resident': True,
                'multiple_matches': True,
                'count': len(district_matches),
                'address': address,
                'matches': district_matches,
                'warning': 'CRITICAL: MULTIPLE ADDRESSES FOUND - DO NOT PICK JUST ONE',
                'note': f'Found {len(district_matches)} DIFFERENT addresses in {CITY_NAME}. ALL {len(district_matches)} addresses must be presented to the user. User must specify quadrant (NW, NE, SE, or SW).'
            }
        
        log_interaction(name, arguments, result)
        return [TextContent(type="text", text=json.dumps(result, indent=2))]

async def main():
    from mcp.server.stdio import stdio_server
    async with stdio_server() as (read_stream, write_stream):
        await app.run(read_stream, write_stream, app.create_initialization_options())

if __name__ == "__main__":
    asyncio.run(main())