Karan6124 commited on
Commit
8419ecc
·
1 Parent(s): b27ea06

feat: implement google gemini llm helper service with structured json configuration

Browse files
Files changed (1) hide show
  1. 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