Skip to content

Guardrails pricing calculation #289

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rjayar3 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Guardrails pricing calculation #289

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
251 changes: 171 additions & 80 deletions retriever/retriever.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,40 @@
from typing import List, Dict, Any, Optional, Tuple, Union
from dataclasses import asdict

def calculate_guardrail_units(text: str,
use_content_filters: bool = False,
use_denied_topics: bool = False,
use_contextual_grounding: bool = False,
source: str = "",
query: str = "") -> dict:
"""
Calculate text units for enabled guardrail policies.
1000 characters = 1 text unit. Partial units round up.
"""
CHARS_PER_TEXT_UNIT = 1000
result = {'total_text_units': 0, 'breakdown': {}}

if not text:
return result

text_units = (len(text) + CHARS_PER_TEXT_UNIT - 1) // CHARS_PER_TEXT_UNIT

if use_content_filters:
result['breakdown']['content_filters'] = text_units
result['total_text_units'] += text_units

if use_denied_topics:
result['breakdown']['denied_topics'] = text_units
result['total_text_units'] += text_units

if use_contextual_grounding:
grounding_units = (len(source) + len(query) + len(text) +
CHARS_PER_TEXT_UNIT - 1) // CHARS_PER_TEXT_UNIT
result['breakdown']['contextual_grounding'] = grounding_units
result['total_text_units'] += grounding_units

return result

def retrieve(config: Config, experimentalConfig: ExperimentalConfig) -> None:
"""
Execute the retrieval process for question answering experiments.
Expand Down Expand Up @@ -88,8 +122,8 @@ def initialize_components(config: Config, experimentalConfig: ExperimentalConfig
"""Initialize all required components for the retrieval process."""
try:
# Initialize embedding processor if required
if experimentalConfig.bedrock_knowledge_base or not experimentalConfig.knowledge_base:
logger.info("Skipping embed processor initialization")
if experimentalConfig.bedrock_knowledge_base:
logger.info("Setting up knowledge base retriever")
embed_processor = None

else:
Expand All @@ -101,20 +135,19 @@ def initialize_components(config: Config, experimentalConfig: ExperimentalConfig
inference_processor = InferenceProcessor(experimentalConfig)

# Initialize vector database
vector_database = None
if experimentalConfig.knowledge_base:
if experimentalConfig.bedrock_knowledge_base:
logger.info("Connecting to Knowledge base")
vector_database = KnowledgeBaseVectorDatabase(region=experimentalConfig.aws_region)
else:
logger.info(f"Connecting to OpenSearch at {config.opensearch_host}")
vector_database = OpenSearchVectorDatabase(
host=config.opensearch_host,
is_serverless=config.opensearch_serverless,
region=config.aws_region,
username=config.opensearch_username,
password=config.opensearch_password
)

if experimentalConfig.bedrock_knowledge_base:
logger.info("Connecting to Knowledge base")
vector_database = KnowledgeBaseVectorDatabase(region=experimentalConfig.aws_region)
else:
logger.info(f"Connecting to OpenSearch at {config.opensearch_host}")
vector_database = OpenSearchVectorDatabase(
host=config.opensearch_host,
is_serverless=config.opensearch_serverless,
region=config.aws_region,
username=config.opensearch_username,
password=config.opensearch_password
)

# Initialize DynamoDB connections
logger.info("Initializing DynamoDB connections")
Expand Down Expand Up @@ -195,24 +228,71 @@ def process_questions(
logger.debug(f"Processing question {idx+1}: {question}")

# Generate embeddings
if experimentalConfig.bedrock_knowledge_base or not experimentalConfig.knowledge_base:
query_metadata, query_embedding = {'inputTokens': '0', 'latencyMs': '0'}, None
else:
if not experimentalConfig.bedrock_knowledge_base:
logger.info("Generating embeddings for the question using provided embedder")
query_metadata, query_embedding = components["embed_processor"].embed_text(
question
)
else:
query_metadata, query_embedding = {'inputTokens': '0', 'latencyMs': '0'}, None


query_results=None
guardrail_input_assessment = None
guardrail_output_assessment = None
guardrail_context_assessment = None
guardrail_id = None
guardrail_blocked = None
guardrail_input_usage = None
guardrail_output_usage = None
guardrail_context_usage = None
has_content_filters = None
has_denied_topics = None
has_contextual_grounding = None

answer_metadata = {}
answer = ""

# Get guardrail configuration
try:
guardrail_client = components['guardrails']['client']
guardrail_response = guardrail_client.get_guardrail(
guardrailIdentifier=guardrail_id,
guardrailVersion=components['guardrails'].get('version', 'latest')
)

# Check if guardrail is in READY state
if guardrail_response.get('status') != 'READY':
logger.warning(f"Guardrail {guardrail_id} is not in READY state. Current status: {guardrail_response.get('status')}")

# Determine which types of checks are enabled
has_content_filters = bool(guardrail_response.get('contentPolicy', {}).get('filters', []))
has_denied_topics = bool(guardrail_response.get('topicPolicy', {}).get('topics', []))
has_contextual_grounding = bool(guardrail_response.get('contextualGroundingPolicy', {}).get('filters', []))

# Calculate usage based on actual guardrail settings
guardrail_usage = calculate_guardrail_units(
text=question,
use_content_filters=has_content_filters,
use_denied_topics=has_denied_topics,
use_contextual_grounding=has_contextual_grounding,
source=' '.join(record['text'] for record in query_results) if query_results else "",
query=question
)

logger.info(f"Guardrail usage for question {idx+1}: {guardrail_usage}")

except Exception as e:
logger.error(f"Error getting guardrail configuration: {str(e)}", exc_info=True)
guardrail_usage = calculate_guardrail_units(
text=question,
use_content_filters=False,
use_denied_topics=False,
use_contextual_grounding=False
)
guardrail_metadata = {'guardrail_usage': guardrail_usage}


# Retrieval query embed is not provided by knowledge base
retrieval_query_embed_tokens += int(query_metadata.get("inputTokens", 0) if query_embedding else 0)

Expand All @@ -225,6 +305,14 @@ def process_questions(

# Apply INPUT guardrails
if experimentalConfig.enable_prompt_guardrails:
guardrail_input_usage = calculate_guardrail_units(
text=question,
use_content_filters=has_content_filters,
use_denied_topics=has_denied_topics,
use_contextual_grounding=has_contextual_grounding,
source="",
query=question
)
blocked, modified_question, guardrail_input_assessment = apply_guardrail_check(
components,
guardrail_id,
Expand All @@ -238,26 +326,33 @@ def process_questions(

# Apply CONTEXT guardrails if not already blocked
if experimentalConfig.enable_context_guardrails and guardrail_blocked == 'NONE':
if experimentalConfig.knowledge_base:
# Search for relevant context once
if isinstance(components["vector_database"], OpenSearchVectorDatabase):
query_results = components["vector_database"].search(
experimentalConfig.index_id, query_embedding, experimentalConfig.knn_num
)
elif isinstance(components["vector_database"], KnowledgeBaseVectorDatabase):
query_results = components["vector_database"].search(
question, experimentalConfig.kb_data, experimentalConfig.knn_num
)
# Search for relevant context once
if isinstance(components["vector_database"], OpenSearchVectorDatabase):
query_results = components["vector_database"].search(
experimentalConfig.index_id, query_embedding, experimentalConfig.knn_num
)
elif isinstance(components["vector_database"], KnowledgeBaseVectorDatabase):
query_results = components["vector_database"].search(
question, experimentalConfig.kb_data, experimentalConfig.knn_num
)

if experimentalConfig.chunking_strategy.lower() == 'hierarchical':
query_results = __duplicate_removal_for_heirarchical_config(query_results)
if experimentalConfig.chunking_strategy.lower() == 'hierarchical':
query_results = __duplicate_removal_for_heirarchical_config(query_results)

if experimentalConfig.rerank_model_id and experimentalConfig.rerank_model_id.lower() != 'none':
#Rerank the query results
query_results = __rerank_query_result(query_results, question, experimentalConfig, idx)
if experimentalConfig.rerank_model_id and experimentalConfig.rerank_model_id.lower() != 'none':
#Rerank the query results
query_results = __rerank_query_result(query_results, question, experimentalConfig, idx)


if query_results:
guardrail_context_usage = calculate_guardrail_units(
text=query_results,
use_content_filters=has_content_filters,
use_denied_topics=has_denied_topics,
use_contextual_grounding=has_contextual_grounding,
source="",
query=query_results
)
context = ' '.join(record['text'] for record in query_results)
blocked, modified_context, guardrail_context_assessment = apply_guardrail_check(
components,
Expand All @@ -274,38 +369,41 @@ def process_questions(
if guardrail_blocked == 'NONE':
# Fetch context if not already done
if query_results is None:
if experimentalConfig.knowledge_base:
if isinstance(components["vector_database"], OpenSearchVectorDatabase):
query_results = components["vector_database"].search(
experimentalConfig.index_id, query_embedding, experimentalConfig.knn_num
)
elif isinstance(components["vector_database"], KnowledgeBaseVectorDatabase):
query_results = components["vector_database"].search(
question, experimentalConfig.kb_data, experimentalConfig.knn_num
)
if experimentalConfig.chunking_strategy.lower() == 'hierarchical':
query_results = __duplicate_removal_for_heirarchical_config(query_results)
if experimentalConfig.rerank_model_id and experimentalConfig.rerank_model_id.lower() != 'none':
#Rerank the query results
query_results = __rerank_query_result(query_results, question, experimentalConfig, idx)

# Generate answer
if experimentalConfig.knowledge_base:
answer_metadata, answer = components["inference_processor"].generate_text(
if isinstance(components["vector_database"], OpenSearchVectorDatabase):
query_results = components["vector_database"].search(
experimentalConfig.index_id, query_embedding, experimentalConfig.knn_num
)
elif isinstance(components["vector_database"], KnowledgeBaseVectorDatabase):
query_results = components["vector_database"].search(
question, experimentalConfig.kb_data, experimentalConfig.knn_num
)

if experimentalConfig.chunking_strategy.lower() == 'hierarchical':
query_results = __duplicate_removal_for_heirarchical_config(query_results)

if experimentalConfig.rerank_model_id and experimentalConfig.rerank_model_id.lower() != 'none':
#Rerank the query results
query_results = __rerank_query_result(query_results, question, experimentalConfig, idx)

# Generate answer
answer_metadata, answer = components["inference_processor"].generate_text(
user_query=question,
context=query_results,
default_prompt=config.inference_system_prompt,
)
else:
answer_metadata, answer = components["inference_processor"].generate_text(
user_query=question,
default_prompt=config.inference_system_prompt,
)
retrieval_input_tokens += int(answer_metadata["inputTokens"])
retrieval_output_tokens += int(answer_metadata["outputTokens"])

# Apply OUTPUT guardrails if enabled
if experimentalConfig.enable_response_guardrails:
guardrail_output_usage = calculate_guardrail_units(
text=answer,
use_content_filters=has_content_filters,
use_denied_topics=has_denied_topics,
use_contextual_grounding=has_contextual_grounding,
source=query_results,
query=question
)
blocked, modified_answer, guardrail_output_assessment = apply_guardrail_check(
components,
guardrail_id,
Expand All @@ -317,36 +415,29 @@ def process_questions(
answer = modified_answer
guardrail_blocked = 'OUTPUT'
else:
if experimentalConfig.knowledge_base:
# Search for relevant context
if isinstance(components["vector_database"], OpenSearchVectorDatabase):
# Search for relevant context
if isinstance(components["vector_database"], OpenSearchVectorDatabase):
query_results = components["vector_database"].search(
experimentalConfig.index_id, query_embedding, experimentalConfig.knn_num
)
elif isinstance(components["vector_database"], KnowledgeBaseVectorDatabase):
query_results = components["vector_database"].search(
question, experimentalConfig.kb_data, experimentalConfig.knn_num
)
elif isinstance(components["vector_database"], KnowledgeBaseVectorDatabase):
query_results = components["vector_database"].search(
question, experimentalConfig.kb_data, experimentalConfig.knn_num
)

if experimentalConfig.chunking_strategy.lower() == 'hierarchical':
query_results = __duplicate_removal_for_heirarchical_config(query_results)
if experimentalConfig.chunking_strategy.lower() == 'hierarchical':
query_results = __duplicate_removal_for_heirarchical_config(query_results)

if experimentalConfig.rerank_model_id and experimentalConfig.rerank_model_id.lower() != 'none':
#Rerank the query results
query_results = __rerank_query_result(query_results, question, experimentalConfig, idx)
if experimentalConfig.rerank_model_id and experimentalConfig.rerank_model_id.lower() != 'none':
#Rerank the query results
query_results = __rerank_query_result(query_results, question, experimentalConfig, idx)

# Generate answer
if experimentalConfig.knowledge_base:
answer_metadata, answer = components["inference_processor"].generate_text(
user_query=question,
context=query_results,
default_prompt=config.inference_system_prompt,
)
else:
answer_metadata, answer = components["inference_processor"].generate_text(
user_query=question,
default_prompt=config.inference_system_prompt
)
answer_metadata, answer = components["inference_processor"].generate_text(
user_query=question,
context=query_results,
default_prompt=config.inference_system_prompt,
)
retrieval_input_tokens += int(answer_metadata["inputTokens"])
retrieval_output_tokens += int(answer_metadata["outputTokens"])

Expand Down