339 lines
11 KiB
Python
339 lines
11 KiB
Python
"""LangSmith profile configuration and auth helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import datetime
|
|
import json
|
|
import os
|
|
import threading
|
|
from collections.abc import Mapping
|
|
from pathlib import Path
|
|
from typing import Any, NamedTuple, Optional, TypedDict, cast
|
|
|
|
import requests
|
|
|
|
_OAUTH_CLIENT_ID = "langsmith-cli"
|
|
_TOKEN_REFRESH_LEEWAY = datetime.timedelta(minutes=1)
|
|
_TOKEN_REFRESH_TIMEOUT = 10
|
|
|
|
|
|
class ProfileOAuth(TypedDict, total=False):
|
|
access_token: str
|
|
refresh_token: str
|
|
expires_at: str
|
|
|
|
|
|
class ProfileConfig(TypedDict, total=False):
|
|
api_key: str
|
|
api_url: str
|
|
workspace_id: str
|
|
oauth: ProfileOAuth
|
|
|
|
|
|
class ProfileConfigFile(TypedDict, total=False):
|
|
current_profile: str
|
|
profiles: dict[str, ProfileConfig]
|
|
|
|
|
|
class ProfileState(NamedTuple):
|
|
path: Path
|
|
config: ProfileConfigFile
|
|
profile_name: str
|
|
|
|
|
|
class ProfileClientConfig(NamedTuple):
|
|
api_url: Optional[str] = None
|
|
api_key: Optional[str] = None
|
|
workspace_id: Optional[str] = None
|
|
oauth_access_token: Optional[str] = None
|
|
oauth_refresh_token: Optional[str] = None
|
|
oauth_expires_at: Optional[str] = None
|
|
profile_state: Optional[ProfileState] = None
|
|
|
|
@property
|
|
def has_oauth(self) -> bool:
|
|
return bool(self.oauth_access_token or self.oauth_refresh_token)
|
|
|
|
|
|
def trim_auth_value(value: Optional[str]) -> Optional[str]:
|
|
if not value:
|
|
return None
|
|
trimmed = value.strip().strip('"').strip("'")
|
|
return trimmed or None
|
|
|
|
|
|
def _profile_config_path() -> Optional[Path]:
|
|
if config_file := os.environ.get("LANGSMITH_CONFIG_FILE"):
|
|
return Path(config_file)
|
|
try:
|
|
return Path.home() / ".langsmith" / "config.json"
|
|
except RuntimeError:
|
|
return None
|
|
|
|
|
|
def _load_profile_state() -> Optional[ProfileState]:
|
|
path = _profile_config_path()
|
|
if path is None or not path.exists():
|
|
return None
|
|
try:
|
|
raw = json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
return None
|
|
if not isinstance(raw, dict):
|
|
return None
|
|
profiles = raw.get("profiles")
|
|
if not isinstance(profiles, dict):
|
|
return None
|
|
profile_name = os.environ.get("LANGSMITH_PROFILE")
|
|
if not profile_name:
|
|
current_profile = raw.get("current_profile")
|
|
if isinstance(current_profile, str) and current_profile:
|
|
profile_name = current_profile
|
|
elif "default" in profiles:
|
|
profile_name = "default"
|
|
if not profile_name or not isinstance(profiles.get(profile_name), dict):
|
|
return None
|
|
return ProfileState(path, cast(ProfileConfigFile, raw), profile_name)
|
|
|
|
|
|
def _profile_from_state(state: ProfileState) -> Optional[ProfileConfig]:
|
|
profiles = state.config.get("profiles") or {}
|
|
profile = profiles.get(state.profile_name)
|
|
if not isinstance(profile, dict):
|
|
return None
|
|
return cast(ProfileConfig, profile)
|
|
|
|
|
|
def load_profile_client_config() -> ProfileClientConfig:
|
|
state = _load_profile_state()
|
|
if state is None:
|
|
return ProfileClientConfig()
|
|
profile = _profile_from_state(state)
|
|
if profile is None:
|
|
return ProfileClientConfig()
|
|
oauth = profile.get("oauth") or {}
|
|
return ProfileClientConfig(
|
|
api_url=profile.get("api_url"),
|
|
api_key=trim_auth_value(profile.get("api_key")),
|
|
workspace_id=profile.get("workspace_id"),
|
|
oauth_access_token=trim_auth_value(oauth.get("access_token")),
|
|
oauth_refresh_token=trim_auth_value(oauth.get("refresh_token")),
|
|
oauth_expires_at=oauth.get("expires_at"),
|
|
profile_state=state,
|
|
)
|
|
|
|
|
|
def _normalize_profile_api_url(api_url: str) -> str:
|
|
while api_url.endswith("/"):
|
|
api_url = api_url[:-1]
|
|
suffix = "/api/v1"
|
|
if api_url.endswith(suffix):
|
|
return api_url[: -len(suffix)]
|
|
return api_url
|
|
|
|
|
|
def _parse_profile_expires_at(expires_at: str) -> Optional[datetime.datetime]:
|
|
try:
|
|
parsed = datetime.datetime.fromisoformat(expires_at.replace("Z", "+00:00"))
|
|
except ValueError:
|
|
return None
|
|
if parsed.tzinfo is None:
|
|
parsed = parsed.replace(tzinfo=datetime.timezone.utc)
|
|
return parsed
|
|
|
|
|
|
def should_refresh_profile_token(profile: ProfileConfig) -> bool:
|
|
oauth = profile.get("oauth") or {}
|
|
if not oauth.get("refresh_token"):
|
|
return False
|
|
if not oauth.get("access_token"):
|
|
return True
|
|
expires_at = oauth.get("expires_at")
|
|
if not expires_at:
|
|
return False
|
|
parsed = _parse_profile_expires_at(expires_at)
|
|
if parsed is None:
|
|
return False
|
|
return (
|
|
parsed <= datetime.datetime.now(datetime.timezone.utc) + _TOKEN_REFRESH_LEEWAY
|
|
)
|
|
|
|
|
|
def _refresh_profile_oauth_token(
|
|
api_url: Optional[str], refresh_token: str
|
|
) -> Optional[dict[str, Any]]:
|
|
refresh_url = _normalize_profile_api_url(
|
|
api_url or "https://api.smith.langchain.com"
|
|
)
|
|
try:
|
|
response = requests.post(
|
|
f"{refresh_url}/oauth/token",
|
|
data={
|
|
"grant_type": "refresh_token",
|
|
"client_id": _OAUTH_CLIENT_ID,
|
|
"refresh_token": refresh_token,
|
|
},
|
|
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
|
timeout=_TOKEN_REFRESH_TIMEOUT,
|
|
)
|
|
except requests.RequestException:
|
|
return None
|
|
if response.status_code < 200 or response.status_code >= 300:
|
|
return None
|
|
try:
|
|
token = response.json()
|
|
except ValueError:
|
|
return None
|
|
if not isinstance(token, dict) or not token.get("access_token"):
|
|
return None
|
|
return token
|
|
|
|
|
|
def _apply_profile_token_response(
|
|
profile: ProfileConfig, token: Mapping[str, Any]
|
|
) -> None:
|
|
oauth = profile.setdefault("oauth", {})
|
|
access_token = token.get("access_token")
|
|
if isinstance(access_token, str) and access_token:
|
|
oauth["access_token"] = access_token
|
|
refresh_token = token.get("refresh_token")
|
|
if isinstance(refresh_token, str) and refresh_token:
|
|
oauth["refresh_token"] = refresh_token
|
|
expires_in = token.get("expires_in")
|
|
if isinstance(expires_in, (int, float)) and expires_in > 0:
|
|
expires_at = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
|
|
seconds=expires_in
|
|
)
|
|
oauth["expires_at"] = expires_at.isoformat().replace("+00:00", "Z")
|
|
|
|
|
|
def _save_profile_config(path: Path, config: ProfileConfigFile) -> None:
|
|
try:
|
|
path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
|
|
temp_path = path.with_name(f"{path.name}.tmp")
|
|
temp_path.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")
|
|
os.chmod(temp_path, 0o600)
|
|
os.replace(temp_path, path)
|
|
os.chmod(path, 0o600)
|
|
except OSError:
|
|
return
|
|
|
|
|
|
class ProfileAuth:
|
|
def __init__(
|
|
self,
|
|
config: ProfileClientConfig,
|
|
*,
|
|
api_key_header: str,
|
|
) -> None:
|
|
self._state = config.profile_state
|
|
self._api_key_header = api_key_header
|
|
self._lock = threading.Lock()
|
|
self._managed_auth_headers: set[tuple[str, str]] = set()
|
|
self._remember_auth_headers(self._auth_headers(refresh=False))
|
|
|
|
@property
|
|
def has_auth(self) -> bool:
|
|
profile = self._profile()
|
|
if profile is None:
|
|
return False
|
|
oauth = profile.get("oauth") or {}
|
|
return bool(
|
|
trim_auth_value(oauth.get("access_token"))
|
|
or trim_auth_value(oauth.get("refresh_token"))
|
|
or trim_auth_value(profile.get("api_key"))
|
|
)
|
|
|
|
@property
|
|
def oauth_access_token(self) -> Optional[str]:
|
|
profile = self._profile()
|
|
if profile is None:
|
|
return None
|
|
return trim_auth_value((profile.get("oauth") or {}).get("access_token"))
|
|
|
|
def needs_refresh(self) -> bool:
|
|
profile = self._profile()
|
|
return profile is not None and should_refresh_profile_token(profile)
|
|
|
|
def current_auth_headers(self) -> dict[str, str]:
|
|
headers = self._auth_headers(refresh=False)
|
|
self._remember_auth_headers(headers)
|
|
return headers
|
|
|
|
def get_auth_headers(self) -> dict[str, str]:
|
|
headers = self._auth_headers(refresh=True)
|
|
self._remember_auth_headers(headers)
|
|
return headers
|
|
|
|
def prepare_request_headers(self, headers: Mapping[str, str]) -> dict[str, str]:
|
|
"""Replace stale profile-managed auth while preserving explicit auth."""
|
|
request_headers = dict(headers)
|
|
for key, value in list(request_headers.items()):
|
|
if self._is_profile_auth_header(key, value):
|
|
del request_headers[key]
|
|
if not self._has_auth_header(request_headers):
|
|
request_headers.update(self.current_auth_headers())
|
|
return request_headers
|
|
|
|
def _profile(self) -> Optional[ProfileConfig]:
|
|
if self._state is None:
|
|
return None
|
|
return _profile_from_state(self._state)
|
|
|
|
def _auth_headers(self, *, refresh: bool) -> dict[str, str]:
|
|
profile = self._profile()
|
|
if profile is None:
|
|
return {}
|
|
if refresh and should_refresh_profile_token(profile):
|
|
with self._lock:
|
|
profile = self._profile()
|
|
if profile is not None and should_refresh_profile_token(profile):
|
|
self._refresh(profile)
|
|
return self._headers_from_profile(profile)
|
|
|
|
def _refresh(self, profile: ProfileConfig) -> None:
|
|
refresh_token = trim_auth_value(
|
|
(profile.get("oauth") or {}).get("refresh_token")
|
|
)
|
|
if refresh_token is None or self._state is None:
|
|
return
|
|
api_url = profile.get("api_url")
|
|
token = _refresh_profile_oauth_token(api_url, refresh_token)
|
|
if token is None:
|
|
return
|
|
_apply_profile_token_response(profile, token)
|
|
profiles = self._state.config.get("profiles") or {}
|
|
profiles[self._state.profile_name] = profile
|
|
self._state.config["profiles"] = profiles
|
|
_save_profile_config(self._state.path, self._state.config)
|
|
|
|
def _headers_from_profile(self, profile: Optional[ProfileConfig]) -> dict[str, str]:
|
|
if profile is None:
|
|
return {}
|
|
oauth_access_token = trim_auth_value(
|
|
(profile.get("oauth") or {}).get("access_token")
|
|
)
|
|
if oauth_access_token:
|
|
return {"Authorization": f"Bearer {oauth_access_token}"}
|
|
api_key = trim_auth_value(profile.get("api_key"))
|
|
if api_key:
|
|
return {self._api_key_header: api_key}
|
|
return {}
|
|
|
|
def _remember_auth_headers(self, headers: Mapping[str, str]) -> None:
|
|
for name, value in headers.items():
|
|
if self._is_auth_header_name(name) and value:
|
|
self._managed_auth_headers.add((name.lower(), value))
|
|
|
|
def _is_profile_auth_header(self, name: str, value: str) -> bool:
|
|
return (name.lower(), value) in self._managed_auth_headers
|
|
|
|
def _has_auth_header(self, headers: Mapping[str, str]) -> bool:
|
|
return any(
|
|
self._is_auth_header_name(name) and bool(value)
|
|
for name, value in headers.items()
|
|
)
|
|
|
|
def _is_auth_header_name(self, name: str) -> bool:
|
|
return name.lower() in {"authorization", self._api_key_header.lower()}
|