alschameri/helping-source

1"""2Gemini API client for embeddings and text generation.3"""4 5import os6import time7import logging8from typing import List, Dict, Optional9import numpy as np10import requests11import json12 13logger = logging.getLogger(__name__)14 15class GeminiClient:16    """Client for interacting with Google Gemini API."""17    18    def __init__(self):19        self.embedding_dim = 768  # Default dimension for Gemini embeddings20        self.api_key = os.getenv('GEMINI_API_KEY')21        if not self.api_key:22            raise ValueError("GEMINI_API_KEY environment variable is required")23        24        self.project_id = os.getenv('GEMINI_PROJECT', '')25        self.base_url = "https://generativelanguage.googleapis.com/v1beta"26        27        # Rate limiting28        self.last_request_time = 029        self.min_request_interval = 1.0  # seconds30        31        logger.info("Gemini client initialized")32    33    def _wait_for_rate_limit(self):34        """Simple rate limiting to avoid hitting API limits."""35        current_time = time.time()36        time_since_last = current_time - self.last_request_time37        38        if time_since_last < self.min_request_interval:39            sleep_time = self.min_request_interval - time_since_last40            time.sleep(sleep_time)41        42        self.last_request_time = time.time()43    44    def _make_request(self, url: str, payload: Dict, retries: int = 3) -> Dict:45        """Make HTTP request to Gemini API with retry logic."""46        for attempt in range(retries):47            try:48                self._wait_for_rate_limit()49                50                headers = {51                    'Content-Type': 'application/json'52                }53                54                response = requests.post(55                    f"{url}?key={self.api_key}",56                    headers=headers,57                    json=payload,58                    timeout=3059                )60                61                if response.status_code == 200:62                    return response.json()63                elif response.status_code == 429:  # Rate limit64                    wait_time = (2 ** attempt) * 2  # Exponential backoff65                    logger.warning(f"Rate limited, waiting {wait_time}s before retry {attempt + 1}")66                    time.sleep(wait_time)67                    continue68                else:69                    logger.error(f"API request failed: {response.status_code} - {response.text}")70                    response.raise_for_status()71                    72            except requests.exceptions.RequestException as e:73                logger.error(f"Request attempt {attempt + 1} failed: {e}")74                if attempt == retries - 1:75                    raise76                time.sleep(2 ** attempt)77        78        raise Exception("All retry attempts failed")79    80    def embed_texts(self, texts: List[str]) -> List[np.ndarray]:81        """Generate embeddings for a list of texts using Gemini."""82        if not texts:83            return []84        85        try:86            # Gemini embedding endpoint87            url = f"{self.base_url}/models/text-embedding-004:embedContent"88            89            embeddings = []90            91            # Process texts in batches to avoid hitting limits92            batch_size = 1093            for i in range(0, len(texts), batch_size):94                batch_texts = texts[i:i + batch_size]95                96                for text in batch_texts:97                    payload = {98                        "model": "models/text-embedding-004",99                        "content": {100                            "parts": [{101                                "text": text102                            }]103                        }104                    }105                    106                    response_data = self._make_request(url, payload)107                    108                    if 'embedding' in response_data and 'values' in response_data['embedding']:109                        embedding = np.array(response_data['embedding']['values'], dtype=np.float32)110                        embeddings.append(embedding)111                    else:112                        logger.error(f"Unexpected embedding response: {response_data}")113                        # Fallback to random embedding114                        embeddings.append(np.random.rand(self.embedding_dim).astype(np.float32))115            116            logger.info(f"Generated {len(embeddings)} embeddings")117            return embeddings118            119        except Exception as e:120            logger.error(f"Error generating embeddings: {e}")121            # Fallback to random embeddings for development122            logger.warning("Using random embeddings as fallback")123            return [np.random.rand(self.embedding_dim).astype(np.float32) for _ in texts]124    125    def generate_with_context(self, system_prompt: str, user_message: str, contexts: List[str], conversation_history: List[Dict] = None) -> Dict:126        """Generate response using Gemini with provided context and conversation history."""127        try:128            # Build the complete prompt with conversation history129            context_section = ""130            if contexts:131                context_section = "\n\nالسياق المتاح:\n" + "\n---\n".join(contexts)132            133            # Add conversation history if available134            history_section = ""135            if conversation_history:136                history_section = "\n\nالمحادثة السابقة:\n"137                for i, entry in enumerate(conversation_history[-3:]):  # Last 3 exchanges138                    history_section += f"المستخدم: {entry.get('user_message', '')}\n"139                    history_section += f"المساعد: {entry.get('assistant_response', '')}\n---\n"140            141            full_prompt = f"""{system_prompt}142 143{history_section}144 145{context_section}146 147سؤال المستخدم الحالي: {user_message}148 149يرجى الإجابة باللغة العربية مع مراعاة سياق المحادثة السابقة. اجعل إجابتك مفيدة ومختصرة (2-4 جمل) واقترح 2-4 أسئلة متابعة مفيدة."""150 151            # Gemini generation endpoint152            url = f"{self.base_url}/models/gemini-2.0-flash:generateContent"153            154            payload = {155                "contents": [{156                    "parts": [{157                        "text": full_prompt158                    }]159                }],160                "generationConfig": {161                    "temperature": 0.7,162                    "topK": 40,163                    "topP": 0.95,164                    "maxOutputTokens": 512,165                    "stopSequences": []166                }167            }168            169            response_data = self._make_request(url, payload)170            171            # Extract generated text172            if ('candidates' in response_data and 173                len(response_data['candidates']) > 0 and174                'content' in response_data['candidates'][0] and175                'parts' in response_data['candidates'][0]['content']):176                177                generated_text = response_data['candidates'][0]['content']['parts'][0]['text']178                179                # Try to extract suggested questions from the response180                suggested_questions = self._extract_suggested_questions(generated_text)181                182                return {183                    'text': generated_text,184                    'suggested_questions': suggested_questions,185                    'usage': response_data.get('usageMetadata', {})186                }187            else:188                logger.error(f"Unexpected generation response: {response_data}")189                return {190                    'text': 'عذراً، حدث خطأ في توليد الإجابة.',191                    'suggested_questions': ["ما هي عروض السفر؟", "عنّا", "التاشيرات"]192                }193                194        except Exception as e:195            logger.error(f"Error generating response: {e}")196            return {197                'text': 'عذراً، حدث خطأ مؤقت. يرجى المحاولة مرة أخرى.',198                'suggested_questions': ["ما هي عروض السفر؟", "عنّا", "التاشيرات"]199            }200    201    def _extract_suggested_questions(self, text: str) -> List[str]:202        """Extract suggested questions from generated text."""203        # Default suggestions204        default_suggestions = [205            "ما هي عروض السفر؟",206            "عنّا", 207            "التاشيرات",208            "احجز رحلة"209        ]210        211        # Simple heuristic to find questions in the response212        lines = text.split('\n')213        questions = []214        215        for line in lines:216            line = line.strip()217            if line.endswith('؟') and len(line) < 100:  # Arabic question mark218                questions.append(line)219        220        # Return found questions or defaults221        if questions and len(questions) <= 6:222            return questions[:4]  # Max 4 suggestions223        else:224            return default_suggestions225    226    def test_connection(self) -> bool:227        """Test if the Gemini API connection is working."""228        try:229            test_response = self.embed_texts(["تجربة الاتصال"])230            return len(test_response) > 0231        except Exception as e:232            logger.error(f"Connection test failed: {e}")233            return False234    235    def get_available_models(self) -> List[str]:236        """Get list of available Gemini models."""237        try:238            url = f"{self.base_url}/models"239            headers = {'Content-Type': 'application/json'}240            response = requests.get(f"{url}?key={self.api_key}", headers=headers, timeout=10)241            242            if response.status_code == 200:243                data = response.json()244                models = [model['name'] for model in data.get('models', [])]245                return models246            else:247                logger.error(f"Failed to get models: {response.status_code}")248                return []249        except Exception as e:250            logger.error(f"Error getting models: {e}")251            return []