Upload 5 files

gradio_smol.py

@@ -0,0 +1,24 @@
+import gradio as gr
+from tools import FreightAgent, EXAMPLE_QUERIES
+from utils import initialize_database
+from smolagents import CodeAgent, HfApiModel, GradioUI
+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 = "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"
+model = HfApiModel(model_id=model_id, token=os.environ["HF_API_KEY"])
+
+agent = CodeAgent(tools=[sql_query, get_schema, get_csv_as_dataframe], model=model)
+
+if __name__ == "__main__":
+    GradioUI(agent).launch()

requirements.txt

@@ -0,0 +1,5 @@
+pandas
+sqlalchemy
+smolagents
+python-dotenv
+gradio

sql_data.py

@@ -0,0 +1,80 @@
+from sqlalchemy import create_engine, text
+from smolagents import tool
+import os
+import pandas as pd
+
+# Create database engine
+engine = create_engine("sqlite:///freights.db")
+
+
+@tool
+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)}"
+
+
+@tool
+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)
+    """
+
+
+@tool
+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)

tools.py

@@ -0,0 +1,43 @@
+from smolagents import CodeAgent, HfApiModel
+from sql_data import sql_query, get_schema
+from sqlalchemy import create_engine, inspect, text
+import os
+from dotenv import load_dotenv
+from typing import Dict, List, Any
+import json
+
+# Load environment variables
+load_dotenv()
+
+# Example queries that the agent can handle
+EXAMPLE_QUERIES = [
+    "Quels sont les tarifs moyens des conteneurs 20ft et 40ft entre tous les ports ?",
+    "Quels sont les ports d'origine les plus fréquents ?",
+    "Montre-moi les routes avec des tarifs élevés pour les conteneurs 40ft",
+    "Quelle est l'évolution des prix au fil du temps pour la route Surabaya vers Nansha ?",
+    "Quelles sont les destinations disponibles depuis Shanghai ?",
+]
+
+
+class FreightAgent:
+    def __init__(self):
+        self.setup_agent()
+
+    def setup_agent(self) -> None:
+        """
+        Initialize the CodeAgent with SQL tools.
+
+        Create a CodeAgent with two tools: `sql_query` and `get_schema`.
+        `sql_query` allows to perform SQL queries on the freights table.
+        `get_schema` returns the schema of the freights table.
+        """
+        self.agent = CodeAgent(
+            tools=[sql_query, get_schema],
+            model=HfApiModel("meta-llama/Llama-3.1-8B-Instruct"),
+        )
+
+    def query(self, question: str) -> str:
+        """
+        Ask a question about the freight data in natural language
+        """
+        return self.agent.run(question)

utils.py

@@ -0,0 +1,80 @@
+import pandas as pd
+import sqlite3
+import requests
+from typing import List, Dict, Any
+import os
+from sqlalchemy import create_engine, Column, Float, String, Integer, DateTime
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+# Create base class for declarative models
+Base = declarative_base()
+
+
+class Freight(Base):
+    """SQLAlchemy model for freight data"""
+
+    __tablename__ = "freights"
+
+    id = Column(Integer, primary_key=True)
+    departure = Column(DateTime)
+    origin_port_locode = Column(String)
+    origin_port_name = Column(String)
+    destination_port = Column(String)
+    destination_port_name = Column(String)
+    dv20rate = Column(Float)
+    dv40rate = Column(Float)
+    currency = Column(String)
+    inserted_on = Column(DateTime)
+
+
+def download_csv(url: str, local_path: str = "freights.csv") -> str:
+    """
+    Download CSV file from Hugging Face and save it locally
+    """
+    response = requests.get(url)
+    with open(local_path, "wb") as f:
+        f.write(response.content)
+    return local_path
+
+
+def create_database(db_name: str = "freights.db") -> None:
+    """
+    Create SQLite database and necessary tables
+    """
+    engine = create_engine(f"sqlite:///{db_name}")
+    Base.metadata.create_all(engine)
+
+
+def load_csv_to_db(csv_path: str, db_name: str = "freights.db") -> None:
+    """
+    Load CSV data into SQLite database
+    """
+    # Read CSV
+    df = pd.read_csv(csv_path, parse_dates=["departure", "inserted_on"])
+
+    # Connect to database
+    engine = create_engine(f"sqlite:///{db_name}")
+
+    # Save to database
+    df.to_sql("freights", engine, if_exists="replace", index=False)
+
+
+def initialize_database(csv_url: str) -> None:
+    """
+    Initialize the database by downloading CSV and loading data.
+
+    Args:
+        csv_url: URL of the CSV file to download and load.
+    """
+    # Download CSV
+    csv_path = download_csv(csv_url)
+
+    # Create and load database
+    create_database()
+    load_csv_to_db(csv_path)
+    print("Database initialized.")
+
+    # Clean up CSV file
+    if os.path.exists(csv_path):
+        os.remove(csv_path)