search.py

from dotenv import load_dotenv
import os
import aiohttp


class GoogleSearch:
    def __init__(self):
        load_dotenv()
        self.counter = 0
        self.api_key = os.environ["GOOGLE_API_KEY"]
        self.cse_id = os.getenv("GOOGLE_CSE_ID")

    async def google_search(self, query: str, num_results: int = 5) -> str:
        """
        Args:
            query: Search query
            num_results: Max results to return
        Returns:
            dict: JSON response from Google API.
        """
        if self.counter > 1:
            return "No more searches, move on"
        self.counter += 1

        if not self.api_key or not self.cse_id:
            raise ValueError(
                "GOOGLE_API_KEY and GOOGLE_CSE_ID must be set in environment variables."
            )

        url = "https://www.googleapis.com/customsearch/v1"
        params = {"key": self.api_key, "cx": self.cse_id, "q": query}

        async with aiohttp.ClientSession() as session:
            async with session.get(url, params=params) as response:
                if response.status == 200:
                    data = await response.json()
                    results = "Web Search results:\n\n" + "\n\n".join(
                        [
                            f"Link:{result['link']}\nTitle:{result['title']}\nSnippet:{result['snippet']}"
                            for result in data["items"][:num_results]
                        ]
                    )
                    return results
                else:
                    return f"Search failed with status {response.status}"

    async def google_image_search(self, query: str) -> str:
        """
        Args:
            query: Search query
        Returns:
            dict: JSON response from Google API.
        """
        if self.counter > 2:
            return "No more searches, move on"
        self.counter += 1

        if not self.api_key or not self.cse_id:
            raise ValueError(
                "GOOGLE_API_KEY and GOOGLE_CSE_ID must be set in environment variables."
            )

        url = "https://www.googleapis.com/customsearch/v1"
        params = {"key": self.api_key, "cx": self.cse_id, "q": query}

        async with aiohttp.ClientSession() as session:
            async with session.get(url, params=params) as response:
                if response.status == 200:
                    data = await response.json()
                    results = "Web Search results:\n\n" + "\n\n".join(
                        [
                            f"Link:{result['link']}\nTitle:{result['title']}"
                            for result in data["items"][:4]
                        ]
                    )
                    return results
                else:
                    return f"Search failed with status {response.status}"


async def main():
    google_search = GoogleSearch()
    results = await google_search.google_search("Fresident of Prance?")
    print(results)


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())