feat: add Space integration for user authentication and model management with OAuth support

2026-06-11 00:06:04 +00:00 · 2025-12-26 00:35:47 +08:00
parent 7479545339
commit 8caab43b00
27 changed files with 5214 additions and 6156 deletions
--- a/src/langbot/pkg/persistence/migrations/dbm014_space_account_support.py
+++ b/src/langbot/pkg/persistence/migrations/dbm014_space_account_support.py
@@ -0,0 +1,82 @@
+import sqlalchemy
+from .. import migration
+
+
+@migration.migration_class(14)
+class DBMigrateSpaceAccountSupport(migration.DBMigration):
+    """Add Space account support fields to users table"""
+
+    async def upgrade(self):
+        """Upgrade"""
+        # Get all column names from the users table
+        columns = []
+
+        if self.ap.persistence_mgr.db.name == 'postgresql':
+            result = await self.ap.persistence_mgr.execute_async(
+                sqlalchemy.text("SELECT column_name FROM information_schema.columns WHERE table_name = 'users';")
+            )
+            all_result = result.fetchall()
+            columns = [row[0] for row in all_result]
+        else:
+            result = await self.ap.persistence_mgr.execute_async(sqlalchemy.text('PRAGMA table_info(users);'))
+            all_result = result.fetchall()
+            columns = [row[1] for row in all_result]
+
+        # Add account_type column
+        if 'account_type' not in columns:
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text("ALTER TABLE users ADD COLUMN account_type VARCHAR(32) DEFAULT 'local' NOT NULL")
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text("ALTER TABLE users ADD COLUMN account_type VARCHAR(32) DEFAULT 'local' NOT NULL")
+                )
+
+        # Add space_account_uuid column
+        if 'space_account_uuid' not in columns:
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE users ADD COLUMN space_account_uuid VARCHAR(255)')
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE users ADD COLUMN space_account_uuid VARCHAR(255)')
+                )
+
+        # Add space_access_token column
+        if 'space_access_token' not in columns:
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE users ADD COLUMN space_access_token TEXT')
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE users ADD COLUMN space_access_token TEXT')
+                )
+
+        # Add space_refresh_token column
+        if 'space_refresh_token' not in columns:
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE users ADD COLUMN space_refresh_token TEXT')
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE users ADD COLUMN space_refresh_token TEXT')
+                )
+
+        # Add space_api_key column
+        if 'space_api_key' not in columns:
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE users ADD COLUMN space_api_key VARCHAR(255)')
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE users ADD COLUMN space_api_key VARCHAR(255)')
+                )
+
+    async def downgrade(self):
+        """Downgrade"""
+        pass
--- a/src/langbot/pkg/persistence/migrations/dbm015_model_source_tracking.py
+++ b/src/langbot/pkg/persistence/migrations/dbm015_model_source_tracking.py
@@ -0,0 +1,78 @@
+import sqlalchemy
+from .. import migration
+
+
+@migration.migration_class(15)
+class DBMigrateModelSourceTracking(migration.DBMigration):
+    """Add source tracking fields to models tables for Space integration"""
+
+    async def upgrade(self):
+        """Upgrade"""
+        # Add source column to llm_models table
+        llm_columns = await self._get_columns('llm_models')
+
+        if 'source' not in llm_columns:
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text("ALTER TABLE llm_models ADD COLUMN source VARCHAR(32) DEFAULT 'local' NOT NULL")
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text("ALTER TABLE llm_models ADD COLUMN source VARCHAR(32) DEFAULT 'local' NOT NULL")
+                )
+
+        if 'space_model_id' not in llm_columns:
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE llm_models ADD COLUMN space_model_id VARCHAR(255)')
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE llm_models ADD COLUMN space_model_id VARCHAR(255)')
+                )
+
+        # Add source column to embedding_models table
+        embedding_columns = await self._get_columns('embedding_models')
+
+        if 'source' not in embedding_columns:
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text(
+                        "ALTER TABLE embedding_models ADD COLUMN source VARCHAR(32) DEFAULT 'local' NOT NULL"
+                    )
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text(
+                        "ALTER TABLE embedding_models ADD COLUMN source VARCHAR(32) DEFAULT 'local' NOT NULL"
+                    )
+                )
+
+        if 'space_model_id' not in embedding_columns:
+            if self.ap.persistence_mgr.db.name == 'postgresql':
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE embedding_models ADD COLUMN space_model_id VARCHAR(255)')
+                )
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.text('ALTER TABLE embedding_models ADD COLUMN space_model_id VARCHAR(255)')
+                )
+
+    async def _get_columns(self, table_name: str) -> list:
+        """Get column names for a table"""
+        if self.ap.persistence_mgr.db.name == 'postgresql':
+            result = await self.ap.persistence_mgr.execute_async(
+                sqlalchemy.text(
+                    f"SELECT column_name FROM information_schema.columns WHERE table_name = '{table_name}';"
+                )
+            )
+            all_result = result.fetchall()
+            return [row[0] for row in all_result]
+        else:
+            result = await self.ap.persistence_mgr.execute_async(sqlalchemy.text(f'PRAGMA table_info({table_name});'))
+            all_result = result.fetchall()
+            return [row[1] for row in all_result]
+
+    async def downgrade(self):
+        """Downgrade"""
+        pass