From 03f8a36065c44a0cfa8fdd96cd47c98d74274ac2 Mon Sep 17 00:00:00 2001 From: rex <1073853456@qq.com> Date: Mon, 5 Aug 2024 03:06:39 +0800 Subject: [PATCH] add main scripts --- .env.example | 6 ++ .github/workflows/test.yml | 11 +-- README.md | 18 +++- demo/one-api.ipynb | 171 ++++++++++++++++++++++++++++++++++++ pyproject.toml | 54 ++++++++++++ requirements_dev.txt | 3 +- setup.py | 44 ---------- src/one-api-cli/__init__.py | 5 -- src/one_api_cli/__init__.py | 8 ++ src/one_api_cli/account.py | 142 ++++++++++++++++++++++++++++++ src/one_api_cli/channel.py | 151 +++++++++++++++++++++++++++++++ src/one_api_cli/constant.py | 23 +++++ 12 files changed, 577 insertions(+), 59 deletions(-) create mode 100644 .env.example create mode 100644 demo/one-api.ipynb create mode 100644 pyproject.toml delete mode 100644 setup.py delete mode 100644 src/one-api-cli/__init__.py create mode 100644 src/one_api_cli/__init__.py create mode 100644 src/one_api_cli/account.py create mode 100644 src/one_api_cli/channel.py create mode 100644 src/one_api_cli/constant.py diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..ddaf188 --- /dev/null +++ b/.env.example @@ -0,0 +1,6 @@ +# ONE API URL +ONE_API_BASE_URL= + +# ACCESS TOKEN at https://{one-api-url}/panel/profile +ONE_API_ACCESS_TOKEN= + diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 8f6f824..a421956 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -32,13 +32,10 @@ jobs: python -m pip install --upgrade pip python -m pip install .[test] python -m pip install -r requirements_dev.txt - - name: Test with pytest and coverage - env: - OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} - OPENAI_API_BASE_URL: ${{ secrets.OPENAI_API_BASE_URL }} - run: | - pip install coverage - coverage run -m pytest tests/ + # - name: Test with pytest and coverage + # run: | + # pip install coverage + # coverage run -m pytest tests/ - name: Upload coverage to Codecov uses: codecov/codecov-action@v4 with: diff --git a/README.md b/README.md index 638b4fb..ec566aa 100644 --- a/README.md +++ b/README.md @@ -15,8 +15,24 @@ -# Installation +A Python CLI for the [One API](https://github.com/songquanpeng/one-api) project. + +## Setup + +Install the package using pip: ```bash pip install one-api-cli ``` + +Setup environment variables: + +```bash +export ONE_API_BASE_URL=https://your_base_url +export ONE_API_ACCESS_TOKEN=your_access_token # Optional +export ONE_API_SESSION_TOKEN=your_session_token # Optional +``` + +## Usage + +See examples in [Jupyter notebooks](demo/one-api.ipynb). \ No newline at end of file diff --git a/demo/one-api.ipynb b/demo/one-api.ipynb new file mode 100644 index 0000000..cf1f7ef --- /dev/null +++ b/demo/one-api.ipynb @@ -0,0 +1,171 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 导入变量，或者从环境变量中加载\n", + "from dotenv import load_dotenv\n", + "load_dotenv(override=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 用户管理" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from one_api_cli import get_users, get_user, update_user, delete_user, create_user" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 查看用户\n", + "users = get_users()\n", + "for user in users:\n", + " print(user['id'], user['username'])\n", + "print(get_user(1))\n", + "print(get_user(100)) # 不存在的用户" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 新增用户\n", + "username = \"test2\"\n", + "display_name = \"test\"\n", + "password = \"complicated_password\"\n", + "create_user(username, display_name, password)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 修改用户信息\n", + "username = \"test2\"\n", + "new_password = \"new_password_233\"\n", + "update_user(username, password=new_password) # 也可以指定 ID" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 删除用户\n", + "delete_user(username) # 或使用 ID" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 渠道管理" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from one_api_cli import get_channels, get_channel, create_channel, update_channel, delete_channel" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 查看频道\n", + "channels = get_channels()\n", + "for channel in channels:\n", + " print(channel['id'], channel['name'])\n", + "print(get_channel(1))\n", + "print(get_channel(100)) # 不存在的频道" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 新增频道\n", + "name = \"test_channel\"\n", + "create_channel(name, 'sk-123', 'https://api.openai.com', 'gpt-test')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 修改频道信息\n", + "channel_id = [channel['id'] for channel in channels if channel['name'] == name][0]\n", + "new_name = \"new_channel\"\n", + "update_channel(channel_id, name=new_name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 删除频道\n", + "delete_channel(channel_id)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "base", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.13" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..e086ccb --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,54 @@ +[project] +name = "one-api-cli" +version = "0.1.0" +authors = [ + { name="Rex Wang", email="1073853456@qq.com" }, +] +description = "A CLI tool for the One API project." +keywords = ["oneapi", "cli", "tool"] +readme = "README.md" +license = { file = "LICENSE" } +requires-python = ">=3.9,<3.12" +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: MacOS", + "Operating System :: POSIX :: Linux", +] +dependencies = [ + "loguru", + "click", + "requests" +] + +[project.optional-dependencies] +dev = [ + "ipython", + "notebook", +] +test = [ + "pytest", +] +docs = [ + "sphinx", + "sphinx-rtd-theme", +] +lint = [ + "mypy", +] +all = [ + "ipython", + "notebook", + "pytest", + "sphinx", + "sphinx-rtd-theme", + "mypy", +] + +[project.urls] +"Bug Tracker" = "https://github.com/rexwzh/one-api-cli/issues" + +[tool.pytest.ini_options] +addopts = [ + "--import-mode=importlib", +] diff --git a/requirements_dev.txt b/requirements_dev.txt index 22e5671..0507590 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -1,10 +1,9 @@ -pip==19.2.3 bump2version==0.5.11 wheel==0.33.6 watchdog==0.9.0 flake8==3.7.8 tox==3.14.0 -coverage==4.5.4 +coverage Sphinx==1.8.5 twine==1.14.0 pytest==6.2.4 diff --git a/setup.py b/setup.py deleted file mode 100644 index 62b10e8..0000000 --- a/setup.py +++ /dev/null @@ -1,44 +0,0 @@ -#!/usr/bin/env python - -"""The setup script.""" - -from setuptools import setup, find_packages - -VERSION = '0.1.0' - -with open('README.md') as readme_file: - readme = readme_file.read() - -requirements = ['requests', 'click'] - -test_requirements = ['pytest>=3'] - -setup( - author="Rex Wang", - author_email='1073853456@qq.com', - python_requires='>=3.6', - classifiers=[ - 'Development Status :: 2 - Pre-Alpha', - 'Intended Audience :: Developers', - 'License :: OSI Approved :: MIT License', - 'Natural Language :: English', - 'Programming Language :: Python :: 3', - 'Programming Language :: Python :: 3.6', - 'Programming Language :: Python :: 3.7', - 'Programming Language :: Python :: 3.8', - 'Programming Language :: Python :: 3.9', - ], - description="A CLI for the One API project", - install_requires=requirements, - license="MIT license", - # long_description=readme + '\n\n' + history , - include_package_data=True, - keywords='One API CLI', - name='one-api-cli', - packages=find_packages(include=['src', 'src.*']), - test_suite='tests', - tests_require=test_requirements, - url='https://github.com/RexWzh/one-api-cli', - version=VERSION, - zip_safe=False, -) diff --git a/src/one-api-cli/__init__.py b/src/one-api-cli/__init__.py deleted file mode 100644 index 41bc4d1..0000000 --- a/src/one-api-cli/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Top-level package for one-api-cli.""" - -__author__ = """Rex Wang""" -__email__ = '1073853456@qq.com' -__version__ = '0.1.0' diff --git a/src/one_api_cli/__init__.py b/src/one_api_cli/__init__.py new file mode 100644 index 0000000..023a407 --- /dev/null +++ b/src/one_api_cli/__init__.py @@ -0,0 +1,8 @@ +"""Top-level package for one-api-cli.""" + +__author__ = """Rex Wang""" +__email__ = '1073853456@qq.com' +__version__ = '0.1.0' + +from .account import get_users, update_user, get_user, create_user, delete_user +from .channel import get_channels, update_channel, delete_channel, create_channel, get_channel \ No newline at end of file diff --git a/src/one_api_cli/account.py b/src/one_api_cli/account.py new file mode 100644 index 0000000..dfc829d --- /dev/null +++ b/src/one_api_cli/account.py @@ -0,0 +1,142 @@ +import requests +from .constant import base_url, headers +from loguru import logger +from typing import List, Union + +def get_users() -> List[dict]: + """ + Retrieve a list of users. + + Returns: + list: A list of user dictionaries. + """ + user_url = f"{base_url}/api/user" + try: + response = requests.get(user_url, headers=headers) + response.raise_for_status() + msg = response.json() + if not msg['success']: + logger.error(msg['message']) + return {} + return msg['data'] + except requests.RequestException as e: + logger.error(f"Error fetching users: {e}") + return [] + +def get_user(ind:int) -> dict: + """ + Retrieve the data of a user. + + Returns: + dict: A user dictionary. + """ + user_url = f"{base_url}/api/user/{ind}" + try: + response = requests.get(user_url, headers=headers) + response.raise_for_status() + msg = response.json() + if not msg['success']: + logger.error(msg['message']) + return {} + return msg['data'] + except requests.RequestException as e: + logger.error(f"Error fetching user: {e}") + return {} + +def create_user( + username:str, + display_name:str, + password:str, + group:str='default', + quota:int=0, + is_edit:bool=False + ) -> bool: + """ + Create a new user. + + Args: + username (str): The username of the user. + display_name (str): The display name of the user. + password (str): The password of the user. + group (str): The group of the user. + quota (int): The quota of the user. + is_edit (bool): Whether the user can edit. + + Returns: + bool: Whether the user was created successfully. + """ + user_url = f"{base_url}/api/user" + user_data = { + 'username': username, + 'display_name': display_name, + 'password': password, + 'group': group, + 'quota': quota, + 'is_edit': is_edit + } + try: + response = requests.post(user_url, headers=headers, json=user_data) + response.raise_for_status() + msg = response.json() + if not msg['success']: + logger.error(msg['message']) + return False + return True + except requests.RequestException as e: + logger.error(f"Error creating user: {e}") + return False + +def update_user(user_id:int, **options) -> bool: + """ + Update a user's data. + + Args: + user_id (int): The ID of the user. + **options: The data to update. + + Returns: + bool: Whether the user was updated successfully. + """ + user_url = f"{base_url}/api/user" + user_data = get_user(user_id) + if not user_data: + logger.error(f"User with ID {user_id} not found.") + return False + + user_data.update(options) + try: + response = requests.put(user_url, headers=headers, json=user_data) + response.raise_for_status() + msg = response.json() + if not msg['success']: + logger.error(msg['message']) + return False + return True + except requests.RequestException as e: + logger.error(f"Error updating user: {e}") + return False + +def delete_user(user:Union[int, str]) -> bool: + """ + Delete a user. + + Args: + user (int, str): The ID or username of the user. + + Returns: + str: The response message. + """ + delete_url = f"{base_url}/api/user/manage" + try: + username = get_user(user)['username'] if isinstance(user, int) else user + data = {"username":username, "action":"delete"} + response = requests.post(delete_url, headers=headers, json=data) + response.raise_for_status() + msg = response.json() + if not msg['success']: + logger.error(msg['message']) + return False + return True + except requests.RequestException as e: + logger.error(f"Error deleting user: {e}") + return False \ No newline at end of file diff --git a/src/one_api_cli/channel.py b/src/one_api_cli/channel.py new file mode 100644 index 0000000..a15950f --- /dev/null +++ b/src/one_api_cli/channel.py @@ -0,0 +1,151 @@ +import requests +from .constant import base_url, headers +from loguru import logger + +channel_url = f"{base_url}/api/channel" + +def get_channels(): + """ + Retrieve a list of channels. + + Returns: + list: A list of channel dictionaries. + """ + try: + response = requests.get(channel_url, headers=headers) + response.raise_for_status() + msg = response.json() + if not msg['success']: + logger.error(msg['message']) + return {} + return msg['data'] + except requests.RequestException as e: + logger.error(f"Error fetching channels: {e}") + return [] + +def get_channel(channel_id:int)->dict: + """ + Retrieve the data of a channel. + + Returns: + dict: A channel dictionary. + """ + channel_id_url = f"{channel_url}/{channel_id}" + try: + response = requests.get(channel_id_url, headers=headers) + response.raise_for_status() + msg = response.json() + if not msg['success']: + logger.error(msg['message']) + return {} + return msg['data'] + except requests.RequestException as e: + logger.error(f"Error fetching channel: {e}") + return {} + +def update_channel(channel_id, **options) -> bool: + """ + Update a channel's data. + + Args: + channel_id (int): The ID of the channel. + **options: The data to update. + + Returns: + bool: True if the channel is updated successfully, False otherwise. + """ + + try: + channel_data = get_channel(channel_id) + if not channel_data: + logger.error(f"Channel with ID {channel_id} not found.") + return False + channel_data.update(options) + response = requests.put(channel_url, headers=headers, json=channel_data) + response.raise_for_status() + msg = response.json() + if not msg['success']: + logger.error(msg['message']) + return False + return True + except requests.RequestException as e: + logger.error(f"Error updating channel: {e}") + return False + +def delete_channel(channel_id) -> bool: + """ + Delete a channel. + + Args: + channel_id (int): The ID of the channel. + + Returns: + bool: True if the channel is deleted successfully, False otherwise. + """ + channel_id_url = f"{channel_url}/{channel_id}" + try: + response = requests.delete(channel_id_url, headers=headers) + response.raise_for_status() + msg = response.json() + if not msg['success']: + logger.error(msg['message']) + return False + return True + except requests.RequestException as e: + logger.error(f"Error deleting channel: {e}") + return False + +def create_channel( + name, key, base_url, models, + type: int = 1, + other: str = '', + model_mapping: str = '', + groups: list = ['default'], + config: str = '{}', + is_edit: bool = False, + group: str = 'default' +) -> bool: + """ + Create a new channel. + + Args: + name (str): The name of the channel. + key (str): The key of the channel. + base_url (str): The base URL of the channel. + models (list): The models of the channel. + type (int): The type of the channel. Default to OpenAI. + other (str): Other information of the channel. + model_mapping (str): The model mapping of the channel. + groups (list): The groups of the channel. + config (str): The config of the channel. + is_edit (bool): Whether the channel can be edited. + group (str): The group of the channel. + + Returns: + bool: True if the channel is created successfully, False otherwise. + """ + + data = { + 'name': name, + 'key': key, + 'base_url': base_url, + 'models': models, + 'type': type, + 'other': other, + 'model_mapping': model_mapping, + 'groups': groups, + 'config': config, + 'is_edit': is_edit, + 'group': group + } + try: + response = requests.post(channel_url, headers=headers, json=data) + response.raise_for_status() + msg = response.json() + if not msg['success']: + logger.error(msg['message']) + return False + return True + except requests.RequestException as e: + logger.error(f"Error creating channel: {e}") + return False \ No newline at end of file diff --git a/src/one_api_cli/constant.py b/src/one_api_cli/constant.py new file mode 100644 index 0000000..752b999 --- /dev/null +++ b/src/one_api_cli/constant.py @@ -0,0 +1,23 @@ +import os + +# Environment variables +base_url = os.getenv("ONE_API_BASE_URL") +access_token = os.getenv("ONE_API_ACCESS_TOKEN") +section_token = os.getenv("ONE_API_SECTION_TOKEN") + + +assert base_url, "ONE_API_BASE_URL is not set" +assert access_token or section_token, "Either ONE_API_ACCESS_TOKEN or ONE_API_SECTION_TOKEN must be set" + +# Headers +if access_token: + headers = { + "Accept": "application/json, text/plain, */*", + "Authorization": f"Bearer {access_token}", + } +else: + section_token = section_token.lstrip("section=").strip("=") + headers = { + "Accept": "application/json, text/plain, */*", + "Cookie": f"section={section_token}=", + } \ No newline at end of file