Spaces:
Running
Running
feat: implement google gemini llm helper service with structured json configuration
Browse files- app/services/llm.py +83 -0
app/services/llm.py
CHANGED
|
@@ -0,0 +1,83 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import json
|
| 3 |
+
from typing import Dict, Any, Optional
|
| 4 |
+
import google.generativeai as genai
|
| 5 |
+
from dotenv import load_dotenv
|
| 6 |
+
|
| 7 |
+
load_dotenv()
|
| 8 |
+
|
| 9 |
+
# Load and validate key
|
| 10 |
+
api_key = os.getenv("GEMINI_API_KEY")
|
| 11 |
+
if not api_key:
|
| 12 |
+
raise ValueError("GEMINI_API_KEY environment variable is missing from .env")
|
| 13 |
+
|
| 14 |
+
# Configure Google GenAI
|
| 15 |
+
genai.configure(api_key=api_key)
|
| 16 |
+
|
| 17 |
+
class LLMService:
|
| 18 |
+
@staticmethod
|
| 19 |
+
def call_gemini(
|
| 20 |
+
prompt: str,
|
| 21 |
+
system_instruction: Optional[str] = None,
|
| 22 |
+
model_name: str = "gemini-2.0-flash",
|
| 23 |
+
temperature: float = 0.2,
|
| 24 |
+
json_output: bool = False
|
| 25 |
+
) -> str:
|
| 26 |
+
"""
|
| 27 |
+
Invokes Gemini LLM model and returns the text response.
|
| 28 |
+
If json_output is True, configures the request to return application/json.
|
| 29 |
+
"""
|
| 30 |
+
# Define generation config
|
| 31 |
+
generation_config = {
|
| 32 |
+
"temperature": temperature,
|
| 33 |
+
}
|
| 34 |
+
|
| 35 |
+
if json_output:
|
| 36 |
+
generation_config["response_mime_type"] = "application/json"
|
| 37 |
+
|
| 38 |
+
# Initialize the model
|
| 39 |
+
model = genai.GenerativeModel(
|
| 40 |
+
model_name=model_name,
|
| 41 |
+
generation_config=generation_config,
|
| 42 |
+
system_instruction=system_instruction
|
| 43 |
+
)
|
| 44 |
+
|
| 45 |
+
# Generate response
|
| 46 |
+
response = model.generate_content(prompt)
|
| 47 |
+
|
| 48 |
+
if not response.text:
|
| 49 |
+
raise RuntimeError("Gemini model returned an empty response.")
|
| 50 |
+
|
| 51 |
+
return response.text
|
| 52 |
+
|
| 53 |
+
@staticmethod
|
| 54 |
+
def call_gemini_json(
|
| 55 |
+
prompt: str,
|
| 56 |
+
system_instruction: Optional[str] = None,
|
| 57 |
+
model_name: str = "gemini-2.0-flash",
|
| 58 |
+
temperature: float = 0.2
|
| 59 |
+
) -> Dict[str, Any]:
|
| 60 |
+
"""
|
| 61 |
+
Calls Gemini forcing a JSON structure and parses the result into a python dictionary.
|
| 62 |
+
"""
|
| 63 |
+
raw_response = LLMService.call_gemini(
|
| 64 |
+
prompt=prompt,
|
| 65 |
+
system_instruction=system_instruction,
|
| 66 |
+
model_name=model_name,
|
| 67 |
+
temperature=temperature,
|
| 68 |
+
json_output=True
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
try:
|
| 72 |
+
return json.loads(raw_response)
|
| 73 |
+
except json.JSONDecodeError as e:
|
| 74 |
+
# Simple fallback if JSON parsing fails
|
| 75 |
+
clean_str = raw_response.strip()
|
| 76 |
+
if clean_str.startswith("```json"):
|
| 77 |
+
clean_str = clean_str.split("```json")[1].split("```")[0].strip()
|
| 78 |
+
elif clean_str.startswith("```"):
|
| 79 |
+
clean_str = clean_str.split("```")[1].split("```")[0].strip()
|
| 80 |
+
try:
|
| 81 |
+
return json.loads(clean_str)
|
| 82 |
+
except Exception:
|
| 83 |
+
raise ValueError(f"Failed to parse JSON response from Gemini. Raw output: {raw_response}") from e
|