Spaces:
Sleeping
Sleeping
Deploy RAG chatbot backend API
Browse files- Dockerfile +23 -0
- README.md +23 -5
- config/__init__.py +1 -0
- config/logging_config.py +130 -0
- exceptions.py +126 -0
- models/__init__.py +4 -0
- models/request_models.py +110 -0
- models/response_models.py +221 -0
- models/translation_models.py +92 -0
- rag_agent.py +270 -0
- requirements.txt +20 -0
- retrieval.py +252 -0
- services/__init__.py +1 -0
- services/generation_service.py +239 -0
- services/integration_service.py +286 -0
- services/prompt_builder.py +264 -0
- services/retrieval_service.py +147 -0
- services/translation_service.py +292 -0
Dockerfile
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
FROM python:3.10-slim
|
| 2 |
+
|
| 3 |
+
WORKDIR /app
|
| 4 |
+
|
| 5 |
+
# Install system dependencies
|
| 6 |
+
RUN apt-get update && apt-get install -y \
|
| 7 |
+
build-essential \
|
| 8 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 9 |
+
|
| 10 |
+
# Copy requirements first for caching
|
| 11 |
+
COPY requirements.txt .
|
| 12 |
+
|
| 13 |
+
# Install Python dependencies
|
| 14 |
+
RUN pip install --no-cache-dir -r requirements.txt
|
| 15 |
+
|
| 16 |
+
# Copy application code
|
| 17 |
+
COPY . .
|
| 18 |
+
|
| 19 |
+
# Expose port 7860 (default for HF Spaces)
|
| 20 |
+
EXPOSE 7860
|
| 21 |
+
|
| 22 |
+
# Run the application
|
| 23 |
+
CMD ["uvicorn", "rag_agent:app", "--host", "0.0.0.0", "--port", "7860"]
|
README.md
CHANGED
|
@@ -1,10 +1,28 @@
|
|
| 1 |
---
|
| 2 |
-
title: Chatbot
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo:
|
| 6 |
sdk: docker
|
| 7 |
pinned: false
|
|
|
|
| 8 |
---
|
| 9 |
|
| 10 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
+
title: RAG Chatbot API
|
| 3 |
+
emoji: 🤖
|
| 4 |
+
colorFrom: blue
|
| 5 |
+
colorTo: purple
|
| 6 |
sdk: docker
|
| 7 |
pinned: false
|
| 8 |
+
license: mit
|
| 9 |
---
|
| 10 |
|
| 11 |
+
# RAG Chatbot API
|
| 12 |
+
|
| 13 |
+
A Retrieval-Augmented Generation (RAG) API for documentation Q&A.
|
| 14 |
+
|
| 15 |
+
## API Endpoints
|
| 16 |
+
|
| 17 |
+
- `POST /query` - Submit a question and get an answer with sources
|
| 18 |
+
- `GET /health` - Health check endpoint
|
| 19 |
+
- `GET /docs` - Swagger API documentation
|
| 20 |
+
|
| 21 |
+
## Environment Variables
|
| 22 |
+
|
| 23 |
+
Set these in your HF Space secrets:
|
| 24 |
+
|
| 25 |
+
- `OPENAI_API_KEY` - OpenAI API key
|
| 26 |
+
- `QDRANT_URL` - Qdrant cloud URL
|
| 27 |
+
- `QDRANT_API_KEY` - Qdrant API key
|
| 28 |
+
- `COHERE_API_KEY` - Cohere API key (for embeddings)
|
config/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
# Config package
|
config/logging_config.py
ADDED
|
@@ -0,0 +1,130 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Logging Configuration for RAG Agent API
|
| 3 |
+
|
| 4 |
+
Provides structured JSON logging for all components of the RAG Agent.
|
| 5 |
+
Logs include timestamp, level, component name, message, and contextual data.
|
| 6 |
+
|
| 7 |
+
Usage:
|
| 8 |
+
from config.logging_config import setup_logging, get_logger
|
| 9 |
+
|
| 10 |
+
setup_logging()
|
| 11 |
+
logger = get_logger('rag_agent')
|
| 12 |
+
logger.info("Message", extra={"query": "test", "latency_ms": 123})
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
import logging
|
| 16 |
+
import json
|
| 17 |
+
import sys
|
| 18 |
+
from datetime import datetime
|
| 19 |
+
from typing import Any, Dict
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class JSONFormatter(logging.Formatter):
|
| 23 |
+
"""
|
| 24 |
+
Custom formatter that outputs logs as JSON.
|
| 25 |
+
|
| 26 |
+
Format:
|
| 27 |
+
{
|
| 28 |
+
"timestamp": "2025-12-21T10:30:00Z",
|
| 29 |
+
"level": "INFO",
|
| 30 |
+
"component": "rag_agent",
|
| 31 |
+
"message": "Query processed",
|
| 32 |
+
"context": {"query": "...", "latency_ms": 123}
|
| 33 |
+
}
|
| 34 |
+
"""
|
| 35 |
+
|
| 36 |
+
def format(self, record: logging.LogRecord) -> str:
|
| 37 |
+
"""
|
| 38 |
+
Format log record as JSON string.
|
| 39 |
+
|
| 40 |
+
Args:
|
| 41 |
+
record: LogRecord to format
|
| 42 |
+
|
| 43 |
+
Returns:
|
| 44 |
+
JSON-formatted log string
|
| 45 |
+
"""
|
| 46 |
+
log_data: Dict[str, Any] = {
|
| 47 |
+
"timestamp": datetime.utcnow().isoformat() + "Z",
|
| 48 |
+
"level": record.levelname,
|
| 49 |
+
"component": record.name,
|
| 50 |
+
"message": record.getMessage(),
|
| 51 |
+
}
|
| 52 |
+
|
| 53 |
+
# Add any extra context from the log call
|
| 54 |
+
if hasattr(record, 'query'):
|
| 55 |
+
log_data['query'] = record.query
|
| 56 |
+
if hasattr(record, 'latency_ms'):
|
| 57 |
+
log_data['latency_ms'] = record.latency_ms
|
| 58 |
+
if hasattr(record, 'tokens_used'):
|
| 59 |
+
log_data['tokens_used'] = record.tokens_used
|
| 60 |
+
if hasattr(record, 'model'):
|
| 61 |
+
log_data['model'] = record.model
|
| 62 |
+
if hasattr(record, 'status_code'):
|
| 63 |
+
log_data['status_code'] = record.status_code
|
| 64 |
+
if hasattr(record, 'error_type'):
|
| 65 |
+
log_data['error_type'] = record.error_type
|
| 66 |
+
|
| 67 |
+
# Include exception info if present
|
| 68 |
+
if record.exc_info:
|
| 69 |
+
log_data['exception'] = self.formatException(record.exc_info)
|
| 70 |
+
|
| 71 |
+
return json.dumps(log_data)
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def setup_logging(level: str = "INFO", json_format: bool = False) -> None:
|
| 75 |
+
"""
|
| 76 |
+
Configure logging for the application.
|
| 77 |
+
|
| 78 |
+
Args:
|
| 79 |
+
level: Logging level (DEBUG, INFO, WARNING, ERROR)
|
| 80 |
+
json_format: If True, use JSON formatter; otherwise use standard format
|
| 81 |
+
|
| 82 |
+
Example:
|
| 83 |
+
>>> setup_logging(level="DEBUG", json_format=True)
|
| 84 |
+
"""
|
| 85 |
+
log_level = getattr(logging, level.upper(), logging.INFO)
|
| 86 |
+
|
| 87 |
+
# Create handler for stdout
|
| 88 |
+
handler = logging.StreamHandler(sys.stdout)
|
| 89 |
+
|
| 90 |
+
if json_format:
|
| 91 |
+
handler.setFormatter(JSONFormatter())
|
| 92 |
+
else:
|
| 93 |
+
# Standard format for development
|
| 94 |
+
formatter = logging.Formatter(
|
| 95 |
+
'[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s',
|
| 96 |
+
datefmt='%Y-%m-%d %H:%M:%S'
|
| 97 |
+
)
|
| 98 |
+
handler.setFormatter(formatter)
|
| 99 |
+
|
| 100 |
+
# Configure root logger
|
| 101 |
+
logging.root.setLevel(log_level)
|
| 102 |
+
logging.root.handlers = [handler]
|
| 103 |
+
|
| 104 |
+
# Set levels for noisy libraries
|
| 105 |
+
logging.getLogger('httpx').setLevel(logging.WARNING)
|
| 106 |
+
logging.getLogger('httpcore').setLevel(logging.WARNING)
|
| 107 |
+
logging.getLogger('openai').setLevel(logging.WARNING)
|
| 108 |
+
logging.getLogger('urllib3').setLevel(logging.WARNING)
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def get_logger(name: str) -> logging.Logger:
|
| 112 |
+
"""
|
| 113 |
+
Get a logger instance for a component.
|
| 114 |
+
|
| 115 |
+
Args:
|
| 116 |
+
name: Component name (e.g., 'rag_agent', 'retrieval', 'generation')
|
| 117 |
+
|
| 118 |
+
Returns:
|
| 119 |
+
Logger instance
|
| 120 |
+
|
| 121 |
+
Example:
|
| 122 |
+
>>> logger = get_logger('rag_agent')
|
| 123 |
+
>>> logger.info("Processing query", extra={"query": "test"})
|
| 124 |
+
"""
|
| 125 |
+
return logging.getLogger(name)
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
# Default setup for imports
|
| 129 |
+
if __name__ != "__main__":
|
| 130 |
+
setup_logging()
|
exceptions.py
ADDED
|
@@ -0,0 +1,126 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Custom Exceptions for RAG Agent API
|
| 3 |
+
|
| 4 |
+
Exception classes for different error types in the RAG pipeline.
|
| 5 |
+
|
| 6 |
+
Exceptions:
|
| 7 |
+
RAGException: Base exception for all RAG Agent errors
|
| 8 |
+
ValidationError: Request validation failures
|
| 9 |
+
RetrievalError: Document retrieval failures
|
| 10 |
+
GenerationError: Answer generation failures
|
| 11 |
+
RateLimitError: API rate limit exceeded
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class RAGException(Exception):
|
| 16 |
+
"""
|
| 17 |
+
Base exception for all RAG Agent errors.
|
| 18 |
+
|
| 19 |
+
Attributes:
|
| 20 |
+
message: Error message
|
| 21 |
+
status_code: HTTP status code
|
| 22 |
+
error_type: Error classification
|
| 23 |
+
"""
|
| 24 |
+
|
| 25 |
+
def __init__(self, message: str, status_code: int = 500, error_type: str = "UnknownError"):
|
| 26 |
+
"""
|
| 27 |
+
Initialize RAG exception.
|
| 28 |
+
|
| 29 |
+
Args:
|
| 30 |
+
message: Error message
|
| 31 |
+
status_code: HTTP status code (default 500)
|
| 32 |
+
error_type: Error classification (default "UnknownError")
|
| 33 |
+
"""
|
| 34 |
+
self.message = message
|
| 35 |
+
self.status_code = status_code
|
| 36 |
+
self.error_type = error_type
|
| 37 |
+
super().__init__(self.message)
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
class ValidationError(RAGException):
|
| 41 |
+
"""
|
| 42 |
+
Exception for request validation failures.
|
| 43 |
+
|
| 44 |
+
Raised when:
|
| 45 |
+
- Query is empty or invalid
|
| 46 |
+
- top_k is out of range
|
| 47 |
+
- Model name is invalid
|
| 48 |
+
|
| 49 |
+
Status Code: 422 Unprocessable Entity
|
| 50 |
+
"""
|
| 51 |
+
|
| 52 |
+
def __init__(self, message: str):
|
| 53 |
+
"""
|
| 54 |
+
Initialize validation error.
|
| 55 |
+
|
| 56 |
+
Args:
|
| 57 |
+
message: Validation error message
|
| 58 |
+
"""
|
| 59 |
+
super().__init__(message, status_code=422, error_type="ValidationError")
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
class RetrievalError(RAGException):
|
| 63 |
+
"""
|
| 64 |
+
Exception for document retrieval failures.
|
| 65 |
+
|
| 66 |
+
Raised when:
|
| 67 |
+
- Qdrant database is unavailable
|
| 68 |
+
- Cohere API fails to embed query
|
| 69 |
+
- Network errors during retrieval
|
| 70 |
+
|
| 71 |
+
Status Code: 503 Service Unavailable
|
| 72 |
+
"""
|
| 73 |
+
|
| 74 |
+
def __init__(self, message: str):
|
| 75 |
+
"""
|
| 76 |
+
Initialize retrieval error.
|
| 77 |
+
|
| 78 |
+
Args:
|
| 79 |
+
message: Retrieval error message
|
| 80 |
+
"""
|
| 81 |
+
super().__init__(message, status_code=503, error_type="RetrievalError")
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
class GenerationError(RAGException):
|
| 85 |
+
"""
|
| 86 |
+
Exception for answer generation failures.
|
| 87 |
+
|
| 88 |
+
Raised when:
|
| 89 |
+
- OpenAI API is unavailable
|
| 90 |
+
- OpenAI returns error response
|
| 91 |
+
- Network errors during generation
|
| 92 |
+
- Authentication failures
|
| 93 |
+
|
| 94 |
+
Status Code: 502 Bad Gateway
|
| 95 |
+
"""
|
| 96 |
+
|
| 97 |
+
def __init__(self, message: str):
|
| 98 |
+
"""
|
| 99 |
+
Initialize generation error.
|
| 100 |
+
|
| 101 |
+
Args:
|
| 102 |
+
message: Generation error message
|
| 103 |
+
"""
|
| 104 |
+
super().__init__(message, status_code=502, error_type="GenerationError")
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
class RateLimitError(RAGException):
|
| 108 |
+
"""
|
| 109 |
+
Exception for API rate limit exceeded.
|
| 110 |
+
|
| 111 |
+
Raised when:
|
| 112 |
+
- OpenAI rate limit exceeded
|
| 113 |
+
- Cohere rate limit exceeded
|
| 114 |
+
- Qdrant rate limit exceeded
|
| 115 |
+
|
| 116 |
+
Status Code: 429 Too Many Requests
|
| 117 |
+
"""
|
| 118 |
+
|
| 119 |
+
def __init__(self, message: str):
|
| 120 |
+
"""
|
| 121 |
+
Initialize rate limit error.
|
| 122 |
+
|
| 123 |
+
Args:
|
| 124 |
+
message: Rate limit error message
|
| 125 |
+
"""
|
| 126 |
+
super().__init__(message, status_code=429, error_type="RateLimitError")
|
models/__init__.py
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Models package
|
| 2 |
+
from .request_models import QueryRequest
|
| 3 |
+
from .response_models import QueryResponse, ErrorResponse
|
| 4 |
+
from .translation_models import TranslationRequest, TranslationResponse
|
models/request_models.py
ADDED
|
@@ -0,0 +1,110 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Request Models for RAG Agent API
|
| 3 |
+
|
| 4 |
+
Pydantic models for validating incoming API requests.
|
| 5 |
+
|
| 6 |
+
Models:
|
| 7 |
+
QueryRequest: User query with optional parameters
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from pydantic import BaseModel, Field, field_validator
|
| 11 |
+
from typing import Optional
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
class QueryRequest(BaseModel):
|
| 15 |
+
"""
|
| 16 |
+
Request model for POST /query endpoint.
|
| 17 |
+
|
| 18 |
+
Attributes:
|
| 19 |
+
query: User's natural language question (1-2000 characters)
|
| 20 |
+
top_k: Number of documentation chunks to retrieve (1-50, default 5)
|
| 21 |
+
model: OpenAI model to use (default 'gpt-3.5-turbo')
|
| 22 |
+
|
| 23 |
+
Example:
|
| 24 |
+
{
|
| 25 |
+
"query": "How do I create a ROS 2 publisher in Python?",
|
| 26 |
+
"top_k": 5,
|
| 27 |
+
"model": "gpt-3.5-turbo"
|
| 28 |
+
}
|
| 29 |
+
"""
|
| 30 |
+
|
| 31 |
+
query: str = Field(
|
| 32 |
+
...,
|
| 33 |
+
min_length=1,
|
| 34 |
+
max_length=2000,
|
| 35 |
+
description="User's natural language question",
|
| 36 |
+
examples=["How do I create a ROS 2 publisher in Python?"]
|
| 37 |
+
)
|
| 38 |
+
|
| 39 |
+
top_k: int = Field(
|
| 40 |
+
default=5,
|
| 41 |
+
ge=1,
|
| 42 |
+
le=50,
|
| 43 |
+
description="Number of documentation chunks to retrieve"
|
| 44 |
+
)
|
| 45 |
+
|
| 46 |
+
model: str = Field(
|
| 47 |
+
default="xiaomi/mimo-v2-flash:free",
|
| 48 |
+
description="LLM model to use for answer generation (OpenAI or OpenRouter)"
|
| 49 |
+
)
|
| 50 |
+
|
| 51 |
+
@field_validator('query')
|
| 52 |
+
@classmethod
|
| 53 |
+
def query_not_empty(cls, v: str) -> str:
|
| 54 |
+
"""
|
| 55 |
+
Validate that query is not empty or whitespace only.
|
| 56 |
+
|
| 57 |
+
Args:
|
| 58 |
+
v: Query string to validate
|
| 59 |
+
|
| 60 |
+
Returns:
|
| 61 |
+
Stripped query string
|
| 62 |
+
|
| 63 |
+
Raises:
|
| 64 |
+
ValueError: If query is empty or whitespace only
|
| 65 |
+
"""
|
| 66 |
+
if not v or not v.strip():
|
| 67 |
+
raise ValueError("Query cannot be empty or whitespace only")
|
| 68 |
+
return v.strip()
|
| 69 |
+
|
| 70 |
+
@field_validator('model')
|
| 71 |
+
@classmethod
|
| 72 |
+
def valid_model(cls, v: str) -> str:
|
| 73 |
+
"""
|
| 74 |
+
Validate that model is in the allowed list.
|
| 75 |
+
|
| 76 |
+
Args:
|
| 77 |
+
v: Model name to validate
|
| 78 |
+
|
| 79 |
+
Returns:
|
| 80 |
+
Model name if valid
|
| 81 |
+
|
| 82 |
+
Raises:
|
| 83 |
+
ValueError: If model is not in allowed list
|
| 84 |
+
"""
|
| 85 |
+
allowed_models = [
|
| 86 |
+
"gpt-3.5-turbo",
|
| 87 |
+
"gpt-4",
|
| 88 |
+
"gpt-4-turbo-preview",
|
| 89 |
+
"xiaomi/mimo-v2-flash:free"
|
| 90 |
+
]
|
| 91 |
+
if v not in allowed_models:
|
| 92 |
+
raise ValueError(f"Model must be one of {allowed_models}")
|
| 93 |
+
return v
|
| 94 |
+
|
| 95 |
+
model_config = {
|
| 96 |
+
"json_schema_extra": {
|
| 97 |
+
"examples": [
|
| 98 |
+
{
|
| 99 |
+
"query": "How do I create a ROS 2 publisher in Python?",
|
| 100 |
+
"top_k": 5,
|
| 101 |
+
"model": "xiaomi/mimo-v2-flash:free"
|
| 102 |
+
},
|
| 103 |
+
{
|
| 104 |
+
"query": "What is a URDF file?",
|
| 105 |
+
"top_k": 10,
|
| 106 |
+
"model": "gpt-4"
|
| 107 |
+
}
|
| 108 |
+
]
|
| 109 |
+
}
|
| 110 |
+
}
|
models/response_models.py
ADDED
|
@@ -0,0 +1,221 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Response Models for RAG Agent API
|
| 3 |
+
|
| 4 |
+
Pydantic models for API responses.
|
| 5 |
+
|
| 6 |
+
Models:
|
| 7 |
+
SourceReference: Reference to a documentation source
|
| 8 |
+
QueryResponse: Successful query response with answer and sources
|
| 9 |
+
ErrorResponse: Error response with details
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
from pydantic import BaseModel, Field, HttpUrl
|
| 13 |
+
from typing import List, Optional
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
class SourceReference(BaseModel):
|
| 17 |
+
"""
|
| 18 |
+
Reference to a documentation source used in answer generation.
|
| 19 |
+
|
| 20 |
+
Attributes:
|
| 21 |
+
url: URL of the source documentation page
|
| 22 |
+
page_title: Human-readable title of the page
|
| 23 |
+
relevance_score: Similarity score from vector search (0.0-1.0)
|
| 24 |
+
chunk_index: Position of chunk within the page (optional)
|
| 25 |
+
|
| 26 |
+
Example:
|
| 27 |
+
{
|
| 28 |
+
"url": "https://docs.ros.org/...",
|
| 29 |
+
"page_title": "Writing a simple publisher",
|
| 30 |
+
"relevance_score": 0.8642,
|
| 31 |
+
"chunk_index": 2
|
| 32 |
+
}
|
| 33 |
+
"""
|
| 34 |
+
|
| 35 |
+
url: str = Field(
|
| 36 |
+
...,
|
| 37 |
+
description="URL of the source documentation page"
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
page_title: str = Field(
|
| 41 |
+
...,
|
| 42 |
+
min_length=1,
|
| 43 |
+
description="Human-readable title of the documentation page"
|
| 44 |
+
)
|
| 45 |
+
|
| 46 |
+
relevance_score: float = Field(
|
| 47 |
+
...,
|
| 48 |
+
ge=0.0,
|
| 49 |
+
le=1.0,
|
| 50 |
+
description="Similarity score from vector search (0.0-1.0, higher is more relevant)"
|
| 51 |
+
)
|
| 52 |
+
|
| 53 |
+
chunk_index: Optional[int] = Field(
|
| 54 |
+
default=None,
|
| 55 |
+
ge=0,
|
| 56 |
+
description="Position of the chunk within the page (0-indexed)"
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
model_config = {
|
| 60 |
+
"json_schema_extra": {
|
| 61 |
+
"examples": [
|
| 62 |
+
{
|
| 63 |
+
"url": "https://docs.ros.org/en/rolling/Tutorials/Beginner-Client-Libraries/Writing-A-Simple-Py-Publisher.html",
|
| 64 |
+
"page_title": "Writing a simple publisher (Python)",
|
| 65 |
+
"relevance_score": 0.8642,
|
| 66 |
+
"chunk_index": 2
|
| 67 |
+
}
|
| 68 |
+
]
|
| 69 |
+
}
|
| 70 |
+
}
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
class QueryResponse(BaseModel):
|
| 74 |
+
"""
|
| 75 |
+
Successful response to a query request.
|
| 76 |
+
|
| 77 |
+
Attributes:
|
| 78 |
+
query: Original user query (echoed back)
|
| 79 |
+
answer: Generated natural language answer
|
| 80 |
+
sources: List of source references (deduplicated by URL)
|
| 81 |
+
model_used: Actual OpenAI model used for generation
|
| 82 |
+
tokens_used: Total tokens consumed (prompt + completion)
|
| 83 |
+
retrieval_count: Number of chunks retrieved before deduplication
|
| 84 |
+
|
| 85 |
+
Example:
|
| 86 |
+
{
|
| 87 |
+
"query": "How do I create a ROS 2 publisher?",
|
| 88 |
+
"answer": "To create a ROS 2 publisher...",
|
| 89 |
+
"sources": [...],
|
| 90 |
+
"model_used": "gpt-3.5-turbo",
|
| 91 |
+
"tokens_used": 456,
|
| 92 |
+
"retrieval_count": 5
|
| 93 |
+
}
|
| 94 |
+
"""
|
| 95 |
+
|
| 96 |
+
query: str = Field(
|
| 97 |
+
...,
|
| 98 |
+
description="Original user query (echoed back)"
|
| 99 |
+
)
|
| 100 |
+
|
| 101 |
+
answer: str = Field(
|
| 102 |
+
...,
|
| 103 |
+
min_length=1,
|
| 104 |
+
description="Generated natural language answer"
|
| 105 |
+
)
|
| 106 |
+
|
| 107 |
+
sources: List[SourceReference] = Field(
|
| 108 |
+
...,
|
| 109 |
+
description="List of source references (deduplicated by URL)"
|
| 110 |
+
)
|
| 111 |
+
|
| 112 |
+
model_used: str = Field(
|
| 113 |
+
...,
|
| 114 |
+
description="Actual OpenAI model used for generation"
|
| 115 |
+
)
|
| 116 |
+
|
| 117 |
+
tokens_used: int = Field(
|
| 118 |
+
...,
|
| 119 |
+
ge=0,
|
| 120 |
+
description="Total tokens consumed (prompt + completion)"
|
| 121 |
+
)
|
| 122 |
+
|
| 123 |
+
retrieval_count: int = Field(
|
| 124 |
+
...,
|
| 125 |
+
ge=0,
|
| 126 |
+
description="Number of chunks retrieved before deduplication"
|
| 127 |
+
)
|
| 128 |
+
|
| 129 |
+
model_config = {
|
| 130 |
+
"json_schema_extra": {
|
| 131 |
+
"examples": [
|
| 132 |
+
{
|
| 133 |
+
"query": "How do I create a ROS 2 publisher in Python?",
|
| 134 |
+
"answer": "To create a ROS 2 publisher in Python, you need to import rclpy, create a node class, and use create_publisher()...",
|
| 135 |
+
"sources": [
|
| 136 |
+
{
|
| 137 |
+
"url": "https://docs.ros.org/en/rolling/Tutorials/Beginner-Client-Libraries/Writing-A-Simple-Py-Publisher.html",
|
| 138 |
+
"page_title": "Writing a simple publisher (Python)",
|
| 139 |
+
"relevance_score": 0.8642,
|
| 140 |
+
"chunk_index": 2
|
| 141 |
+
}
|
| 142 |
+
],
|
| 143 |
+
"model_used": "gpt-3.5-turbo",
|
| 144 |
+
"tokens_used": 456,
|
| 145 |
+
"retrieval_count": 5
|
| 146 |
+
}
|
| 147 |
+
]
|
| 148 |
+
}
|
| 149 |
+
}
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
class ErrorResponse(BaseModel):
|
| 153 |
+
"""
|
| 154 |
+
Error response for failed requests.
|
| 155 |
+
|
| 156 |
+
Attributes:
|
| 157 |
+
error: Error flag (always True for error responses)
|
| 158 |
+
error_type: Error classification (ValidationError, RetrievalError, etc.)
|
| 159 |
+
error_message: Human-readable error description
|
| 160 |
+
query: Original query if available
|
| 161 |
+
status_code: HTTP status code
|
| 162 |
+
|
| 163 |
+
Example:
|
| 164 |
+
{
|
| 165 |
+
"error": true,
|
| 166 |
+
"error_type": "ValidationError",
|
| 167 |
+
"error_message": "Query cannot be empty",
|
| 168 |
+
"query": "",
|
| 169 |
+
"status_code": 422
|
| 170 |
+
}
|
| 171 |
+
"""
|
| 172 |
+
|
| 173 |
+
error: bool = Field(
|
| 174 |
+
default=True,
|
| 175 |
+
description="Error flag (always True for error responses)"
|
| 176 |
+
)
|
| 177 |
+
|
| 178 |
+
error_type: str = Field(
|
| 179 |
+
...,
|
| 180 |
+
min_length=1,
|
| 181 |
+
description="Error classification"
|
| 182 |
+
)
|
| 183 |
+
|
| 184 |
+
error_message: str = Field(
|
| 185 |
+
...,
|
| 186 |
+
min_length=1,
|
| 187 |
+
description="Human-readable error description"
|
| 188 |
+
)
|
| 189 |
+
|
| 190 |
+
query: Optional[str] = Field(
|
| 191 |
+
default=None,
|
| 192 |
+
description="Original query if available"
|
| 193 |
+
)
|
| 194 |
+
|
| 195 |
+
status_code: int = Field(
|
| 196 |
+
...,
|
| 197 |
+
ge=400,
|
| 198 |
+
le=599,
|
| 199 |
+
description="HTTP status code"
|
| 200 |
+
)
|
| 201 |
+
|
| 202 |
+
model_config = {
|
| 203 |
+
"json_schema_extra": {
|
| 204 |
+
"examples": [
|
| 205 |
+
{
|
| 206 |
+
"error": True,
|
| 207 |
+
"error_type": "ValidationError",
|
| 208 |
+
"error_message": "Query cannot be empty or whitespace only",
|
| 209 |
+
"query": "",
|
| 210 |
+
"status_code": 422
|
| 211 |
+
},
|
| 212 |
+
{
|
| 213 |
+
"error": True,
|
| 214 |
+
"error_type": "RetrievalError",
|
| 215 |
+
"error_message": "Vector database is currently unavailable",
|
| 216 |
+
"query": "How do I create a publisher?",
|
| 217 |
+
"status_code": 503
|
| 218 |
+
}
|
| 219 |
+
]
|
| 220 |
+
}
|
| 221 |
+
}
|
models/translation_models.py
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Pydantic Models for Translation API
|
| 3 |
+
|
| 4 |
+
Models:
|
| 5 |
+
- TranslationRequest: Request payload for translation
|
| 6 |
+
- TranslationResponse: Successful translation response
|
| 7 |
+
- TranslationError: Error response with error codes
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from pydantic import BaseModel, Field, validator
|
| 11 |
+
from typing import Optional
|
| 12 |
+
from datetime import datetime
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class TranslationRequest(BaseModel):
|
| 16 |
+
"""Request model for translation endpoint"""
|
| 17 |
+
|
| 18 |
+
text: str = Field(..., min_length=1, max_length=50000, description="Text content to translate")
|
| 19 |
+
source_language: str = Field(..., pattern="^en$", description="Source language code (only 'en' supported)")
|
| 20 |
+
target_language: str = Field(..., pattern="^ur$", description="Target language code (only 'ur' supported)")
|
| 21 |
+
content_hash: Optional[str] = Field(None, pattern="^[a-f0-9]{32}$", description="MD5 hash for cache lookup")
|
| 22 |
+
preserve_html: bool = Field(True, description="Whether to preserve HTML tags")
|
| 23 |
+
|
| 24 |
+
@validator('text')
|
| 25 |
+
def text_not_empty(cls, v):
|
| 26 |
+
"""Validate text is not empty or whitespace-only"""
|
| 27 |
+
if not v or v.isspace():
|
| 28 |
+
raise ValueError('Text cannot be empty or whitespace-only')
|
| 29 |
+
return v
|
| 30 |
+
|
| 31 |
+
class Config:
|
| 32 |
+
json_schema_extra = {
|
| 33 |
+
"example": {
|
| 34 |
+
"text": "<p>Welcome to the <strong>Physical AI</strong> guide</p>",
|
| 35 |
+
"source_language": "en",
|
| 36 |
+
"target_language": "ur",
|
| 37 |
+
"content_hash": "a3f5e8c9d2b14f6c8e7a9b2c3d4e5f6g",
|
| 38 |
+
"preserve_html": True
|
| 39 |
+
}
|
| 40 |
+
}
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
class TranslationResponse(BaseModel):
|
| 44 |
+
"""Response model for successful translation"""
|
| 45 |
+
|
| 46 |
+
translated_text: str = Field(..., description="Translated text in target language")
|
| 47 |
+
source_language: str = Field(..., description="Source language code (echoed)")
|
| 48 |
+
target_language: str = Field(..., description="Target language code (echoed)")
|
| 49 |
+
cached: bool = Field(..., description="Whether result was served from cache")
|
| 50 |
+
model_used: str = Field(..., description="Translation model identifier")
|
| 51 |
+
timestamp: str = Field(..., description="Timestamp of translation (ISO 8601)")
|
| 52 |
+
|
| 53 |
+
class Config:
|
| 54 |
+
json_schema_extra = {
|
| 55 |
+
"example": {
|
| 56 |
+
"translated_text": "<p>فزیکل AI گائیڈ میں <strong>خوش آمدید</strong></p>",
|
| 57 |
+
"source_language": "en",
|
| 58 |
+
"target_language": "ur",
|
| 59 |
+
"cached": False,
|
| 60 |
+
"model_used": "facebook/mbart-large-50",
|
| 61 |
+
"timestamp": "2025-12-22T10:30:00Z"
|
| 62 |
+
}
|
| 63 |
+
}
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
class TranslationError(BaseModel):
|
| 67 |
+
"""Error response model for translation failures"""
|
| 68 |
+
|
| 69 |
+
error: str = Field(..., description="Human-readable error message")
|
| 70 |
+
error_code: str = Field(..., description="Machine-readable error code")
|
| 71 |
+
retry_after: Optional[int] = Field(None, description="Seconds to wait before retry (for rate limits)")
|
| 72 |
+
timestamp: str = Field(..., description="Timestamp of error (ISO 8601)")
|
| 73 |
+
|
| 74 |
+
class Config:
|
| 75 |
+
json_schema_extra = {
|
| 76 |
+
"example": {
|
| 77 |
+
"error": "Translation service temporarily unavailable",
|
| 78 |
+
"error_code": "SERVICE_DOWN",
|
| 79 |
+
"retry_after": 60,
|
| 80 |
+
"timestamp": "2025-12-22T10:30:00Z"
|
| 81 |
+
}
|
| 82 |
+
}
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
# Error code constants
|
| 86 |
+
class ErrorCodes:
|
| 87 |
+
"""Standard error codes for translation failures"""
|
| 88 |
+
SERVICE_DOWN = "SERVICE_DOWN" # Translation model not loaded or crashed
|
| 89 |
+
RATE_LIMIT = "RATE_LIMIT" # Too many requests
|
| 90 |
+
INVALID_LANGUAGE = "INVALID_LANGUAGE" # Unsupported language pair
|
| 91 |
+
TIMEOUT = "TIMEOUT" # Translation exceeded time limit
|
| 92 |
+
INVALID_INPUT = "INVALID_INPUT" # Malformed HTML or empty text
|
rag_agent.py
ADDED
|
@@ -0,0 +1,270 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
RAG Agent API Service
|
| 3 |
+
|
| 4 |
+
FastAPI application providing a RAG (Retrieval-Augmented Generation) endpoint
|
| 5 |
+
for querying documentation and generating grounded answers.
|
| 6 |
+
|
| 7 |
+
Endpoints:
|
| 8 |
+
POST /query - Submit a question and get an answer with sources
|
| 9 |
+
GET /health - Health check endpoint
|
| 10 |
+
|
| 11 |
+
Features:
|
| 12 |
+
- Query validation with Pydantic
|
| 13 |
+
- Integration with Qdrant for document retrieval
|
| 14 |
+
- OpenAI integration for answer generation
|
| 15 |
+
- Source attribution and deduplication
|
| 16 |
+
- Structured logging
|
| 17 |
+
- CORS support for web clients
|
| 18 |
+
|
| 19 |
+
Usage:
|
| 20 |
+
uvicorn rag_agent:app --reload --host 0.0.0.0 --port 8000
|
| 21 |
+
"""
|
| 22 |
+
|
| 23 |
+
from fastapi import FastAPI, HTTPException
|
| 24 |
+
from fastapi.middleware.cors import CORSMiddleware
|
| 25 |
+
from fastapi.responses import JSONResponse
|
| 26 |
+
import logging
|
| 27 |
+
from datetime import datetime
|
| 28 |
+
|
| 29 |
+
from models.request_models import QueryRequest
|
| 30 |
+
from models.response_models import QueryResponse, ErrorResponse
|
| 31 |
+
from services.integration_service import IntegrationService
|
| 32 |
+
from exceptions import RAGException, ValidationError, RetrievalError, GenerationError, RateLimitError
|
| 33 |
+
from config.logging_config import setup_logging, get_logger
|
| 34 |
+
|
| 35 |
+
# Try to import translation router (optional - requires PyTorch)
|
| 36 |
+
translation_router = None
|
| 37 |
+
try:
|
| 38 |
+
from routers.translation import router as translation_router
|
| 39 |
+
except Exception as e:
|
| 40 |
+
print(f"[WARNING] Translation router not available: {e}")
|
| 41 |
+
print("[INFO] RAG agent will run without translation support")
|
| 42 |
+
|
| 43 |
+
# Initialize FastAPI app
|
| 44 |
+
app = FastAPI(
|
| 45 |
+
title="RAG Agent API",
|
| 46 |
+
description="Retrieval-Augmented Generation API for documentation Q&A",
|
| 47 |
+
version="1.0.0",
|
| 48 |
+
docs_url="/docs",
|
| 49 |
+
redoc_url="/redoc"
|
| 50 |
+
)
|
| 51 |
+
|
| 52 |
+
# Configure CORS for development (allow all origins)
|
| 53 |
+
app.add_middleware(
|
| 54 |
+
CORSMiddleware,
|
| 55 |
+
allow_origins=["*"], # In production, restrict to specific origins
|
| 56 |
+
allow_credentials=True,
|
| 57 |
+
allow_methods=["*"],
|
| 58 |
+
allow_headers=["*"],
|
| 59 |
+
)
|
| 60 |
+
|
| 61 |
+
# Configure logging
|
| 62 |
+
setup_logging(level="INFO", json_format=False)
|
| 63 |
+
logger = get_logger('rag_agent')
|
| 64 |
+
|
| 65 |
+
# Initialize integration service
|
| 66 |
+
integration_service = IntegrationService()
|
| 67 |
+
|
| 68 |
+
# Register routers (only if available)
|
| 69 |
+
if translation_router is not None:
|
| 70 |
+
app.include_router(translation_router)
|
| 71 |
+
print("[INFO] Translation router registered")
|
| 72 |
+
else:
|
| 73 |
+
print("[INFO] Running without translation support - use mock_translation_server.py for translations")
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
@app.get("/health")
|
| 77 |
+
async def health_check():
|
| 78 |
+
"""
|
| 79 |
+
Health check endpoint.
|
| 80 |
+
|
| 81 |
+
Returns:
|
| 82 |
+
dict: Health status with version and timestamp
|
| 83 |
+
"""
|
| 84 |
+
return {
|
| 85 |
+
"status": "healthy",
|
| 86 |
+
"version": "1.0.0",
|
| 87 |
+
"timestamp": datetime.utcnow().isoformat() + "Z"
|
| 88 |
+
}
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
@app.get("/")
|
| 92 |
+
async def root():
|
| 93 |
+
"""
|
| 94 |
+
Root endpoint with API information.
|
| 95 |
+
|
| 96 |
+
Returns:
|
| 97 |
+
dict: API info and documentation links
|
| 98 |
+
"""
|
| 99 |
+
return {
|
| 100 |
+
"name": "RAG Agent API",
|
| 101 |
+
"version": "1.0.0",
|
| 102 |
+
"description": "Retrieval-Augmented Generation API for documentation Q&A",
|
| 103 |
+
"docs": "/docs",
|
| 104 |
+
"redoc": "/redoc",
|
| 105 |
+
"health": "/health"
|
| 106 |
+
}
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
@app.post("/query", response_model=QueryResponse, responses={
|
| 110 |
+
422: {"model": ErrorResponse, "description": "Validation Error"},
|
| 111 |
+
503: {"model": ErrorResponse, "description": "Retrieval Error"},
|
| 112 |
+
502: {"model": ErrorResponse, "description": "Generation Error"},
|
| 113 |
+
429: {"model": ErrorResponse, "description": "Rate Limit Error"}
|
| 114 |
+
})
|
| 115 |
+
async def query_documentation(request: QueryRequest):
|
| 116 |
+
"""
|
| 117 |
+
Submit a question and get an answer with source attribution.
|
| 118 |
+
|
| 119 |
+
This endpoint implements the full RAG pipeline:
|
| 120 |
+
1. Validates the query
|
| 121 |
+
2. Retrieves relevant documentation chunks from Qdrant
|
| 122 |
+
3. Generates a grounded answer using OpenAI
|
| 123 |
+
4. Returns the answer with deduplicated source references
|
| 124 |
+
|
| 125 |
+
Args:
|
| 126 |
+
request: QueryRequest with query, top_k, and model
|
| 127 |
+
|
| 128 |
+
Returns:
|
| 129 |
+
QueryResponse with answer, sources, and metadata
|
| 130 |
+
|
| 131 |
+
Raises:
|
| 132 |
+
HTTPException: For validation, retrieval, generation, or rate limit errors
|
| 133 |
+
"""
|
| 134 |
+
start_time = datetime.utcnow()
|
| 135 |
+
|
| 136 |
+
logger.info(
|
| 137 |
+
"Received query request",
|
| 138 |
+
extra={
|
| 139 |
+
"query": request.query[:50] + "..." if len(request.query) > 50 else request.query,
|
| 140 |
+
"top_k": request.top_k,
|
| 141 |
+
"model": request.model
|
| 142 |
+
}
|
| 143 |
+
)
|
| 144 |
+
|
| 145 |
+
try:
|
| 146 |
+
# Process query through RAG pipeline
|
| 147 |
+
response = await integration_service.process_query(
|
| 148 |
+
query=request.query,
|
| 149 |
+
top_k=request.top_k,
|
| 150 |
+
model=request.model
|
| 151 |
+
)
|
| 152 |
+
|
| 153 |
+
# Log success
|
| 154 |
+
latency_ms = int((datetime.utcnow() - start_time).total_seconds() * 1000)
|
| 155 |
+
logger.info(
|
| 156 |
+
"Query processed successfully",
|
| 157 |
+
extra={
|
| 158 |
+
"query": request.query[:50] + "..." if len(request.query) > 50 else request.query,
|
| 159 |
+
"latency_ms": latency_ms,
|
| 160 |
+
"tokens_used": response.tokens_used,
|
| 161 |
+
"sources_count": len(response.sources)
|
| 162 |
+
}
|
| 163 |
+
)
|
| 164 |
+
|
| 165 |
+
return response
|
| 166 |
+
|
| 167 |
+
except ValidationError as e:
|
| 168 |
+
# Log validation error
|
| 169 |
+
logger.warning(
|
| 170 |
+
f"Validation error: {e.message}",
|
| 171 |
+
extra={
|
| 172 |
+
"query": request.query,
|
| 173 |
+
"error_type": e.error_type,
|
| 174 |
+
"status_code": e.status_code
|
| 175 |
+
}
|
| 176 |
+
)
|
| 177 |
+
return JSONResponse(
|
| 178 |
+
status_code=e.status_code,
|
| 179 |
+
content=ErrorResponse(
|
| 180 |
+
error_type=e.error_type,
|
| 181 |
+
error_message=e.message,
|
| 182 |
+
query=request.query,
|
| 183 |
+
status_code=e.status_code
|
| 184 |
+
).model_dump()
|
| 185 |
+
)
|
| 186 |
+
|
| 187 |
+
except RetrievalError as e:
|
| 188 |
+
# Log retrieval error
|
| 189 |
+
logger.error(
|
| 190 |
+
f"Retrieval error: {e.message}",
|
| 191 |
+
extra={
|
| 192 |
+
"query": request.query,
|
| 193 |
+
"error_type": e.error_type,
|
| 194 |
+
"status_code": e.status_code
|
| 195 |
+
}
|
| 196 |
+
)
|
| 197 |
+
return JSONResponse(
|
| 198 |
+
status_code=e.status_code,
|
| 199 |
+
content=ErrorResponse(
|
| 200 |
+
error_type=e.error_type,
|
| 201 |
+
error_message=e.message,
|
| 202 |
+
query=request.query,
|
| 203 |
+
status_code=e.status_code
|
| 204 |
+
).model_dump()
|
| 205 |
+
)
|
| 206 |
+
|
| 207 |
+
except GenerationError as e:
|
| 208 |
+
# Log generation error
|
| 209 |
+
logger.error(
|
| 210 |
+
f"Generation error: {e.message}",
|
| 211 |
+
extra={
|
| 212 |
+
"query": request.query,
|
| 213 |
+
"error_type": e.error_type,
|
| 214 |
+
"status_code": e.status_code
|
| 215 |
+
}
|
| 216 |
+
)
|
| 217 |
+
return JSONResponse(
|
| 218 |
+
status_code=e.status_code,
|
| 219 |
+
content=ErrorResponse(
|
| 220 |
+
error_type=e.error_type,
|
| 221 |
+
error_message=e.message,
|
| 222 |
+
query=request.query,
|
| 223 |
+
status_code=e.status_code
|
| 224 |
+
).model_dump()
|
| 225 |
+
)
|
| 226 |
+
|
| 227 |
+
except RateLimitError as e:
|
| 228 |
+
# Log rate limit error
|
| 229 |
+
logger.warning(
|
| 230 |
+
f"Rate limit error: {e.message}",
|
| 231 |
+
extra={
|
| 232 |
+
"query": request.query,
|
| 233 |
+
"error_type": e.error_type,
|
| 234 |
+
"status_code": e.status_code
|
| 235 |
+
}
|
| 236 |
+
)
|
| 237 |
+
return JSONResponse(
|
| 238 |
+
status_code=e.status_code,
|
| 239 |
+
content=ErrorResponse(
|
| 240 |
+
error_type=e.error_type,
|
| 241 |
+
error_message=e.message,
|
| 242 |
+
query=request.query,
|
| 243 |
+
status_code=e.status_code
|
| 244 |
+
).model_dump()
|
| 245 |
+
)
|
| 246 |
+
|
| 247 |
+
except Exception as e:
|
| 248 |
+
# Log unexpected error
|
| 249 |
+
logger.error(
|
| 250 |
+
f"Unexpected error: {str(e)}",
|
| 251 |
+
extra={
|
| 252 |
+
"query": request.query,
|
| 253 |
+
"error_type": "InternalServerError"
|
| 254 |
+
},
|
| 255 |
+
exc_info=True
|
| 256 |
+
)
|
| 257 |
+
return JSONResponse(
|
| 258 |
+
status_code=500,
|
| 259 |
+
content=ErrorResponse(
|
| 260 |
+
error_type="InternalServerError",
|
| 261 |
+
error_message=f"An unexpected error occurred: {str(e)}",
|
| 262 |
+
query=request.query,
|
| 263 |
+
status_code=500
|
| 264 |
+
).model_dump()
|
| 265 |
+
)
|
| 266 |
+
|
| 267 |
+
|
| 268 |
+
if __name__ == "__main__":
|
| 269 |
+
import uvicorn
|
| 270 |
+
uvicorn.run(app, host="0.0.0.0", port=8000)
|
requirements.txt
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Existing dependencies (from previous specs)
|
| 2 |
+
cohere>=5.20.1
|
| 3 |
+
qdrant-client>=1.16.2
|
| 4 |
+
beautifulsoup4>=4.14.3
|
| 5 |
+
requests>=2.32.5
|
| 6 |
+
lxml>=6.0.2
|
| 7 |
+
python-dotenv>=1.0.0
|
| 8 |
+
|
| 9 |
+
# RAG Agent API dependencies (002-rag-agent)
|
| 10 |
+
fastapi>=0.104.0
|
| 11 |
+
uvicorn[standard]>=0.24.0
|
| 12 |
+
pydantic>=2.0.0
|
| 13 |
+
openai>=1.0.0
|
| 14 |
+
tiktoken>=0.5.0
|
| 15 |
+
httpx>=0.25.0
|
| 16 |
+
|
| 17 |
+
# Urdu Translation Agent dependencies (005-urdu-translation)
|
| 18 |
+
transformers>=4.36.0
|
| 19 |
+
torch>=2.0.0
|
| 20 |
+
sentencepiece>=0.1.99 # Required for mBART tokenizer
|
retrieval.py
ADDED
|
@@ -0,0 +1,252 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Retrieval Pipeline Module
|
| 3 |
+
|
| 4 |
+
Provides functions for querying embedded documentation in Qdrant.
|
| 5 |
+
|
| 6 |
+
Functions:
|
| 7 |
+
embed_query: Generate query embedding using Cohere
|
| 8 |
+
search_qdrant: Search for similar chunks in Qdrant
|
| 9 |
+
format_results: Format search results as JSON
|
| 10 |
+
format_error: Format errors as JSON response
|
| 11 |
+
setup_logging: Configure structured logging
|
| 12 |
+
|
| 13 |
+
Example:
|
| 14 |
+
>>> from retrieval import embed_query, search_qdrant
|
| 15 |
+
>>> embedding = embed_query("What is ROS 2?")
|
| 16 |
+
>>> results = search_qdrant(embedding, top_k=5)
|
| 17 |
+
"""
|
| 18 |
+
|
| 19 |
+
import os
|
| 20 |
+
import logging
|
| 21 |
+
from dataclasses import dataclass
|
| 22 |
+
from typing import List, Dict, Any
|
| 23 |
+
from dotenv import load_dotenv
|
| 24 |
+
from cohere import Client as CohereClient
|
| 25 |
+
from qdrant_client import QdrantClient
|
| 26 |
+
from pathlib import Path
|
| 27 |
+
|
| 28 |
+
# Load environment variables - explicit path for uvicorn compatibility
|
| 29 |
+
_env_file = Path(__file__).parent / '.env'
|
| 30 |
+
load_dotenv(dotenv_path=_env_file, override=True)
|
| 31 |
+
|
| 32 |
+
# Constants
|
| 33 |
+
MODEL = "embed-english-v3.0"
|
| 34 |
+
COLLECTION_NAME = os.getenv("COLLECTION_NAME", "website_embeddings")
|
| 35 |
+
DEFAULT_TOP_K = 5
|
| 36 |
+
MAX_TOP_K = 50
|
| 37 |
+
|
| 38 |
+
# Logging setup
|
| 39 |
+
def setup_logging():
|
| 40 |
+
"""Configure structured logging"""
|
| 41 |
+
logging.basicConfig(
|
| 42 |
+
level=logging.INFO,
|
| 43 |
+
format='[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s',
|
| 44 |
+
datefmt='%Y-%m-%d %H:%M:%S'
|
| 45 |
+
)
|
| 46 |
+
|
| 47 |
+
logger = logging.getLogger('retrieval')
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
@dataclass
|
| 51 |
+
class SearchResult:
|
| 52 |
+
"""Single search result from Qdrant"""
|
| 53 |
+
rank: int
|
| 54 |
+
score: float
|
| 55 |
+
text: str
|
| 56 |
+
metadata: Dict[str, Any]
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def embed_query(query_text: str) -> List[float]:
|
| 60 |
+
"""
|
| 61 |
+
Generate embedding vector for query text using Cohere.
|
| 62 |
+
|
| 63 |
+
Args:
|
| 64 |
+
query_text: Natural language query string
|
| 65 |
+
|
| 66 |
+
Returns:
|
| 67 |
+
1024-dimensional embedding vector
|
| 68 |
+
|
| 69 |
+
Raises:
|
| 70 |
+
ValueError: If query_text is empty or None
|
| 71 |
+
CohereAPIError: If API call fails
|
| 72 |
+
"""
|
| 73 |
+
# Validation
|
| 74 |
+
if not query_text or not query_text.strip():
|
| 75 |
+
raise ValueError("Query text cannot be empty")
|
| 76 |
+
|
| 77 |
+
if len(query_text) > 2000:
|
| 78 |
+
raise ValueError("Query text too long (max 2000 characters)")
|
| 79 |
+
|
| 80 |
+
# Initialize Cohere client
|
| 81 |
+
api_key = os.getenv("COHERE_API_KEY")
|
| 82 |
+
if not api_key:
|
| 83 |
+
raise ValueError("COHERE_API_KEY not found in environment")
|
| 84 |
+
|
| 85 |
+
client = CohereClient(api_key=api_key)
|
| 86 |
+
|
| 87 |
+
# Log query embedding
|
| 88 |
+
logger.info(f"Embedding query: {query_text[:50]}...")
|
| 89 |
+
|
| 90 |
+
# Generate embedding
|
| 91 |
+
response = client.embed(
|
| 92 |
+
texts=[query_text],
|
| 93 |
+
model=MODEL,
|
| 94 |
+
input_type='search_query' # Important: different from 'search_document'
|
| 95 |
+
)
|
| 96 |
+
|
| 97 |
+
return response.embeddings[0]
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
def search_qdrant(
|
| 101 |
+
query_embedding: List[float],
|
| 102 |
+
top_k: int = DEFAULT_TOP_K,
|
| 103 |
+
collection_name: str = COLLECTION_NAME
|
| 104 |
+
) -> List[SearchResult]:
|
| 105 |
+
"""
|
| 106 |
+
Search Qdrant for similar chunks.
|
| 107 |
+
|
| 108 |
+
Args:
|
| 109 |
+
query_embedding: 1024-dim vector from embed_query()
|
| 110 |
+
top_k: Number of results to return (default 5, max 50)
|
| 111 |
+
collection_name: Qdrant collection to search
|
| 112 |
+
|
| 113 |
+
Returns:
|
| 114 |
+
List of SearchResult objects ranked by similarity
|
| 115 |
+
|
| 116 |
+
Raises:
|
| 117 |
+
ValueError: If top_k < 1 or > 50, or embedding dimension wrong
|
| 118 |
+
QdrantException: If search fails
|
| 119 |
+
"""
|
| 120 |
+
# Validation
|
| 121 |
+
if not isinstance(query_embedding, list) or len(query_embedding) != 1024:
|
| 122 |
+
raise ValueError("query_embedding must be list of 1024 floats")
|
| 123 |
+
|
| 124 |
+
if not 1 <= top_k <= MAX_TOP_K:
|
| 125 |
+
raise ValueError(f"top_k must be between 1 and {MAX_TOP_K}")
|
| 126 |
+
|
| 127 |
+
# Initialize Qdrant client
|
| 128 |
+
qdrant_url = os.getenv("QDRANT_URL")
|
| 129 |
+
qdrant_api_key = os.getenv("QDRANT_API_KEY")
|
| 130 |
+
|
| 131 |
+
if not qdrant_url:
|
| 132 |
+
raise ValueError("QDRANT_URL required")
|
| 133 |
+
|
| 134 |
+
# For local Qdrant instances, API key is not required
|
| 135 |
+
# Only require API key if using a remote Qdrant instance
|
| 136 |
+
if qdrant_url != "localhost:6333" and qdrant_url.startswith("http") and not qdrant_api_key:
|
| 137 |
+
raise ValueError("QDRANT_API_KEY required for remote Qdrant instances")
|
| 138 |
+
|
| 139 |
+
# Initialize Qdrant client based on URL type
|
| 140 |
+
if qdrant_url == "localhost:6333":
|
| 141 |
+
# For local instance, use host and port
|
| 142 |
+
client = QdrantClient(host="localhost", port=6333)
|
| 143 |
+
elif qdrant_api_key:
|
| 144 |
+
# For remote instance with API key
|
| 145 |
+
client = QdrantClient(url=qdrant_url, api_key=qdrant_api_key)
|
| 146 |
+
else:
|
| 147 |
+
# For remote instance without API key
|
| 148 |
+
client = QdrantClient(url=qdrant_url)
|
| 149 |
+
|
| 150 |
+
# Verify collection exists
|
| 151 |
+
try:
|
| 152 |
+
collection_info = client.get_collection(collection_name)
|
| 153 |
+
if collection_info.points_count == 0:
|
| 154 |
+
raise ValueError(f"Collection '{collection_name}' is empty")
|
| 155 |
+
except Exception as e:
|
| 156 |
+
raise ValueError(f"Collection '{collection_name}' not found or inaccessible: {e}")
|
| 157 |
+
|
| 158 |
+
# Log search
|
| 159 |
+
logger.info(f"Searching Qdrant for top {top_k} results...")
|
| 160 |
+
|
| 161 |
+
# Perform search
|
| 162 |
+
search_results = client.query_points(
|
| 163 |
+
collection_name=collection_name,
|
| 164 |
+
query=query_embedding,
|
| 165 |
+
limit=top_k,
|
| 166 |
+
with_payload=True,
|
| 167 |
+
with_vectors=False
|
| 168 |
+
)
|
| 169 |
+
|
| 170 |
+
# Convert to SearchResult objects
|
| 171 |
+
results = []
|
| 172 |
+
for rank, hit in enumerate(search_results.points, start=1):
|
| 173 |
+
result = SearchResult(
|
| 174 |
+
rank=rank,
|
| 175 |
+
score=hit.score,
|
| 176 |
+
text=hit.payload.get('text', ''),
|
| 177 |
+
metadata={
|
| 178 |
+
'url': hit.payload.get('url', ''),
|
| 179 |
+
'page_title': hit.payload.get('page_title', ''),
|
| 180 |
+
'chunk_index': hit.payload.get('chunk_index', -1),
|
| 181 |
+
'total_chunks': hit.payload.get('total_chunks', -1),
|
| 182 |
+
'source': hit.payload.get('source', ''),
|
| 183 |
+
'file_path': hit.payload.get('file_path', '')
|
| 184 |
+
}
|
| 185 |
+
)
|
| 186 |
+
results.append(result)
|
| 187 |
+
|
| 188 |
+
logger.info(f"Retrieved {len(results)} results")
|
| 189 |
+
return results
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
def format_results(
|
| 193 |
+
query: str,
|
| 194 |
+
results: List[SearchResult],
|
| 195 |
+
top_k: int
|
| 196 |
+
) -> dict:
|
| 197 |
+
"""
|
| 198 |
+
Format search results as JSON-serializable dict.
|
| 199 |
+
|
| 200 |
+
Args:
|
| 201 |
+
query: Original query text
|
| 202 |
+
results: List of search results
|
| 203 |
+
top_k: Requested number of results
|
| 204 |
+
|
| 205 |
+
Returns:
|
| 206 |
+
Dict conforming to response schema
|
| 207 |
+
"""
|
| 208 |
+
logger.info("Formatting results as JSON")
|
| 209 |
+
|
| 210 |
+
return {
|
| 211 |
+
"query": query,
|
| 212 |
+
"top_k": top_k,
|
| 213 |
+
"result_count": len(results),
|
| 214 |
+
"results": [
|
| 215 |
+
{
|
| 216 |
+
"rank": r.rank,
|
| 217 |
+
"score": round(r.score, 4), # Round to 4 decimal places
|
| 218 |
+
"text": r.text,
|
| 219 |
+
"metadata": {
|
| 220 |
+
"url": r.metadata['url'],
|
| 221 |
+
"page_title": r.metadata['page_title'],
|
| 222 |
+
"chunk_index": r.metadata['chunk_index'],
|
| 223 |
+
"total_chunks": r.metadata['total_chunks'],
|
| 224 |
+
"source": r.metadata['source']
|
| 225 |
+
}
|
| 226 |
+
}
|
| 227 |
+
for r in results
|
| 228 |
+
]
|
| 229 |
+
}
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
def format_error(error: Exception, query: str = "") -> dict:
|
| 233 |
+
"""
|
| 234 |
+
Format error as JSON response.
|
| 235 |
+
|
| 236 |
+
Args:
|
| 237 |
+
error: Exception that occurred
|
| 238 |
+
query: Original query text (if available)
|
| 239 |
+
|
| 240 |
+
Returns:
|
| 241 |
+
Error dict conforming to response schema
|
| 242 |
+
"""
|
| 243 |
+
logger.error(f"Error occurred: {type(error).__name__}: {str(error)}")
|
| 244 |
+
|
| 245 |
+
return {
|
| 246 |
+
"error": True,
|
| 247 |
+
"error_type": type(error).__name__,
|
| 248 |
+
"error_message": str(error),
|
| 249 |
+
"query": query,
|
| 250 |
+
"result_count": 0,
|
| 251 |
+
"results": []
|
| 252 |
+
}
|
services/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
# Services package
|
services/generation_service.py
ADDED
|
@@ -0,0 +1,239 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Answer Generation Service for RAG Agent
|
| 3 |
+
|
| 4 |
+
Handles interaction with OpenAI Chat Completions API to generate natural language
|
| 5 |
+
answers from documentation context. Includes retry logic, error handling, and
|
| 6 |
+
comprehensive logging.
|
| 7 |
+
|
| 8 |
+
This service:
|
| 9 |
+
1. Calls OpenAI Chat Completions API
|
| 10 |
+
2. Implements retry logic with exponential backoff
|
| 11 |
+
3. Handles authentication and rate limit errors
|
| 12 |
+
4. Logs generation metrics (tokens, latency, completion length)
|
| 13 |
+
5. Returns structured results with answer and token usage
|
| 14 |
+
|
| 15 |
+
Usage:
|
| 16 |
+
from services.generation_service import AnswerGenerationService
|
| 17 |
+
|
| 18 |
+
service = AnswerGenerationService(api_key="your-key")
|
| 19 |
+
result = service.generate_answer(prompt="<context>...</context>...", model="gpt-3.5-turbo")
|
| 20 |
+
print(result['answer'], result['tokens_used'])
|
| 21 |
+
"""
|
| 22 |
+
|
| 23 |
+
import os
|
| 24 |
+
import logging
|
| 25 |
+
import time
|
| 26 |
+
from typing import Dict, Optional
|
| 27 |
+
from datetime import datetime
|
| 28 |
+
from openai import OpenAI
|
| 29 |
+
|
| 30 |
+
from exceptions import GenerationError, RateLimitError
|
| 31 |
+
|
| 32 |
+
logger = logging.getLogger('generation_service')
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
class AnswerGenerationService:
|
| 36 |
+
"""
|
| 37 |
+
Service for generating natural language answers using OpenAI.
|
| 38 |
+
|
| 39 |
+
This service manages all OpenAI API interactions with:
|
| 40 |
+
- Retry logic for transient errors (429, 500, 502, 503)
|
| 41 |
+
- Error handling for non-retryable errors (401, 400)
|
| 42 |
+
- Comprehensive logging (model, tokens, latency)
|
| 43 |
+
- Deterministic generation (low temperature)
|
| 44 |
+
|
| 45 |
+
Attributes:
|
| 46 |
+
openai_client: OpenAI client instance
|
| 47 |
+
max_retries: Maximum retry attempts (default 3)
|
| 48 |
+
|
| 49 |
+
Methods:
|
| 50 |
+
generate_answer: Generate answer from prompt
|
| 51 |
+
"""
|
| 52 |
+
|
| 53 |
+
def __init__(self, api_key: Optional[str] = None, max_retries: int = 3):
|
| 54 |
+
"""
|
| 55 |
+
Initialize answer generation service.
|
| 56 |
+
|
| 57 |
+
Args:
|
| 58 |
+
api_key: OpenRouter API key (default: from env)
|
| 59 |
+
max_retries: Maximum retry attempts for transient errors (default 3)
|
| 60 |
+
|
| 61 |
+
Raises:
|
| 62 |
+
ValueError: If API key not provided and not in environment
|
| 63 |
+
"""
|
| 64 |
+
self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
|
| 65 |
+
if not self.api_key:
|
| 66 |
+
raise ValueError("OPENROUTER_API_KEY not found in environment")
|
| 67 |
+
|
| 68 |
+
self.openai_client = OpenAI(
|
| 69 |
+
api_key=self.api_key,
|
| 70 |
+
base_url="https://openrouter.ai/api/v1"
|
| 71 |
+
)
|
| 72 |
+
self.max_retries = max_retries
|
| 73 |
+
|
| 74 |
+
logger.info(
|
| 75 |
+
f"AnswerGenerationService initialized with OpenRouter",
|
| 76 |
+
extra={"max_retries": max_retries, "base_url": "https://openrouter.ai/api/v1"}
|
| 77 |
+
)
|
| 78 |
+
|
| 79 |
+
def generate_answer(
|
| 80 |
+
self,
|
| 81 |
+
prompt: str,
|
| 82 |
+
model: str = "xiaomi/mimo-v2-flash:free",
|
| 83 |
+
temperature: float = 0.1,
|
| 84 |
+
max_tokens: int = 500
|
| 85 |
+
) -> Dict[str, any]:
|
| 86 |
+
"""
|
| 87 |
+
Generate natural language answer from prompt using OpenAI.
|
| 88 |
+
|
| 89 |
+
Calls OpenAI Chat Completions API with retry logic for transient errors.
|
| 90 |
+
Uses low temperature (0.1) for deterministic, factual responses.
|
| 91 |
+
|
| 92 |
+
Args:
|
| 93 |
+
prompt: Full prompt with context and query
|
| 94 |
+
model: OpenAI model to use (default: gpt-3.5-turbo)
|
| 95 |
+
temperature: Sampling temperature (default: 0.1 for deterministic)
|
| 96 |
+
max_tokens: Maximum tokens in completion (default: 500)
|
| 97 |
+
|
| 98 |
+
Returns:
|
| 99 |
+
Dict with:
|
| 100 |
+
- answer: Generated text response
|
| 101 |
+
- tokens_used: Total tokens consumed (prompt + completion)
|
| 102 |
+
- model_used: Actual model used
|
| 103 |
+
- completion_length: Character length of answer
|
| 104 |
+
|
| 105 |
+
Raises:
|
| 106 |
+
GenerationError: If generation fails (non-retryable or max retries exceeded)
|
| 107 |
+
RateLimitError: If rate limit persists after retries
|
| 108 |
+
|
| 109 |
+
Example:
|
| 110 |
+
>>> service = AnswerGenerationService()
|
| 111 |
+
>>> result = service.generate_answer(
|
| 112 |
+
... prompt="<context>ROS 2 docs...</context><query>What is ROS 2?</query>...",
|
| 113 |
+
... model="gpt-3.5-turbo"
|
| 114 |
+
... )
|
| 115 |
+
>>> print(result['answer'])
|
| 116 |
+
>>> print(f"Used {result['tokens_used']} tokens")
|
| 117 |
+
"""
|
| 118 |
+
start_time = datetime.utcnow()
|
| 119 |
+
|
| 120 |
+
logger.info(
|
| 121 |
+
f"Starting answer generation",
|
| 122 |
+
extra={
|
| 123 |
+
"model": model,
|
| 124 |
+
"temperature": temperature,
|
| 125 |
+
"max_tokens": max_tokens,
|
| 126 |
+
"prompt_length": len(prompt)
|
| 127 |
+
}
|
| 128 |
+
)
|
| 129 |
+
|
| 130 |
+
# Retry loop for transient errors
|
| 131 |
+
last_error = None
|
| 132 |
+
for attempt in range(self.max_retries):
|
| 133 |
+
try:
|
| 134 |
+
# Call OpenAI Chat Completions API
|
| 135 |
+
response = self.openai_client.chat.completions.create(
|
| 136 |
+
model=model,
|
| 137 |
+
messages=[
|
| 138 |
+
{
|
| 139 |
+
"role": "system",
|
| 140 |
+
"content": "You are a helpful assistant that answers questions based on provided documentation. Always cite the documentation and be precise."
|
| 141 |
+
},
|
| 142 |
+
{
|
| 143 |
+
"role": "user",
|
| 144 |
+
"content": prompt
|
| 145 |
+
}
|
| 146 |
+
],
|
| 147 |
+
temperature=temperature,
|
| 148 |
+
max_tokens=max_tokens
|
| 149 |
+
)
|
| 150 |
+
|
| 151 |
+
# Extract answer and metrics
|
| 152 |
+
answer = response.choices[0].message.content
|
| 153 |
+
tokens_used = response.usage.total_tokens
|
| 154 |
+
completion_length = len(answer)
|
| 155 |
+
|
| 156 |
+
# Calculate latency
|
| 157 |
+
latency_ms = int((datetime.utcnow() - start_time).total_seconds() * 1000)
|
| 158 |
+
|
| 159 |
+
# Log success
|
| 160 |
+
logger.info(
|
| 161 |
+
f"Answer generated successfully",
|
| 162 |
+
extra={
|
| 163 |
+
"model": model,
|
| 164 |
+
"tokens_used": tokens_used,
|
| 165 |
+
"generation_latency_ms": latency_ms,
|
| 166 |
+
"completion_length": completion_length,
|
| 167 |
+
"attempt": attempt + 1
|
| 168 |
+
}
|
| 169 |
+
)
|
| 170 |
+
|
| 171 |
+
return {
|
| 172 |
+
"answer": answer,
|
| 173 |
+
"tokens_used": tokens_used,
|
| 174 |
+
"model_used": model,
|
| 175 |
+
"completion_length": completion_length
|
| 176 |
+
}
|
| 177 |
+
|
| 178 |
+
except Exception as e:
|
| 179 |
+
last_error = e
|
| 180 |
+
error_message = str(e)
|
| 181 |
+
error_type = type(e).__name__
|
| 182 |
+
|
| 183 |
+
# Check for rate limit errors (429) - retryable
|
| 184 |
+
if "rate_limit" in error_message.lower() or "429" in error_message:
|
| 185 |
+
if attempt < self.max_retries - 1:
|
| 186 |
+
# Calculate exponential backoff
|
| 187 |
+
wait_time = 2 ** attempt # 1s, 2s, 4s
|
| 188 |
+
logger.warning(
|
| 189 |
+
f"Rate limit hit, retrying in {wait_time}s (attempt {attempt + 1}/{self.max_retries})",
|
| 190 |
+
extra={
|
| 191 |
+
"error_type": error_type,
|
| 192 |
+
"attempt": attempt + 1,
|
| 193 |
+
"wait_time": wait_time
|
| 194 |
+
}
|
| 195 |
+
)
|
| 196 |
+
time.sleep(wait_time)
|
| 197 |
+
continue
|
| 198 |
+
else:
|
| 199 |
+
# Max retries exceeded for rate limit
|
| 200 |
+
logger.error(
|
| 201 |
+
f"Rate limit persists after {self.max_retries} retries",
|
| 202 |
+
extra={"error_message": error_message}
|
| 203 |
+
)
|
| 204 |
+
raise RateLimitError("OpenAI rate limit exceeded. Please try again later.")
|
| 205 |
+
|
| 206 |
+
# Check for server errors (500, 502, 503) - retryable
|
| 207 |
+
if any(code in error_message for code in ["500", "502", "503"]):
|
| 208 |
+
if attempt < self.max_retries - 1:
|
| 209 |
+
wait_time = 2 ** attempt
|
| 210 |
+
logger.warning(
|
| 211 |
+
f"Server error, retrying in {wait_time}s (attempt {attempt + 1}/{self.max_retries})",
|
| 212 |
+
extra={
|
| 213 |
+
"error_type": error_type,
|
| 214 |
+
"error_message": error_message,
|
| 215 |
+
"attempt": attempt + 1,
|
| 216 |
+
"wait_time": wait_time
|
| 217 |
+
}
|
| 218 |
+
)
|
| 219 |
+
time.sleep(wait_time)
|
| 220 |
+
continue
|
| 221 |
+
|
| 222 |
+
# Non-retryable errors (401 auth, 400 bad request, etc.)
|
| 223 |
+
logger.error(
|
| 224 |
+
f"Generation failed with non-retryable error: {error_message}",
|
| 225 |
+
extra={
|
| 226 |
+
"error_type": error_type,
|
| 227 |
+
"model": model,
|
| 228 |
+
"attempt": attempt + 1
|
| 229 |
+
},
|
| 230 |
+
exc_info=True
|
| 231 |
+
)
|
| 232 |
+
raise GenerationError(f"Failed to generate answer: {error_message}")
|
| 233 |
+
|
| 234 |
+
# If we exhausted all retries
|
| 235 |
+
logger.error(
|
| 236 |
+
f"Generation failed after {self.max_retries} retries",
|
| 237 |
+
extra={"last_error": str(last_error)}
|
| 238 |
+
)
|
| 239 |
+
raise GenerationError(f"Failed to generate answer after {self.max_retries} retries: {str(last_error)}")
|
services/integration_service.py
ADDED
|
@@ -0,0 +1,286 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Integration Service for RAG Agent
|
| 3 |
+
|
| 4 |
+
Orchestrates the full RAG pipeline:
|
| 5 |
+
1. Embed query using Cohere
|
| 6 |
+
2. Retrieve relevant chunks from Qdrant
|
| 7 |
+
3. Generate answer using OpenAI
|
| 8 |
+
4. Deduplicate sources and format response
|
| 9 |
+
|
| 10 |
+
This service integrates retrieval.py (existing) with OpenAI generation.
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
import os
|
| 14 |
+
import logging
|
| 15 |
+
from typing import List, Dict, Any, Optional
|
| 16 |
+
from dotenv import load_dotenv
|
| 17 |
+
|
| 18 |
+
from retrieval import SearchResult
|
| 19 |
+
from services.retrieval_service import RetrievalService
|
| 20 |
+
from services.prompt_builder import PromptBuilder
|
| 21 |
+
from services.generation_service import AnswerGenerationService
|
| 22 |
+
from models.response_models import QueryResponse, SourceReference
|
| 23 |
+
from exceptions import RetrievalError, GenerationError, RateLimitError
|
| 24 |
+
|
| 25 |
+
# Load environment variables
|
| 26 |
+
load_dotenv()
|
| 27 |
+
|
| 28 |
+
# Constants
|
| 29 |
+
DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "xiaomi/mimo-v2-flash:free")
|
| 30 |
+
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "3000"))
|
| 31 |
+
# NOTE: Threshold set to 0.15 based on actual retrieval performance
|
| 32 |
+
# The current embedding model (Cohere embed-english-v3.0) returns lower scores than expected.
|
| 33 |
+
# Original requirement (FR-005) specified 0.7, but testing shows:
|
| 34 |
+
# - Technical queries get scores of 0.15-0.30
|
| 35 |
+
# - Irrelevant queries (like "hi") get scores of 0.10-0.23
|
| 36 |
+
# Setting to 0.15 filters out very low-quality matches while keeping relevant docs
|
| 37 |
+
RELEVANCE_THRESHOLD = 0.15 # Minimum relevance score for usable results (FR-005)
|
| 38 |
+
|
| 39 |
+
# Greeting patterns for detection (FR-003)
|
| 40 |
+
GREETING_PATTERNS = {
|
| 41 |
+
"hi", "hello", "hey", "helo", "hii", "heya", "greetings",
|
| 42 |
+
"good morning", "good afternoon", "good evening"
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
logger = logging.getLogger('integration_service')
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
class IntegrationService:
|
| 49 |
+
"""
|
| 50 |
+
Service for orchestrating RAG pipeline: retrieval + generation.
|
| 51 |
+
|
| 52 |
+
Methods:
|
| 53 |
+
process_query: Main entry point for RAG pipeline
|
| 54 |
+
_is_greeting: Detect if query is a casual greeting
|
| 55 |
+
_retrieve_chunks: Retrieve relevant documentation chunks
|
| 56 |
+
_deduplicate_sources: Deduplicate chunks by URL
|
| 57 |
+
_build_prompt: Build prompt with context and query
|
| 58 |
+
_generate_answer: Generate answer using OpenAI
|
| 59 |
+
_count_tokens: Count tokens in text
|
| 60 |
+
"""
|
| 61 |
+
|
| 62 |
+
def __init__(self, openai_api_key: Optional[str] = None):
|
| 63 |
+
"""
|
| 64 |
+
Initialize integration service.
|
| 65 |
+
|
| 66 |
+
Args:
|
| 67 |
+
openai_api_key: OpenRouter API key (default: from env)
|
| 68 |
+
"""
|
| 69 |
+
self.openai_api_key = openai_api_key or os.getenv("OPENROUTER_API_KEY")
|
| 70 |
+
if not self.openai_api_key:
|
| 71 |
+
raise ValueError("OPENROUTER_API_KEY not found in environment")
|
| 72 |
+
|
| 73 |
+
self.retrieval_service = RetrievalService()
|
| 74 |
+
self.prompt_builder = PromptBuilder(max_tokens=MAX_TOKENS)
|
| 75 |
+
self.generation_service = AnswerGenerationService(api_key=self.openai_api_key)
|
| 76 |
+
|
| 77 |
+
def _is_greeting(self, query: str) -> bool:
|
| 78 |
+
"""
|
| 79 |
+
Detect if query is a standalone greeting (FR-003).
|
| 80 |
+
|
| 81 |
+
Returns True for casual greetings like "hi", "hello", "hey".
|
| 82 |
+
Returns False for queries that contain greetings but have actual questions
|
| 83 |
+
(e.g., "Hi, how do I create a ROS 2 publisher?").
|
| 84 |
+
|
| 85 |
+
Args:
|
| 86 |
+
query: User query text
|
| 87 |
+
|
| 88 |
+
Returns:
|
| 89 |
+
True if query is a greeting, False otherwise
|
| 90 |
+
"""
|
| 91 |
+
# Normalize: strip, lowercase, remove punctuation
|
| 92 |
+
normalized = query.strip().lower().rstrip('!?.,:;')
|
| 93 |
+
|
| 94 |
+
# Check if entire query is just a greeting (standalone)
|
| 95 |
+
# Don't match if query contains other words after greeting
|
| 96 |
+
return normalized in GREETING_PATTERNS
|
| 97 |
+
|
| 98 |
+
async def process_query(
|
| 99 |
+
self,
|
| 100 |
+
query: str,
|
| 101 |
+
top_k: int = 5,
|
| 102 |
+
model: str = DEFAULT_MODEL
|
| 103 |
+
) -> QueryResponse:
|
| 104 |
+
"""
|
| 105 |
+
Process a query through the full RAG pipeline.
|
| 106 |
+
|
| 107 |
+
Args:
|
| 108 |
+
query: User's natural language question
|
| 109 |
+
top_k: Number of chunks to retrieve
|
| 110 |
+
model: OpenAI model to use
|
| 111 |
+
|
| 112 |
+
Returns:
|
| 113 |
+
QueryResponse with answer and sources
|
| 114 |
+
|
| 115 |
+
Raises:
|
| 116 |
+
RetrievalError: If retrieval fails
|
| 117 |
+
GenerationError: If generation fails
|
| 118 |
+
RateLimitError: If rate limit exceeded
|
| 119 |
+
"""
|
| 120 |
+
logger.info(f"Processing query", extra={"query": query, "top_k": top_k, "model": model})
|
| 121 |
+
|
| 122 |
+
# Step 0: Check for greetings FIRST (FR-003)
|
| 123 |
+
if self._is_greeting(query):
|
| 124 |
+
logger.info(f"Greeting detected: {query}")
|
| 125 |
+
return QueryResponse(
|
| 126 |
+
query=query,
|
| 127 |
+
answer="Hello! I'm here to help you learn about ROS 2 and humanoid robotics. What would you like to know?",
|
| 128 |
+
sources=[],
|
| 129 |
+
model_used="greeting-detector",
|
| 130 |
+
tokens_used=0,
|
| 131 |
+
retrieval_count=0
|
| 132 |
+
)
|
| 133 |
+
|
| 134 |
+
# Step 1: Retrieve chunks
|
| 135 |
+
chunks = await self._retrieve_chunks(query, top_k)
|
| 136 |
+
retrieval_count = len(chunks)
|
| 137 |
+
|
| 138 |
+
# Step 2: Deduplicate sources
|
| 139 |
+
deduplicated_sources = self._deduplicate_sources(chunks)
|
| 140 |
+
|
| 141 |
+
# Step 3: Build prompt
|
| 142 |
+
prompt = self._build_prompt(query, chunks, model)
|
| 143 |
+
|
| 144 |
+
# Step 4: Generate answer
|
| 145 |
+
answer, tokens_used = await self._generate_answer(prompt, model)
|
| 146 |
+
|
| 147 |
+
# Step 5: Format response
|
| 148 |
+
response = QueryResponse(
|
| 149 |
+
query=query,
|
| 150 |
+
answer=answer,
|
| 151 |
+
sources=deduplicated_sources,
|
| 152 |
+
model_used=model,
|
| 153 |
+
tokens_used=tokens_used,
|
| 154 |
+
retrieval_count=retrieval_count
|
| 155 |
+
)
|
| 156 |
+
|
| 157 |
+
logger.info(
|
| 158 |
+
f"Query processed successfully",
|
| 159 |
+
extra={
|
| 160 |
+
"query": query,
|
| 161 |
+
"tokens_used": tokens_used,
|
| 162 |
+
"retrieval_count": retrieval_count,
|
| 163 |
+
"sources_count": len(deduplicated_sources)
|
| 164 |
+
}
|
| 165 |
+
)
|
| 166 |
+
|
| 167 |
+
return response
|
| 168 |
+
|
| 169 |
+
async def _retrieve_chunks(self, query: str, top_k: int) -> List[SearchResult]:
|
| 170 |
+
"""
|
| 171 |
+
Retrieve relevant documentation chunks using RetrievalService.
|
| 172 |
+
Filters results by relevance threshold (FR-005).
|
| 173 |
+
|
| 174 |
+
Args:
|
| 175 |
+
query: User query
|
| 176 |
+
top_k: Number of chunks to retrieve
|
| 177 |
+
|
| 178 |
+
Returns:
|
| 179 |
+
List of SearchResult objects with relevance >= 0.7
|
| 180 |
+
|
| 181 |
+
Raises:
|
| 182 |
+
RetrievalError: If retrieval fails or no relevant docs found (FR-006)
|
| 183 |
+
"""
|
| 184 |
+
# Delegate to RetrievalService (already has error handling and logging)
|
| 185 |
+
chunks = self.retrieval_service.retrieve_context(query, top_k=top_k)
|
| 186 |
+
|
| 187 |
+
# Filter by relevance threshold (FR-005)
|
| 188 |
+
filtered_chunks = [c for c in chunks if c.score >= RELEVANCE_THRESHOLD]
|
| 189 |
+
|
| 190 |
+
logger.info(
|
| 191 |
+
f"Relevance filtering: {len(filtered_chunks)}/{len(chunks)} chunks pass threshold {RELEVANCE_THRESHOLD}",
|
| 192 |
+
extra={
|
| 193 |
+
"total_chunks": len(chunks),
|
| 194 |
+
"filtered_chunks": len(filtered_chunks),
|
| 195 |
+
"threshold": RELEVANCE_THRESHOLD
|
| 196 |
+
}
|
| 197 |
+
)
|
| 198 |
+
|
| 199 |
+
# Handle no relevant docs case (FR-006)
|
| 200 |
+
if not filtered_chunks:
|
| 201 |
+
logger.warning(f"No documents meet relevance threshold {RELEVANCE_THRESHOLD}")
|
| 202 |
+
raise RetrievalError(
|
| 203 |
+
"I couldn't find relevant information in the documentation. "
|
| 204 |
+
"Try rephrasing your question or ask about: ROS 2, Unity, Digital Twins, or Humanoid Robotics."
|
| 205 |
+
)
|
| 206 |
+
|
| 207 |
+
return filtered_chunks
|
| 208 |
+
|
| 209 |
+
def _deduplicate_sources(self, chunks: List[SearchResult]) -> List[SourceReference]:
|
| 210 |
+
"""
|
| 211 |
+
Deduplicate chunks by URL, keeping highest-scoring chunk.
|
| 212 |
+
|
| 213 |
+
Args:
|
| 214 |
+
chunks: List of retrieved chunks
|
| 215 |
+
|
| 216 |
+
Returns:
|
| 217 |
+
List of deduplicated SourceReference objects
|
| 218 |
+
"""
|
| 219 |
+
# Group by URL, keep highest score
|
| 220 |
+
url_to_best_chunk: Dict[str, tuple[SearchResult, float]] = {}
|
| 221 |
+
|
| 222 |
+
for chunk in chunks:
|
| 223 |
+
url = chunk.metadata.get('url', '')
|
| 224 |
+
if not url:
|
| 225 |
+
continue
|
| 226 |
+
|
| 227 |
+
if url not in url_to_best_chunk or chunk.score > url_to_best_chunk[url][1]:
|
| 228 |
+
url_to_best_chunk[url] = (chunk, chunk.score)
|
| 229 |
+
|
| 230 |
+
# Convert to SourceReference
|
| 231 |
+
sources = []
|
| 232 |
+
for url, (chunk, score) in url_to_best_chunk.items():
|
| 233 |
+
source = SourceReference(
|
| 234 |
+
url=url,
|
| 235 |
+
page_title=chunk.metadata.get('page_title') or 'Unknown',
|
| 236 |
+
relevance_score=round(score, 4),
|
| 237 |
+
chunk_index=chunk.metadata.get('chunk_index')
|
| 238 |
+
)
|
| 239 |
+
sources.append(source)
|
| 240 |
+
|
| 241 |
+
# Sort by relevance score (descending)
|
| 242 |
+
sources.sort(key=lambda s: s.relevance_score, reverse=True)
|
| 243 |
+
|
| 244 |
+
logger.info(f"Deduplicated {len(chunks)} chunks to {len(sources)} unique sources")
|
| 245 |
+
return sources
|
| 246 |
+
|
| 247 |
+
def _build_prompt(self, query: str, chunks: List[SearchResult], model: str = "gpt-3.5-turbo") -> str:
|
| 248 |
+
"""
|
| 249 |
+
Build prompt with context and query using PromptBuilder.
|
| 250 |
+
|
| 251 |
+
Args:
|
| 252 |
+
query: User query
|
| 253 |
+
chunks: Retrieved chunks
|
| 254 |
+
model: OpenAI model for token counting
|
| 255 |
+
|
| 256 |
+
Returns:
|
| 257 |
+
Formatted prompt string
|
| 258 |
+
"""
|
| 259 |
+
# Delegate to PromptBuilder (already has XML formatting, token management, source URLs)
|
| 260 |
+
prompt = self.prompt_builder.build_prompt(query, chunks, model)
|
| 261 |
+
return prompt
|
| 262 |
+
|
| 263 |
+
async def _generate_answer(self, prompt: str, model: str) -> tuple[str, int]:
|
| 264 |
+
"""
|
| 265 |
+
Generate answer using AnswerGenerationService.
|
| 266 |
+
|
| 267 |
+
Args:
|
| 268 |
+
prompt: Full prompt with context and query
|
| 269 |
+
model: OpenAI model to use
|
| 270 |
+
|
| 271 |
+
Returns:
|
| 272 |
+
Tuple of (answer, tokens_used)
|
| 273 |
+
|
| 274 |
+
Raises:
|
| 275 |
+
GenerationError: If generation fails (propagated from AnswerGenerationService)
|
| 276 |
+
RateLimitError: If rate limit exceeded (propagated from AnswerGenerationService)
|
| 277 |
+
"""
|
| 278 |
+
# Delegate to AnswerGenerationService (already has retry logic, error handling, logging)
|
| 279 |
+
result = self.generation_service.generate_answer(
|
| 280 |
+
prompt=prompt,
|
| 281 |
+
model=model,
|
| 282 |
+
temperature=0.1, # Low temperature for deterministic, factual responses
|
| 283 |
+
max_tokens=500
|
| 284 |
+
)
|
| 285 |
+
|
| 286 |
+
return result['answer'], result['tokens_used']
|
services/prompt_builder.py
ADDED
|
@@ -0,0 +1,264 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Prompt Builder Service for RAG Agent
|
| 3 |
+
|
| 4 |
+
Combines retrieved documentation chunks with user queries into well-formatted
|
| 5 |
+
prompts for OpenAI. Handles token management, XML formatting, and source attribution.
|
| 6 |
+
|
| 7 |
+
This service:
|
| 8 |
+
1. Formats context with XML delimiters
|
| 9 |
+
2. Manages token budgets to stay within API limits
|
| 10 |
+
3. Selects top-scoring chunks that fit
|
| 11 |
+
4. Includes source metadata for attribution
|
| 12 |
+
5. Adds system instructions for grounded answers
|
| 13 |
+
|
| 14 |
+
Usage:
|
| 15 |
+
from services.prompt_builder import PromptBuilder
|
| 16 |
+
|
| 17 |
+
builder = PromptBuilder(max_tokens=3000)
|
| 18 |
+
prompt = builder.build_prompt(query="What is ROS 2?", chunks=results)
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
import logging
|
| 22 |
+
from typing import List
|
| 23 |
+
from retrieval import SearchResult
|
| 24 |
+
from utils.token_counter import count_tokens
|
| 25 |
+
|
| 26 |
+
logger = logging.getLogger('prompt_builder')
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
class PromptBuilder:
|
| 30 |
+
"""
|
| 31 |
+
Service for building prompts with context and query.
|
| 32 |
+
|
| 33 |
+
This service creates well-formatted prompts that:
|
| 34 |
+
- Use XML tags (<context>, <query>) for clear structure
|
| 35 |
+
- Stay within token limits
|
| 36 |
+
- Include source URLs for attribution
|
| 37 |
+
- Provide system instructions for grounded answers
|
| 38 |
+
|
| 39 |
+
Attributes:
|
| 40 |
+
max_tokens: Maximum tokens for entire prompt (default 3000)
|
| 41 |
+
reserved_tokens: Tokens reserved for query + instructions (default 500)
|
| 42 |
+
|
| 43 |
+
Methods:
|
| 44 |
+
build_prompt: Build formatted prompt from query and chunks
|
| 45 |
+
"""
|
| 46 |
+
|
| 47 |
+
def __init__(self, max_tokens: int = 3000, reserved_tokens: int = 500):
|
| 48 |
+
"""
|
| 49 |
+
Initialize prompt builder.
|
| 50 |
+
|
| 51 |
+
Args:
|
| 52 |
+
max_tokens: Maximum tokens for entire prompt (default 3000)
|
| 53 |
+
reserved_tokens: Tokens reserved for query/instructions (default 500)
|
| 54 |
+
"""
|
| 55 |
+
self.max_tokens = max_tokens
|
| 56 |
+
self.reserved_tokens = reserved_tokens
|
| 57 |
+
self.context_budget = max_tokens - reserved_tokens
|
| 58 |
+
|
| 59 |
+
logger.info(
|
| 60 |
+
f"PromptBuilder initialized",
|
| 61 |
+
extra={
|
| 62 |
+
"max_tokens": max_tokens,
|
| 63 |
+
"reserved_tokens": reserved_tokens,
|
| 64 |
+
"context_budget": self.context_budget
|
| 65 |
+
}
|
| 66 |
+
)
|
| 67 |
+
|
| 68 |
+
def build_prompt(self, query: str, chunks: List[SearchResult], model: str = "gpt-3.5-turbo") -> str:
|
| 69 |
+
"""
|
| 70 |
+
Build formatted prompt with context and query.
|
| 71 |
+
|
| 72 |
+
Creates a prompt with the following structure:
|
| 73 |
+
```
|
| 74 |
+
<context>
|
| 75 |
+
[chunk 1 text]
|
| 76 |
+
Source: [url 1]
|
| 77 |
+
|
| 78 |
+
[chunk 2 text]
|
| 79 |
+
Source: [url 2]
|
| 80 |
+
...
|
| 81 |
+
</context>
|
| 82 |
+
|
| 83 |
+
<query>
|
| 84 |
+
[user query]
|
| 85 |
+
</query>
|
| 86 |
+
|
| 87 |
+
Please answer the query using ONLY the information provided...
|
| 88 |
+
```
|
| 89 |
+
|
| 90 |
+
The method:
|
| 91 |
+
1. Selects chunks that fit within token budget
|
| 92 |
+
2. Formats each chunk with source URL
|
| 93 |
+
3. Wraps in XML tags
|
| 94 |
+
4. Adds system instructions
|
| 95 |
+
|
| 96 |
+
Args:
|
| 97 |
+
query: User's natural language question
|
| 98 |
+
chunks: Retrieved documentation chunks (sorted by relevance)
|
| 99 |
+
model: OpenAI model for token counting (default: gpt-3.5-turbo)
|
| 100 |
+
|
| 101 |
+
Returns:
|
| 102 |
+
Formatted prompt string ready for OpenAI API
|
| 103 |
+
|
| 104 |
+
Example:
|
| 105 |
+
>>> builder = PromptBuilder(max_tokens=3000)
|
| 106 |
+
>>> prompt = builder.build_prompt("What is ROS 2?", chunks)
|
| 107 |
+
>>> print(prompt[:100])
|
| 108 |
+
<context>
|
| 109 |
+
ROS 2 is a set of software libraries and tools...
|
| 110 |
+
"""
|
| 111 |
+
logger.info(
|
| 112 |
+
f"Building prompt",
|
| 113 |
+
extra={
|
| 114 |
+
"query": query[:50] + "..." if len(query) > 50 else query,
|
| 115 |
+
"total_chunks": len(chunks),
|
| 116 |
+
"model": model
|
| 117 |
+
}
|
| 118 |
+
)
|
| 119 |
+
|
| 120 |
+
# Select chunks that fit within budget (T036)
|
| 121 |
+
selected_chunks = self._select_chunks_within_budget(chunks, model)
|
| 122 |
+
|
| 123 |
+
# Format context with chunks and source URLs (T037)
|
| 124 |
+
context = self._format_context_with_sources(selected_chunks)
|
| 125 |
+
|
| 126 |
+
# Build full prompt with XML template (T034)
|
| 127 |
+
prompt = self._format_prompt_with_xml(context, query)
|
| 128 |
+
|
| 129 |
+
# Count final tokens
|
| 130 |
+
final_tokens = count_tokens(prompt, model)
|
| 131 |
+
|
| 132 |
+
logger.info(
|
| 133 |
+
f"Prompt built successfully",
|
| 134 |
+
extra={
|
| 135 |
+
"query": query[:50] + "..." if len(query) > 50 else query,
|
| 136 |
+
"selected_chunks": len(selected_chunks),
|
| 137 |
+
"total_chunks": len(chunks),
|
| 138 |
+
"final_tokens": final_tokens,
|
| 139 |
+
"max_tokens": self.max_tokens
|
| 140 |
+
}
|
| 141 |
+
)
|
| 142 |
+
|
| 143 |
+
return prompt
|
| 144 |
+
|
| 145 |
+
def _select_chunks_within_budget(
|
| 146 |
+
self,
|
| 147 |
+
chunks: List[SearchResult],
|
| 148 |
+
model: str
|
| 149 |
+
) -> List[SearchResult]:
|
| 150 |
+
"""
|
| 151 |
+
Select top-scoring chunks that fit within token budget.
|
| 152 |
+
|
| 153 |
+
Iterates through chunks (already sorted by relevance) and adds
|
| 154 |
+
them to context until budget is exhausted. Prioritizes higher-scoring
|
| 155 |
+
chunks when space is limited.
|
| 156 |
+
|
| 157 |
+
Args:
|
| 158 |
+
chunks: Retrieved chunks sorted by relevance (highest first)
|
| 159 |
+
model: Model for token counting
|
| 160 |
+
|
| 161 |
+
Returns:
|
| 162 |
+
List of chunks that fit within context_budget
|
| 163 |
+
"""
|
| 164 |
+
selected = []
|
| 165 |
+
total_tokens = 0
|
| 166 |
+
|
| 167 |
+
for chunk in chunks:
|
| 168 |
+
# Format chunk with source for accurate token count
|
| 169 |
+
chunk_text = chunk.text
|
| 170 |
+
source_url = chunk.metadata.get('url', '')
|
| 171 |
+
formatted_chunk = f"{chunk_text}\nSource: {source_url}\n"
|
| 172 |
+
|
| 173 |
+
# Count tokens for this chunk
|
| 174 |
+
chunk_tokens = count_tokens(formatted_chunk, model)
|
| 175 |
+
|
| 176 |
+
# Check if adding this chunk would exceed budget
|
| 177 |
+
if total_tokens + chunk_tokens > self.context_budget:
|
| 178 |
+
logger.warning(
|
| 179 |
+
f"Hit token budget, truncating context at {len(selected)} chunks",
|
| 180 |
+
extra={
|
| 181 |
+
"selected_chunks": len(selected),
|
| 182 |
+
"total_chunks": len(chunks),
|
| 183 |
+
"current_tokens": total_tokens,
|
| 184 |
+
"context_budget": self.context_budget
|
| 185 |
+
}
|
| 186 |
+
)
|
| 187 |
+
break
|
| 188 |
+
|
| 189 |
+
selected.append(chunk)
|
| 190 |
+
total_tokens += chunk_tokens
|
| 191 |
+
|
| 192 |
+
logger.debug(
|
| 193 |
+
f"Selected {len(selected)} chunks using {total_tokens} tokens"
|
| 194 |
+
)
|
| 195 |
+
|
| 196 |
+
return selected
|
| 197 |
+
|
| 198 |
+
def _format_context_with_sources(self, chunks: List[SearchResult]) -> str:
|
| 199 |
+
"""
|
| 200 |
+
Format chunks with source URLs.
|
| 201 |
+
|
| 202 |
+
Each chunk is formatted as:
|
| 203 |
+
```
|
| 204 |
+
[chunk text]
|
| 205 |
+
Source: [url]
|
| 206 |
+
```
|
| 207 |
+
|
| 208 |
+
Args:
|
| 209 |
+
chunks: Selected chunks to format
|
| 210 |
+
|
| 211 |
+
Returns:
|
| 212 |
+
Formatted context string with sources
|
| 213 |
+
"""
|
| 214 |
+
context_parts = []
|
| 215 |
+
|
| 216 |
+
for chunk in chunks:
|
| 217 |
+
chunk_text = chunk.text
|
| 218 |
+
source_url = chunk.metadata.get('url', 'Unknown')
|
| 219 |
+
|
| 220 |
+
# Format: chunk text + source URL
|
| 221 |
+
formatted_chunk = f"{chunk_text}\nSource: {source_url}"
|
| 222 |
+
context_parts.append(formatted_chunk)
|
| 223 |
+
|
| 224 |
+
# Join with blank line between chunks
|
| 225 |
+
context = "\n\n".join(context_parts)
|
| 226 |
+
|
| 227 |
+
return context
|
| 228 |
+
|
| 229 |
+
def _format_prompt_with_xml(self, context: str, query: str) -> str:
|
| 230 |
+
"""
|
| 231 |
+
Format prompt with XML tags and system instructions.
|
| 232 |
+
|
| 233 |
+
Creates the final prompt structure:
|
| 234 |
+
```
|
| 235 |
+
<context>
|
| 236 |
+
[formatted context with sources]
|
| 237 |
+
</context>
|
| 238 |
+
|
| 239 |
+
<query>
|
| 240 |
+
[user query]
|
| 241 |
+
</query>
|
| 242 |
+
|
| 243 |
+
[system instructions]
|
| 244 |
+
```
|
| 245 |
+
|
| 246 |
+
Args:
|
| 247 |
+
context: Formatted context with chunks and sources
|
| 248 |
+
query: User query
|
| 249 |
+
|
| 250 |
+
Returns:
|
| 251 |
+
Complete formatted prompt
|
| 252 |
+
"""
|
| 253 |
+
# Build prompt with XML delimiters (T034)
|
| 254 |
+
prompt = f"""<context>
|
| 255 |
+
{context}
|
| 256 |
+
</context>
|
| 257 |
+
|
| 258 |
+
<query>
|
| 259 |
+
{query}
|
| 260 |
+
</query>
|
| 261 |
+
|
| 262 |
+
Please answer the query using ONLY the information provided in the context above. If the context doesn't contain enough information to answer the query, say so. Be concise and cite specific details from the context."""
|
| 263 |
+
|
| 264 |
+
return prompt
|
services/retrieval_service.py
ADDED
|
@@ -0,0 +1,147 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Retrieval Service for RAG Agent
|
| 3 |
+
|
| 4 |
+
Provides a clean service layer abstraction over the retrieval pipeline.
|
| 5 |
+
Wraps retrieval.py functions with enhanced error handling and logging.
|
| 6 |
+
|
| 7 |
+
This service:
|
| 8 |
+
1. Embeds user queries using Cohere
|
| 9 |
+
2. Searches Qdrant for relevant documentation chunks
|
| 10 |
+
3. Handles errors gracefully
|
| 11 |
+
4. Provides detailed logging and metrics
|
| 12 |
+
|
| 13 |
+
Usage:
|
| 14 |
+
from services.retrieval_service import RetrievalService
|
| 15 |
+
|
| 16 |
+
retrieval_service = RetrievalService()
|
| 17 |
+
results = retrieval_service.retrieve_context("How do I create a publisher?", top_k=5)
|
| 18 |
+
"""
|
| 19 |
+
|
| 20 |
+
import logging
|
| 21 |
+
from typing import List
|
| 22 |
+
from datetime import datetime
|
| 23 |
+
|
| 24 |
+
from retrieval import embed_query, search_qdrant, SearchResult
|
| 25 |
+
from exceptions import RetrievalError
|
| 26 |
+
|
| 27 |
+
logger = logging.getLogger('retrieval_service')
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
class RetrievalService:
|
| 31 |
+
"""
|
| 32 |
+
Service for retrieving relevant documentation chunks.
|
| 33 |
+
|
| 34 |
+
This service wraps the lower-level retrieval.py functions and provides:
|
| 35 |
+
- Enhanced error handling with custom exceptions
|
| 36 |
+
- Detailed logging with metrics (latency, relevance scores)
|
| 37 |
+
- Clean abstraction for integration layer
|
| 38 |
+
|
| 39 |
+
Methods:
|
| 40 |
+
retrieve_context: Retrieve relevant chunks for a query
|
| 41 |
+
"""
|
| 42 |
+
|
| 43 |
+
def __init__(self):
|
| 44 |
+
"""Initialize retrieval service."""
|
| 45 |
+
logger.info("RetrievalService initialized")
|
| 46 |
+
|
| 47 |
+
def retrieve_context(self, query: str, top_k: int = 5) -> List[SearchResult]:
|
| 48 |
+
"""
|
| 49 |
+
Retrieve relevant documentation chunks for a query.
|
| 50 |
+
|
| 51 |
+
This method:
|
| 52 |
+
1. Embeds the query using Cohere (via retrieval.embed_query)
|
| 53 |
+
2. Searches Qdrant for similar chunks (via retrieval.search_qdrant)
|
| 54 |
+
3. Calculates and logs relevance metrics
|
| 55 |
+
4. Handles errors with specific error messages
|
| 56 |
+
|
| 57 |
+
Args:
|
| 58 |
+
query: User's natural language question
|
| 59 |
+
top_k: Number of chunks to retrieve (default 5)
|
| 60 |
+
|
| 61 |
+
Returns:
|
| 62 |
+
List of SearchResult objects with text, metadata, and relevance scores
|
| 63 |
+
|
| 64 |
+
Raises:
|
| 65 |
+
RetrievalError: If embedding fails (Cohere issue)
|
| 66 |
+
RetrievalError: If vector search fails (Qdrant issue)
|
| 67 |
+
|
| 68 |
+
Example:
|
| 69 |
+
>>> service = RetrievalService()
|
| 70 |
+
>>> results = service.retrieve_context("What is ROS 2?", top_k=5)
|
| 71 |
+
>>> print(f"Retrieved {len(results)} chunks")
|
| 72 |
+
>>> print(f"Top score: {results[0].score}")
|
| 73 |
+
"""
|
| 74 |
+
start_time = datetime.utcnow()
|
| 75 |
+
|
| 76 |
+
logger.info(
|
| 77 |
+
f"Starting retrieval",
|
| 78 |
+
extra={
|
| 79 |
+
"query": query[:50] + "..." if len(query) > 50 else query,
|
| 80 |
+
"top_k": top_k
|
| 81 |
+
}
|
| 82 |
+
)
|
| 83 |
+
|
| 84 |
+
try:
|
| 85 |
+
# Step 1: Embed query using Cohere
|
| 86 |
+
logger.debug(f"Embedding query with Cohere")
|
| 87 |
+
query_embedding = embed_query(query)
|
| 88 |
+
|
| 89 |
+
except ValueError as e:
|
| 90 |
+
# Handle validation errors from retrieval.py
|
| 91 |
+
error_message = f"Query embedding failed: {str(e)}"
|
| 92 |
+
logger.error(error_message, exc_info=True)
|
| 93 |
+
raise RetrievalError(error_message)
|
| 94 |
+
|
| 95 |
+
except Exception as e:
|
| 96 |
+
# Handle Cohere API failures
|
| 97 |
+
error_message = f"Query embedding failed: {str(e)}"
|
| 98 |
+
logger.error(
|
| 99 |
+
error_message,
|
| 100 |
+
extra={"query": query, "error_type": type(e).__name__},
|
| 101 |
+
exc_info=True
|
| 102 |
+
)
|
| 103 |
+
raise RetrievalError(error_message)
|
| 104 |
+
|
| 105 |
+
try:
|
| 106 |
+
# Step 2: Search Qdrant for similar chunks
|
| 107 |
+
logger.debug(f"Searching Qdrant with top_k={top_k}")
|
| 108 |
+
results = search_qdrant(query_embedding, top_k=top_k)
|
| 109 |
+
|
| 110 |
+
except ValueError as e:
|
| 111 |
+
# Handle validation errors from retrieval.py
|
| 112 |
+
error_message = f"Vector database search failed: {str(e)}"
|
| 113 |
+
logger.error(error_message, exc_info=True)
|
| 114 |
+
raise RetrievalError(error_message)
|
| 115 |
+
|
| 116 |
+
except Exception as e:
|
| 117 |
+
# Handle Qdrant connection/query failures
|
| 118 |
+
error_message = f"Vector database unavailable: {str(e)}"
|
| 119 |
+
logger.error(
|
| 120 |
+
error_message,
|
| 121 |
+
extra={"query": query, "top_k": top_k, "error_type": type(e).__name__},
|
| 122 |
+
exc_info=True
|
| 123 |
+
)
|
| 124 |
+
raise RetrievalError("Vector database unavailable")
|
| 125 |
+
|
| 126 |
+
# Calculate metrics
|
| 127 |
+
latency_ms = int((datetime.utcnow() - start_time).total_seconds() * 1000)
|
| 128 |
+
retrieval_count = len(results)
|
| 129 |
+
|
| 130 |
+
# Calculate average relevance score
|
| 131 |
+
average_relevance_score = 0.0
|
| 132 |
+
if retrieval_count > 0:
|
| 133 |
+
average_relevance_score = sum(r.score for r in results) / retrieval_count
|
| 134 |
+
|
| 135 |
+
# Log success with metrics
|
| 136 |
+
logger.info(
|
| 137 |
+
f"Retrieval completed successfully",
|
| 138 |
+
extra={
|
| 139 |
+
"query": query[:50] + "..." if len(query) > 50 else query,
|
| 140 |
+
"top_k": top_k,
|
| 141 |
+
"retrieval_count": retrieval_count,
|
| 142 |
+
"average_relevance_score": round(average_relevance_score, 4),
|
| 143 |
+
"retrieval_latency_ms": latency_ms
|
| 144 |
+
}
|
| 145 |
+
)
|
| 146 |
+
|
| 147 |
+
return results
|
services/translation_service.py
ADDED
|
@@ -0,0 +1,292 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Translation Service for English to Urdu translation using mBART model.
|
| 3 |
+
|
| 4 |
+
This service handles:
|
| 5 |
+
- Loading and managing the mBART translation model
|
| 6 |
+
- Translating text content from English to Urdu
|
| 7 |
+
- Content extraction (excluding code blocks)
|
| 8 |
+
- Server-side caching with TTL and LRU eviction
|
| 9 |
+
- Error handling and timeout management
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
import asyncio
|
| 13 |
+
import hashlib
|
| 14 |
+
import logging
|
| 15 |
+
from datetime import datetime, timedelta
|
| 16 |
+
from typing import Dict, Optional, Tuple
|
| 17 |
+
from bs4 import BeautifulSoup
|
| 18 |
+
import time
|
| 19 |
+
|
| 20 |
+
# Translation model imports (will be loaded lazily)
|
| 21 |
+
try:
|
| 22 |
+
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
|
| 23 |
+
TRANSFORMERS_AVAILABLE = True
|
| 24 |
+
except ImportError:
|
| 25 |
+
TRANSFORMERS_AVAILABLE = False
|
| 26 |
+
logging.warning("Transformers library not available. Translation service will be disabled.")
|
| 27 |
+
|
| 28 |
+
logger = logging.getLogger(__name__)
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
class TranslationCache:
|
| 32 |
+
"""Server-side cache for translations with TTL and LRU eviction"""
|
| 33 |
+
|
| 34 |
+
def __init__(self, ttl_seconds: int = 86400, max_size_bytes: int = 1_000_000_000):
|
| 35 |
+
"""
|
| 36 |
+
Initialize translation cache
|
| 37 |
+
|
| 38 |
+
Args:
|
| 39 |
+
ttl_seconds: Time-to-live for cache entries (default: 24 hours)
|
| 40 |
+
max_size_bytes: Maximum cache size in bytes (default: 1GB)
|
| 41 |
+
"""
|
| 42 |
+
self.cache: Dict[str, Dict] = {}
|
| 43 |
+
self.ttl_seconds = ttl_seconds
|
| 44 |
+
self.max_size_bytes = max_size_bytes
|
| 45 |
+
self.access_times: Dict[str, float] = {}
|
| 46 |
+
|
| 47 |
+
def get(self, key: str) -> Optional[str]:
|
| 48 |
+
"""
|
| 49 |
+
Get cached translation if not expired
|
| 50 |
+
|
| 51 |
+
Args:
|
| 52 |
+
key: Cache key (hash of content)
|
| 53 |
+
|
| 54 |
+
Returns:
|
| 55 |
+
Translated text if found and not expired, None otherwise
|
| 56 |
+
"""
|
| 57 |
+
if key not in self.cache:
|
| 58 |
+
return None
|
| 59 |
+
|
| 60 |
+
entry = self.cache[key]
|
| 61 |
+
created_at = datetime.fromisoformat(entry["timestamp"])
|
| 62 |
+
age = datetime.now() - created_at
|
| 63 |
+
|
| 64 |
+
# Check if expired
|
| 65 |
+
if age.total_seconds() > self.ttl_seconds:
|
| 66 |
+
del self.cache[key]
|
| 67 |
+
del self.access_times[key]
|
| 68 |
+
return None
|
| 69 |
+
|
| 70 |
+
# Update access time for LRU
|
| 71 |
+
self.access_times[key] = time.time()
|
| 72 |
+
return entry["translated"]
|
| 73 |
+
|
| 74 |
+
def set(self, key: str, translated_text: str):
|
| 75 |
+
"""
|
| 76 |
+
Store translation in cache
|
| 77 |
+
|
| 78 |
+
Args:
|
| 79 |
+
key: Cache key (hash of content)
|
| 80 |
+
translated_text: Translated text to cache
|
| 81 |
+
"""
|
| 82 |
+
# Evict if cache is too large
|
| 83 |
+
self._evict_if_needed()
|
| 84 |
+
|
| 85 |
+
self.cache[key] = {
|
| 86 |
+
"translated": translated_text,
|
| 87 |
+
"timestamp": datetime.now().isoformat(),
|
| 88 |
+
"access_count": 1
|
| 89 |
+
}
|
| 90 |
+
self.access_times[key] = time.time()
|
| 91 |
+
|
| 92 |
+
def _evict_if_needed(self):
|
| 93 |
+
"""Evict least recently used entries if cache exceeds size limit"""
|
| 94 |
+
current_size = self._estimate_size()
|
| 95 |
+
|
| 96 |
+
while current_size > self.max_size_bytes and self.cache:
|
| 97 |
+
# Find least recently used entry
|
| 98 |
+
lru_key = min(self.access_times, key=self.access_times.get)
|
| 99 |
+
del self.cache[lru_key]
|
| 100 |
+
del self.access_times[lru_key]
|
| 101 |
+
current_size = self._estimate_size()
|
| 102 |
+
|
| 103 |
+
def _estimate_size(self) -> int:
|
| 104 |
+
"""Estimate cache size in bytes"""
|
| 105 |
+
total_size = 0
|
| 106 |
+
for entry in self.cache.values():
|
| 107 |
+
total_size += len(entry["translated"].encode('utf-8'))
|
| 108 |
+
return total_size
|
| 109 |
+
|
| 110 |
+
def clear(self) -> int:
|
| 111 |
+
"""Clear all cache entries. Returns number of entries cleared."""
|
| 112 |
+
count = len(self.cache)
|
| 113 |
+
self.cache.clear()
|
| 114 |
+
self.access_times.clear()
|
| 115 |
+
return count
|
| 116 |
+
|
| 117 |
+
def get_stats(self) -> Dict:
|
| 118 |
+
"""Get cache statistics"""
|
| 119 |
+
total_entries = len(self.cache)
|
| 120 |
+
total_size_mb = self._estimate_size() / (1024 * 1024)
|
| 121 |
+
|
| 122 |
+
oldest_age_hours = 0
|
| 123 |
+
if self.cache:
|
| 124 |
+
oldest_timestamp = min(
|
| 125 |
+
datetime.fromisoformat(entry["timestamp"])
|
| 126 |
+
for entry in self.cache.values()
|
| 127 |
+
)
|
| 128 |
+
oldest_age_hours = (datetime.now() - oldest_timestamp).total_seconds() / 3600
|
| 129 |
+
|
| 130 |
+
return {
|
| 131 |
+
"total_entries": total_entries,
|
| 132 |
+
"total_size_mb": round(total_size_mb, 2),
|
| 133 |
+
"oldest_entry_age_hours": round(oldest_age_hours, 2)
|
| 134 |
+
}
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
class TranslationService:
|
| 138 |
+
"""Service for translating English text to Urdu using mBART"""
|
| 139 |
+
|
| 140 |
+
def __init__(self):
|
| 141 |
+
self.model = None
|
| 142 |
+
self.tokenizer = None
|
| 143 |
+
self.model_name = "facebook/mbart-large-50-many-to-many-mmt"
|
| 144 |
+
self.cache = TranslationCache()
|
| 145 |
+
self._model_loading = False
|
| 146 |
+
self._model_loaded = False
|
| 147 |
+
|
| 148 |
+
async def load_model(self):
|
| 149 |
+
"""Load mBART model (lazy loading on first translation request)"""
|
| 150 |
+
if self._model_loaded or not TRANSFORMERS_AVAILABLE:
|
| 151 |
+
return
|
| 152 |
+
|
| 153 |
+
if self._model_loading:
|
| 154 |
+
# Wait for concurrent load to complete
|
| 155 |
+
while self._model_loading:
|
| 156 |
+
await asyncio.sleep(0.1)
|
| 157 |
+
return
|
| 158 |
+
|
| 159 |
+
self._model_loading = True
|
| 160 |
+
try:
|
| 161 |
+
logger.info(f"Loading translation model: {self.model_name}")
|
| 162 |
+
# Note: Actual model loading would happen here
|
| 163 |
+
# For now, we'll mark as loaded but note that actual inference won't work
|
| 164 |
+
# until the model is properly downloaded
|
| 165 |
+
logger.warning(
|
| 166 |
+
"Translation model loading skipped in setup phase. "
|
| 167 |
+
"Run 'python -c \"from transformers import MBartForConditionalGeneration; "
|
| 168 |
+
"MBartForConditionalGeneration.from_pretrained('facebook/mbart-large-50-many-to-many-mmt')\" "
|
| 169 |
+
"to download the model."
|
| 170 |
+
)
|
| 171 |
+
self._model_loaded = True
|
| 172 |
+
logger.info("Translation service initialized (model download pending)")
|
| 173 |
+
except Exception as e:
|
| 174 |
+
logger.error(f"Failed to load translation model: {e}")
|
| 175 |
+
raise
|
| 176 |
+
finally:
|
| 177 |
+
self._model_loading = False
|
| 178 |
+
|
| 179 |
+
def extract_translatable_content(self, html: str) -> Tuple[str, Dict[str, str]]:
|
| 180 |
+
"""
|
| 181 |
+
Extract translatable content from HTML, excluding code blocks
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
html: HTML string to process
|
| 185 |
+
|
| 186 |
+
Returns:
|
| 187 |
+
Tuple of (processable_html, placeholders_dict)
|
| 188 |
+
"""
|
| 189 |
+
soup = BeautifulSoup(html, 'html.parser')
|
| 190 |
+
|
| 191 |
+
# Tags to exclude from translation
|
| 192 |
+
EXCLUDE_TAGS = ['code', 'pre', 'script', 'style']
|
| 193 |
+
|
| 194 |
+
placeholders = {}
|
| 195 |
+
for i, tag in enumerate(soup.find_all(EXCLUDE_TAGS)):
|
| 196 |
+
placeholder = f"__SKIP_{i}__"
|
| 197 |
+
placeholders[placeholder] = str(tag)
|
| 198 |
+
tag.replace_with(placeholder)
|
| 199 |
+
|
| 200 |
+
return str(soup), placeholders
|
| 201 |
+
|
| 202 |
+
def restore_excluded_content(self, translated_html: str, placeholders: Dict[str, str]) -> str:
|
| 203 |
+
"""
|
| 204 |
+
Restore excluded content (code blocks) after translation
|
| 205 |
+
|
| 206 |
+
Args:
|
| 207 |
+
translated_html: Translated HTML with placeholders
|
| 208 |
+
placeholders: Dictionary of placeholders and their original content
|
| 209 |
+
|
| 210 |
+
Returns:
|
| 211 |
+
Final HTML with code blocks restored
|
| 212 |
+
"""
|
| 213 |
+
for placeholder, original_content in placeholders.items():
|
| 214 |
+
translated_html = translated_html.replace(placeholder, original_content)
|
| 215 |
+
|
| 216 |
+
return translated_html
|
| 217 |
+
|
| 218 |
+
async def translate_text(
|
| 219 |
+
self,
|
| 220 |
+
text: str,
|
| 221 |
+
source_lang: str = "en_XX",
|
| 222 |
+
target_lang: str = "ur_PK",
|
| 223 |
+
timeout_seconds: int = 10
|
| 224 |
+
) -> str:
|
| 225 |
+
"""
|
| 226 |
+
Translate text from English to Urdu
|
| 227 |
+
|
| 228 |
+
Args:
|
| 229 |
+
text: Text to translate
|
| 230 |
+
source_lang: Source language code (mBART format)
|
| 231 |
+
target_lang: Target language code (mBART format)
|
| 232 |
+
timeout_seconds: Maximum time for translation
|
| 233 |
+
|
| 234 |
+
Returns:
|
| 235 |
+
Translated text
|
| 236 |
+
|
| 237 |
+
Raises:
|
| 238 |
+
TimeoutError: If translation exceeds timeout
|
| 239 |
+
RuntimeError: If model not loaded or translation fails
|
| 240 |
+
"""
|
| 241 |
+
# Check cache first
|
| 242 |
+
cache_key = hashlib.md5(text.encode('utf-8')).hexdigest()
|
| 243 |
+
cached = self.cache.get(cache_key)
|
| 244 |
+
if cached:
|
| 245 |
+
logger.debug(f"Cache hit for key: {cache_key}")
|
| 246 |
+
return cached
|
| 247 |
+
|
| 248 |
+
# Ensure model is loaded
|
| 249 |
+
await self.load_model()
|
| 250 |
+
|
| 251 |
+
if not self._model_loaded:
|
| 252 |
+
raise RuntimeError("Translation model not loaded")
|
| 253 |
+
|
| 254 |
+
try:
|
| 255 |
+
# For now, return a placeholder since actual model isn't loaded
|
| 256 |
+
# In production, this would call the actual mBART model:
|
| 257 |
+
#
|
| 258 |
+
# self.tokenizer.src_lang = source_lang
|
| 259 |
+
# encoded = self.tokenizer(text, return_tensors="pt")
|
| 260 |
+
# generated = self.model.generate(
|
| 261 |
+
# **encoded,
|
| 262 |
+
# forced_bos_token_id=self.tokenizer.lang_code_to_id[target_lang]
|
| 263 |
+
# )
|
| 264 |
+
# translated = self.tokenizer.batch_decode(generated, skip_special_tokens=True)[0]
|
| 265 |
+
|
| 266 |
+
# Placeholder implementation
|
| 267 |
+
translated = f"[URDU TRANSLATION PLACEHOLDER]: {text}"
|
| 268 |
+
logger.warning("Using placeholder translation (model not actually loaded)")
|
| 269 |
+
|
| 270 |
+
# Cache the result
|
| 271 |
+
self.cache.set(cache_key, translated)
|
| 272 |
+
|
| 273 |
+
return translated
|
| 274 |
+
|
| 275 |
+
except asyncio.TimeoutError:
|
| 276 |
+
logger.error(f"Translation timeout after {timeout_seconds} seconds")
|
| 277 |
+
raise TimeoutError(f"Translation took longer than {timeout_seconds} seconds")
|
| 278 |
+
except Exception as e:
|
| 279 |
+
logger.error(f"Translation failed: {e}")
|
| 280 |
+
raise RuntimeError(f"Translation failed: {str(e)}")
|
| 281 |
+
|
| 282 |
+
|
| 283 |
+
# Global service instance
|
| 284 |
+
_translation_service: Optional[TranslationService] = None
|
| 285 |
+
|
| 286 |
+
|
| 287 |
+
def get_translation_service() -> TranslationService:
|
| 288 |
+
"""Get or create the global translation service instance"""
|
| 289 |
+
global _translation_service
|
| 290 |
+
if _translation_service is None:
|
| 291 |
+
_translation_service = TranslationService()
|
| 292 |
+
return _translation_service
|