import streamlit as st
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, ToolCallingAgent, ManagedAgent, VisitWebpageTool


web_agent = ToolCallingAgent(
    tools=[DuckDuckGoSearchTool(max_results=10), VisitWebpageTool()],
    model=HfApiModel(),
)

managed_web_agent = ManagedAgent(
    agent=web_agent,
    name="search",
    description="Runs veterinarian web searches. Give it your query as an argument. Also, this agent should link to the sources you are using. If any specific sites are provided, this agent will only those sources.",
)

manager_agent = CodeAgent(
    tools=[],
    model=HfApiModel(model_id = "Qwen/Qwen2.5-Coder-32B-Instruct"),
    managed_agents=[managed_web_agent],
    additional_authorized_imports=['pyparsing', 'matplotlib', 'datetime', 'statistics', 'bs4', 'request', 'unicodedata', 'queue', 'time', 'collections', 're', 'math', 'stat', 'random', 
'itertools'],
)

# Function to log agent actions
def log_agent_action(prompt, result, agent_name):
    st.write(f"### Agent Activity ({agent_name}):")
    st.write("**Prompt Sent to Agent:**")
    st.code(prompt, language="text")
    st.write("**Agent Output:**")
    st.code(result, language="text")

# Streamlit app title
st.title("AI Veterinary Assistant Agent researching your query and summarizing it")

# App description
st.write("Generate reports enriched with real-time insights using the AI Veterinary Report Writing Agent powered by SmolAgents and DuckDuckGo.")

# Input blog topic or prompt
prompt = st.text_area("Enter the description of your patient visit:", placeholder="E.g., Today, I helped Max, a 5 year old Chihuahua")

# Button to generate blog content
if st.button("Generate Report"):
    if prompt:
        with st.spinner("Generating blog content..."):
            try:
                # Run the blog agent with the given prompt
                result = manager_agent.run(prompt)
                # Display the generated blog content
                st.subheader("Generated Report:")
                st.write(result)

                # Log backend activity
                log_agent_action(prompt, result, "Veterinary Research Agent with DuckDuckGo")
            except Exception as e:
                st.error(f"An error occurred: {e}")
    else:
        st.warning("Please enter a blog topic or prompt to proceed.")

# Footer
st.markdown("---")
st.caption("Powered by SmolAgents, DuckDuckGo, and Streamlit")