mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-02 12:05:54 +00:00
88132dff8a
* perf: reduce memory usage by ~200MB+ at startup
Two key optimizations:
1. Use importlib.util.find_spec() instead of __import__() in dependency
checking. find_spec() only locates modules without executing them,
avoiding loading all 36 dependencies (~222MB) into memory at startup.
2. Introduce shared aiohttp.ClientSession via httpclient module.
Previously, every HTTP request created a new ClientSession, which
creates a new TCPConnector and SSL context, loading system root
certificates each time (~270MB total allocations observed via memray).
Now all HTTP client code reuses shared sessions.
- satori.py and coze_server_api/client.py are left unchanged as they
create one session per adapter lifecycle (not per-request).
Profiling data (memray):
- Peak memory: 403MB
- SSL context creation: 270MB / 6.7M allocations (67% of total)
- Dependency import: 222MB (55% of peak)
- Expected reduction: 150-350MB at startup
* fix: remove unused aiohttp imports (ruff F401)
* style: ruff format
190 lines
7.6 KiB
Python
190 lines
7.6 KiB
Python
from __future__ import annotations
|
|
|
|
from langbot.pkg.utils import httpclient
|
|
import typing
|
|
import datetime
|
|
import time
|
|
import sqlalchemy
|
|
|
|
from ....core import app
|
|
from ....entity.persistence import user
|
|
from ....entity.dto.space_model import SpaceModel
|
|
|
|
|
|
class SpaceService:
|
|
"""Service for interacting with LangBot Space API"""
|
|
|
|
ap: app.Application
|
|
_credits_cache: typing.Dict[str, typing.Tuple[int, float]] # {user_email: (credits, timestamp)}
|
|
|
|
def __init__(self, ap: app.Application) -> None:
|
|
self.ap = ap
|
|
self._credits_cache = {}
|
|
|
|
def _get_space_config(self) -> typing.Dict[str, str]:
|
|
"""Get Space configuration from config file"""
|
|
space_config = self.ap.instance_config.data.get('space', {})
|
|
return {
|
|
'url': space_config.get('url', 'https://space.langbot.app'),
|
|
'oauth_authorize_url': space_config.get('oauth_authorize_url', 'https://space.langbot.app/auth/authorize'),
|
|
}
|
|
|
|
async def _get_user_by_email(self, user_email: str) -> user.User | None:
|
|
result = await self.ap.persistence_mgr.execute_async(
|
|
sqlalchemy.select(user.User).where(user.User.user == user_email)
|
|
)
|
|
result_list = result.all()
|
|
return result_list[0] if result_list else None
|
|
|
|
async def _ensure_valid_token(self, user_email: str) -> str | None:
|
|
"""Ensure access token is valid, refresh if expired. Returns valid access_token or None."""
|
|
user_obj = await self._get_user_by_email(user_email)
|
|
if not user_obj or user_obj.account_type != 'space':
|
|
return None
|
|
|
|
if not user_obj.space_access_token:
|
|
return None
|
|
|
|
# Check if token is expired (with 60s buffer)
|
|
if user_obj.space_access_token_expires_at:
|
|
if datetime.datetime.now() >= user_obj.space_access_token_expires_at - datetime.timedelta(seconds=60):
|
|
# Token expired, try to refresh
|
|
if user_obj.space_refresh_token:
|
|
try:
|
|
new_token = await self._refresh_and_save_token(user_obj)
|
|
return new_token
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
return user_obj.space_access_token
|
|
|
|
async def _refresh_and_save_token(self, user_obj: user.User) -> str:
|
|
"""Refresh token and save to database"""
|
|
token_data = await self.refresh_token(user_obj.space_refresh_token)
|
|
access_token = token_data.get('access_token')
|
|
expires_in = token_data.get('expires_in', 0)
|
|
|
|
if not access_token:
|
|
raise ValueError('Failed to refresh token')
|
|
|
|
expires_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in) if expires_in > 0 else None
|
|
|
|
await self.ap.persistence_mgr.execute_async(
|
|
sqlalchemy.update(user.User)
|
|
.where(user.User.user == user_obj.user)
|
|
.values(
|
|
space_access_token=access_token,
|
|
space_access_token_expires_at=expires_at,
|
|
)
|
|
)
|
|
|
|
return access_token
|
|
|
|
# === Raw API calls (no token validation) ===
|
|
|
|
def get_oauth_authorize_url(self, redirect_uri: str, state: str = '') -> str:
|
|
"""Get the Space OAuth authorization URL for redirect"""
|
|
space_config = self._get_space_config()
|
|
authorize_url = space_config['oauth_authorize_url']
|
|
params = f'redirect_uri={redirect_uri}'
|
|
if state:
|
|
params += f'&state={state}'
|
|
return f'{authorize_url}?{params}'
|
|
|
|
async def exchange_oauth_code(self, code: str) -> typing.Dict:
|
|
"""Exchange OAuth authorization code for tokens"""
|
|
from langbot.pkg.utils import constants
|
|
|
|
space_config = self._get_space_config()
|
|
space_url = space_config['url']
|
|
|
|
session = httpclient.get_session()
|
|
async with session.post(
|
|
f'{space_url}/api/v1/accounts/oauth/token',
|
|
json={'code': code, 'instance_id': constants.instance_id},
|
|
) as response:
|
|
if response.status != 200:
|
|
raise ValueError(f'Failed to exchange OAuth code: {await response.text()}')
|
|
data = await response.json()
|
|
if data.get('code') != 0:
|
|
raise ValueError(f'Failed to exchange OAuth code: {data.get("msg")}')
|
|
return data.get('data', {})
|
|
|
|
async def refresh_token(self, refresh_token: str) -> typing.Dict:
|
|
"""Refresh Space access token"""
|
|
space_config = self._get_space_config()
|
|
space_url = space_config['url']
|
|
|
|
session = httpclient.get_session()
|
|
async with session.post(
|
|
f'{space_url}/api/v1/accounts/token/refresh', json={'refresh_token': refresh_token}
|
|
) as response:
|
|
if response.status != 200:
|
|
raise ValueError(f'Failed to refresh token: {await response.text()}')
|
|
data = await response.json()
|
|
if data.get('code') != 0:
|
|
raise ValueError(f'Failed to refresh token: {data.get("msg")}')
|
|
return data.get('data', {})
|
|
|
|
async def get_user_info_raw(self, access_token: str) -> typing.Dict:
|
|
"""Get user info from Space using access token (no validation)"""
|
|
space_config = self._get_space_config()
|
|
space_url = space_config['url']
|
|
|
|
session = httpclient.get_session()
|
|
async with session.get(
|
|
f'{space_url}/api/v1/accounts/me', headers={'Authorization': f'Bearer {access_token}'}
|
|
) as response:
|
|
if response.status != 200:
|
|
raise ValueError(f'Failed to get user info: {await response.text()}')
|
|
data = await response.json()
|
|
if data.get('code') != 0:
|
|
raise ValueError(f'Failed to get user info: {data.get("msg")}')
|
|
return data.get('data', {})
|
|
|
|
# === API calls with token validation ===
|
|
|
|
async def get_user_info(self, user_email: str) -> typing.Dict | None:
|
|
"""Get user info from Space (with token validation)"""
|
|
access_token = await self._ensure_valid_token(user_email)
|
|
if not access_token:
|
|
return None
|
|
return await self.get_user_info_raw(access_token)
|
|
|
|
async def get_credits(self, user_email: str, force_refresh: bool = False) -> int | None:
|
|
"""Get Space credits for user with caching (60s TTL)"""
|
|
cache_ttl = 60
|
|
|
|
if not force_refresh and user_email in self._credits_cache:
|
|
credits, ts = self._credits_cache[user_email]
|
|
if time.time() - ts < cache_ttl:
|
|
return credits
|
|
|
|
try:
|
|
info = await self.get_user_info(user_email)
|
|
if info is None:
|
|
return None
|
|
credits = info.get('credits')
|
|
if credits is not None:
|
|
self._credits_cache[user_email] = (credits, time.time())
|
|
return credits
|
|
except Exception:
|
|
return self._credits_cache.get(user_email, (None, 0))[0]
|
|
|
|
async def get_models(self) -> typing.List[SpaceModel]:
|
|
"""Get models from Space"""
|
|
|
|
space_config = self._get_space_config()
|
|
space_url = space_config['url']
|
|
|
|
session = httpclient.get_session()
|
|
async with session.get(f'{space_url}/api/v1/models') as response:
|
|
if response.status != 200:
|
|
raise ValueError(f'Failed to get models: {await response.text()}')
|
|
data = await response.json()
|
|
if data.get('code') != 0:
|
|
raise ValueError(f'Failed to get models: {data.get("msg")}')
|
|
models_data = data.get('data', {}).get('models', [])
|
|
return [SpaceModel.model_validate(model_dict) for model_dict in models_data]
|