Update app.py

app.py

@@ -1,79 +1,76 @@
-from smolagents import CodeAgent, DuckDuckGoSearchTool, load_tool, tool
+from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
 import datetime
 import requests
 import pytz
 import yaml
 import os
-from huggingface_hub import login
 from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
 import fitz  # PyMuPDF
 from sentence_transformers import SentenceTransformer, util
 from transformers import pipeline
-from diffusers import StableDiffusionPipeline
-import torch
 
 # API Key for weather
 API_KEY = os.getenv("Weather_Token")
 
-hf_token = os.getenv("HF_TOKEN")
-login(token=hf_token)
-
-sd_pipe = StableDiffusionPipeline.from_pretrained(
-    "rupeshs/LCM-runwayml-stable-diffusion-v1-5",
-    use_auth_token=hf_token,
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32
-).to(device)
-
 # -------------------- TOOL 1: Get Weather --------------------
 @tool
 def get_current_weather(place: str) -> str:
     """
-    Get the current weather for a given location.
+    A tool that fetches the current weather of a particular place.
     Args:
-        place (str): Name of the city or location (e.g., "London" or "New York").
+        place (str): A string representing a valid place (e.g., 'London/Paris').
     Returns:
-        str: Weather condition, temperature, humidity, and wind speed.
+        str: Weather description including condition, temperature, humidity, and wind speed.
     """
+    api_key = API_KEY
     url = "https://api.openweathermap.org/data/2.5/weather"
     params = {
         "q": place,
-        "appid": API_KEY,
+        "appid": api_key,
         "units": "metric"
     }
+
     try:
         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: {data['weather'][0]['description']}\n"
-                f"- Temperature: {data['main']['temp']}°C\n"
-                f"- Humidity: {data['main']['humidity']}%\n"
-                f"- Wind Speed: {data['wind']['speed']} m/s"
+                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.get('message', 'Unknown error')}"
+            return f"Error: {data['message']}"
     except Exception as e:
-        return f"Error fetching weather data: {str(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:
     """
-    Get the current local time in a given timezone.
+    A tool that fetches the current local time in a specified timezone.
     Args:
-        timezone (str): Timezone string in the format 'Region/City', 
-                        e.g., "America/New_York".
+        timezone (str): A string representing a valid timezone (e.g., 'America/New_York').
     Returns:
-        str: Formatted local time string.
+        str: The current local time formatted as a string.
     """
     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: {str(e)}"
+        return f"Error fetching time for timezone '{timezone}': {str(e)}"
+
 
 # -------------------- TOOL 3: Document QnA --------------------
 embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
@@ -82,104 +79,99 @@ qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
 @tool
 def document_qna_tool(pdf_path: str, question: str) -> str:
     """
-    Answer a natural language question based on the content of a PDF document.
+    A tool that answers natural language questions about a given PDF document.
     Args:
         pdf_path (str): Path to the local PDF file.
-        question (str): Natural language question to ask about the PDF.
+        question (str): Question about the content of the PDF.
     Returns:
-        str: Answer generated from the most relevant section of the document.
+        str: Answer to the question based on the content.
     """
+    import os, fitz, traceback
+    from sentence_transformers import SentenceTransformer, util
+    from transformers import pipeline
+
     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()]
+
+        print("[DEBUG] Opening PDF...")
+        try:
+            doc = fitz.open(pdf_path)
+        except RuntimeError as e:
+            return f"[ERROR] Could not open PDF. It may be corrupted or encrypted. Details: {str(e)}"
+
+        text_chunks = []
+        for page in doc:
+            text = page.get_text()
+            if text.strip():
+                text_chunks.append(text)
         doc.close()
+
         if not text_chunks:
             return "[ERROR] No readable text in the PDF."
 
+        print(f"[DEBUG] Extracted {len(text_chunks)} text chunks.")
+
+        embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
         embeddings = embedding_model.encode(text_chunks, convert_to_tensor=True)
         question_embedding = embedding_model.encode(question, convert_to_tensor=True)
+
+        print("[DEBUG] Performing semantic search...")
         scores = util.pytorch_cos_sim(question_embedding, embeddings)[0]
-        best_context = text_chunks[scores.argmax().item()]
+        best_match_idx = scores.argmax().item()
+        best_context = text_chunks[best_match_idx]
 
+        qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
         prompt = f"Context: {best_context}\nQuestion: {question}"
+        print("[DEBUG] Calling QA model...")
         answer = qa_pipeline(prompt, max_new_tokens=500)[0]['generated_text']
+
         return f"Answer: {answer.strip()}"
+
     except Exception as e:
-        return f"[EXCEPTION] {type(e).__name__}: {str(e)}"
+        return f"[EXCEPTION] {type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
 
-# -------------------- TOOL 4: Local Image Generation --------------------
-@tool
-def image_generator(prompt: str) -> str:
-    """
-    Generate an image from a given text prompt using Stable Diffusion.
-    Args:
-        prompt (str): Description of the image to generate.
-    Returns:
-        str: Path to the saved generated image.
-    """
-    image = sd_pipe(prompt).images[0]
-    output_path = "generated_image.png"
-    image.save(output_path)
-    return f"Image saved at {output_path}"
-
-# -------------------- Local LLM (Replaces HfApiModel) --------------------
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-class LocalModel:
-    def __init__(self):
-        model_name = "openlm-research/open_llama_3b"
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="auto" if torch.cuda.is_available() else None,
-        )
-
-    def generate(self, prompt, max_new_tokens=500, **kwargs):
-        """
-        Generate text from the given prompt.
-        Extra kwargs like 'stop_sequences' are accepted for compatibility.
-        """
-        stop_sequences = kwargs.pop("stop_sequences", None)
-
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-        output = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
-        text = self.tokenizer.decode(output[0], skip_special_tokens=True)
-
-        # If stop_sequences provided, truncate output
-        if stop_sequences:
-            for stop in stop_sequences:
-                if stop in text:
-                    text = text.split(stop)[0]
-                    break
-
-        return text
-
-    def __call__(self, prompt, **kwargs):
-        return self.generate(prompt, **kwargs)
-
-# -------------------- Agent Setup --------------------
+# -------------------- 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,
+)
+
+from smolagents import Tool
+
+image_generation_tool = Tool.from_space(
+    "black-forest-labs/FLUX.1-schnell",
+    name="image_generator",  # You can name it whatever makes sense for your agent
+    description="Generate an image from a prompt"
+)
+
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
 
-model = LocalModel()
 agent = CodeAgent(
     model=model,
     tools=[
         get_current_time_in_timezone,
         get_current_weather,
-        image_generator,
+        image_generation_tool,
         search_tool,
-        document_qna_tool,
+        document_qna_tool,  # ← New Tool Added
         final_answer
     ],
     max_steps=6,
     verbosity_level=1,
+    grammar=None,
+    planning_interval=None,
+    name=None,
+    description=None,
     prompt_templates=prompt_templates
 )