Fix Gemini LLM implementation and imports

modules/llm/__init__.py

@@ -2,10 +2,11 @@ from .base import AbstractLLMModel
 from .registry import LLM_MODEL_REGISTRY, get_llm_model, register_llm_model
 from .hf_pipeline import HFTextGenerationLLM
 from .qwen import QwenLLM
+from .gemini import GeminiLLM
 
 __all__ = [
     "AbstractLLMModel",
     "get_llm_model",
     "register_llm_model",
     "LLM_MODEL_REGISTRY",
-]
+]

modules/llm/gemini.py

@@ -1,23 +1,31 @@
+import os
+
+os.environ["XDG_CACHE_HOME"] = "./.cache" # must be set before importing google.generativeai
+
 import google.generativeai as genai
-from abc import ABC, abstractmethod
+
 from .base import AbstractLLMModel
 from .registry import register_llm_model
 
 
-
-GEMINI_TOKEN = os.getenv("GEMINI_API_KEY")
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 
 
 @register_llm_model("gemini-2.5-flash")
-class GeminiModel(AbstractLLMModel):
+class GeminiLLM(AbstractLLMModel):
     def __init__(
         self, model_id: str, device: str = "cpu", cache_dir: str = "cache", **kwargs
     ):
+        if os.environ.get("XDG_CACHE_HOME") != cache_dir:
+            raise RuntimeError(
+                f"XDG_CACHE_HOME must be set to '{cache_dir}' before importing this module."
+            )
+        if not GOOGLE_API_KEY:
+            raise ValueError("Please set the GOOGLE_API_KEY environment variable to use Gemini.")
         super().__init__(model_id=model_id, **kwargs)
-        genai.configure(api_key=GEMINI_API_KEY)
+        genai.configure(api_key=GOOGLE_API_KEY)
         self.model = genai.GenerativeModel(model_id)
 
     def generate(self, prompt: str, **kwargs) -> str:
         response = self.model.generate_content(prompt, **kwargs)
         return response.text
-

tests/test_llm_infer.py

@@ -11,6 +11,7 @@ if __name__ == "__main__":
         # "deepseek-ai/DeepSeek-R1-0528",
         # "openai-community/gpt2-xl",
         # "google/gemma-2-2b",
+        # "gemini-2.5-flash",
     ]
     for model_id in supported_llms:
         try: