Update nova_agent.py

nova_agent.py

@@ -5,9 +5,9 @@ from dotenv import load_dotenv
 
 load_dotenv()
 
-class NovaLiteAgent:
+class NovaProAgent:
     def __init__(self):
-        print("NovaLiteAgent initialized.")
+        print("NovaProAgent initialized.")
         
         # Get AWS credentials from environment variables
         aws_access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
@@ -31,17 +31,31 @@ class NovaLiteAgent:
         self.accept = "application/json"
         
     async def __call__(self, question: str) -> str:
-        print(f"NovaLiteAgent received question (first 50 chars): {question[:50]}...")
+        print(f"NovaProAgent received question (first 50 chars): {question[:50]}...")
         
         try:
-            # Create a more focused prompt for concise answers
-            prompt = f"""Answer this question directly and concisely. Provide only the essential information requested, not explanations or step-by-step reasoning unless specifically asked.
+            # 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 Lite
+            # Prepare the request payload for Nova Pro
             payload = {
                 "messages": [
                     {
@@ -52,8 +66,8 @@ Answer:"""
                     }
                 ],
                 "inferenceConfig": {
-                    "max_new_tokens": 150,
-                    "temperature": 0.0
+                    "max_new_tokens": max_tokens,
+                    "temperature": 0.1 if needs_reasoning else 0.0
                 }
             }
             
@@ -69,30 +83,46 @@ Answer:"""
             response_body = json.loads(response['body'].read())
             answer = response_body['output']['message']['content'][0]['text']
             
-            # Clean up the answer - remove verbose patterns
+            # Clean up the answer
             answer = answer.strip()
             
-            # Remove common verbose beginnings
+            # 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"
+                "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()) > 10:
+                        if len(sentence.strip()) > 5:
                             answer = sentence.strip()
                             break
             
-            # Limit length
-            if len(answer) > 200:
-                sentences = answer.split('. ')
-                answer = sentences[0] + '.'
+            # Clean up common patterns
+            answer = answer.replace('**', '').replace('*', '')
             
             return answer