refactor: create_engine helper

Author: naaive

openagent/core/database/engine.py

@@ -1,38 +1,22 @@
 import os
-from typing import Optional
+from typing import Optional, Literal
 from urllib.parse import urlparse
 
 from sqlalchemy import create_engine as sa_create_engine, text as sa_text, Engine
 from loguru import logger
 
 
-def _create_sqlite_engine(db_url: Optional[str] = None, db_name: Optional[str] = None, storage_dir: str = "storage") -> Engine:
+def _create_sqlite_engine(db_url: str) -> Engine:
     """
-    Create a SQLite engine from a URL or create a default URL if not provided.
+    Create a SQLite engine from a URL.
     
     Args:
         db_url: SQLite database URL (sqlite:///path/to/file.db)
-        db_name: Database name to use if db_url is not provided
-        storage_dir: Directory to store SQLite databases (default: 'storage')
         
     Returns:
         SQLAlchemy engine instance
-        
-    Raises:
-        ValueError: If neither db_url nor db_name is provided
     """
-    if db_url:
-        return sa_create_engine(db_url)
-    
-    if not db_name:
-        raise ValueError("Either db_url or db_name must be provided for SQLite")
-    
-    # Create default SQLite path
-    db_path = os.path.join(os.getcwd(), storage_dir, f"{db_name}.db")
-    if not os.path.exists(os.path.dirname(db_path)):
-        os.makedirs(os.path.dirname(db_path))
-    
-    return sa_create_engine(f"sqlite:///{db_path}")
+    return sa_create_engine(db_url)
 
 
 def _ensure_postgres_database_exists(db_url: str) -> None:
@@ -86,28 +70,27 @@ def _create_postgres_engine(db_url: str) -> Engine:
     return sa_create_engine(db_url)
 
 
-def create_engine(db_type: str = "sqlite", db_url: Optional[str] = None, db_name: Optional[str] = None, storage_dir: str = "storage") -> Engine:
+def create_engine(db_type: Literal["sqlite", "postgres"] = "sqlite", db_url: str = None) -> Engine:
     """
     Create a database engine based on the provided configuration.
     
     Args:
         db_type: Type of database ('sqlite' or 'postgres')
         db_url: Database URL. For postgres: postgresql://user:password@host:port/database, 
                 for sqlite: sqlite:///path/to/file.db
-        db_name: Database name to use if db_url is not provided (SQLite only)
-        storage_dir: Directory to store SQLite databases (default: 'storage')
             
     Returns:
         SQLAlchemy engine instance
     
     Raises:
-        ValueError: If an unsupported database type is specified or if required parameters are missing
+        ValueError: If an unsupported database type is specified or if db_url is missing
     """
+    if not db_url:
+        raise ValueError("Database URL is required")
+        
     if db_type == "sqlite":
-        return _create_sqlite_engine(db_url, db_name, storage_dir)
+        return _create_sqlite_engine(db_url)
     elif db_type == "postgres":
-        if not db_url:
-            raise ValueError("Database URL is required for PostgreSQL")
         return _create_postgres_engine(db_url)
     else:
         raise ValueError(f"Unsupported database type: {db_type}") 

openagent/tools/pendle/market_analysis.py

@@ -91,7 +91,7 @@ def _init_database(self, config: Optional[DatabaseConfig] = None) -> None:
         """Initialize database connection based on configuration"""
         # Set default configuration if not provided
         db_type = "sqlite"
-        db_url = None
+        db_url = 'sqlite:///' + os.path.join(os.getcwd(), "storage", f"{self.name}.db")
 
         if config:
             db_type = config.type
@@ -101,7 +101,6 @@ def _init_database(self, config: Optional[DatabaseConfig] = None) -> None:
         engine = create_engine(
             db_type=db_type,
             db_url=db_url,
-            db_name=self.name
         )
 
         # Create tables and initialize session factory
@@ -266,10 +265,10 @@ def get_top_markets(markets: list[PendleMarketData], key_attr: str, n: int = 3)
 
         # Combine all relevant symbols
         relevant_symbols = (
-            liquidity_increase_top_symbols
-            | new_market_liquidity_increase_symbols
-            | apy_increase_symbols
-            | new_market_apy_increase_symbols
+                liquidity_increase_top_symbols
+                | new_market_liquidity_increase_symbols
+                | apy_increase_symbols
+                | new_market_apy_increase_symbols
         )
 
         # Filter markets to include only those in the relevant symbols set

openagent/tools/pendle/voter_apy_analysis.py

@@ -14,7 +14,7 @@
 from sqlalchemy.orm import sessionmaker
 
 from openagent.agent.config import ModelConfig
-from openagent.core.database import sqlite
+from openagent.core.database.engine import create_engine
 from openagent.core.tool import Tool
 from openagent.core.utils.fetch_json import fetch_json
 from openagent.core.utils.json_equal import json_equal
@@ -47,8 +47,8 @@ def __init__(self, core_model=None):
         self.core_model = core_model
         self.tool_model = None
         self.tool_prompt = None
-        db_path = os.path.join(os.getcwd(), "storage", f"{self.name}.db")
-        self.engine = sqlite.create_engine(db_path)
+        db_path = 'sqlite:///' + os.path.join(os.getcwd(), "storage", f"{self.name}.db")
+        self.engine = create_engine("sqlite", db_path)
         Base.metadata.create_all(self.engine)
         session = sessionmaker(bind=self.engine)
         self.session = session()

openagent/tools/twitter/feed.py

@@ -10,7 +10,7 @@
 from sqlalchemy import Column, Integer, String, DateTime
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker
-from openagent.core.database import sqlite
+from openagent.core.database.engine import create_engine
 from openagent.core.tool import Tool
 
 Base = declarative_base()
@@ -80,8 +80,8 @@ def __init__(self):
         self.retry_delay = 1
 
         # Initialize database
-        db_path = os.path.join(os.getcwd(), "storage", f"{self.name}.db")
-        self.engine = sqlite.create_engine(db_path)
+        db_path = 'sqlite:///'+os.path.join(os.getcwd(), "storage", f"{self.name}.db")
+        self.engine = create_engine('sqlite',db_path)
         Base.metadata.create_all(self.engine)
         Session = sessionmaker(bind=self.engine)
         self.session = Session()