Spaces:
Sleeping
Sleeping
| """ | |
| Answer Generation Service for RAG Agent | |
| Handles interaction with OpenAI Chat Completions API to generate natural language | |
| answers from documentation context. Includes retry logic, error handling, and | |
| comprehensive logging. | |
| This service: | |
| 1. Calls OpenAI Chat Completions API | |
| 2. Implements retry logic with exponential backoff | |
| 3. Handles authentication and rate limit errors | |
| 4. Logs generation metrics (tokens, latency, completion length) | |
| 5. Returns structured results with answer and token usage | |
| Usage: | |
| from services.generation_service import AnswerGenerationService | |
| service = AnswerGenerationService(api_key="your-key") | |
| result = service.generate_answer(prompt="<context>...</context>...", model="gpt-3.5-turbo") | |
| print(result['answer'], result['tokens_used']) | |
| """ | |
| import os | |
| import logging | |
| import time | |
| from typing import Dict, Optional | |
| from datetime import datetime | |
| from openai import OpenAI | |
| from exceptions import GenerationError, RateLimitError | |
| logger = logging.getLogger('generation_service') | |
| class AnswerGenerationService: | |
| """ | |
| Service for generating natural language answers using OpenAI. | |
| This service manages all OpenAI API interactions with: | |
| - Retry logic for transient errors (429, 500, 502, 503) | |
| - Error handling for non-retryable errors (401, 400) | |
| - Comprehensive logging (model, tokens, latency) | |
| - Deterministic generation (low temperature) | |
| Attributes: | |
| openai_client: OpenAI client instance | |
| max_retries: Maximum retry attempts (default 3) | |
| Methods: | |
| generate_answer: Generate answer from prompt | |
| """ | |
| def __init__(self, api_key: Optional[str] = None, max_retries: int = 3): | |
| """ | |
| Initialize answer generation service. | |
| Args: | |
| api_key: OpenRouter API key (default: from env) | |
| max_retries: Maximum retry attempts for transient errors (default 3) | |
| Raises: | |
| ValueError: If API key not provided and not in environment | |
| """ | |
| self.api_key = api_key or os.getenv("OPENROUTER_API_KEY") | |
| if not self.api_key: | |
| raise ValueError("OPENROUTER_API_KEY not found in environment") | |
| self.openai_client = OpenAI( | |
| api_key=self.api_key, | |
| base_url="https://openrouter.ai/api/v1" | |
| ) | |
| self.max_retries = max_retries | |
| logger.info( | |
| f"AnswerGenerationService initialized with OpenRouter", | |
| extra={"max_retries": max_retries, "base_url": "https://openrouter.ai/api/v1"} | |
| ) | |
| def generate_answer( | |
| self, | |
| prompt: str, | |
| model: str = "xiaomi/mimo-v2-flash:free", | |
| temperature: float = 0.1, | |
| max_tokens: int = 500 | |
| ) -> Dict[str, any]: | |
| """ | |
| Generate natural language answer from prompt using OpenAI. | |
| Calls OpenAI Chat Completions API with retry logic for transient errors. | |
| Uses low temperature (0.1) for deterministic, factual responses. | |
| Args: | |
| prompt: Full prompt with context and query | |
| model: OpenAI model to use (default: gpt-3.5-turbo) | |
| temperature: Sampling temperature (default: 0.1 for deterministic) | |
| max_tokens: Maximum tokens in completion (default: 500) | |
| Returns: | |
| Dict with: | |
| - answer: Generated text response | |
| - tokens_used: Total tokens consumed (prompt + completion) | |
| - model_used: Actual model used | |
| - completion_length: Character length of answer | |
| Raises: | |
| GenerationError: If generation fails (non-retryable or max retries exceeded) | |
| RateLimitError: If rate limit persists after retries | |
| Example: | |
| >>> service = AnswerGenerationService() | |
| >>> result = service.generate_answer( | |
| ... prompt="<context>ROS 2 docs...</context><query>What is ROS 2?</query>...", | |
| ... model="gpt-3.5-turbo" | |
| ... ) | |
| >>> print(result['answer']) | |
| >>> print(f"Used {result['tokens_used']} tokens") | |
| """ | |
| start_time = datetime.utcnow() | |
| logger.info( | |
| f"Starting answer generation", | |
| extra={ | |
| "model": model, | |
| "temperature": temperature, | |
| "max_tokens": max_tokens, | |
| "prompt_length": len(prompt) | |
| } | |
| ) | |
| # Retry loop for transient errors | |
| last_error = None | |
| for attempt in range(self.max_retries): | |
| try: | |
| # Call OpenAI Chat Completions API | |
| response = self.openai_client.chat.completions.create( | |
| model=model, | |
| messages=[ | |
| { | |
| "role": "system", | |
| "content": "You are a helpful assistant that answers questions based on provided documentation. Always cite the documentation and be precise." | |
| }, | |
| { | |
| "role": "user", | |
| "content": prompt | |
| } | |
| ], | |
| temperature=temperature, | |
| max_tokens=max_tokens | |
| ) | |
| # Extract answer and metrics | |
| answer = response.choices[0].message.content | |
| tokens_used = response.usage.total_tokens | |
| completion_length = len(answer) | |
| # Calculate latency | |
| latency_ms = int((datetime.utcnow() - start_time).total_seconds() * 1000) | |
| # Log success | |
| logger.info( | |
| f"Answer generated successfully", | |
| extra={ | |
| "model": model, | |
| "tokens_used": tokens_used, | |
| "generation_latency_ms": latency_ms, | |
| "completion_length": completion_length, | |
| "attempt": attempt + 1 | |
| } | |
| ) | |
| return { | |
| "answer": answer, | |
| "tokens_used": tokens_used, | |
| "model_used": model, | |
| "completion_length": completion_length | |
| } | |
| except Exception as e: | |
| last_error = e | |
| error_message = str(e) | |
| error_type = type(e).__name__ | |
| # Check for rate limit errors (429) - retryable | |
| if "rate_limit" in error_message.lower() or "429" in error_message: | |
| if attempt < self.max_retries - 1: | |
| # Calculate exponential backoff | |
| wait_time = 2 ** attempt # 1s, 2s, 4s | |
| logger.warning( | |
| f"Rate limit hit, retrying in {wait_time}s (attempt {attempt + 1}/{self.max_retries})", | |
| extra={ | |
| "error_type": error_type, | |
| "attempt": attempt + 1, | |
| "wait_time": wait_time | |
| } | |
| ) | |
| time.sleep(wait_time) | |
| continue | |
| else: | |
| # Max retries exceeded for rate limit | |
| logger.error( | |
| f"Rate limit persists after {self.max_retries} retries", | |
| extra={"error_message": error_message} | |
| ) | |
| raise RateLimitError("OpenAI rate limit exceeded. Please try again later.") | |
| # Check for server errors (500, 502, 503) - retryable | |
| if any(code in error_message for code in ["500", "502", "503"]): | |
| if attempt < self.max_retries - 1: | |
| wait_time = 2 ** attempt | |
| logger.warning( | |
| f"Server error, retrying in {wait_time}s (attempt {attempt + 1}/{self.max_retries})", | |
| extra={ | |
| "error_type": error_type, | |
| "error_message": error_message, | |
| "attempt": attempt + 1, | |
| "wait_time": wait_time | |
| } | |
| ) | |
| time.sleep(wait_time) | |
| continue | |
| # Non-retryable errors (401 auth, 400 bad request, etc.) | |
| logger.error( | |
| f"Generation failed with non-retryable error: {error_message}", | |
| extra={ | |
| "error_type": error_type, | |
| "model": model, | |
| "attempt": attempt + 1 | |
| }, | |
| exc_info=True | |
| ) | |
| raise GenerationError(f"Failed to generate answer: {error_message}") | |
| # If we exhausted all retries | |
| logger.error( | |
| f"Generation failed after {self.max_retries} retries", | |
| extra={"last_error": str(last_error)} | |
| ) | |
| raise GenerationError(f"Failed to generate answer after {self.max_retries} retries: {str(last_error)}") | |