app.py

import streamlit as st

from typing import Optional
from smolagents import (
    CodeAgent,
    ToolCallingAgent,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    HfApiModel
)


system_prompt = """
You are a specialist tasked with finding films and TV shows using a web search engine, similar to an IMDB expert.
Your job is to determine the best match based on the user’s provided description, which may also include cast details and/or the year of release.
You have access to several search teams:
- One team searches by description only.
- A second team searches using description plus approximate cast details.
- A third team uses description with approximate year of release.
- A fourth team combines description, cast, and year of release for a comprehensive search.

Follow these steps:
1. Identify which details the user has provided (description, cast, release year).
2. Direct the query to the corresponding search team to obtain the best results.
3. If only a description is given, focus solely on that; if additional details are available, refine the search accordingly.
4. Retrieve the film or TV show title that best matches the query, or a list of the most relevant titles if multiple strong candidates exist.
5. Finally, output only the title(s) of the film or TV show as your answer.
"""

def agent_init(
    together_api_token: str, 
    description: str, 
    content_type: Optional[str] = None, 
    additional_info:  Optional[str] = None, 
    cast_or_year_of_release: Optional[str] = None
):
    model = HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", provider="together", max_tokens=4096, token=together_api_token)

    agent_vox114 = ToolCallingAgent(
        model=model,
        tools=[VisitWebpageTool(), DuckDuckGoSearchTool()],
        max_steps=10,
        planning_interval=4,
        name="film_seracher",
        description="""
        This agent specializes in finding films or TV shows based on a provided description, which may include additional details like cast names and/or the year of release.
        If the user supplies extra context (e.g., specific actors or release year), the agent will refine the search accordingly.
        In cases where an exact match isn’t found, it will return a list of the most relevant results.
        Ensure that the user’s query is a full sentence with clear details to maximize search accuracy.
        """
    )
    
    agent_vox114.prompt_templates['managed_agent']['task'] += system_prompt
    
    main_agent = CodeAgent(
        model=model,
        tools=[],
        additional_authorized_imports=["re", "datetime", "json"],
        managed_agents=[agent_vox114],
        max_steps=10,
        planning_interval=4,
    )

    if content_type is None:
        content_type = "film or TV show"

    if additional_info is None:
        additional_info = ""

    if additional_info is None and cast_or_year_of_release is None:
        cast_or_year_of_release = ""

    if (additional_info and cast_or_year_of_release is None) or (additional_info is None and cast_or_year_of_release):
        st.error("additional_info or cast_or_year_of_release is None, you cannot use one of them they are work in pair.")
        return None # Return None to indicate an error

    vox114 = main_agent.run(f"Find the {content_type} by description {additional_info}: {description}. {cast_or_year_of_release}")

    return vox114


def main():
    st.set_page_config(
        page_title="Film Finder AI Agent",
        page_icon="🎬",
        layout="wide",
        initial_sidebar_state="expanded",
    )

    st.title("🎬 Film & TV Show Finder AI Agent")
    st.markdown("Enter your API token and film/show details below to find what you're looking for!")

    # Sidebar for API Token
    with st.sidebar:
        st.header("API Token")
        together_api_token = st.text_input("Together API Token", type="password", placeholder="Enter your API token here")
        st.markdown("Get your API token from [Together AI](https://together.ai/)", unsafe_allow_html=True)

    # Main content area
    with st.container():
        col1, col2 = st.columns([3, 2]) # Adjust column ratio for better layout

        with col1:
            st.header("Search Criteria")
            description = st.text_area("Film/Show Description", placeholder="e.g., A group of wizards fighting a dark lord...", height=150)
            content_type = st.selectbox("Content Type", ["film", "TV Show"], index=0) # Film is default
            additional_info_expander = st.expander("Optional Details: Cast & Year of Release", expanded=False)
            with additional_info_expander:
                additional_info = st.text_input("Additional Information (e.g., plot points, genre)", placeholder="e.g., set in space, sci-fi")
                cast_or_year_of_release = st.text_input("Cast or Year of Release (e.g., starring Tom Hanks, released in 2010)", placeholder="e.g., starring Leonardo DiCaprio")

            search_button = st.button("🔍 Search for Film/Show", use_container_width=True)

        with col2:
            st.header("Search Results")
            if search_button:
                if not together_api_token:
                    st.error("Please enter your Together API token in the sidebar.")
                elif not description:
                    st.warning("Please provide a description of the film or TV show you are looking for.")
                else:
                    with st.spinner("Agent is searching for the best match..."):
                        result = agent_init(
                            together_api_token=together_api_token,
                            description=description,
                            content_type=content_type.lower(),
                            additional_info=additional_info if additional_info_expander.expanded else None, # Only pass if expander is opened
                            cast_or_year_of_release=cast_or_year_of_release if additional_info_expander.expanded else None # Only pass if expander is opened
                        )

                    if result:
                        st.success("Search Results Found!")
                        st.write(result)
                    else:
                        st.warning("No results found or an error occurred. Please check your input and API token.")
            else:
                st.info("Enter your search criteria and click 'Search for Film/Show' to begin.")

    # Footer with Agent Details (optional, for extra polish)
    st.markdown("---")
    st.caption("Powered by Smolagents by Hugging Face")
    # st.caption("System Prompt (for reference):")
    # st.code(system_prompt, language='python') # Show system prompt in a code block for transparency


if __name__ == "__main__":
    main()