agent.py

import os
import tempfile
import time
import re
import json
from typing import List, Optional, Dict, Any
from urllib.parse import urlparse
import requests
import yt_dlp
from bs4 import BeautifulSoup
from difflib import SequenceMatcher

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper, WikipediaAPIWrapper
from langchain.agents import Tool, AgentExecutor, ConversationalAgent, initialize_agent, AgentType
from langchain.memory import ConversationBufferMemory
from langchain.prompts import MessagesPlaceholder
from langchain.tools import BaseTool, Tool, tool
from google.generativeai.types import HarmCategory, HarmBlockThreshold
from PIL import Image
import google.generativeai as genai
from pydantic import Field

from smolagents import WikipediaSearchTool

def invoke_with_retry(
    llm: ChatGoogleGenerativeAI,
    prompt: str,
    max_retries: int = 5,
    initial_delay: int = 60
):
    """
    Google Generative AIへのAPI呼び出しを、`ResourceExhausted`エラー時に再試行する関数。
    
    Args:
        llm: ChatGoogleGenerativeAIのインスタンス。
        prompt: ユーザーからのプロンプト文字列。
        max_retries: 最大再試行回数。デフォルトは5。
        initial_delay: 最初の再試行までの待機時間（秒）。デフォルトは60。
        
    Returns:
        成功した場合のAPIレスポンス、失敗した場合はNone。
    """
    retries = 0
    delay = initial_delay
    
    while retries < max_retries:
        try:
            messages = [HumanMessage(content=prompt)]
            response = llm.invoke(messages)
            return response
        except ResourceExhausted as e:
            print(f"APIアクセス上限を超えました。待機して再試行します。({retries + 1}/{max_retries})")
            print(f"エラー詳細: {e}")
            time.sleep(delay)
            delay *= 2  # 指数バックオフ
            retries += 1
        except Exception as e:
            # 他の予期せぬエラーに対する処理
            print(f"予期せぬエラーが発生しました: {e}")
            break
            
    print("最大再試行回数に達しました。API呼び出しに失敗しました。")
    return None

class SmolagentToolWrapper(BaseTool):
    """Wrapper for smolagents tools to make them compatible with LangChain."""
    
    wrapped_tool: object = Field(description="The wrapped smolagents tool")
    
    def __init__(self, tool):
        """Initialize the wrapper with a smolagents tool."""
        super().__init__(
            name=tool.name,
            description=tool.description,
            return_direct=False,
            wrapped_tool=tool
        )

    def _run(self, query: str) -> str:
        """Use the wrapped tool to execute the query."""
        try:
            # For WikipediaSearchTool
            if hasattr(self.wrapped_tool, 'search'):
                return self.wrapped_tool.search(query)
            # For DuckDuckGoSearchTool and others
            return self.wrapped_tool(query)
        except Exception as e:
            return f"Error using tool: {str(e)}"
    
    def _arun(self, query: str) -> str:
        """Async version - just calls sync version since smolagents tools don't support async."""
        return self._run(query)



class Agent:
    def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"):
        # Suppress warnings
        import warnings
        warnings.filterwarnings("ignore", category=UserWarning)
        warnings.filterwarnings("ignore", category=DeprecationWarning)
        warnings.filterwarnings("ignore", message=".*will be deprecated.*")
        warnings.filterwarnings("ignore", "LangChain.*")
        
        self.api_key = api_key
        self.model_name = model_name
        
        # Configure Gemini
        genai.configure(api_key=api_key)
        
        # Initialize the LLM
        self.llm = self._setup_llm()
        
        # Setup tools
        self.tools = [
            SmolagentToolWrapper(WikipediaSearchTool()),
            #Tool(
            #    name="analyze_video",
            #    func=self._analyze_video,
            #    description="Analyze YouTube video content directly"
            #),
            #Tool(
            #    name="analyze_image",
            #    func=self._analyze_image,
            #    description="Analyze image content"
            #),
            Tool(
                name="analyze_table",
                func=self._analyze_table,
                description="Analyze table or matrix data"
            ),
            Tool(
                name="analyze_list",
                func=self._analyze_list,
                description="Analyze and categorize list items"
            ),
            Tool(
                name="web_search",
                func=self._web_search,
                description="Search the web for information"
            )
        ]
        
        # Setup memory
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True
        )
        
        # Initialize agent
        self.agent = self._setup_agent()
       

    def run(self, query: str) -> str:
        """Run the agent on a query with incremental retries."""
        max_retries = 3
        base_sleep = 1  # Start with 1 second sleep
        
        for attempt in range(max_retries):
            try:

                # If no match found in answer bank, use the agent
                response = self.agent.run(query)
                return response

            except Exception as e:
                sleep_time = base_sleep * (attempt + 1)  # Incremental sleep: 1s, 2s, 3s
                if attempt < max_retries - 1:
                    print(f"Attempt {attempt + 1} failed. Retrying in {sleep_time} seconds...")
                    time.sleep(sleep_time)
                    continue
                return f"Error processing query after {max_retries} attempts: {str(e)}"

        print("Agent processed all queries!")

    def _clean_response(self, response: str) -> str:
        """Clean up the response from the agent."""
        # Remove any tool invocation artifacts
        cleaned = re.sub(r'> Entering new AgentExecutor chain...|> Finished chain.', '', response)
        cleaned = re.sub(r'Thought:.*?Action:.*?Action Input:.*?Observation:.*?\n', '', cleaned, flags=re.DOTALL)
        return cleaned.strip()

    def run_interactive(self):
        print("AI Assistant Ready! (Type 'exit' to quit)")
        
        while True:
            query = input("You: ").strip()
            if query.lower() == 'exit':
                print("Goodbye!")
                break
            
            print("Assistant:", self.run(query))

    def _web_search(self, query: str, domain: Optional[str] = None) -> str:
        """Perform web search with rate limiting and retries."""
        try:
            # Use DuckDuckGo API wrapper for more reliable results
            search = DuckDuckGoSearchAPIWrapper(max_results=5)
            results = search.run(f"{query} {f'site:{domain}' if domain else ''}")
            
            if not results or results.strip() == "":
                return "No search results found."
                
            return results

        except Exception as e:
            return f"Search error: {str(e)}"

    def _analyze_video(self, url: str) -> str:
        """Analyze video content using Gemini's video understanding capabilities."""
        try:
            # Validate URL
            parsed_url = urlparse(url)
            if not all([parsed_url.scheme, parsed_url.netloc]):
                return "Please provide a valid video URL with http:// or https:// prefix."
            
            # Check if it's a YouTube URL
            if 'youtube.com' not in url and 'youtu.be' not in url:
                return "Only YouTube videos are supported at this time."

            try:
                # Configure yt-dlp with minimal extraction
                ydl_opts = {
                    'quiet': True,
                    'no_warnings': True,
                    'extract_flat': True,
                    'no_playlist': True,
                    'youtube_include_dash_manifest': False
                }

                with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                    try:
                        # Try basic info extraction
                        info = ydl.extract_info(url, download=False, process=False)
                        if not info:
                            return "Could not extract video information."

                        title = info.get('title', 'Unknown')
                        description = info.get('description', '')
                        
                        # Create a detailed prompt with available metadata
                        prompt = f"""Please analyze this YouTube video:
Title: {title}
URL: {url}
Description: {description}
Please provide a detailed analysis focusing on:
1. Main topic and key points from the title and description
2. Expected visual elements and scenes
3. Overall message or purpose
4. Target audience"""

                        # Use the LLM with proper message format
                        #messages = [HumanMessage(content=prompt)]
                        #response = self.llm.invoke(messages)
                        response = invoke_with_retry(self.llm, prompt)
                        return response.content if hasattr(response, 'content') else str(response)

                    except Exception as e:
                        if 'Sign in to confirm' in str(e):
                            return "This video requires age verification or sign-in. Please provide a different video URL."
                        return f"Error accessing video: {str(e)}"

            except Exception as e:
                return f"Error extracting video info: {str(e)}"

        except Exception as e:
            return f"Error analyzing video: {str(e)}"

    def _analyze_table(self, table_data: str) -> str:
        """Analyze table or matrix data."""
        try:
            if not table_data or not isinstance(table_data, str):
                return "Please provide valid table data for analysis."

            prompt = f"""Please analyze this table:
{table_data}
Provide a detailed analysis including:
1. Structure and format
2. Key patterns or relationships
3. Notable findings
4. Any mathematical properties (if applicable)"""

            #messages = [HumanMessage(content=prompt)]
            #response = self.llm.invoke(messages)
            response = invoke_with_retry(self.llm, prompt)
            return response.content if hasattr(response, 'content') else str(response)

        except Exception as e:
            return f"Error analyzing table: {str(e)}"

    def _analyze_image(self, image_data: str) -> str:
        """Analyze image content."""
        try:
            if not image_data or not isinstance(image_data, str):
                return "Please provide a valid image for analysis."

            prompt = f"""Please analyze this image:
{image_data}
Focus on:
1. Visual elements and objects
2. Colors and composition
3. Text or numbers (if present)
4. Overall context and meaning"""

            #messages = [HumanMessage(content=prompt)]
            #response = self.llm.invoke(messages)
            response = invoke_with_retry(self.llm, prompt)
            return response.content if hasattr(response, 'content') else str(response)

        except Exception as e:
            return f"Error analyzing image: {str(e)}"

    def _analyze_list(self, list_data: str) -> str:
        """Analyze and categorize list items."""
        if not list_data:
            return "No list data provided."
        try:
            items = [x.strip() for x in list_data.split(',')]
            if not items:
                return "Please provide a comma-separated list of items."
            # Add list analysis logic here
            return "Please provide the list items for analysis."
        except Exception as e:
            return f"Error analyzing list: {str(e)}"

    def _setup_llm(self):
        """Set up the language model."""
        # Set up model with video capabilities
        generation_config = {
            "temperature": 0.0,
            "max_output_tokens": 2000,
            "candidate_count": 1,
        }
        
        safety_settings = {
            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
        }
        
        return ChatGoogleGenerativeAI(
            model="gemini-2.0-flash",
            google_api_key=self.api_key,
            temperature=0,
            max_output_tokens=2000,
            generation_config=generation_config,
            safety_settings=safety_settings,
            system_message=SystemMessage(content=(
                "You are a precise AI assistant that helps users find information and analyze content. "
                "You can directly understand and analyze YouTube videos, images, and other content. "
                "When analyzing videos, focus on relevant details like dialogue, text, and key visual elements. "
                "For lists, tables, and structured data, ensure proper formatting and organization. "
                "If you need additional context, clearly explain what is needed."
            ))
        )
        
    def _setup_agent(self) -> AgentExecutor:
        """Set up the agent with tools and system message."""
        
        # Define the system message template
        PREFIX = """You are a helpful AI assistant that can use various tools to answer questions and analyze content. You have access to tools for web search, Wikipedia lookup, and multimedia analysis.
TOOLS:
------
You have access to the following tools:"""

        FORMAT_INSTRUCTIONS = """To use a tool, use the following format:
Thought: Do I need to use a tool? Yes
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:
Thought: Do I need to use a tool? No
Final Answer: [your response here]
Begin! Remember to ALWAYS include 'Thought:', 'Action:', 'Action Input:', and 'Final Answer:' in your responses."""

        SUFFIX = """Previous conversation history:
{chat_history}
New question: {input}
{agent_scratchpad}"""

        # Create the base agent
        agent = ConversationalAgent.from_llm_and_tools(
            llm=self.llm,
            tools=self.tools,
            prefix=PREFIX,
            format_instructions=FORMAT_INSTRUCTIONS,
            suffix=SUFFIX,
            input_variables=["input", "chat_history", "agent_scratchpad", "tool_names"],
            handle_parsing_errors=True
        )

        # Initialize agent executor with custom output handling
        return AgentExecutor.from_agent_and_tools(
            agent=agent,
            tools=self.tools,
            memory=self.memory,
            max_iterations=5,
            verbose=True,
            handle_parsing_errors=True,
            return_only_outputs=True  # This ensures we only get the final output
        )