post_processing.py

# post_processing.py
# No changes needed, but ensured absolute imports.

import re
from retrieval import google_search  # Absolute import

def ground_statutes(response, jurisdiction, GOOGLE_CUSTOM_SEARCH_API_KEY, GOOGLE_SEARCH_API, ask_gpt4o):
    # Extract cited statutes from response (e.g., KRS 403.270)
    statute_pattern = r'KRS\s(\d+\.\d+)'
    cited_statutes = re.findall(statute_pattern, response)
    
    for statute_code in cited_statutes:
        # Search web for the statute text
        query = f"{jurisdiction} {statute_code} statute text"
        search_result = google_search(query, GOOGLE_CUSTOM_SEARCH_API_KEY, GOOGLE_SEARCH_API)
        
        # Use GPT to verify if the cited description matches the search result
        verification_prompt = f"Does the following statute description match the searched text?\nDescription in response: {response}\nSearched text: {search_result}"
        verification = ask_gpt4o(verification_prompt)
        
        if "no" in verification.lower() or "does not match" in verification.lower():
            # Correct the response
            correction_prompt = f"Correct the following response based on the accurate statute text: {search_result}\nOriginal response: {response}"
            response = ask_gpt4o(correction_prompt)
    
    return response