Blame for app.py

import os
import gradio as gr
import requests
import pandas as pd
from dotenv import load_dotenv
from duckduckgo_search import DDGS
import yt_dlp
from openai import OpenAI
import base64

from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage
from langchain_openai import ChatOpenAI
from typing import TypedDict, Dict, Any, List, Optional, Annotated, Sequence, Union, Literal
from langchain_core.tools import tool

load_dotenv()

# Environment setup
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") or os.environ.get("OPEN_AI_KEY")
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

# Type definitions for state schema
class AgentState(TypedDict):
    question: str
    messages: List[Union[HumanMessage, AIMessage, ToolMessage, SystemMessage]]
    steps_remaining: int
    answer: str
    intermediate_steps: List[Dict[str, Any]]

# --------------------- IMPROVED TOOLS ---------------------
# Helper functions for the enhanced web search

def refine_search_query(query: str) -> str:
    """
    Refine the search query by adding strategic terms or removing noise.
    
    Args:
        query: Original search query
        
    Returns:
        Refined search query
    """
    # Simple refinements
    # 1. Convert questions to keyword format
    question_words = ['what', 'when', 'where', 'who', 'why', 'how']
    query_lower = query.lower()
    
    for word in question_words:
        if query_lower.startswith(word):
            # Remove the question word and question mark
            query = re.sub(r'^' + word + r'\s+is\s+', '', query, flags=re.IGNORECASE)
            query = re.sub(r'^' + word + r'\s+are\s+', '', query, flags=re.IGNORECASE)
            query = re.sub(r'^' + word + r'\s+', '', query, flags=re.IGNORECASE)
            query = query.rstrip('?')
            break
    
    # 2. Add specific qualifiers for more accurate results
    if re.search(r'\d{4}', query):  # Contains a year
        # Don't add "latest" for historical queries
        pass
    elif any(word in query_lower for word in ['history', 'historical', 'ancient', 'past']):
        # Don't add "latest" for historical queries
        pass
    else:
        # For current topics, prefer recent information
        current_year = datetime.now().year
        query += f" {current_year}"
    
    return query

def get_ddg_results(query: str, max_results: int) -> list:
    """Get search results from DuckDuckGo"""
    results = []
    try:
        with DDGS() as ddgs:
            search_results = list(ddgs.text(query, max_results=max_results))
            results.extend(search_results)
    except Exception as e:
        print(f"DuckDuckGo search error: {e}")
    return results

def get_bing_results(query: str, max_results: int) -> list:
    """Get search results from Bing (example - would need API key in practice)"""
    # This is a placeholder for an actual Bing API implementation
    # In a real implementation, you would use the Bing Web Search API
    return []

def deduplicate_results(results: list) -> list:
    """Remove duplicate search results based on URL or content similarity"""
    seen_urls = set()
    unique_results = []
    
    for result in results:
        url = result.get('href', '')
        # Simple URL normalization
        normalized_url = url.rstrip('/').lower()
        
        if normalized_url and normalized_url not in seen_urls:
            seen_urls.add(normalized_url)
            unique_results.append(result)
    
    return unique_results

def extract_metadata(content: str, url: str) -> dict:
    """Extract additional metadata from content and URL"""
    metadata = {}
    
    # Extract date (simple regex for date patterns)
    date_patterns = [
        r'(\d{1,2}\s+(?:January|February|March|April|May|June|July|August|September|October|November|December)\s+\d{4})',
        r'(\d{1,2}/\d{1,2}/\d{4})',
        r'(\d{4}-\d{2}-\d{2})'
    ]
    
    for pattern in date_patterns:
        date_match = re.search(pattern, content, re.IGNORECASE)
        if date_match:
            metadata['date'] = date_match.group(1)
            break
    
    # Extract author if present
    author_patterns = [
        r'By\s+([A-Z][a-z]+\s+[A-Z][a-z]+)',
        r'Author[:\s]+([A-Z][a-z]+\s+[A-Z][a-z]+)'
    ]
    
    for pattern in author_patterns:
        author_match = re.search(pattern, content)
        if author_match:
            metadata['author'] = author_match.group(1)
            break
    
    return metadata

def clean_content(content: str) -> str:
    """Clean and normalize content"""
    # Remove extra whitespace
    cleaned = re.sub(r'\s+', ' ', content).strip()
    # Truncate if too long
    if len(cleaned) > 300:
        cleaned = cleaned[:297] + "..."
    return cleaned

def highlight_key_info(content: str, query: str) -> str:
    """Highlight key information in the content related to the query"""
    # Extract important keywords from the query
    keywords = [w for w in query.lower().split() if w not in ['the', 'a', 'an', 'is', 'are', 'was', 'were', 'to', 'of', 'in', 'on', 'at', 'by', 'for']]
    
    # Simple highlighting by adding context around keyword matches
    highlighted = content
    for keyword in keywords:
        # Find position of keyword in content
        keyword_pos = highlighted.lower().find(keyword)
        if keyword_pos >= 0:
            # Extract sentence containing the keyword
            sentence_start = max(0, highlighted.rfind('.', 0, keyword_pos) + 1)
            sentence_end = highlighted.find('.', keyword_pos)
            if sentence_end < 0:
                sentence_end = len(highlighted)
                
            # If the sentence is too long, take a window around the keyword
            if sentence_end - sentence_start > 200:
                start_pos = max(0, keyword_pos - 100)
                end_pos = min(len(highlighted), keyword_pos + 100)
                highlighted = highlighted[start_pos:end_pos]
                if start_pos > 0:
                    highlighted = "..." + highlighted
                if end_pos < len(content):
                    highlighted = highlighted + "..."
    
    return highlighted

def extract_domain(url: str) -> str:
    """Extract readable domain name from URL"""
    try:
        domain = urlparse(url).netloc
        # Remove www. prefix if present
        if domain.startswith('www.'):
            domain = domain[4:]
        return domain
    except:
        return "Unknown source"

def organize_results(results: list, query: str) -> list:
    """Organize results by relevance and cluster similar information"""
    # This is a simplified implementation
    # A more advanced version could use NLP to cluster results
    
    # For now, just return the results as is
    return results

def generate_search_summary(results: list, query: str) -> str:
    """Generate a brief summary of search results"""
    # Count sources by domain
    domains = {}
    for result in results:
        url = result.get('href', '')
        domain = extract_domain(url)
        domains[domain] = domains.get(domain, 0) + 1
    
    # Find most common domains
    top_domains = sorted(domains.items(), key=lambda x: x[1], reverse=True)[:3]
    
    # Create summary
    summary_lines = [
        f"Found {len(results)} results for '{query}'",
        f"Top sources: {', '.join(f'{domain} ({count})' for domain, count in top_domains)}"
    ]
    
    # Add a note about result freshness if we have dates
    dates = []
    for result in results:
        body = result.get('body', '')
        metadata = extract_metadata(body, result.get('href', ''))
        if 'date' in metadata:
            dates.append(metadata['date'])
    
    if dates:
        summary_lines.append(f"Results include information from {min(dates)} to {max(dates)}")
    
    return "\n".join(summary_lines)

@tool
def web_search(query: str, max_results: int = 10) -> str:
    """
    Search the web for information about a given query using multiple search engines
    and process the results for better comprehension.
   
    Args:
        query: The search query string.
        max_results: Maximum number of search results to return (default: 10)
       
    Returns:
        Formatted search results as a string with enhanced metadata and summaries.
    """
    try:
        print(f"Searching the web for: {query}")
        
        # Step 1: Query refinement - add strategic terms if needed
        refined_query = refine_search_query(query)
        print(f"Refined query: {refined_query}")
        
        # Step 2: Get results from multiple sources
        results = []
        metadata = {"sources": [], "query": refined_query}
        
        # First try DuckDuckGo
        ddg_results = get_ddg_results(refined_query, max_results)
        results.extend(ddg_results)
        metadata["sources"].append("DuckDuckGo")
        
        # Try Bing if we need more results
        if len(results) < max_results:
            try:
                remaining = max_results - len(results)
                bing_results = get_bing_results(refined_query, remaining)
                results.extend(bing_results)
                metadata["sources"].append("Bing")
            except Exception as e:
                print(f"Bing search error (non-critical): {e}")
        
        # Step 3: Deduplicate results
        results = deduplicate_results(results)
        
        # Step 4: Enhanced metadata extraction
        enhanced_results = []
        for i, result in enumerate(results[:max_results], 1):
            title = result.get('title', 'No title')
            body = result.get('body', 'No content')
            url = result.get('href', 'No URL')
            
            # Extract additional metadata
            additional_metadata = extract_metadata(body, url)
            
            # Clean and format content
            clean_body = clean_content(body)
            
            # Highlight key information
            highlighted_body = highlight_key_info(clean_body, query)
            
            # Structure the result
            formatted_result = f"{i}. {title}\n"
            
            # Add date if available
            if additional_metadata.get('date'):
                formatted_result += f"   Date: {additional_metadata['date']}\n"
                
            # Add summary and source
            formatted_result += f"   Summary: {highlighted_body}\n"
            formatted_result += f"   Source: {extract_domain(url)}\n"
            formatted_result += f"   URL: {url}\n"
            
            enhanced_results.append(formatted_result)
       
        if not enhanced_results:
            return "No search results found."
           
        # Step 5: Results clustering and organization by relevance
        organized_results = organize_results(enhanced_results, query)
        
        # Add search summary
        summary = generate_search_summary(results, query)
        
        # Final formatting
        formatted_results = f"SEARCH SUMMARY:\n{summary}\n\nDETAILED RESULTS:\n" + "\n".join(organized_results)
        print(f"Web search found {len(organized_results)} results")
        
        return formatted_results
    except Exception as e:
        print(f"Error during web search: {e}")
        return f"Error during web search: {str(e)}"

def encode_image(image_url):
    """Helper function to encode an image from URL to base64."""
    try:
        response = requests.get(image_url, stream=True, timeout=10)
        response.raise_for_status()
        image = response.content
        return base64.b64encode(image).decode("utf-8")
    except requests.exceptions.RequestException as e:
        return f"Error accessing image URL: {e}"

@tool
def analyze_image_openai_vision(image_url: str, prompt: str) -> str:
    """
    Analyze an image from a given URL using OpenAI Vision and answer a question about it.
    Provides detailed analysis of visual content.
    
    Args:
        image_url: URL of the image to analyze
        prompt: Specific question or instruction about what to analyze in the image
        
    Returns:
        Analysis of the image based on the prompt
    """
    if not OPENAI_API_KEY:
        return "OpenAI API key not found."
    
    base64_image = encode_image(image_url)
    if isinstance(base64_image, str) and base64_image.startswith("Error"):
        return base64_image
    
    client = OpenAI()
    try:
        response = client.chat.completions.create(
            model="gpt-4-vision-preview",
            messages=[
                {
                    "role": "system", 
                    "content": "You are a precise image analyzer focusing on exact details and specific information. Extract only what's needed to answer the query."
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}},
                    ],
                }
            ],
            max_tokens=500,
        )
        print(f"Image analysis complete: {len(response.choices[0].message.content)} chars")
        return response.choices[0].message.content
    except Exception as e:
        print(f"Error analyzing image with OpenAI Vision: {e}")
        return f"Error analyzing image with OpenAI Vision: {e}"

def get_youtube_transcript(video_url: str) -> str:
    """Extract transcript from a YouTube video with enhanced error handling."""
    try:
        ydl_opts = {
            'writesubtitles': True,
            'writeautomaticsub': True,
            'subtitleslangs': ['en'],
            'skip_download': True,
            'outtmpl': '%(title)s.%(ext)s',
            'quiet': True,
        }
        
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info_dict = ydl.extract_info(video_url, download=False)
            video_title = info_dict.get('title', 'Untitled Video')
            video_duration = info_dict.get('duration', 0)
            
            print(f"Video found: '{video_title}' (Duration: {video_duration}s)")
            
            # Try official subtitles first
            if info_dict and 'subtitles' in info_dict and 'en' in info_dict['subtitles']:
                transcript_url = info_dict['subtitles']['en'][0]['url']
                response = requests.get(transcript_url, timeout=10)
                response.raise_for_status()
                
                # More robust transcript parsing
                transcript_lines = []
                for line in response.text.splitlines():
                    if '-->' not in line and line.strip():
                        parts = line.split(' ', 1)
                        if len(parts) > 1:
                            transcript_lines.append(parts[1])
                
                return f"Video Title: {video_title}\nDuration: {video_duration} seconds\n\nTranscript:\n" + "\n".join(transcript_lines)
            
            # Fall back to automatic captions
            elif info_dict and 'automatic_captions' in info_dict and 'en' in info_dict['automatic_captions']:
                transcript_url = info_dict['automatic_captions']['en'][0]['url']
                response = requests.get(transcript_url, timeout=10)
                response.raise_for_status()
                
                transcript_lines = []
                for line in response.text.splitlines():
                    if '-->' not in line and line.strip():
                        parts = line.split(' ', 1)
                        if len(parts) > 1:
                            transcript_lines.append(parts[1])
                
                return f"Video Title: {video_title}\nDuration: {video_duration} seconds\n\nTranscript (Auto-generated):\n" + "\n".join(transcript_lines)
            else:
                video_info = f"Video Title: {video_title}\nDuration: {video_duration} seconds\n"
                return f"{video_info}\nNo English transcript found for this YouTube video."
    except Exception as e:
        print(f"Error fetching YouTube transcript: {e}")
        return f"Error fetching YouTube transcript: {e}"

@tool
def analyze_youtube_video_openai(video_url: str, prompt: str = "Summarize the key points of this video.") -> str:
    """
    Analyze a YouTube video by summarizing its English transcript using OpenAI.
    Provides comprehensive understanding of video content without watching it.
    
    Args:
        video_url: URL of the YouTube video to analyze
        prompt: Specific question or instruction about what to analyze in the video
        
    Returns:
        Analysis of the video content based on its transcript
    """
    if not OPENAI_API_KEY:
        return "OpenAI API key not found."
    
    transcript = get_youtube_transcript(video_url)
    if not transcript or transcript.startswith("Error"):
        return "Could not retrieve transcript for the YouTube video."
    
    client = OpenAI()
    try:
        response = client.chat.completions.create(
            model="gpt-4-turbo",
            messages=[
                {
                    "role": "system", 
                    "content": "You are a video content analyzer. Extract precise information from video transcripts. Provide specific details and quantitative data when applicable."
                },
                {
                    "role": "user", 
                    "content": f"Here is the transcript of a video:\n\n{transcript[:7000]}\n\n{prompt}"
                }
            ],
            max_tokens=750,
        )
        print(f"YouTube video analysis complete: {len(response.choices[0].message.content)} chars")
        return response.choices[0].message.content
    except Exception as e:
        print(f"Error analyzing video transcript with OpenAI: {e}")
        return f"Error analyzing video transcript with OpenAI: {e}"

@tool
def transcribe_audio_openai(audio_url: str) -> str:
    """
    Transcribe audio from a given URL using OpenAI Whisper API.
    Provides text version of audio content for analysis.
    
    Args:
        audio_url: URL of the audio file to transcribe
        
    Returns:
        Text transcription of the audio content
    """
    if not OPENAI_API_KEY:
        return "OpenAI API key not found."
    
    try:
        print(f"Downloading audio from {audio_url}")
        response = requests.get(audio_url, stream=True, timeout=15)
        response.raise_for_status()
        
        with open("temp_audio.mp3", "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        client = OpenAI()
        print("Transcribing audio with Whisper API...")
        with open("temp_audio.mp3", "rb") as audio_file:
            transcript = client.audio.transcriptions.create(
                model="whisper-1",
                file=audio_file
            )
        
        os.remove("temp_audio.mp3")
        print(f"Audio transcription complete: {len(transcript.text)} chars")
        return transcript.text
    except requests.exceptions.RequestException as e:
        print(f"Error accessing audio URL {audio_url}: {e}")
        return f"Error accessing audio URL {audio_url}: {e}"
    except Exception as e:
        print(f"Error transcribing audio with OpenAI Whisper: {e}")
        return f"Error transcribing audio with OpenAI Whisper: {e}"

@tool
def analyze_audio_transcript_openai(audio_url: str, prompt: str = "Analyze the content of this audio.") -> str:
    """
    Transcribe audio from a URL and then analyze the transcript using OpenAI.
    Provides detailed understanding of audio content without listening to it.
    
    Args:
        audio_url: URL of the audio file to analyze
        prompt: Specific question or instruction about what to analyze in the audio
        
    Returns:
        Analysis of the audio content based on its transcript
    """
    transcript = transcribe_audio_openai(audio_url)
    if isinstance(transcript, str) and transcript.startswith("Error"):
        return transcript
    
    if not transcript:
        return "Could not transcribe the audio."
    
    client = OpenAI()
    try:
        response = client.chat.completions.create(
            model="gpt-4-turbo",
            messages=[
                {
                    "role": "system", 
                    "content": "You are an audio content analyzer. Extract precise information from audio transcripts. Focus on specific details and key points."
                },
                {
                    "role": "user", 
                    "content": f"Here is the transcript of an audio:\n\n{transcript}\n\n{prompt}"
                }
            ],
            max_tokens=750,
        )
        print(f"Audio analysis complete: {len(response.choices[0].message.content)} chars")
        return response.choices[0].message.content
    except Exception as e:
        print(f"Error analyzing audio transcript with OpenAI: {e}")
        return f"Error analyzing audio transcript with OpenAI: {e}"

# Improved Python REPL Tool
from langchain_experimental.tools.python.tool import PythonREPLTool
python_repl = PythonREPLTool()

@tool
def run_python_code(code: str) -> str:
    """
    Execute a Python code snippet and return the output.
    Useful for calculations, data analysis, or any computational tasks.
    
    Args:
        code: Python code to execute
        
    Returns:
        Output of the executed code or error message
    """
    print("Executing Python code...")
    try:
        # Execute the code in a sandbox environment
        result = python_repl.run(code)
        print(f"Python execution successful: {len(str(result))} chars")
        return f"{result}"
    except Exception as e:
        print(f"Error executing Python code: {e}")
        return f"Error executing Python code: {e}"

# --------------------- IMPROVED AGENT ---------------------

class ImprovedLangGraphAgent:
    def __init__(self):
        print("ImprovedLangGraphAgent initializing...")
        
        if not OPENAI_API_KEY:
            raise ValueError("OpenAI API key not found in environment variables")

        print(f"Loaded OpenAI Key: {OPENAI_API_KEY[:3]}..." if OPENAI_API_KEY else "No OpenAI Key found")
        
        # Use gpt-4-turbo for better reasoning capabilities
        self.llm = ChatOpenAI(model="gpt-4-turbo", temperature=0, api_key=OPENAI_API_KEY)
        
        # Register all tools
        self.tools = [
            web_search, 
            analyze_image_openai_vision, 
            analyze_youtube_video_openai, 
            analyze_audio_transcript_openai, 
            run_python_code
        ]
        
        self.tool_map = {
            tool.name: tool for tool in self.tools
        }

        # Improved system prompt with clearer instructions and emphasis on reasoning
        self.sys_prompt = SystemMessage(
            content="""You are a highly capable AI assistant equipped with tools to solve complex problems. Your goal is to determine the correct answer to the given question, using a step-by-step reasoning approach.

AVAILABLE TOOLS:
- 'web_search': Search the internet for up-to-date information.
- 'analyze_image_openai_vision': Analyze images from URLs to extract visual information.
- 'analyze_youtube_video_openai': Analyze YouTube videos by examining their transcripts.
- 'analyze_audio_transcript_openai': Transcribe and analyze audio from URLs.
- 'run_python_code': Execute Python code for calculations, data analysis, or visualization.

REASONING PROCESS:
1. First, understand the question thoroughly.
2. Determine which tools would help gather the necessary information.
3. Execute each step logically, using tools when needed.
4. After collecting information, analyze it critically.
5. Draw conclusions based on evidence collected.
6. Formulate a precise, concise answer following the format requirements.

TOOL USAGE FORMAT:
When you need to use a tool, use the following format:
```tool
tool_name(arg1, arg2, ...)
```

FINAL ANSWER FORMAT:
Your final answer MUST be provided in this format:
<ANSWER>your concise answer here</ANSWER>

ANSWER REQUIREMENTS:
- For numbers: No commas, no units (like $ or %) unless explicitly requested
- For strings: No articles (a, an, the), no abbreviations for places
- For lists: Comma-separated values following above rules

IMPORTANT: 
- Ensure accuracy above all else
- Be extremely precise and concise in your final answer
- For coding questions, execute the code and use the output to determine the answer
- Never include explanations in your final <ANSWER> tag
"""
        )
        
        # Define the various components of our graph
        def setup_agent(state: AgentState) -> Dict[str, Any]:
            """Initialize the agent state with the question."""
            question = state["question"]
            messages = [
                self.sys_prompt,
                HumanMessage(content=question)
            ]
            print(f"[Agent] Setup with question: {question[:50]}...")
            return {
                "question": question,
                "messages": messages,
                "steps_remaining": 5,  # Allow up to 5 reasoning steps
                "answer": "",
                "intermediate_steps": []
            }
        
        def route_agent(state: AgentState) -> Literal["process_agent", "finalize_answer"]:
            """Decide whether to continue reasoning or finalize the answer."""
            if state["steps_remaining"] <= 0:
                print("[Agent] Max steps reached, finalizing...")
                return "finalize_answer"
            
            # Check if an answer has already been determined
            last_message = state["messages"][-1]
            if isinstance(last_message, AIMessage) and "<ANSWER>" in last_message.content:
                print("[Agent] Answer found, finalizing...")
                return "finalize_answer"
                
            print(f"[Agent] Continuing reasoning, {state['steps_remaining']} steps remaining")
            return "process_agent"
            
        def process_agent(state: AgentState) -> Dict[str, Any]:
            """Process the current state and decide the next action."""
            messages = state["messages"]
            
            # Use the LLM to decide the next step
            response = self.llm.invoke(messages)
            print(f"[Agent] LLM response: {response.content[:75]}...")
            
            # Check if the response contains a tool call
            if "I need to use" in response.content.lower() or "let me use" in response.content.lower():
                # Look for tool names in the response
                for tool_name, tool_fn in self.tool_map.items():
                    if tool_name.lower() in response.content.lower():
                        print(f"[Agent] Using tool: {tool_name}")
                        try:
                            # Extract arguments from the response
                            # This is a simple implementation - could be more sophisticated
                            tool_signature = f"{tool_name}("
                            start_idx = response.content.lower().find(tool_name.lower())
                            if start_idx == -1:
                                continue
                                
                            # Find the opening and closing parentheses
                            open_paren = response.content.find("(", start_idx)
                            if open_paren == -1:
                                continue
                                
                            # Extract everything between the first parenthesis and the last one
                            paren_count = 1
                            close_paren = -1
                            for i in range(open_paren + 1, len(response.content)):
                                if response.content[i] == "(":
                                    paren_count += 1
                                elif response.content[i] == ")":
                                    paren_count -= 1
                                    if paren_count == 0:
                                        close_paren = i
                                        break
                            
                            if close_paren == -1:
                                continue
                                
                            args_str = response.content[open_paren + 1:close_paren]
                            
                            # Simple parsing of arguments
                            # This is a basic implementation - a more robust parser would be better
                            args = []
                            kwargs = {}
                            
                            # Split by comma, but respecting quotes
                            in_quotes = False
                            quote_char = None
                            current_arg = ""
                            for char in args_str:
                                if char in ['"', "'"]:
                                    if not in_quotes:
                                        in_quotes = True
                                        quote_char = char
                                    elif char == quote_char:
                                        in_quotes = False
                                        quote_char = None
                                    current_arg += char
                                elif char == ',' and not in_quotes:
                                    args.append(current_arg.strip())
                                    current_arg = ""
                                else:
                                    current_arg += char
                            
                            if current_arg:
                                args.append(current_arg.strip())
                                
                            # Convert string arguments to actual values
                            parsed_args = []
                            for arg in args:
                                arg = arg.strip()
                                if arg.startswith('"') and arg.endswith('"'):
                                    parsed_args.append(arg[1:-1])
                                elif arg.startswith("'") and arg.endswith("'"):
                                    parsed_args.append(arg[1:-1])
                                else:
                                    # Try to parse as other types
                                    try:
                                        parsed_args.append(eval(arg))
                                    except:
                                        parsed_args.append(arg)
                            
                            # Execute the tool
                            tool_result = tool_fn(*parsed_args)
                            print(f"[Agent] Tool result: {tool_result[:75]}...")
                            
                            # Update the state with the tool result
                            messages.append(AIMessage(content=response.content))
                            messages.append(ToolMessage(content=tool_result, tool_call_id=tool_name))
                            
                            # Record the intermediate step
                            state["intermediate_steps"].append({
                                "tool": tool_name,
                                "args": parsed_args,
                                "result": tool_result
                            })
                            
                            # Decrement steps remaining
                            state["steps_remaining"] -= 1
                            
                            return {
                                "question": state["question"],
                                "messages": messages,
                                "steps_remaining": state["steps_remaining"],
                                "answer": state["answer"],
                                "intermediate_steps": state["intermediate_steps"]
                            }
                        except Exception as e:
                            print(f"[Agent] Error executing tool {tool_name}: {e}")
                            messages.append(AIMessage(content=response.content))
                            messages.append(ToolMessage(content=f"Error executing tool {tool_name}: {e}", tool_call_id=tool_name))
                            state["steps_remaining"] -= 1
                            return state
            
            # If no tool call identified, add the response and continue
            messages.append(AIMessage(content=response.content))
            state["steps_remaining"] -= 1
            
            return {
                "question": state["question"],
                "messages": messages,
                "steps_remaining": state["steps_remaining"],
                "answer": state["answer"],
                "intermediate_steps": state["intermediate_steps"]
            }
        
        def finalize_answer(state: AgentState) -> Dict[str, Any]:
            """Extract the final answer from the agent's reasoning."""
            messages = state["messages"]
            answer = ""
            
            # Check if the answer is already in the messages
            for message in reversed(messages):
                if not isinstance(message, HumanMessage) and "<ANSWER>" in message.content:
                    # Extract the answer from between <ANSWER> tags
                    start_idx = message.content.find("<ANSWER>") + len("<ANSWER>")
                    end_idx = message.content.find("</ANSWER>", start_idx)
                    if end_idx > start_idx:
                        answer = message.content[start_idx:end_idx].strip()
                        break
            
            # If no answer found, generate one final answer
            if not answer:
                print("[Agent] No answer found, generating final answer...")
                messages.append(HumanMessage(content="Based on all the information we've gathered, what is the final answer? Remember to format it as <ANSWER>your answer here</ANSWER>"))
                
                final_response = self.llm.invoke(messages)
                messages.append(final_response)
                
                # Extract the answer
                content = final_response.content
                if "<ANSWER>" in content:
                    start_idx = content.find("<ANSWER>") + len("<ANSWER>")
                    end_idx = content.find("</ANSWER>", start_idx)
                    if end_idx > start_idx:
                        answer = content[start_idx:end_idx].strip()
                else:
                    print("[Agent] Warning: Final response didn't contain <ANSWER> tags")
                    answer = content.strip()
            
            print(f"[Agent] Final answer: {answer}")
            return {
                "question": state["question"],
                "messages": messages,
                "steps_remaining": state["steps_remaining"],
                "answer": answer,
                "intermediate_steps": state["intermediate_steps"]
            }
        
        # Build the graph
        workflow = StateGraph(AgentState)
        workflow.add_node("setup", setup_agent)
        workflow.add_node("process_agent", process_agent)
        workflow.add_node("finalize_answer", finalize_answer)
        
        # Add edges
        workflow.set_entry_point("setup")
        workflow.add_conditional_edges(
            "setup",
            route_agent,
            {
                "process_agent": "process_agent",
                "finalize_answer": "finalize_answer"
            }
        )
        workflow.add_conditional_edges(
            "process_agent",
            route_agent,
            {
                "process_agent": "process_agent",
                "finalize_answer": "finalize_answer"
            }
        )
        workflow.add_edge("finalize_answer", END)
        
        self.graph = workflow.compile()
        print("ImprovedLangGraphAgent initialized successfully.")

    def __call__(self, question: str) -> str:
        """Process a question and return the answer."""
        print(f"[ImprovedLangGraphAgent] Processing question: {question[:75]}...")
        output = self.graph.invoke({
            "question": question,
            "messages": [],
            "steps_remaining": 5,
            "answer": "",
            "intermediate_steps": []
        })
        
        answer = output["answer"].strip()
        if not answer:
            print("[ImprovedLangGraphAgent] Warning: Empty answer, falling back to last message content")
            # If answer is empty, try to extract it from the last message
            last_message = output["messages"][-1]
            if isinstance(last_message, AIMessage):
                answer = last_message.content.strip()
        
        return answer

# Update the run_and_submit_all function to use our improved agent
def run_and_submit_all(profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the ImprovedLangGraphAgent on them, submits all answers,
    and displays the results.
    """
    # --- Determine HF Space Runtime URL and Repo URL ---
    space_id = os.getenv("SPACE_ID") # Get the SPACE_ID for sending link to the code

    if profile:
        username = f"{profile.username}"
        print(f"User logged in: {username}")
    else:
        print("User not logged in.")
        return "Please Login to Hugging Face with the button.", None

    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"

    # 1. Instantiate Agent
    try:
        agent = ImprovedLangGraphAgent()
    except Exception as e:
        print(f"Error instantiating agent: {e}")
        return f"Error initializing agent: {e}", None
    
    # In the case of an app running as a hugging Face space, this link points toward your codebase
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
    print(agent_code)

    # 2. Fetch Questions
    print(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
             print("Fetched questions list is empty.")
             return "Fetched questions list is empty or invalid format.", None
        print(f"Fetched {len(questions_data)} questions.")
    except requests.exceptions.RequestException as e:
        print(f"Error fetching questions: {e}")
        return f"Error fetching questions: {e}", None
    except requests.exceptions.JSONDecodeError as e:
         print(f"Error decoding JSON response from questions endpoint: {e}")
         print(f"Response text: {response.text[:500]}")
         return f"Error decoding server response for questions: {e}", None
    except Exception as e:
        print(f"An unexpected error occurred fetching questions: {e}")
        return f"An unexpected error occurred fetching questions: {e}", None

    # 3. Run the Agent
    results_log = []
    answers_payload = []
    print(f"Running agent on {len(questions_data)} questions...")
    for item in questions_data:
        task_id = item.get("task_id")
        question_text = item.get("question")
        if not task_id or question_text is None:
            print(f"Skipping item with missing task_id or question: {item}")
            continue
        try:
            print(f"\n--- Processing Task ID: {task_id} ---")
            submitted_answer = agent(question_text)
            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
            print(f"--- Task ID: {task_id} completed ---\n")
        except Exception as e:
             print(f"Error running agent on task {task_id}: {e}")
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})

    if not answers_payload:
        print("Agent did not produce any answers to submit.")
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    # 4. Prepare Submission 
    submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
    status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
    print(status_update)

    # 5. Submit
    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
    try:
        response = requests.post(submit_url, json=submission_data, timeout=60)
        response.raise_for_status()
        result_data = response.json()
        final_status = (
            f"Submission Successful!\n"
            f"User: {result_data.get('username')}\n"
            f"Overall Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
            f"Message: {result_data.get('message', 'No message received.')}"
        )
        print("Submission successful.")
        results_df = pd.DataFrame(results_log)
        return final_status, results_df
        
    except requests.exceptions.HTTPError as e:
        error_detail = f"Server responded with status {e.response.status_code}."
        try:
            error_json = e.response.json()
            error_detail += f" Detail: {error_json.get('detail', e.response.text)}"
        except requests.exceptions.JSONDecodeError:
            error_detail += f" Response: {e.response.text[:500]}"
        status_message = f"Submission Failed: {error_detail}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
        
    except requests.exceptions.Timeout:
        status_message = "Submission Failed: The request timed out."
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
        
    except requests.exceptions.RequestException as e:
        status_message = f"Submission Failed: Network error - {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
        
    except Exception as e:
        status_message = f"An unexpected error occurred during submission: {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df


# --- Build Gradio Interface using Blocks ---
with gr.Blocks() as demo:
    gr.Markdown("# Basic Agent Evaluation Runner")
    gr.Markdown(
        """
        **Instructions:**
        1.  Please clone this space, then modify the code to define your agent's logic, the tools, the necessary packages, etc ...
        2.  Log in to your Hugging Face account using the button below. This uses your HF username for submission.
        3.  Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.
        ---
        **Disclaimers:**
        Once clicking on the "submit button, it can take quite some time ( this is the time for the agent to go through all the questions).
        This space provides a basic setup and is intentionally sub-optimal to encourage you to develop your own, more robust solution. For instance for the delay process of the submit button, a solution could be to cache the answers and submit in a seperate action or even to answer the questions in async.
        """
    )

    with gr.Row():
        login_button = gr.LoginButton()

    run_button = gr.Button("Run Evaluation & Submit All Answers")

    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)

    run_button.click(
        fn=run_and_submit_all,
        outputs=[status_output, results_table]
    )

if __name__ == "__main__":
    print("\n" + "-"*30 + " App Starting " + "-"*30)
    # Check for SPACE_HOST and SPACE_ID at startup for information
    space_host_startup = os.getenv("SPACE_HOST")
    space_id_startup = os.getenv("SPACE_ID") # Get SPACE_ID at startup

    if space_host_startup:
        print(f"✅ SPACE_HOST found: {space_host_startup}")
        print(f"   Runtime URL should be: https://{space_host_startup}.hf.space")
    else:
        print("ℹ️  SPACE_HOST environment variable not found (running locally?).")

    if space_id_startup: # Print repo URLs if SPACE_ID is found
        print(f"✅ SPACE_ID found: {space_id_startup}")
        print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
        print(f"   Repo Tree URL: https://huggingface.co/spaces/{space_id_startup}/tree/main")
    else:
        print("ℹ️  SPACE_ID environment variable not found (running locally?). Repo URL cannot be determined.")

    print("-"*(60 + len(" App Starting ")) + "\n")

    print("Launching Gradio Interface for Basic Agent Evaluation...")
    demo.launch(debug=True, share=True)