SamarthPujari/First_agent_template

1from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool, Tool2import datetime3import requests4import pytz5import yaml6import os7from tools.final_answer import FinalAnswerTool8from Gradio_UI import GradioUI9import fitz  # PyMuPDF10from sentence_transformers import SentenceTransformer, util11from transformers import pipeline12from PIL import Image13import io14 15# API Key for weather16API_KEY = os.getenv("Weather_Token")17 18# -------------------- TOOL 1: Get Weather --------------------19@tool20def get_current_weather(place: str) -> str:21    """22    A tool that fetches the current weather of a particular place.23    Args:24        place (str): A string representing a valid place (e.g., 'London/Paris').25    Returns:26        str: Weather description including condition, temperature, humidity, and wind speed.27    """28    api_key = API_KEY29    url = "https://api.openweathermap.org/data/2.5/weather"30    params = {31        "q": place,32        "appid": api_key,33        "units": "metric"34    }35 36    try:37        response = requests.get(url, params=params)38        data = response.json()39 40        if response.status_code == 200:41            weather_desc = data["weather"][0]["description"]42            temperature = data["main"]["temp"]43            humidity = data["main"]["humidity"]44            wind_speed = data["wind"]["speed"]45 46            return (47                f"Weather in {place}:\n"48                f"- Condition: {weather_desc}\n"49                f"- Temperature: {temperature}°C\n"50                f"- Humidity: {humidity}%\n"51                f"- Wind Speed: {wind_speed} m/s"52            )53        else:54            return f"Error: {data['message']}"55    except Exception as e:56        return f"Error fetching weather data for '{place}': {str(e)}"57 58 59# -------------------- TOOL 2: Get Time --------------------60@tool61def get_current_time_in_timezone(timezone: str) -> str:62    """63    A tool that fetches the current local time in a specified timezone.64    Args:65        timezone (str): A string representing a valid timezone (e.g., 'America/New_York').66    Returns:67        str: The current local time formatted as a string.68    """69    try:70        tz = pytz.timezone(timezone)71        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")72        return f"The current local time in {timezone} is: {local_time}"73    except Exception as e:74        return f"Error fetching time for timezone '{timezone}': {str(e)}"75 76 77# -------------------- TOOL 3: Document QnA --------------------78embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")79qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")80 81@tool82def document_qna_tool(pdf_path: str, question: str) -> str:83    """84    A tool that answers natural language questions about a given PDF document.85    Args:86        pdf_path (str): Path to the local PDF file.87        question (str): Question about the content of the PDF.88    Returns:89        str: Answer to the question based on the content.90    """91    import os, fitz, traceback92    from sentence_transformers import SentenceTransformer, util93    from transformers import pipeline94 95    try:96        print(f"[DEBUG] PDF Path: {pdf_path}")97        print(f"[DEBUG] Question: {question}")98 99        if not os.path.exists(pdf_path):100            return f"[ERROR] File not found: {pdf_path}"101 102        print("[DEBUG] Opening PDF...")103        try:104            doc = fitz.open(pdf_path)105        except RuntimeError as e:106            return f"[ERROR] Could not open PDF. It may be corrupted or encrypted. Details: {str(e)}"107 108        text_chunks = []109        for page in doc:110            text = page.get_text()111            if text.strip():112                text_chunks.append(text)113        doc.close()114 115        if not text_chunks:116            return "[ERROR] No readable text in the PDF."117 118        print(f"[DEBUG] Extracted {len(text_chunks)} text chunks.")119 120        embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")121        embeddings = embedding_model.encode(text_chunks, convert_to_tensor=True)122        question_embedding = embedding_model.encode(question, convert_to_tensor=True)123 124        print("[DEBUG] Performing semantic search...")125        scores = util.pytorch_cos_sim(question_embedding, embeddings)[0]126        best_match_idx = scores.argmax().item()127        best_context = text_chunks[best_match_idx]128 129        qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")130        prompt = f"Context: {best_context}\nQuestion: {question}"131        print("[DEBUG] Calling QA model...")132        answer = qa_pipeline(prompt, max_new_tokens=500)[0]['generated_text']133 134        return f"Answer: {answer.strip()}"135 136    except Exception as e:137        return f"[EXCEPTION] {type(e).__name__}: {str(e)}\n{traceback.format_exc()}"138 139# -------------------- TOOL 4: Image Generation --------------------140@tool141def generate_image(prompt: str) -> Image.Image:142    """143    A tool that generates an image from a text prompt using a Hugging Face Space.144    Args:145        prompt (str): The text description of the image to generate.146    Returns:147        Image.Image: The generated image as a PIL Image object.148    """149    try:150        # Use the hosted image generation model from Hugging Face Spaces151        image_generator = Tool.from_space(152            "black-forest-labs/FLUX.1-schnell",153            name="image_generator",154            description="Generate an image from a prompt"155        )156 157        # Call the model with the prompt158        result = image_generator(prompt=prompt)159 160        # If the result is bytes → convert to PIL.Image161        if isinstance(result, bytes):162            return Image.open(io.BytesIO(result))163 164        # If the result is already a PIL Image165        if isinstance(result, Image.Image):166            return result167 168        # If the model gave back a path, open it169        if isinstance(result, str):170            return Image.open(result)171 172        raise ValueError("Unexpected output format from image generator.")173 174    except Exception as e:175        raise RuntimeError(f"Error generating image: {str(e)}")176 177# -------------------- Other Components --------------------178final_answer = FinalAnswerTool()179search_tool = DuckDuckGoSearchTool()180 181model = HfApiModel(182    max_tokens=2096,183    temperature=0.5,184    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',185    custom_role_conversions=None,186)187 188with open("prompts.yaml", 'r') as stream:189    prompt_templates = yaml.safe_load(stream)190 191agent = CodeAgent(192    model=model,193    tools=[194        get_current_time_in_timezone,195        get_current_weather,196        generate_image,197        search_tool,198        document_qna_tool,199        final_answer200    ],201    max_steps=6,202    verbosity_level=1,203    grammar=None,204    planning_interval=None,205    name=None,206    description=None,207    prompt_templates=prompt_templates208)209 210GradioUI(agent).launch()