refactor: database URL handling

Author: naaive

openagent/core/database/engine.py

@@ -70,27 +70,29 @@ def _create_postgres_engine(db_url: str) -> Engine:
     return sa_create_engine(db_url)
 
 
-def create_engine(db_type: Literal["sqlite", "postgres"] = "sqlite", db_url: str = None) -> Engine:
+def create_engine(db_url: str) -> Engine:
     """
     Create a database engine based on the provided configuration.
+    Database type is automatically detected from the URL.
     
     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
             
     Returns:
         SQLAlchemy engine instance
     
     Raises:
-        ValueError: If an unsupported database type is specified or if db_url is missing
+        ValueError: If an unsupported database type is detected or if db_url is missing
     """
     if not db_url:
         raise ValueError("Database URL is required")
-        
-    if db_type == "sqlite":
+    
+    # Auto-detect database type from URL
+    if db_url.startswith('sqlite:'):
         return _create_sqlite_engine(db_url)
-    elif db_type == "postgres":
+    elif db_url.startswith('postgresql:'):
         return _create_postgres_engine(db_url)
     else:
-        raise ValueError(f"Unsupported database type: {db_type}") 
+        raise ValueError(f"Could not detect database type from URL: {db_url}. "
+                        "Supported URL formats: 'sqlite:///path/to/file.db' or 'postgresql://user:password@host:port/database'") 

openagent/tools/pendle/market_analysis.py

@@ -71,9 +71,9 @@ class PendleMarketConfig(BaseModel):
         default=None,
         description="Model configuration for LLM. If not provided, will use agent's core model",
     )
-    database: Optional[DatabaseConfig] = Field(
+    db_url: Optional[str] = Field(
         default=None,
-        description="Database configuration. If not provided, will use SQLite with default path",
+        description="Database URL. For postgres: postgresql://user:password@host:port/database, for sqlite: sqlite:///path/to/file.db"
     )
 
 
@@ -87,21 +87,14 @@ def __init__(self, core_model=None):
         self.tool_prompt = None
         self.DBSession = None
 
-    def _init_database(self, config: Optional[DatabaseConfig] = None) -> None:
+    def _init_database(self, db_url: Optional[str]) -> None:
         """Initialize database connection based on configuration"""
         # Set default configuration if not provided
-        db_type = "sqlite"
-        db_url = 'sqlite:///' + os.path.join(os.getcwd(), "storage", f"{self.name}.db")
-
-        if config:
-            db_type = config.type
-            db_url = config.url
+        if not db_url:
+            db_url = 'sqlite:///' + os.path.join(os.getcwd(), "storage", f"{self.name}.db")
 
         # Create engine using the database module's create_engine function
-        engine = create_engine(
-            db_type=db_type,
-            db_url=db_url,
-        )
+        engine = create_engine(db_url)
 
         # Create tables and initialize session factory
         Base.metadata.create_all(engine)
@@ -119,7 +112,7 @@ async def setup(self, config: PendleMarketConfig) -> None:
         """Setup the analysis tool with model and prompt"""
 
         # Initialize database
-        self._init_database(config.database)
+        self._init_database(config.db_url)
 
         # Initialize the model
         model_config = config.model if config.model else self.core_model

openagent/tools/pendle/voter_apy_analysis.py

@@ -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 = 'sqlite:///' + os.path.join(os.getcwd(), "storage", f"{self.name}.db")
-        self.engine = create_engine("sqlite", db_path)
+        db_url = 'sqlite:///' + os.path.join(os.getcwd(), "storage", f"{self.name}.db")
+        self.engine = create_engine(db_url)
         Base.metadata.create_all(self.engine)
         session = sessionmaker(bind=self.engine)
         self.session = session()

openagent/tools/twitter/feed.py

@@ -80,8 +80,8 @@ def __init__(self):
         self.retry_delay = 1
 
         # Initialize database
-        db_path = 'sqlite:///'+os.path.join(os.getcwd(), "storage", f"{self.name}.db")
-        self.engine = create_engine('sqlite',db_path)
+        db_url = 'sqlite:///' + os.path.join(os.getcwd(), "storage", f"{self.name}.db")
+        self.engine = create_engine(db_url)
         Base.metadata.create_all(self.engine)
         Session = sessionmaker(bind=self.engine)
         self.session = Session()
@@ -99,7 +99,7 @@ async def setup(self, config: TwitterFeedConfig) -> None:
         self.config = config
 
     async def _fetch_single_handle(
-        self, client: httpx.AsyncClient, handle: str
+            self, client: httpx.AsyncClient, handle: str
     ) -> List[dict]:
         """Fetch tweets for a single handle with retry logic"""
         params = {"limit": self.config.limit if self.config else 50}

test/tools/test_pendle_market.py

@@ -1,5 +1,7 @@
 import asyncio
 
+from sqlalchemy.testing.config import db_url
+
 from openagent.agent.config import ModelConfig
 from openagent.tools.pendle.market_analysis import PendleMarketTool, PendleMarketConfig, DatabaseConfig
 
@@ -12,10 +14,7 @@ async def test_pendle_market():
             name="gpt-4",
             temperature=0.7
         ),
-        database=DatabaseConfig(
-            type="sqlite",
-            url="sqlite:///storage/test_pendle_market.db"
-        )
+        db_url='sqlite:///storage/test_pendle_market.db'
     )
 
     # Initialize the tool
@@ -36,10 +35,7 @@ async def test_pendle_market_postgres():
             name="gpt-4",
             temperature=0.7
         ),
-        database=DatabaseConfig(
-            type="postgres",
-            url="postgresql://postgres:password@localhost:5434/pendle_market_test"
-        )
+        db_url='postgresql://postgres:password@localhost:5434/pendle_market_test'
     )
 
     # Initialize the tool
@@ -54,5 +50,5 @@ async def test_pendle_market_postgres():
 
 if __name__ == "__main__":
     # Run both tests
-    asyncio.run(test_pendle_market())
-    # asyncio.run(test_pendle_market_postgres())
+    # asyncio.run(test_pendle_market())
+    asyncio.run(test_pendle_market_postgres())