app.py

import gradio as gr
from tools import FreightAgent, EXAMPLE_QUERIES
from utils import initialize_database
from smolagents import CodeAgent, OpenAIServerModel
import os
from dotenv import load_dotenv
from sql_data import sql_query, get_schema, get_csv_as_dataframe

# Load environment variables
load_dotenv()

# Initialize the database if it doesn't exist
if not os.path.exists("freights.db"):
    csv_url = "https://huggingface.co/datasets/sasu-SpidR/fretmaritime/resolve/main/freights.csv"
    initialize_database(csv_url)

# Create the main agent
model_id = "gpt-4.1-mini"
model = OpenAIServerModel(model_id=model_id, api_key=os.environ["OPENAI_API_KEY"])

agent = CodeAgent(tools=[sql_query, get_schema, get_csv_as_dataframe], model=model)

def run_agent(question: str) -> str:
    """
    Run the agent with the given question.
    This ReAct Agent can make request to give you information about the freight data.

    Args:
        question: The question to run the agent with.
    Returns:
        The response of the agent.
    """
    return agent.run(question,max_steps=5)


if __name__ == "__main__":
    gr.Interface(
        fn=run_agent,
        inputs=gr.Textbox(lines=7, label="Question"),
        outputs=gr.Textbox(),
        title="Freight Agent MCP",
        description="Ask a question about the freight data in natural language",
        examples=EXAMPLE_QUERIES if "EXAMPLE_QUERIES" in globals() else None
    ).launch(mcp_server=True)