Function gemini_api

WebAgent/WebWalker/src/rag_system.py:74–146

(ds, output_path, search=False)

    asyncio.run(request_model(prompts))

def gemini_api(ds, output_path, search=False):
    if GEMINI_API_KEY is None or GEMINI_BASE_URL is None:
        print("Please set GEMINI_API_URL and GEMINI_AUTH_TOKEN environment variables.")
        return
    headers = {
        'Authorization': f'Bearer {GEMINI_API_KEY}',
        'Content-Type': 'application/json'
    }
    MAX_CONCURRENT = 16

    async def fetch(session, url, headers, data, semaphore, query_text):
        async with semaphore:
            return query_text, await _fetch_with_retry(session, url, headers, data)

    @retry(stop=stop_after_attempt(10), wait=wait_exponential(min=4, max=60), reraise=True)
    async def _fetch_with_retry(session, url, headers, data):
        async with session.post(url, headers=headers, json=data) as response:
            if response.status != 200:
                print(f"Error: {response.status} - {response.reason}")
            response.raise_for_status()
            return await response.json()

    async def run_gemini_api():
        if not os.path.exists(output_path):
            open(output_path, "w").close()
        with open(output_path, "r", encoding="utf-8") as f:
            visited = [json.loads(line)["question"] for line in f]
        data_list = []
        for item in ds["question"]:
            if item not in visited:
                data_list.append(item)
        semaphore = asyncio.Semaphore(MAX_CONCURRENT)
        async with aiohttp.ClientSession() as session:
            if search:
                tasks = [
                    fetch(session, GEMINI_BASE_URL, headers, {
                        "model": "gemini-1.5-pro",
                        "contents": [
                            {"role": "user", 
                             "parts": 
                                [   
                                    {"text": query}, 
                                    {"tools": {
                                        "google_search_retrieval": {
                                            "dynamic_retrieval_config": {
                                                "mode": "MODE_DYNAMIC",
                                                "dynamic_threshold": 0
                                                } 
                                            }
                                        }
                                    }
                                ]
                            }
                        ]
                    }, semaphore, query)
                    for query in data_list
                ]
            else: