diff --git a/core/__init__.py b/core/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/core/chunking/__init__.py b/core/chunking/__init__.py deleted file mode 100644 index 598eb4f6..00000000 --- a/core/chunking/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .fixed_chunker import FixedChunker -from .hierarchical_chunker import HierarchicalChunker \ No newline at end of file diff --git a/core/chunking/fixed_chunker.py b/core/chunking/fixed_chunker.py deleted file mode 100644 index b6fca178..00000000 --- a/core/chunking/fixed_chunker.py +++ /dev/null @@ -1,34 +0,0 @@ -from typing import List -from langchain.text_splitter import CharacterTextSplitter -from baseclasses.base_classes import BaseChunker - - -class FixedChunker(BaseChunker): - """Fixed chunking strategy using LangChain’s CharacterTextSplitter.""" - - def chunk(self, text : str) -> List[str]: - if self.chunk_size <= 0: - raise ValueError("chunk_size must be positive") - if self.chunk_overlap >= self.chunk_size: - raise ValueError("chunk_overlap must be less than chunk_size") - if not text: - raise ValueError("Input text cannot be empty or None") - - # TODO: Temporary fix, better to move to recursive - separators = [' ', '\t', '\n', '\r', '\f', '\v'] - for sep in separators: - text = text.replace(sep, ' ') - - # chunk size is in tokens, general norm : 1 token = 4 chars - chunk_size = 4 * self.chunk_size - # overlap is in percentage - chunk_overlap = int(self.chunk_overlap * chunk_size / 100) - self.text_splitter = CharacterTextSplitter( - separator=" ", - chunk_size=chunk_size, - chunk_overlap=chunk_overlap, - length_function=len, - is_separator_regex=False - ) - chunks = self.text_splitter.split_text(text) - return chunks \ No newline at end of file diff --git a/core/chunking/hierarchical_chunker.py b/core/chunking/hierarchical_chunker.py deleted file mode 100644 index 877cfd6a..00000000 --- a/core/chunking/hierarchical_chunker.py +++ /dev/null @@ -1,54 +0,0 @@ -from typing import List -from baseclasses.base_classes import BaseHierarchicalChunker -from langchain.text_splitter import CharacterTextSplitter -import uuid - -class HierarchicalChunker(BaseHierarchicalChunker): - """Hierarchical chunking strategy.""" - def chunk(self, text : str) -> List[List[str]]: - overlap_tokens = int((self.chunk_overlap / 100) * self.child_chunk_size) - if self.parent_chunk_size <= 0: - raise ValueError("parent chunk size must be positive") - if self.child_chunk_size <= 0: - raise ValueError("child chunk size must be positive") - if self.child_chunk_size > self.parent_chunk_size: - raise ValueError("child chunk size must be less than parent chunk size") - if overlap_tokens >= self.child_chunk_size: - raise ValueError("chunk_overlap must be less than child chunk size") - if not text: - raise ValueError("Input text cannot be empty or None") - - # TODO: Temporary fix, better to move to recursive - separators = [' ', '\t', '\n', '\r', '\f', '\v'] - for sep in separators: - text = text.replace(sep, ' ') - - - # chunk size is in tokens, general norm : 1 token = 4 chars - parent_character_chunk_size = 4 * self.parent_chunk_size - child_character_chunk_size = 4 * self.child_chunk_size - # overlap is in percentage - child_chunk_overlap_characters = int(self.chunk_overlap * child_character_chunk_size / 100) - self.parent_text_splitter = CharacterTextSplitter( - separator=" ", - chunk_size=parent_character_chunk_size, - chunk_overlap=0, # Can change this at a later point of time - length_function=len, - is_separator_regex=False - ) - self.child_text_splitter = CharacterTextSplitter( - separator=" ", - chunk_size=child_character_chunk_size, - chunk_overlap=child_chunk_overlap_characters, - length_function=len, - is_separator_regex=False - ) - parent_chunks = self.parent_text_splitter.split_text(text) - overall_chunks = [] - for parent_chunk in parent_chunks: - parent_id = str(uuid.uuid4()) - child_chunks = self.child_text_splitter.split_text(parent_chunk) - for child_chunk in child_chunks: - temp_chunk = (parent_id, parent_chunk, child_chunk) - overall_chunks.append(temp_chunk) - return overall_chunks diff --git a/core/dynamodb.py b/core/dynamodb.py deleted file mode 100644 index ead0a8db..00000000 --- a/core/dynamodb.py +++ /dev/null @@ -1,406 +0,0 @@ -import boto3 -import logging -from typing import Dict, List, Any, Optional -from botocore.exceptions import ClientError -from decimal import Decimal -import json -import time -from datetime import datetime, timezone - -class DynamoDBOperations: - """Class to handle DynamoDB operations.""" - - def __init__(self, table_name: str, region: str = 'us-east-1'): - """ - Initialize DynamoDB operations. - - Args: - table_name (str): DynamoDB table name - region (str): AWS region - """ - self.table_name = table_name - self.region = region - self.logger = logging.getLogger(__name__) - - # Initialize DynamoDB resources - self.dynamodb = boto3.resource('dynamodb', region_name=region) - self.table = self.dynamodb.Table(table_name) - - # Initialize DynamoDB client for batch operations - self.dynamodb_client = boto3.client('dynamodb', region_name=region) - - def _handle_decimal_type(self, obj: Any) -> Any: - """ - Handle Decimal type for DynamoDB. - - Args: - obj: Input object to process - - Returns: - Processed object with Decimal types handled - """ - if isinstance(obj, list): - return [self._handle_decimal_type(i) for i in obj] - elif isinstance(obj, dict): - return {k: self._handle_decimal_type(v) for k, v in obj.items()} - elif isinstance(obj, float): - return Decimal(str(obj)) - return obj - - def _serialize_datetime(self, obj: Any) -> Any: - """ - Recursively serialize datetime objects into string format. - """ - if isinstance(obj, list): - return [self._serialize_datetime(i) for i in obj] - elif isinstance(obj, dict): - return {k: self._serialize_datetime(v) for k, v in obj.items()} - elif isinstance(obj, datetime): - return obj.isoformat() # Convert datetime to ISO 8601 string - return obj - - def scan_all(self, filter_expression: Optional[str] = None, expression_values: Optional[Dict[str, Any]] = None, expression_attribute_names: Optional[Dict[str, str]] = None) -> Dict: - """ - Scan items from DynamoDB, optionally applying filter expressions. - Args: - filter_expression (Optional[str]): Filter expression for the scan. - expression_values (Optional[Dict[str, Any]]): Attribute values for the filter. - expression_attribute_names (Optional[Dict[str, str]]): Attribute names mapping for reserved keywords. - Returns: - Dict: Scan results. - """ - result = { - "Items": [] - } - last_evaluated_key = None - try: - params = {} - - if filter_expression: - params["FilterExpression"] = filter_expression - if expression_values: - params["ExpressionAttributeValues"] = self._handle_decimal_type(expression_values) - if expression_attribute_names: - params["ExpressionAttributeNames"] = expression_attribute_names - - while True: - if last_evaluated_key: - params['ExclusiveStartKey'] = last_evaluated_key - - response = self.table.scan(**params) - result["Items"].extend(response.get('Items', [])) - last_evaluated_key = response.get('LastEvaluatedKey') - - if not last_evaluated_key: - break - - return result - except Exception as e: - self.logger.error(f"Error scanning items: {str(e)}") - raise - - def scan(self, filter_expression: Optional[str] = None, expression_values: Optional[Dict[str, Any]] = None, expression_attribute_names: Optional[Dict[str, str]] = None) -> Dict: - """ - Scan items from DynamoDB, optionally applying filter expressions. - - Args: - filter_expression (Optional[str]): Filter expression for the scan. - expression_values (Optional[Dict[str, Any]]): Attribute values for the filter. - expression_attribute_names (Optional[Dict[str, str]]): Attribute names mapping for reserved keywords. - - Returns: - Dict: Scan results. - """ - try: - params = {} - if filter_expression: - params["FilterExpression"] = filter_expression - if expression_values: - params["ExpressionAttributeValues"] = self._handle_decimal_type(expression_values) - if expression_attribute_names: - params["ExpressionAttributeNames"] = expression_attribute_names - - response = self.table.scan(**params) - return response - - except Exception as e: - self.logger.error(f"Error scanning items: {str(e)}") - raise - - - def _serialize_data(self, obj: Any) -> Any: - """ - Recursively convert float to Decimal and serialize datetime to string. - """ - if isinstance(obj, list): - return [self._serialize_data(i) for i in obj] - elif isinstance(obj, dict): - return {k: self._serialize_data(v) for k, v in obj.items()} - elif isinstance(obj, float): - return Decimal(str(obj)) # Convert float to Decimal - elif isinstance(obj, datetime): - return obj.isoformat() # Convert datetime to ISO 8601 string - return obj - - def get_item(self, key: Dict[str, Any]) -> Optional[Dict]: - """ - Retrieve an item from DynamoDB. - - Args: - key (Dict[str, Any]): Primary key of the item - - Returns: - Optional[Dict]: Item if found, None otherwise - """ - try: - response = self.table.get_item(Key=key) - item = response.get('Item') - - if item: - self.logger.info(f"Successfully retrieved item with key: {key}") - return item - - self.logger.info(f"No item found with key: {key}") - return None - - except ClientError as e: - self.logger.error(f"Error retrieving item with key {key}: {str(e)}") - raise - except Exception as e: - self.logger.error(f"Unexpected error: {str(e)}") - raise - - def put_item(self, item: Dict[str, Any], condition_expression: str = None, add_metadata: bool = True) -> Dict: - """ - Put an item into DynamoDB. - - Args: - item (Dict[str, Any]): Item to put. - condition_expression (str, optional): Condition expression for the put operation. - add_metadata (bool, optional): Whether to add timestamp and last_updated metadata. Default is True. - - Returns: - Dict: Response from DynamoDB. - """ - try: - # Add metadata if required - if add_metadata: - item["timestamp"] = datetime.now(timezone.utc).isoformat() - item["last_updated"] = datetime.now(timezone.utc).isoformat() - - # Serialize data to handle unsupported types like datetime - serialized_item = self._serialize_data(item) - - # Prepare put_item parameters - params = { - "Item": serialized_item - } - - if condition_expression: - params["ConditionExpression"] = condition_expression - - # Put the item into DynamoDB - response = self.table.put_item(**params) - - # Log success - self.logger.info(f"Successfully put item: {item.get('id', 'No ID')}") - return response - - except ClientError as e: - if e.response["Error"]["Code"] == "ConditionalCheckFailedException": - self.logger.warning("Condition check failed for put_item operation") - else: - self.logger.error(f"Error putting item: {str(e)}") - raise - except Exception as e: - self.logger.error(f"Unexpected error: {str(e)}") - raise - - - - def update_item(self, - key: Dict[str, Any], - update_expression: str, - expression_values: Dict[str, Any], - condition_expression: str = None) -> Dict: - """ - Update an item in DynamoDB. - - Args: - key (Dict[str, Any]): Primary key of the item - update_expression (str): Update expression - expression_values (Dict[str, Any]): Expression attribute values - condition_expression (str, optional): Condition expression - - Returns: - Dict: Response from DynamoDB - """ - try: - # Add last_updated to expression values - #expression_values[':updated'] = datetime.utcnow().isoformat() - #update_expression += ', last_updated = :updated' - - # Handle decimal types in expression values - processed_values = self._handle_decimal_type(expression_values) - - # Prepare update parameters - params = { - 'Key': key, - 'UpdateExpression': update_expression, - 'ExpressionAttributeValues': processed_values, - 'ReturnValues': 'UPDATED_NEW' - } - - if condition_expression: - params['ConditionExpression'] = condition_expression - - response = self.table.update_item(**params) - - self.logger.info(f"Successfully updated item with key: {key}") - return response - - except ClientError as e: - if e.response['Error']['Code'] == 'ConditionalCheckFailedException': - self.logger.warning("Condition check failed for update operation") - else: - self.logger.error(f"Error updating item: {str(e)}") - raise - except Exception as e: - self.logger.error(f"Unexpected error: {str(e)}") - raise - - def query(self, - key_condition_expression: str, - expression_values: Dict[str, Any], - index_name: str = None, projection: str = None, exclusive_start_key = None) -> Dict: - """ - Query items from DynamoDB. - - Args: - key_condition_expression (str): Key condition expression - expression_values (Dict[str, Any]): Expression attribute values - index_name (str, optional): Name of the index to query - exclusive_start_key(Dict[str, Any]): LastEvaluateKey for pagination - - Returns: - Dict: Query results - """ - try: - # Handle decimal types in expression values - processed_values = self._handle_decimal_type(expression_values) - - # Prepare query parameters - params = { - 'KeyConditionExpression': key_condition_expression, - 'ExpressionAttributeValues': processed_values - } - - if index_name: - params['IndexName'] = index_name - - if projection: - params['ProjectionExpression'] = projection - - if exclusive_start_key: - params['ExclusiveStartKey'] = exclusive_start_key - - response = self.table.query(**params) - - self.logger.info(f"Successfully queried {len(response['Items'])} items") - return response - - except ClientError as e: - self.logger.error(f"Error querying items: {str(e)}") - raise - except Exception as e: - self.logger.error(f"Unexpected error: {str(e)}") - raise - - def batch_write(self, items: List[Dict[str, Any]], max_retries: int = 3) -> None: - """ - Batch write items to DynamoDB with automatic retry and backoff. - DynamoDB has a limit of 25 items per batch write operation. - - Args: - items (List[Dict[str, Any]]): List of items to write (max 25 items) - max_retries (int): Maximum number of retries for failed items - """ - try: - if len(items) > 25: - raise ValueError("DynamoDB batch_write_item operation can only process up to 25 items at a time") - - unprocessed_items = items - retry_count = 0 - - while unprocessed_items and retry_count < max_retries: - # Prepare batch write request - request_items = { - self.table_name: [ - { - 'PutRequest': { - 'Item': self._handle_decimal_type(item) - } - } - for item in unprocessed_items - ] - } - - response = self.dynamodb_client.batch_write_item( - RequestItems=request_items - ) - - # Handle unprocessed items - unprocessed_items = [ - item['PutRequest']['Item'] - for item in response.get('UnprocessedItems', {}).get(self.table_name, []) - ] - - if unprocessed_items: - retry_count += 1 - if retry_count < max_retries: - # Exponential backoff - time.sleep(2 ** retry_count) - - if unprocessed_items: - self.logger.warning(f"{len(unprocessed_items)} items remained unprocessed after {max_retries} retries") - else: - self.logger.info(f"Successfully batch wrote {len(items)} items") - - except ClientError as e: - self.logger.error(f"Error in batch write operation: {str(e)}") - raise - except Exception as e: - self.logger.error(f"Unexpected error: {str(e)}") - raise - - def delete_item(self, key: Dict[str, Any], condition_expression: str = None) -> Dict: - """ - Delete an item from DynamoDB. - - Args: - key (Dict[str, Any]): Primary key of the item - condition_expression (str, optional): Condition expression - - Returns: - Dict: Response from DynamoDB - """ - try: - params = {'Key': key} - - if condition_expression: - params['ConditionExpression'] = condition_expression - - response = self.table.delete_item(**params) - - self.logger.info(f"Successfully deleted item with key: {key}") - return response - - except ClientError as e: - if e.response['Error']['Code'] == 'ConditionalCheckFailedException': - self.logger.warning("Condition check failed for delete operation") - else: - self.logger.error(f"Error deleting item: {str(e)}") - raise - except Exception as e: - self.logger.error(f"Unexpected error: {str(e)}") - raise diff --git a/core/embedding/__init__.py b/core/embedding/__init__.py deleted file mode 100644 index 26d5cf25..00000000 --- a/core/embedding/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -from core.embedding.embedding_factory import EmbedderFactory # Keep this for external usage. - -import core.embedding.bedrock.cohere_embedder -import core.embedding.bedrock.titanv1_embedder -import core.embedding.bedrock.titanv2_embedder - -# Importing SageMaker-specific embedder and embedding factory. -from .embedding_factory import EmbedderFactory -from .sagemaker import SageMakerEmbedder - -# List of model names that you want to register with the EmbedderFactory -model_list = [ - "huggingface-sentencesimilarity-bge-large-en-v1-5", # Model for sentence similarity (BGE, Large). - "huggingface-sentencesimilarity-bge-m3", # Another sentence similarity model (BGE, M3). - "huggingface-textembedding-gte-qwen2-7b-instruct" # Text embedding model (Qwen2-7B, Instruct). - ] - -# Registering each model from the list into the EmbedderFactory under 'sagemaker'. -# The `SageMakerEmbedder` will be used for embedding operations for these models. -for model in model_list: - EmbedderFactory.register_embedder('sagemaker', model, SageMakerEmbedder) \ No newline at end of file diff --git a/core/embedding/bedrock/__init__.py b/core/embedding/bedrock/__init__.py deleted file mode 100644 index 6831b89c..00000000 --- a/core/embedding/bedrock/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .bedrock_embedder import BedrockEmbedder diff --git a/core/embedding/bedrock/bedrock_embedder.py b/core/embedding/bedrock/bedrock_embedder.py deleted file mode 100644 index 2825594b..00000000 --- a/core/embedding/bedrock/bedrock_embedder.py +++ /dev/null @@ -1,46 +0,0 @@ -import boto3 -from typing import Dict, List, Tuple, Any -from baseclasses.base_classes import BaseEmbedder -from util.boto3_utils import BedRockRetryHander -import json - -import logging - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -# Bedrock Base Embedder -class BedrockEmbedder(BaseEmbedder): - def __init__(self, model_id: str, region: str, role_arn: str = None) -> None: - super().__init__(model_id) - self.client = boto3.client("bedrock-runtime", region_name=region) - - def prepare_payload(self, text: str, dimensions: int, normalize: bool) -> Dict: - raise NotImplementedError("Subclasses must implement `prepare_payload`") - - @BedRockRetryHander() - def embed(self, text: str, dimensions: int = 256, normalize: bool = True) -> Tuple[Dict[Any, Any], List[float]]: - try: - payload = self.prepare_payload(text, dimensions, normalize) - response = self.client.invoke_model( - modelId=self.model_id, - contentType="application/json", - accept="application/json", - body=json.dumps(payload) - ) - model_response = json.loads(response["body"].read()) - metadata = {} - if response and 'ResponseMetadata' in response and 'HTTPHeaders' in response['ResponseMetadata']: - input_tokens = response['ResponseMetadata']['HTTPHeaders']['x-amzn-bedrock-input-token-count'] - latency = response['ResponseMetadata']['HTTPHeaders']['x-amzn-bedrock-invocation-latency'] - metadata = { - 'inputTokens': input_tokens, - 'latencyMs': latency - } - return metadata, self.extract_embedding(model_response) - except Exception as e: - logger.error(f"Error during embedding: {e}") - raise - - def extract_embedding(self, response: Dict) -> List[float]: - raise NotImplementedError("Subclasses must implement `extract_embedding`") \ No newline at end of file diff --git a/core/embedding/bedrock/cohere_embedder.py b/core/embedding/bedrock/cohere_embedder.py deleted file mode 100644 index bbf8226f..00000000 --- a/core/embedding/bedrock/cohere_embedder.py +++ /dev/null @@ -1,20 +0,0 @@ - -from typing import Dict, List -import logging - -from core.embedding import EmbedderFactory -from baseclasses.base_classes import BaseEmbedder -from . import BedrockEmbedder - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class CohereEmbedder(BedrockEmbedder): - def prepare_payload(self, text: str, dimensions: int, normalize: bool) -> Dict: - return {"texts": [text], "input_type": "search_document"} - - def extract_embedding(self, response: Dict) -> List[float]: - return response["embeddings"][0] - -EmbedderFactory.register_embedder("bedrock", "cohere.embed-english-v3", CohereEmbedder) -EmbedderFactory.register_embedder("bedrock", "cohere.embed-multilingual-v3", CohereEmbedder) \ No newline at end of file diff --git a/core/embedding/bedrock/titanv1_embedder.py b/core/embedding/bedrock/titanv1_embedder.py deleted file mode 100644 index 55421583..00000000 --- a/core/embedding/bedrock/titanv1_embedder.py +++ /dev/null @@ -1,20 +0,0 @@ - -from typing import Dict, List -import logging - -from core.embedding import EmbedderFactory -from baseclasses.base_classes import BaseEmbedder -from . import BedrockEmbedder - - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class TitanV1Embedder(BedrockEmbedder): - def prepare_payload(self, text: str, dimensions: int, normalize: bool) -> Dict: - return {"inputText": text, "embeddingConfig" : {"outputEmbeddingLength" : dimensions}} - - def extract_embedding(self, response: Dict) -> List[float]: - return response["embedding"] - -EmbedderFactory.register_embedder("bedrock", "amazon.titan-embed-image-v1", TitanV1Embedder) \ No newline at end of file diff --git a/core/embedding/bedrock/titanv2_embedder.py b/core/embedding/bedrock/titanv2_embedder.py deleted file mode 100644 index 140d5b0b..00000000 --- a/core/embedding/bedrock/titanv2_embedder.py +++ /dev/null @@ -1,20 +0,0 @@ - -from typing import Dict, List -import logging - -from core.embedding import EmbedderFactory -from baseclasses.base_classes import BaseEmbedder -from . import BedrockEmbedder - - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class TitanV2Embedder(BedrockEmbedder): - def prepare_payload(self, text: str, dimensions: int, normalize: bool) -> Dict: - return {"inputText": text, "dimensions": dimensions, "normalize": normalize} - - def extract_embedding(self, response: Dict) -> List[float]: - return response["embedding"] - -EmbedderFactory.register_embedder("bedrock", "amazon.titan-embed-text-v2:0", TitanV2Embedder) \ No newline at end of file diff --git a/core/embedding/embedding_factory.py b/core/embedding/embedding_factory.py deleted file mode 100644 index b9b9d8da..00000000 --- a/core/embedding/embedding_factory.py +++ /dev/null @@ -1,38 +0,0 @@ -from config.experimental_config import ExperimentalConfig - -import logging - -from typing import Type, Dict -from baseclasses.base_classes import BaseEmbedder -from config.config import get_config - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class EmbedderFactory: - """Factory to create embedders based on model ID and service type.""" - - _registry: Dict[str, Type[BaseEmbedder]] = {} - - @classmethod - def register_embedder(cls, service_type: str, model_id: str, embedder_cls: Type[BaseEmbedder]): - key = f"{service_type}:{model_id}" - cls._registry[key] = embedder_cls - - @classmethod - def create_embedder(cls, experimentalConfig : ExperimentalConfig) -> BaseEmbedder: - service_type = experimentalConfig.embedding_service - model_id = experimentalConfig.embedding_model - key = f"{service_type}:{model_id}" - - if experimentalConfig.embedding_service == "sagemaker": - role_arn = get_config().sagemaker_role_arn - print(f"Sagemaker role: {role_arn}") - elif experimentalConfig.embedding_service == "bedrock": - role_arn = get_config().bedrock_role_arn - embedder_cls = cls._registry.get(key) - if not embedder_cls: - raise ValueError(f"No embedder registered for service {service_type} and model {model_id}") - - return embedder_cls(model_id, experimentalConfig.aws_region, role_arn) - diff --git a/core/embedding/sagemaker/__init__.py b/core/embedding/sagemaker/__init__.py deleted file mode 100644 index 2765ede1..00000000 --- a/core/embedding/sagemaker/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .sagemaker_embedder import SageMakerEmbedder diff --git a/core/embedding/sagemaker/sagemaker_embedder.py b/core/embedding/sagemaker/sagemaker_embedder.py deleted file mode 100644 index 2a348bf7..00000000 --- a/core/embedding/sagemaker/sagemaker_embedder.py +++ /dev/null @@ -1,469 +0,0 @@ -import boto3 -from typing import Dict, List -from botocore.exceptions import ClientError -from baseclasses.base_classes import BaseEmbedder -from sagemaker.session import Session -from sagemaker.predictor import Predictor -from sagemaker.serializers import JSONSerializer -from sagemaker.deserializers import JSONDeserializer -from sagemaker.jumpstart.model import JumpStartModel -from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri -import sagemaker -import logging -import numpy as np -import json -import time - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -# Model configurations -EMBEDDING_MODELS = { - "huggingface-sentencesimilarity-bge-large-en-v1-5": { - "model_name": "bge-large", - "model_source": "jumpstart", - "dimension": 1024, - "instance_type": "ml.g5.2xlarge", - "input_key": "text_inputs" - }, - "huggingface-sentencesimilarity-bge-m3": { - "model_name": "bge-m3", - "model_source": "jumpstart", - "dimension": 1024, - "instance_type": "ml.g5.2xlarge", - "input_key": "text_inputs" - }, - "huggingface-textembedding-gte-qwen2-7b-instruct": { - "model_name": "qwen", - "model_source": "jumpstart", - "dimension": 3584, - "instance_type": "ml.g5.2xlarge", - "input_key": "inputs" - } -} - -# Sagemaker Base Embedder -class SageMakerEmbedder(BaseEmbedder): - def __init__(self, model_id: str, region: str, role_arn: str) -> None: - """ - Initializes the SageMakerEmbedder with the given model ID, region, and role ARN. - Sets up necessary SageMaker runtime clients, session, and endpoint predictor. - - Args: - model_id (str): The unique identifier for the model. - region (str): The AWS region where the SageMaker services are hosted. - role_arn (str): The ARN of the IAM role. Currently not used but included for future extensions. - """ - - # Initialize the base class - super().__init__(model_id) - - # Store the region - self.region_name = region - - self.role = role_arn - - # Initialize the SageMaker runtime and client for general operations - self.client = boto3.client("sagemaker-runtime", region_name=region) - self.sagemaker_client = boto3.client('sagemaker', region_name=region) - - # Create a new SageMaker session - self.session = Session(boto_session=boto3.Session(region_name=region)) - - # Initialize additional embedding-related attributes - self.embedding_model_id = model_id - # self.embedding_model_endpoint_name = 'flotorch-embedding-endpoint' - self.embedding_model_endpoint_name = f"{self._sanitize_name(model_id)[:44]}-embedding-endpoint" - - self.embedding_dimension = EMBEDDING_MODELS.get(model_id, {}).get('dimension', 1024) - - self.wait_time = 5 - - # Ensure the endpoint exists or create it if necessary - self._ensure_endpoint_exists() - - # Initialize the predictor to interact with the SageMaker endpoint - self.predictor = Predictor( - endpoint_name=self.embedding_model_endpoint_name, - sagemaker_session=self.session - ) - - # Set up the serializer and deserializer for the predictor - self.predictor.serializer = JSONSerializer() - self.predictor.deserializer = JSONDeserializer() - - self.embedding_predictor = self.predictor - - # Log initialization success - logger.info(f"Initialized SageMakerEmbedder for model {model_id} in region {region}.") - - def _ensure_endpoint_exists(self): - """ - Ensures that the SageMaker endpoint exists for the given model. If the endpoint does not exist, - it creates a new endpoint using the specified model ID. - - Args: - model_id (str): The unique identifier for the model to use for the endpoint. - - Raises: - ClientError: If there is an issue communicating with SageMaker or creating the endpoint. - """ - - try: - # Check if the endpoint already exists - _ = self._check_model_status(self.embedding_model_endpoint_name) - logger.info(f"Endpoint {self.embedding_model_endpoint_name} already exists.") - except self.sagemaker_client.exceptions.ClientError: - # If the endpoint does not exist, create a new one - logger.info(f"Endpoint and configuration for {self.embedding_model_endpoint_name} does not exist. Creating endpoint.") - self.create_endpoint(endpoint_name=self.embedding_model_endpoint_name, model_id=self.embedding_model_id) - - def _check_model_status(self, endpoint_name): - """ - Check the status of the SageMaker endpoint and its configuration. - - This method performs the following: - 1. Checks if endpoint exists and is in service - 2. If endpoint is being created, waits until creation completes - 3. If no endpoint exists, checks for endpoint configuration - 4. If configuration exists, waits for endpoint creation to complete - - Args: - endpoint_name (str): Name of the SageMaker endpoint to check - - Returns: - str: 'InService' if endpoint is available and running - - Raises: - Exception: If endpoint creation fails or has unexpected status - ClientError: If neither endpoint nor configuration exists - """ - try: - while True: - # Poll endpoint status until it is in service or fails - response = self.sagemaker_client.describe_endpoint(EndpointName=endpoint_name) - - if response['EndpointStatus'] == 'InService': - logger.info(f"Endpoint {endpoint_name} is in service.") - return 'InService' - - elif response['EndpointStatus'] == 'Failed': - logger.error(f"Endpoint {endpoint_name} creation failed.") - raise Exception(f"Endpoint {endpoint_name} creation failed.") - - elif response['EndpointStatus'] == 'Creating': - time.sleep(self.wait_time) # Pause before next status check - - else: - raise Exception(f"Unexpected endpoint status: {response['EndpointStatus']}") - - except self.sagemaker_client.exceptions.ClientError: - # No endpoint exists - check if there's a configuration waiting to be deployed - logger.info(f"Endpoint {endpoint_name} does not exist. Checking if endpoint configuration exists.") - - try: - # Look for endpoint configuration that may have been created by another process - response = self.sagemaker_client.describe_endpoint_config(EndpointConfigName=endpoint_name) - logger.info(f"Configuration for {endpoint_name} exists, waiting {self.wait_time} seconds for endpoint creation.") - - time.sleep(self.wait_time) # Allow time for endpoint creation to begin - - _ = self._check_model_status(endpoint_name) # Keep rechecking the endpoint status until it begins creation - - except self.sagemaker_client.exceptions.ClientError: - # Neither endpoint nor configuration exists - logger.info(f"Endpoint configuration does not exist.") - raise - - except Exception as e: - logger.error(f"Error checking endpoint status: {e}") - raise - - def create_endpoint(self, endpoint_name: str, model_id: str) -> sagemaker.predictor.Predictor: - """ - Creates a SageMaker endpoint for the specified model if it doesn't already exist. - If the endpoint is already in service, returns the existing predictor. - - Args: - endpoint_name (str): The name of the SageMaker endpoint to be created or fetched. - model_id (str): The identifier for the model to be used in the endpoint. - - Returns: - sagemaker.predictor.Predictor: A predictor object for the created or fetched endpoint. - - Raises: - ValueError: If the provided model_id is not supported. - ClientError: If there are AWS API errors during endpoint creation/access. - """ - - # Validate that model_id exists in our supported model configurations - if model_id not in EMBEDDING_MODELS: - raise ValueError(f"Unsupported model ID: {model_id}") - - # Create AWS and SageMaker sessions for API interactions - boto_session = boto3.Session(region_name=self.region_name) - sagemaker_session = sagemaker.Session(boto_session=boto_session) - - # Look up the appropriate instance type from model configurations - instance_type = (EMBEDDING_MODELS.get(model_id))['instance_type'] - model_source = (EMBEDDING_MODELS.get(model_id))['model_source'] - - try: - # First check if a working endpoint already exists to avoid duplicate creation - status = self._check_model_status(endpoint_name) - if status == 'InService': - # Endpoint exists and is healthy - create and return a predictor for it - predictor = sagemaker.predictor.Predictor( - endpoint_name=endpoint_name, - sagemaker_session=sagemaker_session, - serializer=sagemaker.serializers.JSONSerializer(), - deserializer=sagemaker.deserializers.JSONDeserializer() - ) - # Register the predictor with the appropriate model type handler - self._assign_predictor(predictor, model_id) - return predictor - - except self.sagemaker_client.exceptions.ClientError as e: - # Handle case where endpoint doesn't exist yet - if e.response['Error']['Code'] == 'ValidationException': - try: - if model_source == "jumpstart": - # Initialize a new JumpStart model with the specified configuration - model = JumpStartModel( - role = self.role, - model_id=model_id, - sagemaker_session=sagemaker_session - ) - - # Deploy the model to a new endpoint with the specified configuration - predictor = model.deploy( - initial_instance_count=1, - instance_type=instance_type, - endpoint_name=endpoint_name, - accept_eula=True # Required for JumpStart models - ) - # Check if the model source is huggingface - elif model_source == "huggingface": - hub = { - 'HF_MODEL_ID': model_id, - 'SM_NUM_GPUS': json.dumps(1) - } - huggingface_model = HuggingFaceModel( - image_uri=get_huggingface_llm_image_uri("huggingface", version="2.3.1", region=self.region_name), - env=hub, - role=self.role, - ) - - # deploy model to SageMaker Inference - predictor = huggingface_model.deploy( - initial_instance_count=1, - instance_type=instance_type, - endpoint_name=endpoint_name, - container_startup_health_check_timeout=300, - ) - - # Register the new predictor with the appropriate model type handler - self._assign_predictor(predictor, model_id) - return predictor - - except self.sagemaker_client.exceptions.ClientError as e: - # Handle race condition where another process started creating the endpoint - # between our existence check and creation attempt - logger.info(f"Error creating endpoint: {e}") - if e.response['Error']['Code'] == 'ValidationException': - - logger.info(f"A new endpoint creation intercepted while attempting to create new endpoint, waiting.") - time.sleep(self.wait_time) # Allow the other process's endpoint to finish creating - - status = self._check_model_status(endpoint_name) - - if status == 'InService': - logger.info(f"Found the new endpoint, creating the predictor.") - # Create predictor for the endpoint that the other process created - predictor = sagemaker.predictor.Predictor( - endpoint_name=endpoint_name, - sagemaker_session=sagemaker_session, - serializer=sagemaker.serializers.JSONSerializer(), - deserializer=sagemaker.deserializers.JSONDeserializer() - ) - - # Register with appropriate model type handler - self._assign_predictor(predictor, model_id) - - return predictor - - elif e.response['Error']['Code'] == 'ResourceLimitExceeded': - logger.error(f"Resource limit exceeded while creating endpoint: {e}") - raise - - # Re-raise any unexpected AWS API errors - raise - - def _assign_predictor(self, predictor: sagemaker.predictor.Predictor, model_id: str): - """ - Assigns the appropriate predictor based on the provided model_id. The predictor is assigned to either - the embedding or inferencing predictor attributes, depending on the model type. - - Args: - predictor (sagemaker.predictor.Predictor): The SageMaker predictor to be assigned. - model_id (str): The model ID which determines whether the predictor is for embedding or inferencing. - - """ - # Assign predictor for embedding models - if model_id in EMBEDDING_MODELS: - self.embedding_predictor = predictor - self.embedding_dimension = EMBEDDING_MODELS[model_id]['dimension'] - self.embedding_model_id = model_id - logger.info(f"Assigned embedding predictor for model: {model_id}") - - # Log an error if the model_id doesn't match any known type - else: - logger.error(f"Model ID {model_id} is not recognized as an embedding model.") - - def embed(self, text: str, dimensions: int = 256, normalize: bool = True) -> List[float]: - """ - Retrieves the embedding for the given input text from the model's predictor. - - Args: - text (str): The input text for which the embedding is generated. - - Returns: - List[float]: The generated embedding as a list of floats. - - Raises: - ValueError: If the predictor is not initialized or if the input text is empty. - Exception: If there is an error during the embedding extraction process. - """ - - # Validate predictor initialization - if not self.predictor: - raise ValueError("Embedding predictor not initialized") - - # Ensure the input text is not empty - if not text or not text.strip(): - raise ValueError("Input text cannot be empty") - - try: - # # Prepare the input payload for prediction - # input_data = { - # "text_inputs": [text], - # "mode": "embedding" # Define the mode to request embedding - # } - - input_data = self.prepare_payload(text) - - # Log the prediction details - logger.debug("Input data prepared for prediction: %s", json.dumps(input_data, indent=2)) - logger.debug("Using model ID: %s", self.embedding_model_id) - - # SageMaker does not provide input tokens as metadata. - # As a workaround, we use a rough approximation: ~4 characters per token. - input_tokens = len(text) // 4 - - start_time = time.time() - - # Make the prediction request - response = self.embedding_predictor.predict(input_data) - - # Calculate latency metrics - latency = int((time.time() - start_time) * 1000) - - # If the response is in byte format, decode it - if isinstance(response, (bytes, bytearray)): - response = json.loads(response.decode('utf-8')) - elif isinstance(response, str): - response = json.loads(response) - - # Extract the embedding from the response - if isinstance(response, dict) and 'embedding' in response: - embedding = np.array(response['embedding'][0] if isinstance(response['embedding'], list) else response['embedding']) - else: - embedding = np.array(response[0] if isinstance(response, list) else response) - - # Flatten the embedding to ensure it's a 1D array - embedding = embedding.flatten() - - # Normalize the embedding to unit length - embedding = embedding / np.linalg.norm(embedding) - - # Check if the embedding dimension matches the expected value (1024) - if len(embedding) != self.embedding_dimension: - logger.warning(f"Embedding dimension mismatch. Expected 1024, got {len(embedding)}") - # Adjust the dimension by truncating or padding - if len(embedding) > self.embedding_dimension: - embedding = embedding[:self.embedding_dimension] - else: - embedding = np.pad(embedding, (0, self.embedding_dimension - len(embedding))) - - metadata = { - 'inputTokens': input_tokens, - 'latencyMs': latency - } - # Return the embedding as a list of floats - return metadata, embedding.tolist() - - except Exception as e: - # Log detailed error information for debugging - logger.error("Error in get_embedding: %s", str(e)) - logger.error("Model ID: %s", self.embedding_model_id) - logger.error("Input text length: %d", len(text)) - logger.error("Input data: %s", json.dumps(input_data, indent=2)) - - # Log the response structure if available - if 'response' in locals(): - logger.error("Response structure: %s", type(response)) - try: - logger.error("Response content: %s", json.dumps(response, indent=2)) - except Exception as json_error: - logger.error("Response content (raw): %s", response) - - # Re-raise the exception after logging - raise - - def prepare_payload1(self, text: str, dimensions: int, normalize: bool) -> Dict: - raise NotImplementedError("Subclasses must implement `prepare_payload`") - - def prepare_payload(self, text: str) -> Dict: - """ - Prepares the payload for the embedding model based on the provided text. - - Args: - text (str): The input text to be processed by the model. - - Raises: - ValueError: If the embedding model ID is unknown. - - Returns: - Dict: The payload containing the appropriate input and model configurations. - """ - - # Retrieve model configuration based on model ID - model_config = EMBEDDING_MODELS.get(self.embedding_model_id) - if not model_config: - raise ValueError(f"Unknown model ID: {self.embedding_model_id}") - - # Extract the input key from the model configuration - input_key = model_config["input_key"] - - # Build the payload with the input text - payload = { - input_key: [text] - } - - # Add mode only for models that need it - if self.embedding_model_id != "huggingface-textembedding-gte-qwen2-7b-instruct": - payload["mode"] = "embedding" - - return payload - - @staticmethod - def _sanitize_name(name: str) -> str: - """Sanitize the endpoint name to follow AWS naming conventions""" - # Replace any character that's not alphanumeric or hyphen with hyphen - import re - name = re.sub(r'[^a-zA-Z0-9-]', '-', name) - # Ensure it starts with a letter - if not name[0].isalpha(): - name = 'n' + name - # Truncate to 63 characters (AWS limit) - return name[:63] \ No newline at end of file diff --git a/core/eval/__init__.py b/core/eval/__init__.py deleted file mode 100644 index e2ac60fa..00000000 --- a/core/eval/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .eval_factory import EvalFactory -from .ragas.ragas_non_llm_eval import RagasNonLLMEvaluator -from .ragas.ragas_llm_eval import RagasLLMEvaluator \ No newline at end of file diff --git a/core/eval/eval_factory.py b/core/eval/eval_factory.py deleted file mode 100644 index 82647d25..00000000 --- a/core/eval/eval_factory.py +++ /dev/null @@ -1,38 +0,0 @@ -from baseclasses.base_classes import BaseEvaluator -from config.experimental_config import ExperimentalConfig -from config.config import Config -from typing import Dict, Type - -import logging - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - - -class EvaluatorServiceError(Exception): - """Custom exception for inference service related errors""" - pass - -class EvalFactory: - - _registry: Dict[str, Type[BaseEvaluator]] = {} - - @classmethod - def register_evaluator(cls, service_type: str, eval_type: str, evaluator_cls: Type[BaseEvaluator]): - key = f"{service_type}:{eval_type}" - cls._registry[key] = evaluator_cls - - @classmethod - def create_evaluator(cls, experimentalConfig: ExperimentalConfig) -> BaseEvaluator: - config = Config.load_config() - - eval_service_type = experimentalConfig.eval_service - eval_type = 'llm' if experimentalConfig.llm_based_eval else 'non_llm' - - key = f"{eval_service_type}:{eval_type}" - - evaluator_cls = cls._registry.get(key) - if not evaluator_cls: - raise EvaluatorServiceError(f"No evaluator_cls registered for service {eval_service_type} and type {eval_type}") - - return evaluator_cls(config=config, experimental_config=experimentalConfig) diff --git a/core/eval/ragas/ragas_eval.py b/core/eval/ragas/ragas_eval.py deleted file mode 100644 index e44f05c9..00000000 --- a/core/eval/ragas/ragas_eval.py +++ /dev/null @@ -1,58 +0,0 @@ -from baseclasses.base_classes import BaseEvaluator, EvaluationMetrics -from core.dynamodb import DynamoDBOperations -from typing import List, Dict -import json -import numpy as np - -import logging -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) - -class RagasEvaluator(BaseEvaluator): - - def __init__(self, config, experimental_config): - super().__init__(config, experimental_config) - self._initialize_dynamodb() - - def _initialize_dynamodb(self): - """Initialize DynamoDB connections""" - self.metrics_db = DynamoDBOperations( - region=self.config.aws_region, - table_name=self.config.experiment_question_metrics_table - ) - self.experiment_db = DynamoDBOperations( - region=self.config.aws_region, - table_name=self.config.experiment_table - ) - - def get_all_questions(self, experiment_id: str) -> List[Dict]: - """Fetch all questions for a given experiment""" - expression_values = {":experimentId": experiment_id} - return self.metrics_db.query( - "experiment_id = :experimentId", - expression_values=expression_values, - index_name=self.config.experiment_question_metrics_experimentid_index - ) - - def update_experiment_metrics(self, experiment_id: str, experiment_eval_metrics: Dict[str, float]): - """Update overall experiment metrics""" - try: - if experiment_eval_metrics: - logger.info(f"Updating experiment metrics for experiment {experiment_id}") - self.experiment_db.update_item( - key={'id': experiment_id}, - update_expression="SET eval_metrics = :eval", - expression_values={':eval': {'M' : experiment_eval_metrics}} - ) - except Exception as e: - logger.error(f"Error updating experiment metrics: {e}") - - def calculate_eval_score(self,evaluator,data): - try: - score=evaluator.single_turn_score(data) - return score - except Exception as e: - logger.error(f"Error processing sample : {e}") - return 0.0 - - \ No newline at end of file diff --git a/core/eval/ragas/ragas_llm_eval.py b/core/eval/ragas/ragas_llm_eval.py deleted file mode 100644 index 7485a62d..00000000 --- a/core/eval/ragas/ragas_llm_eval.py +++ /dev/null @@ -1,131 +0,0 @@ -from core.eval.ragas.ragas_eval import RagasEvaluator -from ragas import evaluate -from ragas.dataset_schema import SingleTurnSample, EvaluationDataset -from ragas.metrics._string import NonLLMStringSimilarity -from ragas.metrics import Faithfulness, AspectCritic, LLMContextPrecisionWithoutReference, ResponseRelevancy, LLMContextPrecisionWithReference -from langchain_aws import ChatBedrockConverse, BedrockEmbeddings -from ragas.embeddings import LangchainEmbeddingsWrapper -from ragas.llms import LangchainLLMWrapper -from baseclasses.base_classes import ExperimentQuestionMetrics, EvaluationMetrics -from typing import Optional, List -from core.eval.eval_factory import EvalFactory - -import logging - -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) - -class RagasLLMEvaluator(RagasEvaluator): - def __init__(self, config, experimental_config): - super().__init__(config, experimental_config) - self._initialze_llm() - self._initialize_scorers() - - - def _initialze_llm(self): - self.evaluator_llm = LangchainLLMWrapper(ChatBedrockConverse( - region_name=self.experimental_config.aws_region, - base_url=f"https://bedrock-runtime.{self.experimental_config.aws_region}.amazonaws.com", - model=self.experimental_config.eval_retrieval_model, - temperature=self.experimental_config.eval_retrieval_temperature, - )) - - self.embedding_llm = LangchainEmbeddingsWrapper(BedrockEmbeddings( - region_name=self.experimental_config.aws_region, - model_id=self.experimental_config.eval_embedding_model, - )) - - - def _initialize_scorers(self): - """Initialize all metric scorers""" - self.faithfulness = Faithfulness(llm=self.evaluator_llm) - - self.aspect_critic = AspectCritic(name="maliciousness", - definition="Is the submission intended to harm, deceive, or exploit users?", - llm=self.evaluator_llm) - - self.context_precision = LLMContextPrecisionWithReference(llm=self.evaluator_llm) - - self.answers_relevancy = ResponseRelevancy(llm=self.evaluator_llm, - embeddings=self.embedding_llm) - - def get_questions(self, experiment_id): - return super().get_questions(experiment_id) - - def evaluate(self, experiment_id: str): - """Perform evaluation for all questions in an experiment""" - if not experiment_id: - raise ValueError("Experiment ID cannot be None") - - questions = self.get_all_questions(experiment_id)['Items'] # Contains question, generated answer, GT answer and retrieved context - - metrics_records = [ExperimentQuestionMetrics(**question) for question in questions] - metrics = self.evaluate_bulk_questions(metrics_records) - - experiment_eval_metrics = {} - if metrics: - experiment_eval_metrics = metrics._repr_dict - experiment_eval_metrics = {key: round(value, 2) if isinstance(value, float) else value for key, value in experiment_eval_metrics.items()} - experiment_eval_metrics = EvaluationMetrics().from_dict(experiment_eval_metrics).to_dict() - logger.info(f"Experiment {experiment_id} evaluation metrics: {experiment_eval_metrics}") - - - self.update_experiment_metrics(experiment_id, experiment_eval_metrics) - - def evaluate_bulk_questions(self, metrics_records: List[ExperimentQuestionMetrics]): - """Evaluate a list of metrics records""" - answer_samples = [] - - metrics_to_evaluate = [self.aspect_critic, self.answers_relevancy] - if self.experimental_config.knowledge_base: - metrics_to_evaluate = metrics_to_evaluate + [self.faithfulness, self.context_precision] - - - for metrics_record in metrics_records: - sample_params = { - 'user_input': metrics_record.question, - 'response': metrics_record.generated_answer, - 'reference': metrics_record.gt_answer - } - if self.experimental_config.knowledge_base: - sample_params['retrieved_contexts'] = metrics_record.reference_contexts - - answer_sample = SingleTurnSample(**sample_params) - - answer_samples.append(answer_sample) - - evaluation_dataset = EvaluationDataset(answer_samples) - metrics = evaluate(evaluation_dataset, metrics_to_evaluate) - - return metrics - - - def _evaluate_single_question(self, metrics_record: ExperimentQuestionMetrics) -> Optional[EvaluationMetrics]: - """Evaluate a single question and return its metrics""" - try: - if metrics_record.generated_answer: - answer_sample = SingleTurnSample( - user_input=metrics_record.question, - response=metrics_record.generated_answer, - reference=metrics_record.gt_answer, - retrieved_contexts=metrics_record.reference_contexts - ) - - metrics = EvaluationMetrics( - - faithfulness_score=self.calculate_eval_score(self.faithfulness,answer_sample), - - context_precision_score=self.calculate_eval_score(self.context_precision,answer_sample), - - aspect_critic_score=self.calculate_eval_score(self.aspect_critic,answer_sample), - - answers_relevancy_score=self.calculate_eval_score(self.answers_relevancy,answer_sample) - - ) - return metrics - except Exception as e: - logger.error(f"Error processing sample {metrics_record.id}: {e}") - return {} - - -EvalFactory.register_evaluator('ragas', 'llm', RagasLLMEvaluator) \ No newline at end of file diff --git a/core/eval/ragas/ragas_non_llm_eval.py b/core/eval/ragas/ragas_non_llm_eval.py deleted file mode 100644 index cda755be..00000000 --- a/core/eval/ragas/ragas_non_llm_eval.py +++ /dev/null @@ -1,82 +0,0 @@ -from core.eval.ragas.ragas_eval import RagasEvaluator -from ragas.metrics import NonLLMStringSimilarity, NonLLMContextRecall, NonLLMContextPrecisionWithReference, RougeScore, BleuScore, Faithfulness -from baseclasses.base_classes import ExperimentQuestionMetrics, EvaluationMetrics -from ragas.dataset_schema import SingleTurnSample -from typing import Optional -from core.eval.eval_factory import EvalFactory - -import logging -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) - -class RagasNonLLMEvaluator(RagasEvaluator): - - def __init__(self, config, experimental_config): - super().__init__(config, experimental_config) - self._initialize_scorers() - - - def _initialize_scorers(self): - """Initialize all metric scorers""" - self.str_similar_scorer = NonLLMStringSimilarity() - self.context_recall = NonLLMContextRecall() - self.context_precision = NonLLMContextPrecisionWithReference() - self.rouge_score = RougeScore() - self.bleu_score = BleuScore() - - def get_questions(self, experiment_id): - return super().get_questions(experiment_id) - - def evaluate(self, experiment_id: str): - """Perform evaluation for all questions in an experiment""" - if not experiment_id: - raise ValueError("Experiment ID cannot be None") - - questions = self.get_all_questions(experiment_id)['Items'] - metrics_list = [] - evaluation_dict = {} - - for question in questions: - metrics_record = ExperimentQuestionMetrics(**question) - metrics = self._evaluate_single_question(metrics_record) - - if metrics: - metrics_list.append(metrics) - update_expression = "SET eval_metrics = :metrics" - expression_attribute_values = { - ':metrics': {'M' : metrics.to_dict()} - } - - self.metrics_db.update_item( - key={'id': metrics_record.id}, - update_expression=update_expression, - expression_values=expression_attribute_values - ) - self.update_experiment_metrics(experiment_id, metrics_list) - - def _evaluate_single_question(self, metrics_record: ExperimentQuestionMetrics) -> Optional[EvaluationMetrics]: - """Evaluate a single question and return its metrics""" - try: - answer_sample = SingleTurnSample( - response=metrics_record.generated_answer, - reference=metrics_record.gt_answer - ) - - context_sample = SingleTurnSample( - retrieved_contexts=[metrics_record.generated_answer], - reference_contexts=metrics_record.reference_contexts - ) - - metrics = EvaluationMetrics( - string_similarity=self.calculate_eval_score(self.str_similar_scorer, answer_sample), - context_precision=self.calculate_eval_score(self.context_precision, context_sample), - context_recall=self.calculate_eval_score(self.context_recall, context_sample), - rouge_score=self.calculate_eval_score(self.rouge_score,answer_sample) - ) - - return metrics - except Exception as e: - logger.error(f"Error processing sample {metrics_record.id}: {e}") - return None - -EvalFactory.register_evaluator('ragas', 'non_llm', RagasNonLLMEvaluator) \ No newline at end of file diff --git a/core/guardrails/bedrock_guardrails.py b/core/guardrails/bedrock_guardrails.py deleted file mode 100644 index 9c58e242..00000000 --- a/core/guardrails/bedrock_guardrails.py +++ /dev/null @@ -1,284 +0,0 @@ -from typing import List, Dict, Optional, Any -import boto3, yaml, uuid -from botocore.exceptions import ClientError - -class BedrockGuardrails: - def __init__(self, region: str = 'us-east-1'): - self.bedrock_client = boto3.client('bedrock', region_name = region) - self.runtime_client = boto3.client('bedrock-runtime', region_name = region) - - def create_guardrail( - self, - guardrail_config: dict - ) -> Dict: - """ - Create a comprehensive guardrail using Amazon Bedrock - - Args: - name: Name of the guardrail - description: Description of the guardrail - content_policy: Content policy configuration - topic_policy: Topic policy configuration - word_policy: Word policy configuration - sensitive_info_policy: Sensitive information policy configuration - contextual_grounding_policy: Contextual grounding policy configuration - input_filter: Input filtering configuration - output_filter: Output filtering configuration - """ - try: - - response = self.bedrock_client.create_guardrail(**guardrail_config) - return response - - except ClientError as e: - print(f"Error creating guardrail: {str(e)}") - raise - - def apply_guardrail( - self, - guardrail_id: str, - guardrail_version: str, - content: str, - source: str = 'INPUT', - model_id: str = None, - inference_configuration: Dict = None - ) -> Dict: - """ - Apply a guardrail to content using Amazon Bedrock ApplyGuardrails API - - Args: - guardrail_id (str): The unique identifier of the guardrail - guardrail_version (str): The version of the guardrail to apply - content (str): The content to validate against the guardrail - source (str): The source of the content ('INPUT' or 'OUTPUT') - model_id (str, optional): The model identifier to use for inference - inference_configuration (Dict, optional): Additional inference configuration parameters - - Returns: - Dict: Response from the ApplyGuardrails API - - Example response structure: - { - 'results': [{ - 'status': 'ALLOWED'|'FILTERED'|'DENIED', - 'statusMessage': 'string', - 'violations': [{ - 'policyName': 'string', - 'violationType': 'string', - 'violationMessage': 'string' - }] - }], - 'responseMetadata': { - 'requestId': 'string', - 'attempts': 123, - 'totalRetryDelay': 123.0 - } - } - """ - try: - request_params = { - 'guardrailIdentifier': guardrail_id, - 'guardrailVersion': guardrail_version, - 'source': source, - 'content': content - } - - response = self.runtime_client.apply_guardrail(**request_params) - return response - - except ClientError as e: - print(f"Error applying guardrail: {str(e)}") - raise - - def load_guardrail_config_from_yaml(self, yaml_file_path: str) -> Dict[str, Any]: - """ - Loads guardrail configuration from a YAML file and converts it to Bedrock-compatible format. - - Args: - yaml_file_path (str): Path to the YAML configuration file - - Returns: - Dict[str, Any]: Bedrock-compatible guardrail configuration - """ - try: - with open(yaml_file_path, 'r') as file: - yaml_config = yaml.safe_load(file) - - if not yaml_config or 'guardrails' not in yaml_config: - raise ValueError("Invalid YAML configuration: 'guardrails' section not found") - - config = yaml_config['guardrails'] - - bedrock_config = {} - - if 'name' in config: - random_suffix = str(uuid.uuid4())[:4] - bedrock_config['name'] = f"{config['name']}-{random_suffix}" - - if 'description' in config: - bedrock_config['description'] = config['description'] - - # Convert content policy - if 'content_policy' in config: - bedrock_config['contentPolicyConfig'] = { - 'filtersConfig': config['content_policy']['filtersConfig'] - } - - # Convert topic policy - if 'topic_policy' in config: - bedrock_config['topicPolicyConfig'] = { - 'topicsConfig': config['topic_policy']['topicsConfig'] - } - - # Convert word policy - if 'word_policy' in config: - bedrock_config['wordPolicyConfig'] = { - 'wordsConfig': config['word_policy'].get('wordsConfig', []), - 'managedWordListsConfig': config['word_policy'].get('managedWordListsConfig', []) - } - - # Convert sensitive info policy - if 'sensitive_info_policy' in config: - bedrock_config['sensitiveInformationPolicyConfig'] = { - 'piiEntitiesConfig': config['sensitive_info_policy'].get('piiEntitiesConfig', []), - 'regexesConfig': config['sensitive_info_policy'].get('regexesConfig', []) - } - - # Convert contextual grounding policy - if 'contextual_grounding_policy' in config: - bedrock_config['contextualGroundingPolicyConfig'] = { - 'filtersConfig': config['contextual_grounding_policy']['filtersConfig'] - } - - # Convert filter messages - if 'blocked_input_message' in config: - bedrock_config['blockedInputMessaging'] = config['blocked_input_message'] - - if 'blocked_outputs_message' in config: - bedrock_config['blockedOutputsMessaging'] = config['blocked_outputs_message'] - - return bedrock_config - - except FileNotFoundError: - raise FileNotFoundError(f"Configuration file not found: {yaml_file_path}") - except yaml.YAMLError as e: - raise ValueError(f"Error parsing YAML configuration: {str(e)}") - except Exception as e: - raise Exception(f"Error processing configuration: {str(e)}") - -# Example configurations for different policy types -def get_comprehensive_guardrail_config(): - """ - Complete guardrail configuration with all policy types - """ - # Content Policy Configuration - content_policy = { - 'filtersConfig': [{ - 'type': 'HATE', #'SEXUAL'|'VIOLENCE'|'HATE'|'INSULTS'|'MISCONDUCT'|'PROMPT_ATTACK' - 'inputStrength': 'LOW', #'NONE'|'LOW'|'MEDIUM'|'HIGH' - 'outputStrength': 'LOW', #'NONE'|'LOW'|'MEDIUM'|'HIGH', - 'inputModalities': [ - 'TEXT' #'TEXT'|'IMAGE' - ], - 'outputModalities': [ - 'TEXT' #'TEXT'|'IMAGE' - ] - }] - } - - # Topic Policy Configuration - topic_policy = { - 'topicsConfig': [{ - 'name': 'Financial Advice', - 'definition': 'Providing financial or investment advice', - 'examples': [ - 'What stocks should I invest in?', - 'How should I invest my money?' - ], - 'type': 'DENY' - }, { - 'name': 'Medical Advice', - 'definition': 'Providing medical diagnosis or treatment advice', - 'examples': [ - 'What medication should I take?', - 'How should I treat this condition?' - ], - 'type': 'DENY' - }] - } - - # Word Policy Configuration - word_policy = { - 'wordsConfig': [ - { - 'text': 'crazy' #Any string - }, - ], - 'managedWordListsConfig': [ - { - 'type': 'PROFANITY' - }, - ] - } - - # Sensitive Information Policy Configuration - sensitive_info_policy = { - 'piiEntitiesConfig': [{ - 'type': 'EMAIL_ADDRESS', #'ADDRESS'|'AGE'|'AWS_ACCESS_KEY'|'AWS_SECRET_KEY'|'CA_HEALTH_NUMBER'|'CA_SOCIAL_INSURANCE_NUMBER'|'CREDIT_DEBIT_CARD_CVV'|'CREDIT_DEBIT_CARD_EXPIRY'|'CREDIT_DEBIT_CARD_NUMBER'|'DRIVER_ID'|'EMAIL'|'INTERNATIONAL_BANK_ACCOUNT_NUMBER'|'IP_ADDRESS'|'LICENSE_PLATE'|'MAC_ADDRESS'|'NAME'|'PASSWORD'|'PHONE'|'PIN'|'SWIFT_CODE'|'UK_NATIONAL_HEALTH_SERVICE_NUMBER'|'UK_NATIONAL_INSURANCE_NUMBER'|'UK_UNIQUE_TAXPAYER_REFERENCE_NUMBER'|'URL'|'USERNAME'|'US_BANK_ACCOUNT_NUMBER'|'US_BANK_ROUTING_NUMBER'|'US_INDIVIDUAL_TAX_IDENTIFICATION_NUMBER'|'US_PASSPORT_NUMBER'|'US_SOCIAL_SECURITY_NUMBER'|'VEHICLE_IDENTIFICATION_NUMBER' - 'mode': 'BLOCK' # 'BLOCK'|'ANONYMIZE' - }, { - 'type': 'PHONE_NUMBER', - 'mode': 'BLOCK' - }, { - 'type': 'CREDIT_DEBIT_NUMBER', - 'mode': 'ANONYMIZE' - }, { - 'type': 'SSN', - 'mode': 'ANONYMIZE' - }], - 'regexesConfig': [{ - 'name': 'SSN', #Any string - 'pattern': r'\b\d{3}-\d{2}-\d{4}\b', #Any string - 'description': 'SSN pattern', #Any string - 'action': 'BLOCK' # 'BLOCK'|'ANONYMIZE' - }] - } - - # Contextual Grounding Policy Configuration - contextual_grounding_policy = { - 'filtersConfig': [ - { - 'type': 'GROUNDING', #'GROUNDING'|'RELEVANCE' - 'threshold': 123.0 # Value - }, - ] - } - - input_filter = 'This input contains restricted content.' - - output_filter = 'This output contains restricted content.' - - return { - 'content_policy': content_policy, - 'topic_policy': topic_policy, - 'word_policy': word_policy, - 'sensitive_info_policy': sensitive_info_policy, - 'contextual_grounding_policy': contextual_grounding_policy, - 'input_filter': input_filter, - 'output_filter': output_filter - } - -# Example configurations -def create_domain_specific_guardrail(): - """ - Create a domain-specific guardrail with all policy types - """ - guardrails = BedrockGuardrails() - config = get_comprehensive_guardrail_config() - - response = guardrails.create_guardrail( - name='ComprehensiveGuardrail', - description='Complete guardrail with all policy configurations', - **config - ) - return response diff --git a/core/inference/__init__.py b/core/inference/__init__.py deleted file mode 100644 index 995496f2..00000000 --- a/core/inference/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -from .bedrock.bedrock_inferencer import BedrockInferencer - -# Importing SageMaker-specific inference and inference factory. -from .inference_factory import InferencerFactory -from .sagemaker.sagemaker_inferencer import SageMakerInferencer -from .sagemaker.llama_inferencer import LlamaInferencer - -# List of model names that you want to register with the InferencerFactory -model_list = [ - "meta-textgeneration-llama-3-1-8b-instruct", # Model for text generation (Llama) - "huggingface-llm-falcon-7b-instruct-bf16", # Model for text generation (Falcon) - "meta-vlm-llama-4-scout-17b-16e-instruct", - "meta-textgeneration-llama-3-3-70b-instruct", #Llama model with more parameters for text generation - "deepseek-ai/DeepSeek-R1-Distill-Llama-8B", - "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", - "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", - "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B" - ] - -# Registering each model from the list into the InferencerFactory under 'sagemaker'. -# The `SageMakerInferencer` will be used for inferencing operations for these models. -for model in model_list: - if model.startswith("meta-vlm-llama-4"): - InferencerFactory.register_inferencer('sagemaker', model, LlamaInferencer) - else: - InferencerFactory.register_inferencer('sagemaker', model, SageMakerInferencer) \ No newline at end of file diff --git a/core/inference/bedrock/bedrock_inferencer.py b/core/inference/bedrock/bedrock_inferencer.py deleted file mode 100644 index 248afd66..00000000 --- a/core/inference/bedrock/bedrock_inferencer.py +++ /dev/null @@ -1,167 +0,0 @@ -from baseclasses.base_classes import BaseInferencer -import boto3 -from typing import List, Dict, Any, Union, Tuple -import logging -from config.experimental_config import ExperimentalConfig, NShotPromptGuide -from core.inference.inference_factory import InferencerFactory -from util.boto3_utils import BedRockRetryHander -import random - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -# Class for handling inference using Amazon Bedrock -class BedrockInferencer(BaseInferencer): - """Base class for all Bedrock models since they share the same invocation pattern""" - - # This part was added as part of the SageMaker implementation changes. - # Due to updates in the base class implementation, the invocation point was changed from - # '_initialize_client' to '__init__'. - def __init__(self, model_id: str, experiment_config: ExperimentalConfig, region: str = 'us-east-1', role_arn: str = None): - super().__init__(model_id, experiment_config, region, role_arn) - self._initialize_client() - - def _initialize_client(self) -> None: - self.client = boto3.client( - service_name='bedrock-runtime', - region_name=self.region_name - ) - - def generate_prompt(self, experiment_config: ExperimentalConfig, default_prompt: str, user_query: str, context: List[Dict] = None) -> Tuple[str, List[Dict[str, Any]]]: - # Get n_shot config values first to avoid repeated lookups - n_shot_prompt_guide = experiment_config.n_shot_prompt_guide_obj - n_shot_prompt = experiment_config.n_shot_prompts - - messages = [] - - context_text = "" - if context: - context_text = self._format_context(context) - if context_text: - messages.append(self._prepare_conversation(role="user", message=context_text)) - - # Input validation - if n_shot_prompt < 0: - raise ValueError("n_shot_prompt must be non-negative") - - # Get system prompt - system_prompt = default_prompt if n_shot_prompt_guide is None or n_shot_prompt_guide.system_prompt is None else n_shot_prompt_guide.system_prompt - - base_prompt = n_shot_prompt_guide.user_prompt if n_shot_prompt_guide.user_prompt else "" - messages.append(self._prepare_conversation(role="user", message=base_prompt)) - - # Get examples - examples = n_shot_prompt_guide.examples - - # Format examples - selected_examples = (random.sample(examples, n_shot_prompt) - if len(examples) > n_shot_prompt - else examples) - - # Use string concatenation for example formatting - for example in selected_examples: - if 'example' in example: - messages.append(self._prepare_conversation(role="user", message=example['example'])) - elif 'question' in example and 'answer' in example: - messages.append(self._prepare_conversation(role="user", message=example['question'])) - messages.append(self._prepare_conversation(role="assistant", message=example['answer'])) - - logger.info(f"into {n_shot_prompt} shot prompt with examples {len(selected_examples)}") - - # Add the current user prompt - messages.append(self._prepare_conversation(role="user", message=user_query)) - - return system_prompt, messages - - @BedRockRetryHander() - def generate_text(self, user_query: str, default_prompt: str, context: List[Dict] = None, **kwargs) -> Tuple[Dict[Any, Any], str]: - try: - # Code to generate prompt considering the upload prompt config file - system_prompt, messages = self.generate_prompt(self.experiment_config, default_prompt, user_query, context) - - inference_config = { - "maxTokens": 512, - "temperature": self.experiment_config.temp_retrieval_llm, - "topP": 0.9 - } - - skip_system_param = self.model_id in ("amazon.titan-text-express-v1", "amazon.titan-text-lite-v1", "mistral.mistral-7b-instruct-v0:2") - - request_params = { - "modelId": self.model_id, - "messages": ([self._prepare_conversation(role="user", message=system_prompt)] if skip_system_param else []) + messages, - "inferenceConfig": inference_config - } - - # Add system parameter only for non-Titan-v1 models - #TODO: Short-term fix, will be addressed using inheritence as part of refactoring - if not skip_system_param: - request_params["system"] = [{"text" : system_prompt}] - - response = self.client.converse(**request_params) - - metadata = {} - if 'usage' in response: - for key, value in response['usage'].items(): - metadata[key] = value - if 'metrics' in response: - for key, value in response['metrics'].items(): - metadata[key] = value - return metadata, self._extract_response(response) - except Exception as e: - logger.error(f"Error generating text with Bedrock: {str(e)}") - raise - - def _prepare_conversation(self, message: str, role: str): - # Format message and role into a conversation - if not message or not role: - logger.error(f"Error in parsing message or role") - conversation = { - "role": role, - "content": [{"text" : message}] - } - - return conversation - - def _format_context(self, context: List[Dict[str, str]]) -> str: - """Format context documents into a single string.""" - context_text = "\n".join([ - f"Context {i+1}:\n{doc.get('text', '')}" - for i, doc in enumerate(context) - ]) - logger.debug(f"Formatted context text length: {len(context_text)}") - context_text = "\n" + context_text + "\n" + "\n" - return context_text - - def _extract_response(self, response: Dict) -> str: - """Extract text from Bedrock response.""" - response_text = response["output"]["message"]["content"][0]["text"] - logger.info("Successfully generated response") - logger.debug(f"Response length: {len(response_text)}") - return response_text - - -model_list = ["mistral.mistral-7b-instruct-v0:2", - "mistral.mistral-large-2402-v1:0", - "us.meta.llama3-2-90b-instruct-v1:0", - "us.meta.llama3-2-11b-instruct-v1:0", - "us.meta.llama3-2-3b-instruct-v1:0", - "us.meta.llama3-2-1b-instruct-v1:0", - "cohere.command-r-v1:0", - "cohere.command-r-plus-v1:0", - "anthropic.claude-3-5-sonnet-20240620-v1:0", - "us.anthropic.claude-3-5-sonnet-20241022-v2:0", - "us.anthropic.claude-3-5-haiku-20241022-v1:0", - "amazon.titan-text-express-v1", - "amazon.titan-text-lite-v1", - "us.amazon.nova-lite-v1:0", - "us.amazon.nova-micro-v1:0", - "us.amazon.nova-pro-v1:0", - "us.amazon.nova-premier-v1:0", - "us.anthropic.claude-3-7-sonnet-20250219-v1:0" - ] - -for model in model_list: - InferencerFactory.register_inferencer('bedrock', model, BedrockInferencer) - - \ No newline at end of file diff --git a/core/inference/inference_factory.py b/core/inference/inference_factory.py deleted file mode 100644 index 351dbac1..00000000 --- a/core/inference/inference_factory.py +++ /dev/null @@ -1,49 +0,0 @@ -from config.config import get_config -from config.experimental_config import ExperimentalConfig -import logging -from baseclasses.base_classes import BaseInferencer -from typing import Dict, Type - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class InferenceServiceError(Exception): - """Custom exception for inference service related errors""" - pass - - -class InferencerFactory: - - """Factory to create embedders based on model ID and service type.""" - - _registry: Dict[str, Type[BaseInferencer]] = {} - - @classmethod - def register_inferencer(cls, service_type: str, model_id: str, embedder_cls: Type[BaseInferencer]): - key = f"{service_type}:{model_id}" - cls._registry[key] = embedder_cls - - @classmethod - def create_inferencer(cls, experimentalConfig : ExperimentalConfig) -> BaseInferencer: - service_type = experimentalConfig.retrieval_service - model_id = experimentalConfig.retrieval_model - key = f"{service_type}:{model_id}" - - inferencer_cls = cls._registry.get(key) - if not inferencer_cls: - raise InferenceServiceError(f"No inferencer_cls registered for service {service_type} and model {model_id}") - - if service_type == "sagemaker": - role_arn = get_config().sagemaker_role_arn - elif service_type == "bedrock": - role_arn = get_config().bedrock_role_arn - else: - role_arn = None - - # return inferencer_cls(model_id=model_id, region=experimentalConfig.aws_region, experiment_config=experimentalConfig) - return inferencer_cls( - model_id=model_id, - experiment_config=experimentalConfig, - region=experimentalConfig.aws_region, - role_arn=role_arn - ) \ No newline at end of file diff --git a/core/inference/sagemaker/llama_inferencer.py b/core/inference/sagemaker/llama_inferencer.py deleted file mode 100644 index ff443924..00000000 --- a/core/inference/sagemaker/llama_inferencer.py +++ /dev/null @@ -1,115 +0,0 @@ -from typing import List, Dict -import logging -import time -import random -from config.experimental_config import ExperimentalConfig -from .sagemaker_inferencer import SageMakerInferencer - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class LlamaInferencer(SageMakerInferencer): - def __init__(self, model_id: str, experiment_config: ExperimentalConfig, region: str, role_arn: str): - super().__init__(model_id, experiment_config, region, role_arn) - - def _prepare_conversation(self, message: str, role: str): - # Format message and role into a conversation - if not message or not role: - logger.error(f"Error in parsing message or role") - conversation = { - "role": role, - "content": message - } - return conversation - - def generate_prompt(self, experiment_config: ExperimentalConfig, default_prompt: str, user_query: str, context: List[Dict] = None): - n_shot_prompt_guide = experiment_config.n_shot_prompt_guide_obj - n_shot_prompt = experiment_config.n_shot_prompts - # Input validation - if n_shot_prompt < 0: - raise ValueError("n_shot_prompt must be non-negative") - - # Get system prompt - system_prompt = default_prompt if n_shot_prompt_guide is None or n_shot_prompt_guide.system_prompt is None else n_shot_prompt_guide.system_prompt - - context_text = "" - if context: - context_text = self._format_context(user_query, context) - - base_prompt = n_shot_prompt_guide.user_prompt if n_shot_prompt_guide.user_prompt else "" - - if n_shot_prompt == 0: - logger.info("into zero shot prompt") - - messages = [] - messages.append(self._prepare_conversation(role="user", message=base_prompt)) - if context_text: - messages.append(self._prepare_conversation(role="user", message=context_text)) - messages.append(self._prepare_conversation(role="user", message=user_query)) - - return system_prompt, messages - - # Get examples if nshot is not zero - examples = n_shot_prompt_guide.examples - - # Format examples - selected_examples = (random.sample(examples, n_shot_prompt) - if len(examples) > n_shot_prompt - else examples) - - logger.info(f"into {n_shot_prompt} shot prompt with examples {len(selected_examples)}") - - messages = [] - messages.append(self._prepare_conversation(role="user", message=base_prompt)) - for example in selected_examples: - if 'example' in example: - messages.append(self._prepare_conversation(role="user", message=example['example'])) - elif 'question' in example and 'answer' in example: - messages.append(self._prepare_conversation(role="user", message=example['question'])) - messages.append(self._prepare_conversation(role="assistant", message=example['answer'])) - - if context_text: - messages.append(self._prepare_conversation(role="user", message=context_text)) - - messages.append(self._prepare_conversation(role="user", message=user_query)) - - return system_prompt, messages - - def construct_payload(self, system_prompt: str, prompt: str) -> dict: - """ - Constructs llama 4 payload dictionary for model inference with the given prompts and default parameters. - - Args: - system_prompt (str): The system-level prompt that guides the model's behavior - prompt (str): The actual prompt/query to be sent to the model - - """ - # Define default parameters for the model's generation - default_params = { - "max_new_tokens": 256, - "temperature": self.experiment_config.temp_retrieval_llm, - "top_p": 0.9, - "do_sample": True - } - - # Prepare payload for model inference - payload = { - "system": system_prompt, - "messages": prompt, - "parameters": default_params - } - - return payload - - def parse_response(self, response: dict) -> str: - """ - Parses the response from the model and extracts the generated text. - - Args: - response (dict): The raw response from the model - """ - if "choices" in response and isinstance(response["choices"], list): - return response["choices"][0]["message"]["content"] - else: - raise ValueError(f"Unexpected Llama-4 response format: {response}") - \ No newline at end of file diff --git a/core/inference/sagemaker/sagemaker_inferencer.py b/core/inference/sagemaker/sagemaker_inferencer.py deleted file mode 100644 index 81f89811..00000000 --- a/core/inference/sagemaker/sagemaker_inferencer.py +++ /dev/null @@ -1,611 +0,0 @@ -import boto3 -from typing import List, Dict -from botocore.exceptions import ClientError -from baseclasses.base_classes import BaseInferencer -from config.experimental_config import ExperimentalConfig -from sagemaker.session import Session -from sagemaker.predictor import Predictor -from sagemaker.serializers import JSONSerializer -from sagemaker.deserializers import JSONDeserializer -from sagemaker.jumpstart.model import JumpStartModel -from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri -import sagemaker -import logging -import time -import random -import json - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -# Model configurations - -INFERENCER_MODELS = { - "meta-textgeneration-llama-3-1-8b-instruct": { - "model_source": "jumpstart", - "instance_type": "ml.g5.2xlarge" - }, - "huggingface-llm-falcon-7b-instruct-bf16": { - "model_source": "jumpstart", - "instance_type": "ml.g5.2xlarge" - }, - "meta-textgeneration-llama-3-3-70b-instruct": { - "model_source": "jumpstart", - "instance_type": "ml.p4d.24xlarge" - }, - "meta-vlm-llama-4-scout-17b-16e-instruct": { - "model_source": "jumpstart", - "instance_type": "ml.p4d.24xlarge" - } - , - "deepseek-ai/DeepSeek-R1-Distill-Llama-8B": { - "model_source": "huggingface", - "instance_type": "ml.g5.2xlarge" - } - , - "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B": { - "model_source": "huggingface", - "instance_type": "ml.g5.xlarge" - } - , - "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B": { - "model_source": "huggingface", - "instance_type": "ml.g5.xlarge" - } - , - "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { - "model_source": "huggingface", - "instance_type": "ml.g6e.12xlarge" - } -} - -# Sagemaker Base Inferencer -class SageMakerInferencer(BaseInferencer): - - def __init__(self, model_id: str, experiment_config: ExperimentalConfig, region: str, role_arn: str): - """ - Initializes the SageMakerInferencer with the given model ID, region, and role ARN. - Sets up necessary SageMaker runtime clients, session, and endpoint predictor. - - Args: - model_id (str): The unique identifier for the model. - region (str): The AWS region where the SageMaker services are hosted. - role_arn (str): The ARN of the IAM role. Currently not used but included for future extensions. - """ - - # Store the region - self.region_name = region - - self.role = role_arn - - # Initialize the base class - super().__init__(model_id, experiment_config, region, role_arn) - - logger.info(f"Initializing SageMaker Generator for model: {model_id}") - - # Initialize the SageMaker runtime and client for general operations - self.client = boto3.client("sagemaker-runtime", region_name=region) - self.sagemaker_client = boto3.client('sagemaker', region_name=region) - - # Create a new SageMaker session - self.session = Session(boto_session=boto3.Session(region_name=region)) - - # Initialize additional inferencing-related attributes - self.inferencing_model_id = model_id - self.inferencing_model_endpoint_name = f"{self._sanitize_name(model_id)[:42]}-inferencing-endpoint" - - self.wait_time = 5 - - # Ensure the endpoint exists or create it if necessary - self._ensure_endpoint_exists() - - # Initialize the predictor to interact with the SageMaker endpoint - self.predictor = Predictor( - endpoint_name=self.inferencing_model_endpoint_name, - sagemaker_session=self.session - ) - - # Set up the serializer and deserializer for the predictor - self.predictor.serializer = JSONSerializer() - self.predictor.deserializer = JSONDeserializer() - - self.inferencing_predictor = self.predictor - - # Log initialization success - logger.info(f"Initialized SageMakerInferencer for model {model_id} in region {region}.") - - def _ensure_endpoint_exists(self): - """ - Ensures that the SageMaker endpoint exists for the given model. If the endpoint does not exist, - it creates a new endpoint using the specified model ID. - - Args: - model_id (str): The unique identifier for the model to use for the endpoint. - - Raises: - ClientError: If there is an issue communicating with SageMaker or creating the endpoint. - """ - - try: - # Check if the endpoint already exists - _ = self._check_model_status(self.inferencing_model_endpoint_name) - logger.info(f"Endpoint {self.inferencing_model_endpoint_name} already exists.") - except self.sagemaker_client.exceptions.ClientError: - # If the endpoint does not exist, create a new one - logger.info(f"Endpoint and configuration for {self.inferencing_model_endpoint_name} does not exist. Creating endpoint.") - self.create_endpoint(endpoint_name=self.inferencing_model_endpoint_name, model_id=self.inferencing_model_id) - - def _check_model_status(self, endpoint_name): - """ - Check the status of the SageMaker endpoint and its configuration. - - This method performs the following: - 1. Checks if endpoint exists and is in service - 2. If endpoint is being created, waits until creation completes - 3. If no endpoint exists, checks for endpoint configuration - 4. If configuration exists, waits for endpoint creation to complete - - Args: - endpoint_name (str): Name of the SageMaker endpoint to check - - Returns: - str: 'InService' if endpoint is available and running - - Raises: - Exception: If endpoint creation fails or has unexpected status - ClientError: If neither endpoint nor configuration exists - """ - try: - - while True: - # Poll endpoint status until it is in service or fails - response = self.sagemaker_client.describe_endpoint(EndpointName=endpoint_name) - - if response['EndpointStatus'] == 'InService': - logger.info(f"Endpoint {endpoint_name} is in service.") - return 'InService' - - elif response['EndpointStatus'] == 'Failed': - logger.error(f"Endpoint {endpoint_name} creation failed.") - raise Exception(f"Endpoint {endpoint_name} creation failed.") - - elif response['EndpointStatus'] == 'Creating': - time.sleep(self.wait_time) # Pause before next status check - - else: - raise Exception(f"Unexpected endpoint status: {response['EndpointStatus']}") - except self.sagemaker_client.exceptions.ClientError: - # No endpoint exists - check if there's a configuration waiting to be deployed - logger.info(f"Endpoint {endpoint_name} does not exist. Checking if endpoint configuration exists.") - - try: - # Look for endpoint configuration that may have been created by another process - response = self.sagemaker_client.describe_endpoint_config(EndpointConfigName=endpoint_name) - logger.info(f"Configuration for {endpoint_name} exists, waiting {self.wait_time} seconds for endpoint creation.") - - time.sleep(self.wait_time) # Allow time for endpoint creation to begin - - _ = self._check_model_status(endpoint_name) # Keep rechecking the endpoint status until it begins creation - - except self.sagemaker_client.exceptions.ClientError: - # Neither endpoint nor configuration exists - logger.info(f"Endpoint configuration does not exist.") - raise - - except Exception as e: - logger.error(f"Error checking endpoint status: {e}") - raise - - def create_endpoint(self, endpoint_name: str, model_id: str) -> sagemaker.predictor.Predictor: - """ - Creates a SageMaker endpoint for the specified model if it doesn't already exist. - If the endpoint is already in service, returns the existing predictor. - - Args: - endpoint_name (str): The name of the SageMaker endpoint to be created or fetched. - model_id (str): The identifier for the model to be used in the endpoint. - - Returns: - sagemaker.predictor.Predictor: A predictor object for the created or fetched endpoint. - - Raises: - ValueError: If the provided model_id is not supported. - ClientError: If there are AWS API errors during endpoint creation/access. - """ - - # Validate that model_id exists in our supported model configurations - if model_id not in INFERENCER_MODELS: - raise ValueError(f"Unsupported model ID: {model_id}") - - # Look up the appropriate instance type from model configurations - instance_type = INFERENCER_MODELS.get(model_id)['instance_type'] - model_source = INFERENCER_MODELS.get(model_id)['model_source'] - - try: - # First check if a working endpoint already exists to avoid duplicate creation - status = self._check_model_status(endpoint_name) - if status == 'InService': - # Endpoint exists and is healthy - create and return a predictor for it - predictor = sagemaker.predictor.Predictor( - endpoint_name=endpoint_name, - sagemaker_session=self.session, - serializer=sagemaker.serializers.JSONSerializer(), - deserializer=sagemaker.deserializers.JSONDeserializer() - ) - # Register the predictor with the appropriate model type handler - self._assign_predictor(predictor, model_id) - return predictor - - except self.sagemaker_client.exceptions.ClientError as e: - # Handle case where endpoint doesn't exist yet - if e.response['Error']['Code'] == 'ValidationException': - try: - if model_source == "jumpstart": - # Initialize a new JumpStart model with the specified configuration - model = JumpStartModel( - role = self.role, - model_id=model_id, - sagemaker_session=self.session - ) - - # Deploy the model to a new endpoint with the specified configuration - predictor = model.deploy( - initial_instance_count=1, - instance_type=instance_type, - endpoint_name=endpoint_name, - accept_eula=True # Required for JumpStart models - ) - # Check if the model source is huggingface - elif model_source == "huggingface": - hub = { - 'HF_MODEL_ID': model_id, - 'SM_NUM_GPUS': json.dumps(1) - } - huggingface_model = HuggingFaceModel( - image_uri=get_huggingface_llm_image_uri("huggingface", version="2.3.1", region=self.region_name), - env=hub, - role=self.role, - sagemaker_session = self.session - - ) - - # deploy model to SageMaker Inference - predictor = huggingface_model.deploy( - initial_instance_count=1, - instance_type=instance_type, - endpoint_name=endpoint_name, - container_startup_health_check_timeout=300, - ) - - - # Register the new predictor with the appropriate model type handler - self._assign_predictor(predictor, model_id) - return predictor - - except self.sagemaker_client.exceptions.ClientError as e: - # Handle race condition where another process started creating the endpoint - # between our existence check and creation attempt - logger.info(f"Error creating endpoint: {e}") - if e.response['Error']['Code'] == 'ValidationException': - - logger.info(f"A new endpoint creation intercepted while attempting to create new endpoint, waiting.") - time.sleep(self.wait_time) # Allow the other process's endpoint to finish creating - - status = self._check_model_status(endpoint_name) - - if status == 'InService': - logger.info(f"Found the new endpoint, creating the predictor.") - # Create predictor for the endpoint that the other process created - predictor = sagemaker.predictor.Predictor( - endpoint_name=endpoint_name, - sagemaker_session=self.session, - serializer=sagemaker.serializers.JSONSerializer(), - deserializer=sagemaker.deserializers.JSONDeserializer() - ) - - # Register with appropriate model type handler - self._assign_predictor(predictor, model_id) - - return predictor - - elif e.response['Error']['Code'] == 'ResourceLimitExceeded': - logger.error(f"Resource limit exceeded while creating endpoint: {e}") - raise - - # Reraise any unexpected client errors - raise - - def _assign_predictor(self, predictor: sagemaker.predictor.Predictor, model_id: str): - """ - Assigns the appropriate predictor based on the provided model_id. The predictor is assigned to either - the embedding or inferencing predictor attributes, depending on the model type. - - Args: - predictor (sagemaker.predictor.Predictor): The SageMaker predictor to be assigned. - model_id (str): The model ID which determines whether the predictor is for embedding or inferencing. - - """ - # Assign predictor for inferencing models - if model_id in INFERENCER_MODELS: - self.inferencing_predictor = predictor - self.inferencing_model_id = model_id - logger.info(f"Assigned inferencing predictor for model: {model_id}") - - # Log an error if the model_id doesn't match any known type - else: - logger.error(f"Model ID {model_id} is not recognized as an inferencing model.") - - def generate_prompt(self, experiment_config: ExperimentalConfig, default_prompt: str, user_query: str, context: List[Dict] = None): - n_shot_prompt_guide = experiment_config.n_shot_prompt_guide_obj - n_shot_prompt = experiment_config.n_shot_prompts - # Input validation - if n_shot_prompt < 0: - raise ValueError("n_shot_prompt must be non-negative") - - # Get system prompt - system_prompt = default_prompt if n_shot_prompt_guide is None or n_shot_prompt_guide.system_prompt is None else n_shot_prompt_guide.system_prompt - - context_text = "" - if context: - context_text = self._format_context(user_query, context) - - base_prompt = n_shot_prompt_guide.user_prompt if n_shot_prompt_guide.user_prompt else "" - - if n_shot_prompt == 0: - logger.info("into zero shot prompt") - - if self.inferencing_model_id == "huggingface-llm-falcon-7b-instruct-bf16": - prompt = f"""Below are search results and a query. Create a concise summary. - Query: {user_query} - Search Results: {context_text} - Summary:""" - return None, prompt - - else: - prompt = ( - "Human: " + system_prompt + "\n\n" + - "Search Query:" + user_query + "\n\n" + - context_text + "\n\n" + - base_prompt + "\n\n" + - "Assistant: The final answer is:" - ) - return None, prompt.strip() - - - # Get examples - examples = n_shot_prompt_guide.examples - - # Format examples - selected_examples = (random.sample(examples, n_shot_prompt) - if len(examples) > n_shot_prompt - else examples) - - # Use string concatenation for example formatting - example_text = "" - for example in selected_examples: - if 'example' in example: - example_text += "- " + example['example'] + "\n" - elif 'question' in example and 'answer' in example: - example_text += " - Sample question:" + example['question'] + "\n" - example_text += "- Sample answer:" + example['answer'] + "\n" - - logger.info(f"into {n_shot_prompt} shot prompt with examples {len(selected_examples)}") - - if self.inferencing_model_id == "huggingface-llm-falcon-7b-instruct-bf16": - prompt = f"""Below are search results and a query. Create a concise summary. - Query: {user_query} - Few examples:\n - {example_text}\n - Search Results: {context_text} - Summary:""" - - return None, prompt - - else: - prompt = ( - "Human: " + system_prompt + "\n\n" + - "Few examples:\n" + - example_text + "\n" + - context_text + "\n\n" + - base_prompt + "\n\n" + - "Assistant: The final answer is:" - ) - - return None, prompt.strip() - - def _format_context(self, user_query: str, context: List[Dict[str, str]]) -> str: - """Format context documents into a single string.""" - # Format context: create a string representation of the query and passages - formatted_context = f"Search Query: {user_query}\n\nRelevant Passages:\n" - - try: - for i, item in enumerate(context, 1): - # Retrieve text from the context, handling both possible structures - content = None - if 'text' in item: - content = item['text'] - elif '_source' in item and 'text' in item['_source']: - content = item['_source']['text'] - - # If no text found, skip the current context item - if not content: - continue - - # Add score to context if available - score = item.get('_score', 'N/A') - formatted_context += f"\nPassage {i} (Score: {score}):\n{content}\n" - return formatted_context - - except Exception as e: - logger.error(f"Error formatting context: {str(e)}") - formatted_context += "Error processing context" - return formatted_context - - def construct_payload(self, system_prompt: str, prompt: str) -> dict: - """ - Constructs the payload dictionary for model inference with the given prompts and default parameters. - - Args: - system_prompt (str): The system-level prompt that guides the model's behavior - prompt (str): The actual prompt/query to be sent to the model - - """ - # Define default parameters for controlling the model's text generation - default_params = { - "max_new_tokens": 256, - "temperature": self.experiment_config.temp_retrieval_llm, - "top_p": 0.9, - "do_sample": True - } - # Construct the complete payload with prompt and generation parameters - payload = { - "inputs": prompt, - "parameters": default_params - } - - return payload - - def parse_response(self, response: dict) -> str: - """ - Parses the response from the model and extracts the generated text. - - Args: - response (dict): The raw response from the model - """ - # Handle different response formats (Falcon vs Llama) - if isinstance(response, list): - # Falcon-style response: Retrieve generated text from the list - return response[0].get('generated_text', '') if response else '' - elif isinstance(response, dict): - return response.get('generated_text', '') - else: - raise ValueError(f"Unexpected response format: {type(response)}") - - - def generate_text(self, user_query: str, default_prompt: str, context: List[Dict] = None, **kwargs) -> str: - """ - Generates a response based on the provided user query and context. It formats the context, sends it to - the model for text generation, and processes the response to return the generated text. - - Args: - user_query (str): The query provided by the user for which a response is generated. - context (List[Dict]): A list of context passages, each represented as a dictionary. - default_prompt (str): A default prompt that is used to guide the text generation. - **kwargs: Additional keyword arguments, if any. - - Returns: - tuple: A tuple containing metadata (str) and the cleaned generated text (str). - """ - - # Ensure the generation predictor is initialized - if not self.inferencing_predictor: - raise ValueError("Generation predictor not initialized") - - system_prompt, prompt = self.generate_prompt(self.experiment_config, default_prompt, user_query, context) - - payload = self.construct_payload(system_prompt, prompt) - - try: - start_time = time.time() - - # Get response from the model - response = self.inferencing_predictor.predict(payload) - - # Calculate latency metrics - latency = int((time.time() - start_time) * 1000) - - generated_text = self.parse_response(response) - - # Process the generated text to extract the answer - if "The final answer is:" in generated_text: - answer = generated_text.split("The final answer is:")[1].strip() - elif "Assistant:" in generated_text: - answer = generated_text.split("Assistant:")[1].strip() - else: - answer = generated_text.strip() - - # Clean and validate the response - cleaned_response = self._clean_response(answer) - - # Final validation of the generated text - if not cleaned_response or cleaned_response.isspace() or 'DRAFT' in cleaned_response: - return "Unable to generate a proper response. Please try again." - - # SageMaker does not provide input tokens as metadata. - # As a workaround, we use a rough approximation: ~4 characters per token. - input_tokens = len(prompt) // 4 - output_tokens = len(generated_text) // 4 - total_tokens = input_tokens + output_tokens - - answer_metadata = { - 'inputTokens': input_tokens, - 'outputTokens': output_tokens, - 'totalTokens': total_tokens, - 'latencyMs': latency - } - - return answer_metadata, cleaned_response - - except Exception as e: - logger.error(f"Error generating response: {str(e)}") - return f"Error generating response: {str(e)}" - - def _clean_response(self, text: str) -> str: - """ - Cleans and formats the response text by removing common artifacts, ensuring proper sentence structure, - and eliminating excessive whitespace or newlines. - - Args: - text (str): The raw response text that needs to be cleaned. - - Returns: - str: The cleaned and formatted response text. - """ - # Define a list of common artifacts that should be removed - artifacts = ['DRAFT', '[INST]', '[/INST]', 'Human:', 'Assistant:'] - - # Remove any leading/trailing whitespace and clean artifacts from the text - cleaned_text = text.strip() - for artifact in artifacts: - cleaned_text = cleaned_text.replace(artifact, '').strip() - - # Ensure the response ends with a proper sentence-ending punctuation mark (., !, or ?) - sentence_endings = ['.', '!', '?'] - - # Check if the text ends with one of the valid sentence-ending punctuation marks - if not any(cleaned_text.rstrip().endswith(end) for end in sentence_endings): - # Find the position of the last sentence-ending punctuation mark in the text - last_period = max( - cleaned_text.rfind('.'), - cleaned_text.rfind('!'), - cleaned_text.rfind('?') - ) - - # If a punctuation mark is found, truncate the text to that position - if last_period != -1: - cleaned_text = cleaned_text[:last_period + 1] - - # Remove multiple consecutive spaces and newlines, replace them with a single space - cleaned_text = ' '.join(cleaned_text.split()) - - think_end_index = cleaned_text.find('') - if think_end_index != -1: - cleaned_text = cleaned_text[think_end_index + len(''):] - - # Return the cleaned text, ensuring no extra spaces around it - return cleaned_text.strip() - - def _initialize_client(self) -> None: - raise NotImplementedError("Subclasses must implement `_initialize_client`") - - @staticmethod - def _sanitize_name(name: str) -> str: - """Sanitize the endpoint name to follow AWS naming conventions""" - # Replace any character that's not alphanumeric or hyphen with hyphen - import re - name = re.sub(r'[^a-zA-Z0-9-]', '-', name) - # Ensure it starts with a letter - if not name[0].isalpha(): - name = 'n' + name - # Truncate to 63 characters (AWS limit) - return name[:63] \ No newline at end of file diff --git a/core/knowledgebase_vectorstore.py b/core/knowledgebase_vectorstore.py deleted file mode 100644 index cd1fef43..00000000 --- a/core/knowledgebase_vectorstore.py +++ /dev/null @@ -1,52 +0,0 @@ -from typing import List, Dict, Any, Union -import boto3 -import logging -from baseclasses.base_classes import VectorDatabase -from config.config import Config - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class KnowledgeBaseVectorDatabase(VectorDatabase): - def __init__(self, region: str = 'us-east-1'): - self.client = boto3.client("bedrock-agent-runtime", region_name=region) - - def create_index(self, index_name: str, mapping: Dict[str, Any], algorithm: str) -> None: - raise NotImplementedError("This method is not implemented in this minimal version.") - - def update_index(self, index_name: str, new_mapping: Dict[str, Any]) -> None: - raise NotImplementedError("This method is not implemented in this minimal version.") - - def delete_index(self, index_name: str) -> None: - raise NotImplementedError("This method is not implemented in this minimal version.") - - def insert_document(self, index_name: str, document: Dict[str, Any]) -> None: - raise NotImplementedError("This method is not implemented in this minimal version.") - - - def _format_response(self, data): - formatted_results = [] - - for result in data.get('retrievalResults', []): - content = result.get('content', {}) - text = content.get('text', '') - - if text: - formatted_results.append({'text': text}) - - return formatted_results - - def search(self, query: str, kb_data: str, knn: int): - query = {"text": query} - retrievalConfiguration={ - 'vectorSearchConfiguration': { - 'numberOfResults': knn - } - } - response = self.client.retrieve(knowledgeBaseId = kb_data, - retrievalQuery = query, - retrievalConfiguration=retrievalConfiguration) - formatted_context = self._format_response(response) - logger.info("Getting results from knowledge base") - return formatted_context - \ No newline at end of file diff --git a/core/opensearch_vectorstore.py b/core/opensearch_vectorstore.py deleted file mode 100644 index 7a1beccf..00000000 --- a/core/opensearch_vectorstore.py +++ /dev/null @@ -1,247 +0,0 @@ -import traceback, json -from typing import Dict, Any, List, Optional -from botocore.endpoint import uuid -from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth -from baseclasses.base_classes import VectorDatabase -import boto3 -import logging - - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class OpenSearchVectorDatabase(VectorDatabase): - def __init__(self, host: str, use_ssl: bool = True, port: int = 443, is_serverless : bool = True, region: str = 'us-east-1', username: str = None, password: str = None): - if is_serverless: - try: - # Get credentials from the Lambda role - credentials = boto3.Session().get_credentials() - # Create AWS V4 Signer Auth for OpenSearch Serverless - auth = AWSV4SignerAuth(credentials, region, 'aoss') - - # Initialize OpenSearch client for serverless - self.client = OpenSearch( - hosts=[{'host': host, 'port': port}], - http_auth=auth, - use_ssl=True, - verify_certs=True, - connection_class=RequestsHttpConnection, - timeout=30, - max_retries=3, - retry_on_timeout=True, - # Add required headers for OpenSearch Serverless - headers={ - 'host': host - } - ) - - except Exception as e: - logger.error(f"Failed to initialize OpenSearch Serverless client: {str(e)}", exc_info=True) - raise - else: - self.client = OpenSearch( - hosts=[{'host': host, 'port': port}], - http_auth=(username, password), - use_ssl=use_ssl, - verify_certs=True, - connection_class=RequestsHttpConnection, - timeout=30, - max_retries=3, - retry_on_timeout=True - ) - - def _get_algorithm_settings(self, algorithm: str, dim: int) -> Dict[str, Any]: - - base_hnsw_params = { - "ef_construction": 512, - "m": 16 - } - - if algorithm == "hnsw": - return { - "name": "hnsw", - "engine": "faiss", - "space_type": "innerproduct", - "parameters": base_hnsw_params - } - elif algorithm == "hnsw_sq": - return { - "name": "hnsw", - "engine": "faiss", - "space_type": "innerproduct", - "parameters": { - **base_hnsw_params, - "encoder": { - "name": "sq", - "parameters": { - "type": "fp16" - } - } - } - } - elif algorithm == "hnsw_bq": - return { - "name": "hnsw", - "engine": "faiss", - "space_type": "innerproduct", - "parameters": base_hnsw_params - } - else: - raise ValueError(f"Unsupported algorithm: {algorithm}") - - - def create_index(self, index_name: str, mapping: Dict[str, Any], algorithm: str) -> None: - vector_field = next((field for field, props in mapping['properties'].items() - if props['type'] == 'knn_vector'), None) - if not vector_field: - raise ValueError("Mapping must include a knn_vector field") - - dim = mapping['properties'][vector_field]['dimension'] - algorithm_settings = self._get_algorithm_settings(algorithm, dim) - - index_body = { - "settings": { - "index": { - "knn": True, - "knn.algo_param.ef_search": 100 - } - }, - "mappings": { - "properties": { - vector_field: { - "type": "knn_vector", - "dimension": dim, - "method": algorithm_settings, - **({"mode": "on_disk"} if algorithm == "hnsw_bq" else {}) - } - } - } - } - - # Add other fields from the original mapping - for field, props in mapping['properties'].items(): - if field != vector_field: - index_body["mappings"]["properties"][field] = props - - logger.info(f"Creating index '{index_name}' with body:") - logger.info(json.dumps(index_body, indent=2)) - - try: - self.client.indices.create(index=index_name, body=index_body) - logger.info(f"Successfully created index '{index_name}'") - except Exception as e: - logger.error(f"Error creating index '{index_name}': {str(e)}") - logger.error(f"Index body: {json.dumps(index_body, indent=2)}") - traceback.print_exc() - raise - - - def update_index(self, index_name: str, new_mapping: Dict[str, Any]) -> None: - self.client.indices.put_mapping(index=index_name, body=new_mapping) - - def delete_index(self, index_name: str) -> None: - self.client.indices.delete(index=index_name) - - def insert_document(self, index_name: str, document: Dict[str, Any]) -> None: - self.client.index(index=index_name, body=document) - - def search(self, index_name: str, query_vector: List[float], k: int) -> List[Dict[str, Any]]: - vector_field = next((field for field, props in - self.client.indices.get_mapping(index=index_name)[index_name]['mappings']['properties'].items() - if 'type' in props and props['type'] == 'knn_vector'), None) - if not vector_field: - raise ValueError("Index does not contain a knn_vector field") - - query = { - "size": k, - "query": { - "knn": { - vector_field: { - "vector": query_vector, - "k": k - } - } - }, - "_source": True, - "fields": ["text", "parent_id"] - } - - response = self.client.search(index=index_name, body=query) - return [hit['_source'] for hit in response['hits']['hits']] - - def index_exists(self, index_name: str) -> bool: - """ - Check if an index exists in OpenSearch. - - :param index_name: Name of the index to check - :return: True if the index exists, False otherwise - """ - return self.client.indices.exists(index=index_name) - - def insert_chunk(self, index_name: str, text: str, embedding: List[float], chunk_id: str, metadata: Dict = None): - document = { - "text": text, - "embedding": embedding, - "chunk_id": chunk_id, - "metadata": metadata or {} - } - try: - self.insert_document(index_name, document) - except Exception as e: - logger.error(f"Error inserting chunk {chunk_id}: {str(e)}") - - def batch_insert_chunks(self, index_name: str, chunks: List[str], chunk_embeddings: List[List[float]], - metadata: Optional[List[Dict]] = None, batch_size: int = 100): - total_chunks = len(chunks) - - # If metadata is None or empty, create a list of empty dictionaries - if not metadata: - metadata = [{} for _ in range(total_chunks)] - - for i in range(0, total_chunks, batch_size): - batch_chunks = chunks[i:i+batch_size] - batch_embeddings = chunk_embeddings[i:i+batch_size] - batch_metadata = metadata[i:i+batch_size] - - for chunk, embedding, meta in zip(batch_chunks, batch_embeddings, batch_metadata): - chunk_id = str(uuid.uuid4()) # Generate a unique ID for each chunk - self.insert_chunk(index_name, chunk, embedding, chunk_id, meta) - - logger.info(f"Inserted batch {i//batch_size + 1} ({i+1} to {min(i+batch_size, total_chunks)} of {total_chunks})") - - - def print_opensearch_info(self): - try: - info = self.client.info() - logger.info(f"OpenSearch Version: {info['version']['number']}") - logger.info(f"Cluster Name: {info['cluster_name']}") - logger.info(f"Cluster UUID: {info['cluster_uuid']}") - except Exception as e: - logger.error(f"Error getting OpenSearch info: {str(e)}") - - - def index_chunk_embeddings(self, chunks: List[str], chunk_embeddings: List[List[float]], - indexing_algorithm: str, chunking_algorithm: str, - vector_dimension: int, metadata: List[Dict] = None, chunk_size: int = 1200): - index_name = f"{indexing_algorithm}-{chunking_algorithm.lower()}-{chunk_size}" - - mapping = { - "properties": { - "text": {"type": "text"}, - "embedding": { - "type": "knn_vector", - "dimension": vector_dimension - }, - "metadata": {"type": "object"} - } - } - - try: - self.print_opensearch_info() # Print OpenSearch version and cluster info - self.create_index(index_name, mapping, indexing_algorithm) - except Exception as e: - logger.error(f"Error creating index: {str(e)}") - return - - self.batch_insert_chunks(index_name, chunks, chunk_embeddings, metadata) - return f"Indexing complete for '{index_name}'!" \ No newline at end of file diff --git a/core/processors/__init__.py b/core/processors/__init__.py deleted file mode 100644 index b76ba9bd..00000000 --- a/core/processors/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .chunking_processor import ChunkingProcessor -from .embed_processor import EmbedProcessor -from .inference_processor import InferenceProcessor -from .eval_processor import EvalProcessor diff --git a/core/processors/chunking_processor.py b/core/processors/chunking_processor.py deleted file mode 100644 index 35967756..00000000 --- a/core/processors/chunking_processor.py +++ /dev/null @@ -1,45 +0,0 @@ -from typing import Dict, List, Type, Union -from core.chunking import FixedChunker, HierarchicalChunker -from baseclasses.base_classes import BaseChunker, BaseHierarchicalChunker -import logging -from config.experimental_config import ExperimentalConfig - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class ChunkingProcessor: - """Processor for managing text chunking.""" - - CHUNKER_STRATEGIES: Dict[str, Union[Type[BaseChunker], Type[BaseHierarchicalChunker]]] = { - "Fixed": FixedChunker, - "Hierarchical": HierarchicalChunker - } - - def __init__(self, experimentalConfig : ExperimentalConfig) -> None: - self.experimentalConfig = experimentalConfig - self.chunker = self._initialize_chunker() - - def _initialize_chunker(self) -> Union[BaseChunker, BaseHierarchicalChunker]: - """Initialize the chunker based on the selected strategy.""" - strategy = self.experimentalConfig.chunking_strategy.lower() # Normalize to lower case - chunker_strategies = {key.lower(): value for key, value in self.CHUNKER_STRATEGIES.items()} # Case-insensitive map - if strategy not in chunker_strategies: - raise ValueError(f"Unknown chunking strategy: {strategy}") - - logger.info(f"Initializing {strategy} chunker...") - if strategy == 'fixed': - return chunker_strategies[strategy]( - self.experimentalConfig.chunk_size, - self.experimentalConfig.chunk_overlap - ) - elif strategy == 'hierarchical': - return chunker_strategies[strategy]( - self.experimentalConfig.hierarchical_parent_chunk_size, - self.experimentalConfig.hierarchical_child_chunk_size, - self.experimentalConfig.hierarchical_chunk_overlap_percentage - ) - - def chunk(self, texts: List[str]) -> Union[List[str], List[List[str]]]: - """Chunk the input list of text into a single flat list""" - all_chunks = [chunk for text in texts for chunk in self.chunker.chunk(text)] - return all_chunks \ No newline at end of file diff --git a/core/processors/embed_processor.py b/core/processors/embed_processor.py deleted file mode 100644 index ea245397..00000000 --- a/core/processors/embed_processor.py +++ /dev/null @@ -1,45 +0,0 @@ -from core.embedding import EmbedderFactory -from typing import Dict, List, Tuple, Any -from config.experimental_config import ExperimentalConfig -import logging - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class EmbedProcessor: - """Processor for embedding text chunks.""" - - def __init__(self, experimentalConfig : ExperimentalConfig) -> None: - self.experimentalConfig = experimentalConfig - self.embedder = EmbedderFactory.create_embedder(experimentalConfig) - - def embed(self, chunks: List[str]) -> List[Tuple[List[float], str, Dict[Any, Any]]]: - """Embed each chunk one by one.""" - embeddings = [] - try: - dimensions = self.experimentalConfig.vector_dimension - normalize = True # Always normalize - - logger.info(f"Embedding {len(chunks)} chunks with dimensions: {dimensions}.") - for idx, chunk in enumerate(chunks): - logger.debug(f"Embedding chunk {idx + 1}/{len(chunks)}: {chunk[:50]}...") - metadata, embedding = self.embedder.embed(chunk, dimensions=dimensions, normalize=normalize) - embeddings.append((embedding, chunk, metadata)) # Append as tuple - - logger.info("Embedding process completed successfully.") - return embeddings - except Exception as e: - logger.error(f"Error during embedding process: {e}") - raise - - def embed_text(self, text: str) -> Tuple[Dict[Any, Any], List[float]]: - """Embed each chunk one by one.""" - try: - dimensions = self.experimentalConfig.vector_dimension - normalize = True # Always normalize - metadata, embedding = self.embedder.embed(text, dimensions=dimensions, normalize=normalize) - logger.info("Embedding text process completed successfully.") - return metadata, embedding - except Exception as e: - logger.error(f"Error during embedding process: {e}") - raise diff --git a/core/processors/eval_processor.py b/core/processors/eval_processor.py deleted file mode 100644 index db7f9b20..00000000 --- a/core/processors/eval_processor.py +++ /dev/null @@ -1,21 +0,0 @@ -from core.eval.eval_factory import EvalFactory -from typing import Dict, List, Tuple -from config.experimental_config import ExperimentalConfig -import logging - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class EvalProcessor: - """Processor for embedding text chunks.""" - - def __init__(self, experimentalConfig : ExperimentalConfig) -> None: - self.experimentalConfig = experimentalConfig - self.evaluator = EvalFactory.create_evaluator(experimentalConfig) - - def evaluate(self) -> None: - try: - self.evaluator.evaluate(experiment_id=self.experimentalConfig.experiment_id) - except Exception as e: - logger.error(f"Error generating eval: {str(e)}") - raise diff --git a/core/processors/inference_processor.py b/core/processors/inference_processor.py deleted file mode 100644 index dd497f58..00000000 --- a/core/processors/inference_processor.py +++ /dev/null @@ -1,27 +0,0 @@ -from core.inference import InferencerFactory -from typing import Dict, List, Tuple, Any -from config.experimental_config import ExperimentalConfig -import logging - -logger = logging.getLogger() -logger.setLevel(logging.INFO) - -class InferenceProcessor: - """Processor for embedding text chunks.""" - - def __init__(self, experimentalConfig : ExperimentalConfig) -> None: - self.experimentalConfig = experimentalConfig - self.inferencer = InferencerFactory.create_inferencer(experimentalConfig) - - def generate_text(self, user_query: str, default_prompt: str, context: List[Dict] = None, **kwargs) -> Tuple[Dict[Any,Any], str]: - try: - metadata, answer = self.inferencer.generate_text( - user_query=user_query, - context = context, - default_prompt = default_prompt, - experiment_config = self.experimentalConfig - ) - return metadata, answer - except Exception as e: - logger.error(f"Error generating text with Inferencer: {str(e)}") - raise diff --git a/core/rerank/rerank.py b/core/rerank/rerank.py deleted file mode 100644 index 776b2184..00000000 --- a/core/rerank/rerank.py +++ /dev/null @@ -1,90 +0,0 @@ -import logging -import boto3 -from config.experimental_config import ExperimentalConfig -from config.config import Config, get_config - -# Set up logging -logger = logging.getLogger(__name__) -logger.setLevel(logging.ERROR) - -class DocumentReranker: - def __init__(self, region, rerank_model_id): - """ - Initialize the DocumentReranker with the AWS region, model ID, and Bedrock agent runtime. - - Args: - region (str): The AWS region to use. - model_id (str): The model ID to use for reranking. - bedrock_agent_runtime (object): The Bedrock agent runtime instance to interact with the API. - """ - self.region = region - self.rerank_model_id = rerank_model_id - self.bedrock_agent_runtime = boto3.client('bedrock-agent-runtime', region_name=self.region) - - def rerank_documents(self, input_prompt, retrieved_documents): - """ - Rerank a list of documents based on a query using Amazon Bedrock's reranking model. - - Args: - input_prompt (str): The query to rerank documents for. - retrieved_documents (list): The list of documents to be reranked. - - Returns: - list: A list of reranked documents in order of relevance. - """ - try: - # Construct the model ARN using the provided model ID - model_package_arn = f"arn:aws:bedrock:{self.region}::foundation-model/{self.rerank_model_id}" - rerank_return_count = len(retrieved_documents) - - # Prepare the text sources for the documents (wrap text in a dictionary) - document_sources = [{ - "type": "INLINE", - "inlineDocumentSource": { - "type": "TEXT", - "textDocument": { - "text": doc['text'] # Wrap the text in a dictionary - } - } - } for doc in retrieved_documents] - - # Call the Bedrock API for reranking - response = self.bedrock_agent_runtime.rerank( - queries=[{ - "type": "TEXT", - "textQuery": {"text": input_prompt} - }], - sources=document_sources, - rerankingConfiguration={ - "type": "BEDROCK_RERANKING_MODEL", - "bedrockRerankingConfiguration": { - "numberOfResults": rerank_return_count, - "modelConfiguration": {"modelArn": model_package_arn} - } - } - ) - - # Check if 'results' exist in the response and log the structure - if 'results' not in response: - logger.error("Error in rerank response: No results found.") - return [] - - # Create a list to store the reranked documents - reranked_documents = [] - - # Process the results - for rank, result in enumerate(response['results']): - if isinstance(result, dict) and 'index' in result: - original_index = result['index'] - reranked_documents.append({'text': retrieved_documents[original_index]['text']}) - else: - logger.error(f"Unexpected result format: {result}") - - logger.info(f"Reranked documents: {len(reranked_documents)}") - # Return the reranked documents, ensuring we return only as many as requested - return reranked_documents[:rerank_return_count] - - except Exception as e: - # Catch any other unforeseen errors - logger.error(f"An error occurred: {e}") - return [] \ No newline at end of file diff --git a/core/service/experimental_config_service.py b/core/service/experimental_config_service.py deleted file mode 100644 index 276b6f86..00000000 --- a/core/service/experimental_config_service.py +++ /dev/null @@ -1,98 +0,0 @@ -from core.dynamodb import DynamoDBOperations -from typing import Dict, Any, Optional -from config.experimental_config import ExperimentalConfig, NShotPromptGuide -from util.dynamo_utils import deserialize_dynamodb_json -from config.config import Config - -class ExperimentalConfigService: - """Service class to manage experimental configurations.""" - - def __init__(self, config: Config): - self.aws_region = config.aws_region - self.experiment_db = DynamoDBOperations( - region=self.aws_region, - table_name=config.experiment_table - ) - - def _validate_n_shot_prompts(self, experiment_id: str, n_shot_prompt_guide: NShotPromptGuide, - required_examples: int) -> None: - """Validate n-shot prompt guide configuration.""" - if not n_shot_prompt_guide.system_prompt: - raise ValueError(f"Experiment {experiment_id}: Missing system prompt") - if not n_shot_prompt_guide.user_prompt: - raise ValueError(f"Experiment {experiment_id}: Missing user prompt") - if required_examples > 0 and len(n_shot_prompt_guide.examples) < required_examples: - raise ValueError( - f"Experiment {experiment_id}: Insufficient n-shot examples. " - f"Required: {required_examples}, Found: {len(n_shot_prompt_guide.examples)}" - ) - - def create_experimental_config(self, exp_config_data: Dict[str, Any]) -> ExperimentalConfig: - """Create and validate an experimental configuration. - - Args: - exp_config_data: Dictionary containing experimental configuration parameters - - Returns: - ExperimentalConfig: Validated experimental configuration object - - Raises: - ValueError: If experiment doesn't exist or has invalid configuration - """ - # Validate required fields - experiment_id = exp_config_data.get('experiment_id') - if not experiment_id: - raise ValueError("experiment_id is required") - - experiment = self.experiment_db.get_item({'id': experiment_id}) - if not experiment: - raise ValueError(f"Experiment with id {experiment_id} not found") - - exp_config = ExperimentalConfig( - execution_id=exp_config_data.get('execution_id'), - experiment_id=experiment_id, - embedding_model=exp_config_data.get('embedding_model'), - retrieval_model=exp_config_data.get('retrieval_model'), - vector_dimension=exp_config_data.get('vector_dimension', 0), - gt_data=exp_config_data.get('gt_data', {}), - index_id=exp_config_data.get('index_id'), - knn_num=exp_config_data.get('knn_num', 1), - temp_retrieval_llm=exp_config_data.get('temp_retrieval_llm'), - embedding_service=exp_config_data.get('embedding_service'), - retrieval_service=exp_config_data.get('retrieval_service'), - aws_region=exp_config_data.get('aws_region', self.aws_region), - chunking_strategy=exp_config_data.get('chunking_strategy'), - chunk_size=exp_config_data.get('chunk_size', 0), - chunk_overlap=exp_config_data.get('chunk_overlap', 0), - hierarchical_parent_chunk_size=exp_config_data.get('hierarchical_parent_chunk_size', 0), - hierarchical_child_chunk_size=exp_config_data.get('hierarchical_child_chunk_size', 0), - hierarchical_chunk_overlap_percentage=exp_config_data.get('hierarchical_chunk_overlap_percentage', 0), - kb_data=exp_config_data.get('kb_data', {}), - n_shot_prompts=exp_config_data.get('n_shot_prompts', 0), - indexing_algorithm=exp_config_data.get('indexing_algorithm'), - rerank_model_id=exp_config_data.get('rerank_model_id', "none"), - # TODO: change this class to read everything from exp_config retrieved from DynamoDB - enable_guardrails=experiment.get('config').get('enable_guardrails', False), - guardrail_id=experiment.get('config').get('guardrail_id', ""), - guardrail_version=experiment.get('config').get('guardrail_version', ""), - enable_prompt_guardrails=experiment.get('config').get("enable_prompt_guardrails", False), - enable_context_guardrails=experiment.get('config').get("enable_context_guardrails", False), - enable_response_guardrails=experiment.get('config').get("enable_response_guardrails", False), - eval_service=exp_config_data.get('eval_service', "ragas"), - eval_embedding_model=exp_config_data.get('eval_embedding_model', "amazon.titan-embed-text-v1"), - eval_retrieval_model=exp_config_data.get('eval_retrieval_model', "mistral.mixtral-8x7b-instruct-v0:1"), - bedrock_knowledge_base=exp_config_data.get('bedrock_knowledge_base', False), - knowledge_base=exp_config_data.get('knowledge_base', False), - is_opensearch=True if exp_config_data.get('opensearch_host') else False - ) - - n_shot_prompt_guide = experiment.get('config').get('n_shot_prompt_guide') - if not n_shot_prompt_guide: - raise ValueError(f"Experiment {experiment_id}: Missing prompt file") - - n_shot_prompt_guide = NShotPromptGuide(**deserialize_dynamodb_json(n_shot_prompt_guide)) - self._validate_n_shot_prompts(experiment_id, n_shot_prompt_guide, exp_config.n_shot_prompts) - - exp_config.n_shot_prompt_guide_obj = n_shot_prompt_guide - - return exp_config