Spaces:
Running
Running
File size: 4,044 Bytes
cbb33e5 2bbd95d cbb33e5 2bbd95d cbb33e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
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
import pandas as pd
from sqlalchemy import create_engine, text
# Create database engine
engine = create_engine("sqlite:///freights.db")
# 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 sql_query(query: str) -> str:
"""
Allows you to perform SQL queries on the freights table. Returns a string representation of the result.
The table is named 'freights'. Its description is as follows:
Columns:
- departure: DateTime (Date and time of departure)
- origin_port_locode: String (Origin port code)
- origin_port_name: String (Name of the origin port)
- destination_port: String (Destination port code)
- destination_port_name: String (Name of the destination port)
- dv20rate: Float (Rate for 20ft container in USD)
- dv40rate: Float (Rate for 40ft container in USD)
- currency: String (Currency of the rates)
- inserted_on: DateTime (Date when the rate was inserted)
Args:
query: The query to perform. This should be correct SQL.
Returns:
A string representation of the result of the query.
"""
try:
with engine.connect() as con:
result = con.execute(text(query))
rows = [dict(row._mapping) for row in result]
if not rows:
return "Aucun résultat trouvé."
# Convert to markdown table
headers = list(rows[0].keys())
table = "| " + " | ".join(headers) + " |\n"
table += "| " + " | ".join(["---" for _ in headers]) + " |\n"
for row in rows:
table += "| " + " | ".join(str(row[h]) for h in headers) + " |\n"
return table
except Exception as e:
return f"Error executing query: {str(e)}"
def get_schema() -> str:
"""
Returns the schema of the freights table.
"""
return """
Table: freights
Columns:
- departure: DateTime (Date and time of departure)
- origin_port_locode: String (Origin port code)
- origin_port_name: String (Name of the origin port)
- destination_port: String (Destination port code)
- destination_port_name: String (Name of the destination port)
- dv20rate: Float (Rate for 20ft container in USD)
- dv40rate: Float (Rate for 40ft container in USD)
- currency: String (Currency of the rates)
- inserted_on: DateTime (Date when the rate was inserted)
"""
def get_csv_as_dataframe() -> str:
"""
Returns a string representation of the freights table as a CSV file.
"""
df = pd.read_sql_table("freights", engine)
return df.to_csv(index=False)
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)
|