Update app.py

app.py

@@ -9,23 +9,20 @@ import uuid
 from datetime import datetime, timezone, timedelta
 from dotenv import load_dotenv
 import json
-from huggingface_hub import HfApi
 
-# Load Hugging Face token from the environment variable
+# Load environment variables
+load_dotenv()
+
+# Get Google API key from environment
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 if GOOGLE_API_KEY is None:
     raise ValueError("GOOGLE_API_KEY environment variable is not set. Please set it before running the script.")
 
-
-hf_api = HfApi(
-    token= GOOGLE_API_KEY, # Token is not persisted on the machine.
-)
-
 app = Flask(__name__)
 CORS(app)
 
-# Initialize Gemini client
-client = genai.Client(api_key=hf_api)
+# Initialize Gemini client with correct API key
+client = genai.Client(api_key=GOOGLE_API_KEY)
 
 # In-memory storage for demo (in production, use a database)
 document_caches = {}
@@ -609,60 +606,93 @@ def upload_file():
         if file.filename == '':
             return jsonify({'success': False, 'error': 'No file selected'})
         
+        # Check file size (limit to 10MB for example)
+        file.seek(0, 2)  # Seek to end
+        file_size = file.tell()
+        file.seek(0)  # Reset to beginning
+        
+        if file_size > 10 * 1024 * 1024:  # 10MB limit
+            return jsonify({'success': False, 'error': 'File too large. Maximum size is 10MB.'})
+        
         # Read file content
         file_content = file.read()
+        if not file_content:
+            return jsonify({'success': False, 'error': 'File is empty'})
+        
         file_io = io.BytesIO(file_content)
         
         # Upload to Gemini File API
-        document = client.files.upload(
-            file=file_io,
-            config=dict(mime_type='application/pdf')
-        )
+        try:
+            document = client.files.upload(
+                file=file_io,
+                config=types.FileUploadConfig(
+                    mime_type='application/pdf',
+                    display_name=file.filename
+                )
+            )
+            print(f"Document uploaded successfully: {document.name}")
+        except Exception as upload_error:
+            print(f"Upload error: {upload_error}")
+            return jsonify({'success': False, 'error': f'Failed to upload file to Gemini: {str(upload_error)}'})
         
         # Create cache with system instruction
         try:
             system_instruction = "You are an expert document analyzer. Provide detailed, accurate answers based on the uploaded document content. Always be helpful and thorough in your responses."
             
-            # Use the correct model format as per documentation
-            model = 'models/gemini-2.0-flash-001'
+            # Use the correct model name (without 'models/' prefix)
+            model = 'gemini-2.0-flash-001'
             
             cache = client.caches.create(
                 model=model,
                 config=types.CreateCachedContentConfig(
-                    display_name='pdf document cache',
+                    display_name=f'PDF document cache - {file.filename}',
                     system_instruction=system_instruction,
                     contents=[document],
                     ttl="3600s",  # 1 hour TTL
                 )
             )
             
+            print(f"Cache created successfully: {cache.name}")
+            
             # Store cache info
             cache_id = str(uuid.uuid4())
             document_caches[cache_id] = {
                 'cache_name': cache.name,
                 'document_name': file.filename,
+                'document_file_name': document.name,
                 'created_at': datetime.now().isoformat()
             }
             
+            # Get token count safely
+            token_count = 'Unknown'
+            if hasattr(cache, 'usage_metadata') and cache.usage_metadata:
+                if hasattr(cache.usage_metadata, 'total_token_count'):
+                    token_count = cache.usage_metadata.total_token_count
+                elif hasattr(cache.usage_metadata, 'cached_token_count'):
+                    token_count = cache.usage_metadata.cached_token_count
+            
             return jsonify({
                 'success': True,
                 'cache_id': cache_id,
-                'token_count': getattr(cache.usage_metadata, 'cached_token_count', 'Unknown')
+                'token_count': token_count,
+                'document_name': file.filename
             })
             
         except Exception as cache_error:
+            print(f"Cache error: {cache_error}")
             # If caching fails due to small content, provide alternative approach
-            if "Cached content is too small" in str(cache_error):
+            if "too small" in str(cache_error).lower():
                 return jsonify({
                     'success': False, 
-                    'error': 'PDF is too small for caching. Please upload a larger document (minimum 4,096 tokens required).',
+                    'error': 'PDF content is too small for caching. Please upload a larger document with more text content.',
                     'suggestion': 'Try uploading a longer document or combine multiple documents.'
                 })
             else:
-                raise cache_error
+                return jsonify({'success': False, 'error': f'Failed to create cache: {str(cache_error)}'})
         
     except Exception as e:
-        return jsonify({'success': False, 'error': str(e)})
+        print(f"General error: {e}")
+        return jsonify({'success': False, 'error': f'Server error: {str(e)}'})
 
 @app.route('/upload-url', methods=['POST'])
 def upload_from_url():
@@ -673,62 +703,107 @@ def upload_from_url():
         if not url:
             return jsonify({'success': False, 'error': 'No URL provided'})
         
-        # Download file from URL
-        response = httpx.get(url)
-        response.raise_for_status()
+        # Download file from URL with timeout and size limits
+        try:
+            with httpx.Client(timeout=30.0) as client_http:
+                response = client_http.get(url)
+                response.raise_for_status()
+                
+                # Check content type
+                content_type = response.headers.get('content-type', '').lower()
+                if 'pdf' not in content_type and not url.lower().endswith('.pdf'):
+                    return jsonify({'success': False, 'error': 'URL does not point to a PDF file'})
+                
+                # Check file size
+                content_length = len(response.content)
+                if content_length > 10 * 1024 * 1024:  # 10MB limit
+                    return jsonify({'success': False, 'error': 'File too large. Maximum size is 10MB.'})
+                
+                file_io = io.BytesIO(response.content)
+                
+        except httpx.TimeoutException:
+            return jsonify({'success': False, 'error': 'Request timeout. Please try a different URL.'})
+        except httpx.HTTPError as e:
+            return jsonify({'success': False, 'error': f'Failed to download file: {str(e)}'})
         
-        file_io = io.BytesIO(response.content)
+        # Extract filename from URL
+        filename = url.split('/')[-1]
+        if not filename.endswith('.pdf'):
+            filename += '.pdf'
         
         # Upload to Gemini File API
-        document = client.files.upload(
-            file=file_io,
-            config=dict(mime_type='application/pdf')
-        )
+        try:
+            document = client.files.upload(
+                file=file_io,
+                config=types.FileUploadConfig(
+                    mime_type='application/pdf',
+                    display_name=filename
+                )
+            )
+            print(f"Document uploaded successfully: {document.name}")
+        except Exception as upload_error:
+            print(f"Upload error: {upload_error}")
+            return jsonify({'success': False, 'error': f'Failed to upload file to Gemini: {str(upload_error)}'})
         
         # Create cache with system instruction
         try:
             system_instruction = "You are an expert document analyzer. Provide detailed, accurate answers based on the uploaded document content. Always be helpful and thorough in your responses."
             
-            # Use the correct model format as per documentation
-            model = 'models/gemini-2.0-flash-001'
+            # Use the correct model name (without 'models/' prefix)
+            model = 'gemini-2.0-flash-001'
             
             cache = client.caches.create(
                 model=model,
                 config=types.CreateCachedContentConfig(
-                    display_name='pdf document cache',
+                    display_name=f'PDF document cache - {filename}',
                     system_instruction=system_instruction,
                     contents=[document],
                     ttl="3600s",  # 1 hour TTL
                 )
             )
             
+            print(f"Cache created successfully: {cache.name}")
+            
             # Store cache info
             cache_id = str(uuid.uuid4())
             document_caches[cache_id] = {
                 'cache_name': cache.name,
-                'document_name': url,
+                'document_name': filename,
+                'document_file_name': document.name,
+                'source_url': url,
                 'created_at': datetime.now().isoformat()
             }
             
+            # Get token count safely
+            token_count = 'Unknown'
+            if hasattr(cache, 'usage_metadata') and cache.usage_metadata:
+                if hasattr(cache.usage_metadata, 'total_token_count'):
+                    token_count = cache.usage_metadata.total_token_count
+                elif hasattr(cache.usage_metadata, 'cached_token_count'):
+                    token_count = cache.usage_metadata.cached_token_count
+            
             return jsonify({
                 'success': True,
                 'cache_id': cache_id,
-                'token_count': getattr(cache.usage_metadata, 'cached_token_count', 'Unknown')
+                'token_count': token_count,
+                'document_name': filename
             })
             
         except Exception as cache_error:
+            print(f"Cache error: {cache_error}")
             # If caching fails due to small content, provide alternative approach
-            if "Cached content is too small" in str(cache_error):
+            if "too small" in str(cache_error).lower():
                 return jsonify({
                     'success': False, 
-                    'error': 'PDF is too small for caching. Please upload a larger document (minimum 4,096 tokens required).',
+                    'error': 'PDF content is too small for caching. Please upload a larger document with more text content.',
                     'suggestion': 'Try uploading a longer document or combine multiple documents.'
                 })
             else:
-                raise cache_error
+                return jsonify({'success': False, 'error': f'Failed to create cache: {str(cache_error)}'})
         
     except Exception as e:
-        return jsonify({'success': False, 'error': str(e)})
+        print(f"General error: {e}")
+        return jsonify({'success': False, 'error': f'Server error: {str(e)}'})
 
 @app.route('/ask', methods=['POST'])
 def ask_question():
@@ -741,26 +816,38 @@ def ask_question():
             return jsonify({'success': False, 'error': 'Missing question or cache_id'})
         
         if cache_id not in document_caches:
-            return jsonify({'success': False, 'error': 'Cache not found'})
+            return jsonify({'success': False, 'error': 'Cache not found. Please upload a document first.'})
         
         cache_info = document_caches[cache_id]
         
         # Generate response using cached content with correct model format
-        response = client.models.generate_content(
-            model='models/gemini-2.0-flash-001',
-            contents=question,
-            config=types.GenerateContentConfig(
-                cached_content=cache_info['cache_name']
+        try:
+            response = client.models.generate_content(
+                model='gemini-2.0-flash-001',  # No 'models/' prefix here
+                contents=question,
+                config=types.GenerateContentConfig(
+                    cached_content=cache_info['cache_name']
+                )
             )
-        )
-        
-        return jsonify({
-            'success': True,
-            'answer': response.text
-        })
+            
+            if response and response.text:
+                return jsonify({
+                    'success': True,
+                    'answer': response.text
+                })
+            else:
+                return jsonify({
+                    'success': False,
+                    'error': 'No response generated from the model'
+                })
+                
+        except Exception as gen_error:
+            print(f"Generation error: {gen_error}")
+            return jsonify({'success': False, 'error': f'Failed to generate response: {str(gen_error)}'})
         
     except Exception as e:
-        return jsonify({'success': False, 'error': str(e)})
+        print(f"General error in ask_question: {e}")
+        return jsonify({'success': False, 'error': f'Server error: {str(e)}'})
 
 @app.route('/caches', methods=['GET'])
 def list_caches():
@@ -787,7 +874,11 @@ def delete_cache(cache_id):
         cache_info = document_caches[cache_id]
         
         # Delete from Gemini API
-        client.caches.delete(cache_info['cache_name'])
+        try:
+            client.caches.delete(cache_info['cache_name'])
+        except Exception as delete_error:
+            print(f"Error deleting cache from Gemini API: {delete_error}")
+            # Continue to remove from local storage even if API deletion fails
         
         # Remove from local storage
         del document_caches[cache_id]
@@ -797,7 +888,23 @@ def delete_cache(cache_id):
     except Exception as e:
         return jsonify({'success': False, 'error': str(e)})
 
+# Health check endpoint
+@app.route('/health', methods=['GET'])
+def health_check():
+    return jsonify({'status': 'healthy', 'service': 'Smart Document Analysis Platform'})
+
+# Error handlers
+@app.errorhandler(413)
+def too_large(e):
+    return jsonify({'success': False, 'error': 'File too large'}), 413
+
+@app.errorhandler(500)
+def internal_error(e):
+    return jsonify({'success': False, 'error': 'Internal server error'}), 500
+
 if __name__ == '__main__':
     import os
     port = int(os.environ.get("PORT", 7860))
-    app.run(debug=True, host='0.0.0.0', port=port) 
+    print(f"Starting server on port {port}")
+    print(f"Google API Key configured: {'Yes' if GOOGLE_API_KEY else 'No'}")
+    app.run(debug=False, host='0.0.0.0', port=port)