search_engine.py

from google import genai
from google.genai import types
import config

def initialize_gemini_client():
    """
    Initialize Gemini client.
    """
    return genai.Client(api_key=config.GEMINI_API_KEY)

def search_and_synthesize(query, rag_context):
    """
    Search the web using Gemini's Google Search grounding tool and combine with RAG context.
    Enforces Pakistan-specific content filtering.
    """
    try:
        client = initialize_gemini_client()

        # Create grounding tool with Google Search
        grounding_tool = types.Tool(
            google_search=types.GoogleSearch()
        )

        # Configuration with grounding tool
        # Using the standard configuration approach compatible with google-genai
        generate_config = types.GenerateContentConfig(
            tools=[grounding_tool],
            temperature=0.2,
            system_instruction="""You are a specialized Legal Assistant for Pakistan. 
            Your primary job is to answer the user's legal question by combining:
            1. The User's Question.
            2. The provided 'Legal Context' (which comes from local legal documents like the Constitution and PPC).
            3. Real-time information from Google Search.

            CRITICAL RULES:
            - FILTER: You must ONLY answer questions related to Pakistan Law, the Pakistani Legal System, or general legal queries applicable in Pakistan. 
            - DENIAL: If the user asks about anything else (e.g., "Capital of Peru", "Movie reviews", "Laws of France"), query unrelated to Pakistan, you MUST REFUSE to answer. Say exactly: "I specialize only in Pakistani Law. I cannot assist with this query."
            - SYNTHESIS: Provide a single, cohesive answer. citations are encouraged.
            - Do not treat 'Legal Context' as the only truth if Search reveals it's outdated, but prioritize the Constitution/Acts if they are standard texts.
            - If the user asks for a specific section, quote it if available in Context or Search.
            """
        )

        prompt = f"""
        User Query: {query}
        
        Legal Context from Local Documents:
        {rag_context}
        
        Please provide a comprehensive answer based on the above instructions.
        """

        # Generate response with web search
        response = client.models.generate_content(
            model=config.GEMINI_MODEL_NAME,
            contents=prompt,
            config=generate_config,
        )

        if response and response.text:
            return response.text
        else:
            # Fallback if model returns empty but no error raised
            return "No information could be generated. Please try again."

    except Exception as e:
        print(f"Gemini search error: {e}")
        error_msg = str(e).lower()
        # Check for common "server busy" or quota errors
        if "503" in error_msg or "429" in error_msg or "busy" in error_msg or "quota" in error_msg:
            return "SERVER_BUSY"
        
        # Generic error handling to avoid crashing
        return "SERVER_BUSY"