app.py

from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool
from smolagents.tools import Tool
import datetime
import requests
import pytz
import yaml
import os
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
import fitz  # PyMuPDF
from sentence_transformers import SentenceTransformer, util
from transformers import pipeline

# API Key for weather
API_KEY = os.getenv("Weather_Token")

# -------------------- TOOL 1: Get Weather --------------------
@Tool
def get_current_weather(place: str) -> str:
    try:
        url = "https://api.openweathermap.org/data/2.5/weather"
        params = {
            "q": place,
            "appid": API_KEY,
            "units": "metric"
        }

        response = requests.get(url, params=params)
        data = response.json()

        if response.status_code == 200:
            weather_desc = data["weather"][0]["description"]
            temperature = data["main"]["temp"]
            humidity = data["main"]["humidity"]
            wind_speed = data["wind"]["speed"]

            return (
                f"Weather in {place}:\n"
                f"- Condition: {weather_desc}\n"
                f"- Temperature: {temperature}°C\n"
                f"- Humidity: {humidity}%\n"
                f"- Wind Speed: {wind_speed} m/s"
            )
        else:
            return f"Error: {data['message']}"
    except Exception as e:
        return f"Error fetching weather data for '{place}': {str(e)}"


# -------------------- TOOL 2: Get Time --------------------
@Tool
def get_current_time_in_timezone(timezone: str) -> str:
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


# -------------------- TOOL 3: Document QnA --------------------
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")

def document_qna_tool_fn(pdf_path: str, question: str) -> str:
    import traceback

    try:
        print(f"[DEBUG] PDF Path: {pdf_path}")
        print(f"[DEBUG] Question: {question}")

        if not os.path.exists(pdf_path):
            return f"[ERROR] File not found: {pdf_path}"

        doc = fitz.open(pdf_path)
        text_chunks = [page.get_text() for page in doc if page.get_text().strip()]
        doc.close()

        if not text_chunks:
            return "[ERROR] No readable text in the PDF."

        embeddings = embedding_model.encode(text_chunks, convert_to_tensor=True)
        question_embedding = embedding_model.encode(question, convert_to_tensor=True)

        scores = util.pytorch_cos_sim(question_embedding, embeddings)[0]
        if scores.shape[0] == 0:
            return "[ERROR] No semantic matches found in PDF text."

        best_match_idx = scores.argmax().item()
        best_context = text_chunks[best_match_idx]

        prompt = f"Context: {best_context}\nQuestion: {question}"
        answer = qa_pipeline(prompt, max_new_tokens=100)[0]['generated_text']

        return f"Answer: {answer.strip()}"

    except Exception as e:
        return f"[EXCEPTION] {type(e).__name__}: {str(e)}\n{traceback.format_exc()}"

# Wrap the function using Tool(...)
document_qna_tool = Tool(
    name="document_qna_tool",
    description="Answer questions about a PDF document.",
    func=document_qna_tool_fn
)


# -------------------- Other Components --------------------
final_answer = FinalAnswerTool()
search_tool = DuckDuckGoSearchTool()

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[
        get_current_time_in_timezone,
        get_current_weather,
        image_generation_tool,
        search_tool,
        document_qna_tool,  # Correctly wrapped tool
        final_answer
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

print("[DEBUG] Registered Tools:")
for t in agent.tools:
    print(f" - {getattr(t, 'name', str(t))}")

GradioUI(agent).launch()