switch provider based on environment variables

utils/call_llm.py

@@ -25,7 +25,33 @@ logger.addHandler(file_handler)
 # Simple cache configuration
 cache_file = "llm_cache.json"
 
-def call_llm(prompt, use_cache: bool = True) -> str:
+
+def load_cache():
+    try:
+        with open(cache_file, 'r') as f:
+            return json.load(f)
+    except:
+        logger.warning(f"Failed to load cache.")
+    return {}
+
+
+def save_cache(cache):
+    try:
+        with open(cache_file, 'w') as f:
+            json.dump(cache, f)
+    except:
+        logger.warning(f"Failed to save cache")
+
+
+def get_llm_provider():
+    provider = os.getenv("LLM_PROVIDER")
+    if not provider and (os.getenv("GEMINI_PROJECT_ID") or os.getenv("GEMINI_API_KEY")):
+        provider = "GEMINI"
+    # if necessary, add ANTHROPIC/OPENAI
+    return provider
+
+
+def _call_llm_provider(prompt: str) -> str:
     """
     Call an LLM provider based on environment variables.
     Environment variables:
@@ -59,7 +85,7 @@ def call_llm(prompt, use_cache: bool = True) -> str:
         raise ValueError(f"{base_url_var} environment variable is required")
 
     # Append the endpoint to the base URL
-    url = f"{base_url}/v1/chat/completions"
+    url = f"{base_url.rstrip('/')}/v1/chat/completions"
 
     # Configure headers and payload based on provider
     headers = {
@@ -98,6 +124,58 @@ def call_llm(prompt, use_cache: bool = True) -> str:
     except ValueError:
         raise Exception(f"Failed to parse response as JSON from {provider}. The server might have returned an invalid response.")
 
+# By default, we Google Gemini 2.5 pro, as it shows great performance for code understanding
+def call_llm(prompt: str, use_cache: bool = True) -> str:
+    # Log the prompt
+    logger.info(f"PROMPT: {prompt}")
+
+    # Check cache if enabled
+    if use_cache:
+        # Load cache from disk
+        cache = load_cache()
+        # Return from cache if exists
+        if prompt in cache:
+            logger.info(f"RESPONSE: {cache[prompt]}")
+            return cache[prompt]
+
+    provider = get_llm_provider()
+    if provider == "GEMINI":
+        response_text = _call_llm_gemini(prompt)
+    else:  # generic method using a URL that is OpenAI compatible API (Ollama, ...)
+        response_text = _call_llm_provider(prompt)
+
+    # Log the response
+    logger.info(f"RESPONSE: {response_text}")
+
+    # Update cache if enabled
+    if use_cache:
+        # Load cache again to avoid overwrites
+        cache = load_cache()
+        # Add to cache and save
+        cache[prompt] = response_text
+        save_cache(cache)
+
+    return response_text
+
+
+def _call_llm_gemini(prompt: str) -> str:
+    if os.getenv("GEMINI_PROJECT_ID"):
+        client = genai.Client(
+            vertexai=True,
+            project=os.getenv("GEMINI_PROJECT_ID"),
+            location=os.getenv("GEMINI_LOCATION", "us-central1")
+        )
+    elif os.getenv("GEMINI_API_KEY"):
+        client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+    else:
+        raise ValueError("Either GEMINI_PROJECT_ID or GEMINI_API_KEY must be set in the environment")
+    model = os.getenv("GEMINI_MODEL", "gemini-2.5-pro-exp-03-25")
+    response = client.models.generate_content(
+        model=model,
+        contents=[prompt]
+    )
+    return response.text
+
 if __name__ == "__main__":
     test_prompt = "Hello, how are you?"