refactor: model management

src/langbot/pkg/api/http/controller/groups/user.py

@@ -1,6 +1,7 @@
 import quart
 import argon2
 import asyncio
+import traceback
 
 from .. import group
 
@@ -141,8 +142,10 @@ class UserRouterGroup(group.RouterGroup):
                     }
                 )
             except ValueError as e:
+                traceback.print_exc()
                 return self.fail(1, str(e))
             except Exception as e:
+                traceback.print_exc()
                 return self.fail(2, f'OAuth callback failed: {str(e)}')
 
         @self.route('/info', methods=['GET'], auth_type=group.AuthType.USER_TOKEN)

src/langbot/pkg/api/http/service/model.py

@@ -85,10 +85,17 @@ class LLMModelsService:
 
         await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_model.LLMModel).values(**model_data))
 
-        llm_model = await self.get_llm_model(model_data['uuid'])
-        await self.ap.model_mgr.load_llm_model(llm_model)
+        runtime_provider = self.ap.model_mgr.provider_dict.get(model_data['provider_uuid'])
+        if runtime_provider is None:
+            raise Exception('provider not found')
 
-        # Check if default pipeline has no model bound
+        runtime_llm_model = await self.ap.model_mgr.load_llm_model_with_provider(
+            persistence_model.LLMModel(**model_data),
+            runtime_provider,
+        )
+        self.ap.model_mgr.llm_models.append(runtime_llm_model)
+
+        # set the default pipeline model to this model
         result = await self.ap.persistence_mgr.execute_async(
             sqlalchemy.select(persistence_pipeline.LegacyPipeline).where(
                 persistence_pipeline.LegacyPipeline.is_default == True
@@ -152,8 +159,16 @@ class LLMModelsService:
         )
 
         await self.ap.model_mgr.remove_llm_model(model_uuid)
-        llm_model = await self.get_llm_model(model_uuid)
-        await self.ap.model_mgr.load_llm_model(llm_model)
+
+        runtime_provider = self.ap.model_mgr.provider_dict.get(model_data['provider_uuid'])
+        if runtime_provider is None:
+            raise Exception('provider not found')
+
+        runtime_llm_model = await self.ap.model_mgr.load_llm_model_with_provider(
+            persistence_model.LLMModel(**model_data),
+            runtime_provider,
+        )
+        self.ap.model_mgr.llm_models.append(runtime_llm_model)
 
     async def delete_llm_model(self, model_uuid: str) -> None:
         """Delete an LLM model"""
@@ -174,10 +189,10 @@ class LLMModelsService:
             if runtime_llm_model is None:
                 raise Exception('model not found')
         else:
-            runtime_llm_model = await self.ap.model_mgr.init_runtime_llm_model(model_data)
+            runtime_llm_model = await self.ap.model_mgr.init_temporary_runtime_llm_model(model_data)
 
         extra_args = model_data.get('extra_args', {})
-        await runtime_llm_model.requester.invoke_llm(
+        await runtime_llm_model.provider.requester.invoke_llm(
             query=None,
             model=runtime_llm_model,
             messages=[provider_message.Message(role='user', content='Hello, world! Please just reply a "Hello".')],
@@ -244,8 +259,15 @@ class EmbeddingModelsService:
             sqlalchemy.insert(persistence_model.EmbeddingModel).values(**model_data)
         )
 
-        embedding_model = await self.get_embedding_model(model_data['uuid'])
-        await self.ap.model_mgr.load_embedding_model(embedding_model)
+        runtime_provider = self.ap.model_mgr.provider_dict.get(model_data['provider_uuid'])
+        if runtime_provider is None:
+            raise Exception('provider not found')
+
+        runtime_embedding_model = await self.ap.model_mgr.load_embedding_model_with_provider(
+            persistence_model.EmbeddingModel(**model_data),
+            runtime_provider,
+        )
+        self.ap.model_mgr.embedding_models.append(runtime_embedding_model)
 
         return model_data['uuid']
 
@@ -298,8 +320,16 @@ class EmbeddingModelsService:
         )
 
         await self.ap.model_mgr.remove_embedding_model(model_uuid)
-        embedding_model = await self.get_embedding_model(model_uuid)
-        await self.ap.model_mgr.load_embedding_model(embedding_model)
+
+        runtime_provider = self.ap.model_mgr.provider_dict.get(model_data['provider_uuid'])
+        if runtime_provider is None:
+            raise Exception('provider not found')
+
+        runtime_embedding_model = await self.ap.model_mgr.load_embedding_model_with_provider(
+            persistence_model.EmbeddingModel(**model_data),
+            runtime_provider,
+        )
+        self.ap.model_mgr.embedding_models.append(runtime_embedding_model)
 
     async def delete_embedding_model(self, model_uuid: str) -> None:
         """Delete an embedding model"""
@@ -322,9 +352,9 @@ class EmbeddingModelsService:
             if runtime_embedding_model is None:
                 raise Exception('model not found')
         else:
-            runtime_embedding_model = await self.ap.model_mgr.init_runtime_embedding_model(model_data)
+            runtime_embedding_model = await self.ap.model_mgr.init_temporary_runtime_embedding_model(model_data)
 
-        await runtime_embedding_model.requester.invoke_embedding(
+        await runtime_embedding_model.provider.requester.invoke_embedding(
             model=runtime_embedding_model,
             input_text=['Hello, world!'],
             extra_args={},

src/langbot/pkg/api/http/service/provider.py

@@ -61,6 +61,10 @@ class ModelProviderService:
         await self.ap.persistence_mgr.execute_async(
             sqlalchemy.insert(persistence_model.ModelProvider).values(**provider_data)
         )
+
+        # load to runtime
+        runtime_provider = await self.ap.model_mgr.load_provider(provider_data)
+        self.ap.model_mgr.provider_dict[runtime_provider.provider_entity.uuid] = runtime_provider
         return provider_data['uuid']
 
     async def update_provider(self, provider_uuid: str, provider_data: dict) -> None:
@@ -72,8 +76,7 @@ class ModelProviderService:
             .where(persistence_model.ModelProvider.uuid == provider_uuid)
             .values(**provider_data)
         )
-        # Reload all models using this provider
-        await self.ap.model_mgr.load_models_from_db()
+        await self.ap.model_mgr.reload_provider(provider_uuid)
 
     async def delete_provider(self, provider_uuid: str) -> None:
         """Delete a provider (only if no models reference it)"""
@@ -100,6 +103,8 @@ class ModelProviderService:
             )
         )
 
+        await self.ap.model_mgr.remove_provider(provider_uuid)
+
     async def get_provider_model_counts(self, provider_uuid: str) -> dict:
         """Get count of models using this provider"""
         llm_result = await self.ap.persistence_mgr.execute_async(
@@ -150,3 +155,12 @@ class ModelProviderService:
                 'api_keys': api_keys or [],
             }
         )
+
+    async def update_space_model_provider_api_keys(self, api_key: str) -> None:
+        """Update Space model provider API keys"""
+        await self.ap.persistence_mgr.execute_async(
+            sqlalchemy.update(persistence_model.ModelProvider)
+            .where(persistence_model.ModelProvider.uuid == '00000000-0000-0000-0000-000000000000')
+            .values(api_keys=[api_key])
+        )
+        await self.ap.model_mgr.reload_provider('00000000-0000-0000-0000-000000000000')

src/langbot/pkg/api/http/service/user.py

@@ -149,6 +149,7 @@ class UserService:
                         space_access_token_expires_at=expires_at,
                     )
                 )
+                await self.ap.provider_service.update_space_model_provider_api_keys(api_key)
                 return await self.get_user_by_space_account_uuid(space_account_uuid)
 
             # Check if user with same email exists
@@ -167,6 +168,7 @@ class UserService:
                         space_access_token_expires_at=expires_at,
                     )
                 )
+                await self.ap.provider_service.update_space_model_provider_api_keys(api_key)
                 return await self.get_user_by_email(email)
 
             # Check if system is already initialized
@@ -189,6 +191,7 @@ class UserService:
                     space_access_token_expires_at=expires_at,
                 )
             )
+            await self.ap.provider_service.update_space_model_provider_api_keys(api_key)
 
             return await self.get_user_by_space_account_uuid(space_account_uuid)