85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
# services/llm/groq_service.py
|
|
"""
|
|
Groq service implementation
|
|
"""
|
|
from groq import Groq
|
|
from typing import Dict, List
|
|
import json
|
|
from .base import LLMService
|
|
from config.api_keys import APIKeyManager
|
|
from utils.logger import setup_logger
|
|
|
|
|
|
class GroqService(LLMService):
|
|
def __init__(
|
|
self,
|
|
model: str = "llama3-8b-8192",
|
|
temperature: float = 0.3,
|
|
max_tokens: int = 8000,
|
|
):
|
|
api_key = APIKeyManager.get_groq_key()
|
|
if not api_key:
|
|
raise ValueError(
|
|
"Groq API key not found. Please set the GROQ_API_KEY environment variable."
|
|
)
|
|
|
|
self.client = Groq(api_key=api_key)
|
|
self.model = model
|
|
self.temperature = temperature
|
|
self.max_tokens = max_tokens
|
|
self.logger = setup_logger("groq")
|
|
|
|
def generate_text(self, prompt: str) -> str:
|
|
self.logger.info(f"--- PROMPT ---\n{prompt}")
|
|
try:
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[{"role": "user", "content": prompt}],
|
|
temperature=self.temperature,
|
|
max_tokens=self.max_tokens,
|
|
)
|
|
response_content = response.choices[0].message.content
|
|
self.logger.info(f"--- RESPONSE ---\n{response_content}")
|
|
return response_content
|
|
except Exception as e:
|
|
self.logger.error(f"Error in Groq API call: {e}")
|
|
print(f"Error in Groq API call: {e}")
|
|
return None
|
|
|
|
def get_similarity_scores(self, texts_pairs: Dict[str, List[str]]) -> List[float]:
|
|
system_prompt = (
|
|
"Evaluate the semantic similarity between the following table of pairs of texts in json format on a scale from 0 to 1. "
|
|
"Return the similarity scores for every row in JSON format as a list of numbers, without any additional text or formatting."
|
|
)
|
|
|
|
request_payload = json.dumps(texts_pairs)
|
|
|
|
try:
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": request_payload},
|
|
],
|
|
temperature=self.temperature,
|
|
max_tokens=self.max_tokens,
|
|
response_format={"type": "json_object"},
|
|
)
|
|
|
|
response_content = response.choices[0].message.content
|
|
|
|
try:
|
|
scores = json.loads(response_content)
|
|
if isinstance(scores, dict) and "similarity_scores" in scores:
|
|
return scores["similarity_scores"]
|
|
elif isinstance(scores, list):
|
|
return scores
|
|
else:
|
|
raise ValueError("Unexpected response format")
|
|
except json.JSONDecodeError:
|
|
raise ValueError("Could not decode response as JSON")
|
|
|
|
except Exception as e:
|
|
print(f"Error in Groq similarity calculation: {e}")
|
|
return None
|