# unsloth_doc_app.py
import streamlit as st
import os
from smolagents import (
    CodeAgent,
    ToolCallingAgent,
    HfApiModel,
    DuckDuckGoSearchTool,
    LiteLLMModel,
    tool,
    VisitWebpageTool,
    GradioUI,
    PythonInterpreterTool
)
import requests
import json
from typing import List, Dict, Any

# Configuration
PAGE_CONFIG = {
    "page_title": "Unsloth Documentation Search",
    "page_icon": "📘",
    "layout": "wide"
}

AUTHORIZED_IMPORTS: List[str] = [
    'random', 'collections', 'datetime', 'time', 'queue', 'unicodedata', 
    'os', 'stat', 'json', 'math', 're', 'itertools', 'statistics', 
    'pandas', 'numpy', 'requests'
]

# Agent Initialization
def initialize_unsloth_agent(model=None):
    """Initialize the Unsloth Documentation Agent"""
    tools = [
        DuckDuckGoSearchTool(),
        VisitWebpageTool()
    ]
    if model is None:
        from smolagents import HfApiModel
        model = HfApiModel(
            model_id="meta-llama/Llama-4-Scout-17B-16E-Instruct"
            # provider="together",
        )
    
    return CodeAgent(
        name="unsloth_agent",
        description="Unsloth Documentation AI Agent",
        tools=tools,
        model=model,
        max_steps=12,
        additional_authorized_imports=AUTHORIZED_IMPORTS
    )

# UI Components
def setup_ui():
    """Setup Streamlit UI components"""
    st.set_page_config(**PAGE_CONFIG)
    
    st.title("📘 Unsloth Documentation Search")
    st.markdown("""
    This tool uses an AI agent to help you explore Unsloth documentation.
    Ask questions about Unsloth features, usage, or troubleshooting.
    """)

def display_results(results: str):
    """Display search results in a formatted way"""
    st.markdown("### 📝 Results")
    st.markdown(results, unsafe_allow_html=True)

def setup_sidebar():
    """Setup sidebar content"""
    with st.sidebar:
        st.markdown("### About This Tool")
        st.markdown("""
        This agent specializes in Unsloth documentation, providing:
        - 📚 Quick answers from web searches
        - 🌐 Webpage content retrieval
        - 🤖 AI-powered insights
        """)
        st.markdown("### Tips")
        st.markdown("""
        - Use specific queries like "How to install Unsloth" or "Unsloth fine-tuning guide".
        - Results may include raw data from searches or webpages.
        """)

# Main App
@st.cache_resource
def initialize_app():
    """Initialize the application"""
    return initialize_unsloth_agent()

def main():
    setup_ui()
    setup_sidebar()
    
    if 'agent' not in st.session_state:
        with st.spinner("Initializing Unsloth Documentation Agent..."):
            st.session_state.agent = initialize_app()

    search_query = st.text_input(
        "🔍 Search Unsloth Documentation",
        placeholder="E.g., How to fine-tune with Unsloth"
    )

    # 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")

if __name__ == "__main__":
    main()