tools.py

import os
import base64
import tempfile
import requests

from typing import Type, Dict, Any

from pydantic import BaseModel, Field, PrivateAttr

from youtube_transcript_api import YouTubeTranscriptApi

from langchain.tools import BaseTool, tool
from langchain_community.document_loaders import UnstructuredExcelLoader
from langchain_community.llms import HuggingFaceEndpoint


class ImageRecognitionTool(BaseTool):
    name: str = "image_recognition"
    description: str = (
        "Recognizes the contents of a base64-encoded image using a Qwen-VL model. "
        "Input should be a dictionary with keys: 'image' (base64 string), and optionally 'prompt'."
    )
    _hf_endpoint: HuggingFaceEndpoint = PrivateAttr()

    def __init__(self, hf_endpoint: HuggingFaceEndpoint):
        super().__init__()
        self._hf_endpoint = hf_endpoint

    def _run(self, input: Dict[str, Any]) -> str:
        image = input.get("image")
        if not image:
            raise ValueError("Input must include 'image' as a base64-encoded string.")
        prompt = input.get("prompt", "What is in this image?")

        # The Qwen-VL model expects {"image": base64, "prompt": string}
        payload = {
            "image": image,
            "prompt": prompt
        }

        result = self.hf_endpoint.invoke(payload)

        # Handle different possible formats
        if isinstance(result, dict):
            return result.get("generated_text", str(result))
        return str(result)

@tool
def fetch_and_process_file(task_id: str) -> str:
    """
    Fetches and processes a file from the given task_id URL.

    - .py  → returns the full code as a string.
    - .png → returns base64-encoded string.
    - .xlsx → returns extracted content using UnstructuredExcelLoader.
    - .mp3 → transcribes audio using Hugging Face Whisper model.
    """
    url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
    response = requests.get(url)
    if response.status_code != 200:
        return f"Error fetching file: HTTP {response.status_code}"

    content_disposition = response.headers.get("Content-Disposition", "")
    filename = content_disposition.split("filename=")[-1].strip('"')
    extension = filename.split(".")[-1].lower()

    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{extension}") as tmp:
        tmp.write(response.content)
        tmp_path = tmp.name

    try:
        if extension == "py":
            # Return code as a single string
            with open(tmp_path, "r", encoding="utf-8") as f:
                return "Python file content:\n" + f.read()

        elif extension in ("png", "jpg", "jpeg"):
            # Base64 encode image
            with open(tmp_path, "rb") as f:
                return "Image base64:\n" + base64.b64encode(f.read()).decode("utf-8")

        elif extension == "xlsx":
            loader = UnstructuredExcelLoader(tmp_path)
            docs = loader.load()
            return "Excel file content:\n" + "\n".join(doc.page_content for doc in docs)

        elif extension == "mp3":
            whisper_endpoint = HuggingFaceEndpoint(
                repo_id="openai/whisper-base",
                task="automatic-speech-recognition"
            )

            with open(tmp_path, "rb") as f:
                audio_b64 = base64.b64encode(f.read()).decode()

            # Call the HuggingFaceEndpoint with base64 audio input
            transcription = whisper_endpoint.invoke(audio_b64)
            
            # The response is normally a dict or string, might need adjustment
            if isinstance(transcription, dict) and "text" in transcription:
                return transcription["text"]

            return "MP3 transcription:\n" + str(transcription)

        else:
            return f"Unsupported file type: .{extension}"

    finally:
        os.remove(tmp_path)




class YouTubeTranscriptInput(BaseModel):
    url: str = Field(..., description="A YouTube video URL")

class GetYouTubeTranscriptTool(BaseTool):
    name: str = "get_youtube_transcript"
    description: str = "Fetches transcript from a YouTube video, if available."
    args_schema: Type[BaseModel] = YouTubeTranscriptInput

    def _run(self, url: str) -> str:
        try:
            video_id = url.split("v=")[-1].split("&")[0]
            transcript_list = YouTubeTranscriptApi.get_transcript(video_id)
            transcript = " ".join([t["text"] for t in transcript_list])
            return transcript
        except Exception as e:
            return f"Failed to get transcript: {str(e)}"

    def _arun(self, url: str) -> str:
        raise NotImplementedError("Async not supported yet.")