Update app.py

app.py

@@ -1,7 +1,17 @@
 # unsloth_doc_app.py
 import streamlit as st
 import os
-from smolagents import CodeAgent, Tool
+from smolagents import (
+    CodeAgent,
+    ToolCallingAgent,
+    HfApiModel,
+    DuckDuckGoSearchTool,
+    LiteLLMModel,
+    tool,
+    VisitWebpageTool,
+    GradioUI,
+    PythonInterpreterTool
+)
 import requests
 import json
 from typing import List, Dict, Any
@@ -19,43 +29,6 @@ AUTHORIZED_IMPORTS: List[str] = [
     'pandas', 'numpy', 'requests'
 ]
 
-# Tools
-class DuckDuckGoSearchTool(Tool):
-    name = "duckduckgo_search"
-    description = "Searches DuckDuckGo for Unsloth-related queries."
-    inputs = {"query": {"type": "string", "description": "The search query"}}
-    output_type = "string"
-
-    def __init__(self):
-        super().__init__()
-        self.url = "https://api.duckduckgo.com/"
-        self.headers = {'Content-Type': 'application/json'}
-
-    def forward(self, query: str) -> str:
-        payload = {"q": f"unsloth {query}", "format": "json", "no_html": 1}
-        response = requests.get(self.url, params=payload, headers=self.headers)
-        if response.status_code == 200:
-            return json.dumps(response.json(), indent=2)
-        return "Error: Unable to fetch search results."
-
-class VisitWebpageTool(Tool):
-    name = "visit_webpage"
-    description = "Fetches content from a specific webpage."
-    inputs = {"url": {"type": "string", "description": "The URL to visit"}}
-    output_type = "string"
-
-    def __init__(self):
-        super().__init__()
-        self.headers = {'User-Agent': 'Mozilla/5.0'}
-
-    def forward(self, url: str) -> str:
-        try:
-            response = requests.get(url, headers=self.headers)
-            response.raise_for_status()
-            return response.text[:2000]  # Limit to first 2000 characters
-        except Exception as e:
-            return f"Error visiting webpage: {e}"
-
 # Agent Initialization
 def initialize_unsloth_agent(model=None):
     """Initialize the Unsloth Documentation Agent"""
@@ -65,8 +38,7 @@ def initialize_unsloth_agent(model=None):
     ]
     if model is None:
         from smolagents import HfApiModel
-        model = HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct",provider="together",)
-
+        model = HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", provider="together")
     
     return CodeAgent(
         name="unsloth_agent",
@@ -128,22 +100,24 @@ def main():
         placeholder="E.g., How to fine-tune with Unsloth"
     )
 
-    if st.button("Search", type="primary", key="search_button_with_query") and search_query:
-        with st.spinner("Searching Unsloth documentation..."):
-            try:
-                results = st.session_state.agent.run(search_query)
-                display_results(results)
-                
-                st.markdown("---")
-                st.info("""
-                💡 **How to read the results:**
-                - Results may include JSON from searches or webpage snippets.
-                - Check the sidebar for tips on refining your query.
-                """)
-            except Exception as e:
-                st.error(f"An error occurred: {e}")
-    elif st.button("Search", type="primary", key="search_button_no_query"):
-        st.warning("Please enter a search query.")
+    # Single button to handle both cases
+    if st.button("Search", type="primary", key="search_button"):
+        if search_query:
+            with st.spinner("Searching Unsloth documentation..."):
+                try:
+                    results = st.session_state.agent.run(search_query, additional_args={"sources" : ["https://docs.unsloth.ai/"],"instructions": "Response to the user answer as a structured documentation format."})
+                    display_results(results)
+                    
+                    st.markdown("---")
+                    st.info("""
+                    💡 **How to read the results:**
+                    - Results may include JSON from searches or webpage snippets.
+                    - Check the sidebar for tips on refining your query.
+                    """)
+                except Exception as e:
+                    st.error(f"An error occurred: {e}")
+        else:
+            st.warning("Please enter a search query.")
 
     st.markdown("---")
     st.caption("Powered by SmolAgents and Unsloth")