Update app.py

app.py

@@ -1,4 +1,4 @@
-from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
+from smolagents import CodeAgent, DuckDuckGoSearchTool, load_tool, tool
 import datetime
 import requests
 import pytz
@@ -9,6 +9,8 @@ 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")
@@ -16,61 +18,37 @@ API_KEY = os.getenv("Weather_Token")
 # -------------------- TOOL 1: Get Weather --------------------
 @tool
 def get_current_weather(place: str) -> str:
-    """
-    A tool that fetches the current weather of a particular place.
-    Args:
-        place (str): A string representing a valid place (e.g., 'London/Paris').
-    Returns:
-        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: {weather_desc}\n"
-                f"- Temperature: {temperature}°C\n"
-                f"- Humidity: {humidity}%\n"
-                f"- Wind Speed: {wind_speed} m/s"
+                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"
             )
         else:
-            return f"Error: {data['message']}"
+            return f"Error: {data.get('message', 'Unknown error')}"
     except Exception as e:
-        return f"Error fetching weather data for '{place}': {str(e)}"
-
+        return f"Error fetching weather data: {str(e)}"
 
 # -------------------- TOOL 2: Get Time --------------------
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
-    """
-    A tool that fetches the current local time in a specified timezone.
-    Args:
-        timezone (str): A string representing a valid timezone (e.g., 'America/New_York').
-    Returns:
-        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 for timezone '{timezone}': {str(e)}"
-
+        return f"Error fetching time: {str(e)}"
 
 # -------------------- TOOL 3: Document QnA --------------------
 embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
@@ -78,80 +56,63 @@ qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
 
 @tool
 def document_qna_tool(pdf_path: str, question: str) -> str:
-    """
-    A tool that answers natural language questions about a given PDF document.
-    Args:
-        pdf_path (str): Path to the local PDF file.
-        question (str): Question about the content of the PDF.
-    Returns:
-        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}"
-
-        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 = 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."
 
-        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_match_idx = scores.argmax().item()
-        best_context = text_chunks[best_match_idx]
+        best_context = text_chunks[scores.argmax().item()]
 
-        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)}"
 
+# -------------------- TOOL 4: Local Image Generation --------------------
+@tool
+def image_generator(prompt: str) -> str:
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        pipe = StableDiffusionPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5",
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32
+        ).to(device)
+        image = pipe(prompt).images[0]
+        output_path = "generated_image.png"
+        image.save(output_path)
+        return f"Image saved at {output_path}"
     except Exception as e:
-        return f"[EXCEPTION] {type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
+        return f"Image generation failed: {str(e)}"
 
-# -------------------- Other Components --------------------
-final_answer = FinalAnswerTool()
-search_tool = DuckDuckGoSearchTool()
+# -------------------- Local LLM (Replaces HfApiModel) --------------------
+from smolagents import LocalModel
 
-model = HfApiModel(
-    max_tokens=2096,
-    temperature=0.5,
-    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
-    custom_role_conversions=None,
-)
+class TransformersModel(LocalModel):
+    def __init__(self):
+        self.pipeline = pipeline(
+            "text-generation",
+            model="tiiuae/falcon-7b-instruct",
+            device_map="auto",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        )
 
-from smolagents import Tool
+    def generate(self, prompt, **kwargs):
+        result = self.pipeline(prompt, max_new_tokens=500, do_sample=True)
+        return result[0]['generated_text']
 
-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"
-)
+model = TransformersModel()
+
+# -------------------- Agent Setup --------------------
+final_answer = FinalAnswerTool()
+search_tool = DuckDuckGoSearchTool()
 
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
@@ -161,17 +122,13 @@ agent = CodeAgent(
     tools=[
         get_current_time_in_timezone,
         get_current_weather,
-        image_generation_tool,
+        image_generator,
         search_tool,
-        document_qna_tool,  # ← New Tool Added
+        document_qna_tool,
         final_answer
     ],
     max_steps=6,
     verbosity_level=1,
-    grammar=None,
-    planning_interval=None,
-    name=None,
-    description=None,
     prompt_templates=prompt_templates
 )