nova_agent.py

import os
import boto3
import json
from dotenv import load_dotenv

load_dotenv()

class NovaLiteAgent:
    def __init__(self):
        print("NovaLiteAgent 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-lite-v1:0"
        self.content_type = "application/json"
        self.accept = "application/json"
        
    async def __call__(self, question: str) -> str:
        print(f"NovaLiteAgent received question (first 50 chars): {question[:50]}...")
        
        try:
            # Prepare the request payload for Nova Lite
            payload = {
                "messages": [
                    {
                        "role": "user",
                        "content": [{
                            "text": question
                        }]
                    }
                ],
                "inferenceConfig": {
                    "max_new_tokens": 1000,
                    "temperature": 0.1
                }
            }
            
            # 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']
            
            return answer.strip()
            
        except Exception as e:
            print(f"Error calling Nova Lite: {e}")
            return f"I apologize, but I'm currently experiencing technical difficulties: {e}"