Function run_gemini_api

WebAgent/WebWalker/src/rag_system.py:96–144

()

            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:
                tasks = [
                    fetch(session, GEMINI_BASE_URL, headers, {
                        "model": "gemini-1.5-pro",
                        "contents": [{"role": "user", "parts": [{"text": query}]}],
                        "candidates": 1
                    }, semaphore, query)
                    for query in data_list
                ]
            for future in tqdm.as_completed(tasks, total=len(tasks), desc="Processing queries"):
                query_text, result = await future
                adic = {"question": query_text, "pred": result["candidates"][0]["content"]["parts"][0]["text"]}
                with open(output_path, "a", encoding="utf-8") as f:
                    f.write(json.dumps(adic, ensure_ascii=False) + "\n")

    asyncio.run(run_gemini_api())