84 lines
3.3 KiB
Python
84 lines
3.3 KiB
Python
# services/llm/gemini_service.py
|
|
"""
|
|
Gemini (Google) service implementation
|
|
"""
|
|
import google.generativeai as genai
|
|
from typing import Dict, List
|
|
import json
|
|
from .base import LLMService
|
|
from config.api_keys import APIKeyManager
|
|
from utils.logger import setup_logger
|
|
|
|
|
|
class GeminiService(LLMService):
|
|
def __init__(
|
|
self,
|
|
model: str = "gemini-1.5-flash",
|
|
temperature: float = 0.3,
|
|
max_tokens: int = 16000,
|
|
):
|
|
api_key = APIKeyManager.get_gemini_key()
|
|
if not api_key:
|
|
raise ValueError(
|
|
"Gemini API key not found. Please set the GEMINI_API_KEY environment variable."
|
|
)
|
|
|
|
genai.configure(api_key=api_key)
|
|
self.model = genai.GenerativeModel(model)
|
|
self.temperature = temperature
|
|
self.max_tokens = max_tokens
|
|
self.logger = setup_logger("gemini")
|
|
|
|
def generate_text(self, prompt: str) -> str:
|
|
self.logger.info(f"--- PROMPT ---\n{prompt}")
|
|
try:
|
|
generation_config = genai.types.GenerationConfig(
|
|
max_output_tokens=self.max_tokens, temperature=self.temperature
|
|
)
|
|
response = self.model.generate_content(
|
|
prompt, generation_config=generation_config
|
|
)
|
|
response_content = response.text
|
|
self.logger.info(f"--- RESPONSE ---\n{response_content}")
|
|
return response_content
|
|
except Exception as e:
|
|
self.logger.error(f"Error in Gemini API call: {e}")
|
|
print(f"Error in Gemini API call: {e}")
|
|
return None
|
|
|
|
def get_similarity_scores(self, texts_pairs: Dict[str, List[str]]) -> List[float]:
|
|
system_prompt = (
|
|
"You are an expert in semantic analysis. Evaluate the semantic similarity between the pairs of texts provided. "
|
|
"Return your response ONLY as a JSON object containing a single key 'similarity_scores' with a list of floats from 0.0 to 1.0. "
|
|
"Do not include any other text, explanation, or markdown formatting. The output must be a valid JSON."
|
|
)
|
|
|
|
request_payload = json.dumps(texts_pairs)
|
|
full_prompt = f"{system_prompt}\n\n{request_payload}"
|
|
|
|
try:
|
|
generation_config = genai.types.GenerationConfig(
|
|
max_output_tokens=self.max_tokens,
|
|
temperature=self.temperature,
|
|
response_mime_type="application/json",
|
|
)
|
|
response = self.model.generate_content(
|
|
full_prompt, generation_config=generation_config
|
|
)
|
|
response_content = response.text
|
|
|
|
try:
|
|
scores_data = json.loads(response_content)
|
|
if isinstance(scores_data, dict) and "similarity_scores" in scores_data:
|
|
return scores_data["similarity_scores"]
|
|
else:
|
|
raise ValueError("Unexpected JSON format from Gemini.")
|
|
except (json.JSONDecodeError, ValueError) as e:
|
|
print(f"Error decoding Gemini JSON response: {e}")
|
|
raise ValueError(
|
|
"Could not decode or parse similarity scores from Gemini response."
|
|
)
|
|
except Exception as e:
|
|
print(f"Error in Gemini similarity calculation: {e}")
|
|
return None
|