Fix/google genai api (#124)

* exported namespace management, for use in other libraries.

* prep_namespace exposed for use in other libraries.

* merge conflict in uv.lock

* Updated genai api.

* Fixing uv.lock

* black

---------

Co-authored-by: Niki van Stein <n.van.stein@liacs.leidenuniv.nl>

Author: anantashahane

llamea/llm.py

@@ -11,7 +11,8 @@
 from abc import ABC, abstractmethod
 
 try:
-    import google.generativeai as genai
+    from google import genai
+    from google.genai import types
 except ModuleNotFoundError:  # pragma: no cover - optional dependency
     genai = None
 
@@ -372,22 +373,21 @@ def __init__(self, api_key, model="gemini-2.0-flash", **kwargs):
         """
         if genai is None:  # pragma: no cover - optional dependency
             raise ImportError(
-                "google-generativeai is required to use Gemini_LLM. Install the 'google-generativeai' package."
+                "google-genai is required to use Gemini_LLM. Install the 'google-genai' package."
             )
         super().__init__(api_key, model, None, **kwargs)
-        genai.configure(api_key=api_key)
-        generation_config = {
+
+        self.generation_config = {
+            "system_instruction": "You are a computer scientist and excellent Python programmer.",
             "temperature": 1,
             "top_p": 0.95,
             "top_k": 64,
             "max_output_tokens": 65536,
             "response_mime_type": "text/plain",
         }
-
-        self.client = genai.GenerativeModel(
-            model_name=self.model,  # "gemini-1.5-flash","gemini-2.0-flash",
-            generation_config=generation_config,
-            system_instruction="You are a computer scientist and excellent Python programmer.",
+        self.api_key = api_key
+        self.client = genai.Client(
+            api_key=api_key,
         )
 
     def __getstate__(self):
@@ -399,21 +399,29 @@ def __getstate__(self):
     def __setstate__(self, state):
         """Restore from a pickled state."""
         self.__dict__.update(state)  # put back the simple stuff
-        generation_config = {
-            "temperature": 1,
-            "top_p": 0.95,
-            "top_k": 64,
-            "max_output_tokens": 65536,
-            "response_mime_type": "text/plain",
-        }
 
-        self.client = genai.GenerativeModel(
-            model_name=self.model,  # "gemini-1.5-flash","gemini-2.0-flash",
-            generation_config=generation_config,
-            system_instruction="You are a computer scientist and excellent Python programmer.",
-        )
+        self.client = genai.Client(
+            api_key=self.api_key
+        )  # expecting implicit pull for env var GOOGLE_API_KEY, too risky to pickle.
 
-    def query(self, session_messages, max_retries: int = 5, default_delay: int = 10):
+    def __deepcopy__(self, memo):
+        cls = self.__class__
+        new = cls.__new__(cls)
+        memo[id(self)] = new
+        for k, v in self.__dict__.items():
+            if k == "client":
+                continue
+            setattr(new, k, copy.deepcopy(v, memo))
+        new.client = genai.Client(api_key=new.api_key)
+        return new
+
+    def query(
+        self,
+        session_messages: list[dict[str, str]],
+        max_retries: int = 5,
+        default_delay: int = 10,
+        **kwargs,
+    ):
         """
         Sends the conversation history to Gemini, retrying on 429 ResourceExhausted exceptions.
 
@@ -433,14 +441,18 @@ def query(self, session_messages, max_retries: int = 5, default_delay: int = 10)
         attempt = 0
         while True:
             try:
-                chat = self.client.start_chat(history=history)
+                config = self.generation_config.copy()
+                config.update(**kwargs)
+                chat = self.client.chats.create(
+                    model=self.model, history=history, config=config
+                )
                 response = chat.send_message(last)
                 return response.text
 
             except Exception as err:
                 attempt += 1
                 if attempt > max_retries:
-                    raise  # bubble out after N tries
+                    raise err  # bubble out after N tries
 
                 # Prefer the structured retry_delay field if present
                 delay = getattr(err, "retry_delay", None)

pyproject.toml

@@ -13,7 +13,7 @@ dependencies = [
     "ollama>=0.5.3,<1",
     "jsonlines>=4.0.0,<5",
     "configspace>=1.2.0,<2",
-    "google-generativeai>=0.8.1,<0.9",
+    "google-genai>=1,<2",
     "joblib>=1.4.2,<2",
     "lizard>=1.17.13,<2",
     "networkx>=3.4.2,<4",

tests/test_llm.py

@@ -219,13 +219,13 @@ def test_gemini_llm_retries_then_succeeds(monkeypatch):
     chat_ok.send_message.return_value = type("R", (), {"text": "OK-DONE"})
 
     fake_client = MagicMock()
-    fake_client.start_chat.side_effect = [chat_fail, chat_ok]
+    fake_client.chats.create.side_effect = [chat_fail, chat_ok]
     llm.client = fake_client
 
     reply = llm.query([{"role": "user", "content": "hello"}], max_retries=3)
 
     assert reply == "OK-DONE"
-    assert fake_client.start_chat.call_count == 2  # 1 failure + 1 success
+    assert fake_client.chats.create.call_count == 2  # 1 failure + 1 success
     slept.assert_called_once_with(3)  # 2 s + 1 s safety buffer
 
 
@@ -240,7 +240,7 @@ def test_gemini_llm_gives_up_after_max_retries(monkeypatch):
     chat_fail.send_message.side_effect = _resource_exhausted(1)
 
     fake_client = MagicMock()
-    fake_client.start_chat.return_value = chat_fail
+    fake_client.chats.create.return_value = chat_fail
     llm.client = fake_client
 
     with pytest.raises(Exception):