Update extensions/youtube.py

extensions/youtube.py

@@ -1,7 +1,7 @@
 """
-YouTube Video Summarization Extension
+Enhanced YouTube Video Summarization Extension
 Uses Gemini 2.5 Flash with native YouTube URL support
-Enhanced with caching, metadata, and smart query handling
+Now with better caching, state management, and orchestrator integration
 """
 
 from base_extension import BaseExtension
@@ -30,6 +30,10 @@ class YouTubeExtension(BaseExtension):
     def icon(self) -> str:
         return "📺"
     
+    @property
+    def version(self) -> str:
+        return "2.0.0"
+    
     def get_system_context(self) -> str:
         return """
 You have access to YouTube video analysis using Gemini's native video understanding.
@@ -61,9 +65,33 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
     def _get_default_state(self) -> Dict[str, Any]:
         return {
             "analyzed_videos": [],
-            "cache": {},  # Cache for analysis results
+            "cache": {},
             "cache_hits": 0,
-            "cache_misses": 0
+            "cache_misses": 0,
+            "total_analyses": 0,
+            "created_at": datetime.datetime.now().isoformat(),
+            "last_updated": datetime.datetime.now().isoformat()
+        }
+    
+    def get_state_summary(self, user_id: str) -> Optional[str]:
+        """Provide state summary for system prompt"""
+        state = self.get_state(user_id)
+        video_count = len(state.get("analyzed_videos", []))
+        if video_count > 0:
+            return f"{video_count} videos analyzed"
+        return None
+    
+    def get_metrics(self, user_id: str) -> Dict[str, Any]:
+        """Provide usage metrics"""
+        state = self.get_state(user_id)
+        total_requests = state.get("cache_hits", 0) + state.get("cache_misses", 0)
+        hit_rate = (state.get("cache_hits", 0) / total_requests * 100) if total_requests > 0 else 0
+        
+        return {
+            "total_analyses": state.get("total_analyses", 0),
+            "videos_analyzed": len(state.get("analyzed_videos", [])),
+            "cache_hit_rate": f"{hit_rate:.1f}%",
+            "cached_items": len(state.get("cache", {}))
         }
     
     def get_tools(self) -> List[types.Tool]:
@@ -111,10 +139,17 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
             }
         )
         
+        clear_cache = types.FunctionDeclaration(
+            name="clear_youtube_cache",
+            description="Clear the YouTube analysis cache to free up memory",
+            parameters={"type": "object", "properties": {}}
+        )
+        
         return [types.Tool(function_declarations=[
             analyze_video,
             list_analyzed_videos,
-            get_video_chapters
+            get_video_chapters,
+            clear_cache
         ])]
     
     def _extract_video_id(self, url: str) -> Optional[str]:
@@ -133,7 +168,6 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
     
     def _generate_cache_key(self, video_id: str, query: str) -> str:
         """Generate a cache key based on video ID and query intent"""
-        # Normalize query to catch similar intents
         query_lower = query.lower()
         
         # Categorize queries into types for better cache hits
@@ -146,7 +180,6 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
         elif any(word in query_lower for word in ['key point', 'main point', 'important']):
             query_type = "keypoints"
         else:
-            # For specific questions, use a hash of the query
             query_type = hashlib.md5(query.encode()).hexdigest()[:8]
         
         return f"{video_id}_{query_type}"
@@ -178,18 +211,14 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
             "has_chapters": any(word in analysis.lower() for word in ['chapter', 'section', 'part']),
         }
         
-        # Try to extract duration if mentioned
         duration_match = re.search(r'(\d+)\s*(?:minute|min|hour|hr)', analysis.lower())
         if duration_match:
             metadata["estimated_duration"] = duration_match.group(0)
         
         return metadata
     
-    def handle_tool_call(self, user_id: str, tool_name: str, args: Dict[str, Any]) -> Any:
-        # Ensure state is initialized (user_id is actually the API key in this app)
-        if user_id not in self.state:
-            self.initialize_state(user_id)
-        
+    def _execute_tool(self, user_id: str, tool_name: str, args: Dict[str, Any]) -> Any:
+        """Execute tool logic"""
         state = self.get_state(user_id)
         
         if tool_name == "analyze_youtube_video":
@@ -197,7 +226,6 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
             query = args.get("query", "Provide a detailed transcript of this video with timestamps")
             force_refresh = args.get("force_refresh", False)
             
-            # Validate and extract video ID
             video_id = self._extract_video_id(video_url)
             if not video_id:
                 return {
@@ -205,7 +233,7 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
                     "error": "Invalid YouTube URL. Please provide a valid youtube.com or youtu.be URL."
                 }
             
-            # Check cache first (unless force_refresh)
+            # Check cache
             cache_key = self._generate_cache_key(video_id, query)
             
             if not force_refresh and cache_key in state.get("cache", {}):
@@ -221,22 +249,19 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
             # Cache miss - perform analysis
             print(f"🔄 Cache MISS for {video_id} ({cache_key})")
             state["cache_misses"] = state.get("cache_misses", 0) + 1
+            state["total_analyses"] = state.get("total_analyses", 0) + 1
             
             try:
                 print(f"🎬 Analyzing YouTube video with Gemini: {video_url}")
                 
-                # Enhance query for better results
                 enhanced_query = self._enhance_query(query)
                 print(f"📝 Enhanced query: {enhanced_query[:100]}...")
                 
-                # Import genai client
                 from google import genai
                 from google.genai import types as genai_types
                 
-                # Create Gemini client (user_id is the API key in this app's architecture)
                 client = genai.Client(api_key=user_id)
                 
-                # Create the content with YouTube URL
                 contents = [
                     genai_types.Part(
                         file_data=genai_types.FileData(
@@ -246,7 +271,6 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
                     genai_types.Part(text=enhanced_query)
                 ]
                 
-                # Configure with thinking and grounding
                 config = genai_types.GenerateContentConfig(
                     system_instruction="You are a YouTube video analyst. Extract transcripts with timestamps when requested. Provide detailed, accurate analysis of video content. Format timestamps as [MM:SS].",
                     temperature=0.7,
@@ -259,14 +283,12 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
                 
                 print(f"🤖 Calling Gemini 2.5 Flash with YouTube URL...")
                 
-                # Generate response
                 response = client.models.generate_content(
                     model="gemini-2.5-flash",
                     contents=contents,
                     config=config
                 )
                 
-                # Extract response text
                 result_text = ""
                 if response.candidates and response.candidates[0].content:
                     for part in response.candidates[0].content.parts:
@@ -275,7 +297,6 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
                 
                 print(f"✅ Gemini analysis complete: {len(result_text)} characters")
                 
-                # Extract metadata
                 metadata = self._extract_metadata_from_analysis(result_text, video_id)
                 
                 # Track analyzed video
@@ -287,12 +308,10 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
                     "analysis_length": len(result_text)
                 }
                 
-                # Update or append video entry
                 existing_videos = [v for v in state.get("analyzed_videos", []) if v["video_id"] != video_id]
                 existing_videos.append(video_entry)
-                state["analyzed_videos"] = existing_videos[-20:]  # Keep last 20
+                state["analyzed_videos"] = existing_videos[-20:]
                 
-                # Build result
                 result = {
                     "success": True,
                     "video_url": video_url,
@@ -312,19 +331,31 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
                 
                 # Limit cache size (keep last 50 entries)
                 if len(state["cache"]) > 50:
-                    # Remove oldest entries
                     cache_keys = list(state["cache"].keys())
                     for old_key in cache_keys[:len(cache_keys)-50]:
                         del state["cache"][old_key]
                 
                 self.update_state(user_id, state)
                 
+                # Log activity
+                self.log_activity(user_id, "video_analyzed", {
+                    "video_id": video_id,
+                    "query_type": cache_key.split('_', 1)[1],
+                    "analysis_length": len(result_text)
+                })
+                
                 return result
                 
             except Exception as e:
                 import traceback
                 error_details = traceback.format_exc()
                 print(f"❌ YouTube analysis error: {error_details}")
+                
+                self.log_activity(user_id, "analysis_failed", {
+                    "video_id": video_id,
+                    "error": str(e)
+                })
+                
                 return {
                     "success": False,
                     "error": f"Error analyzing video: {str(e)}",
@@ -340,7 +371,7 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
             
             return {
                 "count": len(videos),
-                "videos": videos[-10:],  # Last 10 videos
+                "videos": videos[-10:],
                 "cache_stats": {
                     "hits": cache_hits,
                     "misses": cache_misses,
@@ -352,24 +383,55 @@ Note: This uses Gemini 2.5 Flash with thinking for high-quality analysis.
         elif tool_name == "get_video_chapters":
             video_url = args["video_url"]
             
-            # Use the analyze function with chapter-specific query
             chapter_query = "Break down this video into main chapters/sections with timestamps. For each chapter provide: timestamp, title, and brief description of what's covered."
             
-            return self.handle_tool_call(user_id, "analyze_youtube_video", {
+            return self._execute_tool(user_id, "analyze_youtube_video", {
                 "video_url": video_url,
                 "query": chapter_query
             })
         
+        elif tool_name == "clear_youtube_cache":
+            cache_size = len(state.get("cache", {}))
+            state["cache"] = {}
+            state["cache_hits"] = 0
+            state["cache_misses"] = 0
+            self.update_state(user_id, state)
+            
+            self.log_activity(user_id, "cache_cleared", {"items_cleared": cache_size})
+            
+            return {
+                "success": True,
+                "message": f"Cleared {cache_size} cached analyses",
+                "items_cleared": cache_size
+            }
+        
         return {"error": f"Unknown tool: {tool_name}"}
     
     def on_enable(self, user_id: str) -> str:
         self.initialize_state(user_id)
-        return "📺 YouTube Summarizer enabled! Share a YouTube URL and I'll analyze it using Gemini's native video understanding. Analysis results are cached for faster follow-up questions. No additional API key needed!"
+        return "📺 YouTube Summarizer enabled! Share a YouTube URL and I'll analyze it using Gemini's native video understanding. Analysis results are cached for faster follow-up questions!"
     
     def on_disable(self, user_id: str) -> str:
-        # Optionally clear cache on disable to save memory
         state = self.get_state(user_id)
         cache_size = len(state.get("cache", {}))
         state["cache"] = {}
         self.update_state(user_id, state)
-        return f"📺 YouTube Summarizer disabled. Cleared {cache_size} cached analyses."
+        return f"📺 YouTube Summarizer disabled. Cleared {cache_size} cached analyses."
+    
+    def health_check(self, user_id: str) -> Dict[str, Any]:
+        """Check extension health"""
+        state = self.get_state(user_id)
+        issues = []
+        
+        # Check cache size
+        cache_size = len(state.get("cache", {}))
+        if cache_size > 100:
+            issues.append(f"Large cache size ({cache_size} items) - consider clearing")
+        
+        return {
+            "healthy": len(issues) == 0,
+            "extension": self.name,
+            "version": self.version,
+            "cache_size": cache_size,
+            "issues": issues if issues else None
+        }