Final_Assignment_Template / nova_agent.py
Kackle's picture
Update nova_agent.py
e96e18c verified
raw
history blame
4.96 kB
import os
import boto3
import json
from dotenv import load_dotenv
load_dotenv()
class NovaProAgent:
def __init__(self):
print("NovaProAgent initialized.")
# Get AWS credentials from environment variables
aws_access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
aws_secret_access_key = os.getenv('AWS_SECRET_ACCESS_KEY')
# Initialize the AWS client
boto3.client(
's3',
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key
)
session = boto3.session.Session()
self.bedrock_client = boto3.client(
service_name='bedrock-runtime',
region_name='us-east-1'
)
self.model_id = "amazon.nova-pro-v1:0"
self.content_type = "application/json"
self.accept = "application/json"
async def __call__(self, question: str) -> str:
print(f"NovaProAgent received question (first 50 chars): {question[:50]}...")
try:
# Detect question type and adjust approach
needs_reasoning = any(keyword in question.lower() for keyword in [
'calculate', 'how many', 'what is the', 'find', 'determine', 'solve',
'table', 'given', 'prove', 'counter-example'
])
if needs_reasoning:
prompt = f"""You are an expert problem solver. Think step by step to solve this question accurately.
Question: {question}
Think through this step by step, then provide your final answer:"""
max_tokens = 300
else:
prompt = f"""Answer this question directly and concisely. Provide only the essential information requested.
Question: {question}
Answer:"""
max_tokens = 150
# Prepare the request payload for Nova Pro
payload = {
"messages": [
{
"role": "user",
"content": [{
"text": prompt
}]
}
],
"inferenceConfig": {
"max_new_tokens": max_tokens,
"temperature": 0.1 if needs_reasoning else 0.0
}
}
# Call Nova Lite model
response = self.bedrock_client.invoke_model(
modelId=self.model_id,
contentType=self.content_type,
accept=self.accept,
body=json.dumps(payload)
)
# Parse response
response_body = json.loads(response['body'].read())
answer = response_body['output']['message']['content'][0]['text']
# Clean up the answer
answer = answer.strip()
# Handle questions requiring external resources
if any(phrase in question.lower() for phrase in [
'attached', 'video', 'image', 'audio', 'file', 'excel', '.mp3', '.jpg', '.png'
]):
if 'video' in question.lower() or 'audio' in question.lower():
return "I cannot access external media files."
elif 'image' in question.lower():
return "I cannot view images."
elif 'excel' in question.lower() or 'file' in question.lower():
return "I cannot access attached files."
# Extract final answer if reasoning was used
if needs_reasoning and 'final answer' in answer.lower():
lines = answer.split('\n')
for line in reversed(lines):
if 'final answer' in line.lower() or 'answer:' in line.lower():
# Extract the part after the colon
if ':' in line:
answer = line.split(':', 1)[1].strip()
break
# Remove verbose beginnings
verbose_starts = [
"To answer this question", "Based on the information", "According to",
"The answer is", "Looking at", "Step by step", "Let me think"
]
for start in verbose_starts:
if answer.lower().startswith(start.lower()):
sentences = answer.split('. ')
for sentence in sentences[1:]:
if len(sentence.strip()) > 5:
answer = sentence.strip()
break
# Clean up common patterns
answer = answer.replace('**', '').replace('*', '')
return answer
except Exception as e:
print(f"Error calling Nova Lite: {e}")
return "Unable to process request."