Update app.py

app.py

@@ -42,6 +42,7 @@ def init_state(question: str):
     }
 
 
+
 # ====== Tool Definitions ======
 @tool
 def duckduckgo_search(query: str) -> str:
@@ -102,6 +103,45 @@ def python_execution(code: str) -> str:
     except Exception as e:
         return f"Error: {str(e)}"
 
+from langchain_core.tools import BaseTool
+from youtube_transcript_api import YouTubeTranscriptApi
+
+class VideoTranscriptionTool(BaseTool):
+    name = "transcript_video"
+    description = "Fetch text transcript from YouTube videos using URL or ID. Optionally include timestamps."
+
+    def _run(self, url: str, include_timestamps: bool = False) -> str:
+        video_id = None
+        if "youtube.com/watch?v=" in url:
+            video_id = url.split("v=")[1].split("&")[0]
+        elif "youtu.be/" in url:
+            video_id = url.split("youtu.be/")[1].split("?")[0]
+        elif len(url.strip()) == 11 and not ("http://" in url or "https://" in url):
+            video_id = url.strip()
+        
+        if not video_id:
+            return f"Invalid or unsupported YouTube URL/ID: {url}"
+
+        try:
+            transcription = YouTubeTranscriptApi.get_transcript(video_id)
+
+            if include_timestamps:
+                formatted = []
+                for part in transcription:
+                    timestamp = f"{int(part['start']//60)}:{int(part['start']%60):02d}"
+                    formatted.append(f"[{timestamp}] {part['text']}")
+                return "\n".join(formatted)
+            else:
+                return " ".join([part['text'] for part in transcription])
+
+        except Exception as e:
+            return f"Error fetching transcript: {str(e)}"
+
+    def _arun(self, *args, **kwargs):
+        raise NotImplementedError("This tool does not support async yet.")
+
+        
+
 def indent_code(code: str) -> str:
     return '\n    '.join(code.splitlines())
 
@@ -149,58 +189,72 @@ def should_continue(state: AgentState) -> str:
 
 def reasoning_node(state: AgentState) -> AgentState:
     import os
+    import time
     from langchain_google_genai import ChatGoogleGenerativeAI
     from langchain.schema import HumanMessage, AIMessage
     from langchain.prompts import ChatPromptTemplate
+    from google.api_core.exceptions import ResourceExhausted
 
-    # Load and verify the Google API key
+    # Load API key
     GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
     if not GOOGLE_API_KEY:
         raise ValueError("GOOGLE_API_KEY not set in environment variables.")
 
-    # Ensure history is initialized and ends with a HumanMessage
+    # Ensure history is well-formed
     if "history" not in state or not isinstance(state["history"], list):
         state["history"] = []
     if not state["history"] or not isinstance(state["history"][-1], HumanMessage):
         state["history"].append(HumanMessage(content="Continue."))
 
-    # Ensure context is a dictionary
-    if not isinstance(state.get("context"), dict):
-        state["context"] = {}
-
-    # Ensure reasoning and iterations keys are present
+    # Ensure context and reasoning fields
+    state.setdefault("context", {})
     state.setdefault("reasoning", "")
     state.setdefault("iterations", 0)
 
-    # Initialize Gemini model via LangChain
+    # Create Gemini model wrapper
     llm = ChatGoogleGenerativeAI(
         model="gemini-1.5-flash",
         temperature=0.1,
         google_api_key=GOOGLE_API_KEY
     )
 
-    # Create prompt
+    # Create prompt chain
     prompt = ChatPromptTemplate.from_messages([
         ("system", (
             "You're an expert problem solver. Analyze the question, select the best tool, "
             "and provide reasoning. Available tools: duckduckgo_search, wikipedia_search, "
             "arxiv_search, document_qa, python_execution.\n\n"
+            "Important: You must select a tool for questions involving video, audio, or code.\n\n"
             "Current Context:\n{context}\n\n"
             "Reasoning Steps:\n{reasoning}\n\n"
-            "Response Format:\nReasoning: [Your analysis]\nAction: [Tool name OR 'Final Answer']\n"
+            "Response Format:\n"
+            "Reasoning: [Your analysis]\n"
+            "Action: [Tool name OR 'Final Answer']\n"
             "Action Input: [Input for tool OR final response]"
         )),
         *state["history"]
     ])
 
-    # Invoke model
     chain = prompt | llm
-    response = chain.invoke({
+
+    # === Add Retry Logic ===
+    def call_with_retry(inputs, retries=3, delay=60):
+        for attempt in range(retries):
+            try:
+                return chain.invoke(inputs)
+            except ResourceExhausted as e:
+                print(f"[Retry {attempt+1}] Gemini rate limit hit. Waiting {delay}s...")
+                time.sleep(delay)
+        raise RuntimeError("Failed after multiple retries due to Gemini quota limit.")
+
+    # Call model with retry protection
+    response = call_with_retry({
         "context": state["context"],
         "reasoning": state["reasoning"],
         "question": state["question"]
     })
 
+    # Parse output
     content = response.content
     reasoning, action, action_input = parse_agent_response(content)
 
@@ -209,7 +263,6 @@ def reasoning_node(state: AgentState) -> AgentState:
     state["reasoning"] += f"\nStep {state['iterations'] + 1}: {reasoning}"
     state["iterations"] += 1
 
-    # Store either final answer or tool to call
     if "final answer" in action.lower():
         state["history"].append(AIMessage(content=f"FINAL ANSWER: {action_input}"))
     else:
@@ -222,6 +275,7 @@ def reasoning_node(state: AgentState) -> AgentState:
 
 
 
+
     
 
 def tool_node(state: AgentState) -> AgentState:
@@ -316,7 +370,8 @@ class BasicAgent:
             wikipedia_search,
             arxiv_search,
             document_qa,
-            python_execution
+            python_execution,
+            VideoTranscriptionTool() 
         ]
 
     def __call__(self, question: str) -> str: