rwkv-api 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,218 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[codz]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ .Python
11
+ build/
12
+ develop-eggs/
13
+ dist/
14
+ downloads/
15
+ eggs/
16
+ .eggs/
17
+ lib/
18
+ lib64/
19
+ parts/
20
+ sdist/
21
+ var/
22
+ wheels/
23
+ share/python-wheels/
24
+ *.egg-info/
25
+ .installed.cfg
26
+ *.egg
27
+ MANIFEST
28
+
29
+ # PyInstaller
30
+ # Usually these files are written by a python script from a template
31
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
32
+ *.manifest
33
+ *.spec
34
+
35
+ # Installer logs
36
+ pip-log.txt
37
+ pip-delete-this-directory.txt
38
+
39
+ # Unit test / coverage reports
40
+ htmlcov/
41
+ .tox/
42
+ .nox/
43
+ .coverage
44
+ .coverage.*
45
+ .cache
46
+ nosetests.xml
47
+ coverage.xml
48
+ *.cover
49
+ *.py.cover
50
+ .hypothesis/
51
+ .pytest_cache/
52
+ cover/
53
+
54
+ # Translations
55
+ *.mo
56
+ *.pot
57
+
58
+ # Django stuff:
59
+ *.log
60
+ local_settings.py
61
+ db.sqlite3
62
+ db.sqlite3-journal
63
+
64
+ # Flask stuff:
65
+ instance/
66
+ .webassets-cache
67
+
68
+ # Scrapy stuff:
69
+ .scrapy
70
+
71
+ # Sphinx documentation
72
+ docs/_build/
73
+
74
+ # PyBuilder
75
+ .pybuilder/
76
+ target/
77
+
78
+ # Jupyter Notebook
79
+ .ipynb_checkpoints
80
+
81
+ # IPython
82
+ profile_default/
83
+ ipython_config.py
84
+
85
+ # pyenv
86
+ # For a library or package, you might want to ignore these files since the code is
87
+ # intended to run in multiple environments; otherwise, check them in:
88
+ # .python-version
89
+
90
+ # pipenv
91
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
92
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
93
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
94
+ # install all needed dependencies.
95
+ # Pipfile.lock
96
+
97
+ # UV
98
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
99
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
100
+ # commonly ignored for libraries.
101
+ # uv.lock
102
+
103
+ # poetry
104
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
105
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
106
+ # commonly ignored for libraries.
107
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
108
+ # poetry.lock
109
+ # poetry.toml
110
+
111
+ # pdm
112
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
113
+ # pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
114
+ # https://pdm-project.org/en/latest/usage/project/#working-with-version-control
115
+ # pdm.lock
116
+ # pdm.toml
117
+ .pdm-python
118
+ .pdm-build/
119
+
120
+ # pixi
121
+ # Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
122
+ # pixi.lock
123
+ # Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
124
+ # in the .venv directory. It is recommended not to include this directory in version control.
125
+ .pixi
126
+
127
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
128
+ __pypackages__/
129
+
130
+ # Celery stuff
131
+ celerybeat-schedule
132
+ celerybeat.pid
133
+
134
+ # Redis
135
+ *.rdb
136
+ *.aof
137
+ *.pid
138
+
139
+ # RabbitMQ
140
+ mnesia/
141
+ rabbitmq/
142
+ rabbitmq-data/
143
+
144
+ # ActiveMQ
145
+ activemq-data/
146
+
147
+ # SageMath parsed files
148
+ *.sage.py
149
+
150
+ # Environments
151
+ .env
152
+ .envrc
153
+ .venv
154
+ env/
155
+ venv/
156
+ ENV/
157
+ env.bak/
158
+ venv.bak/
159
+
160
+ # Spyder project settings
161
+ .spyderproject
162
+ .spyproject
163
+
164
+ # Rope project settings
165
+ .ropeproject
166
+
167
+ # mkdocs documentation
168
+ /site
169
+
170
+ # mypy
171
+ .mypy_cache/
172
+ .dmypy.json
173
+ dmypy.json
174
+
175
+ # Pyre type checker
176
+ .pyre/
177
+
178
+ # pytype static type analyzer
179
+ .pytype/
180
+
181
+ # Cython debug symbols
182
+ cython_debug/
183
+
184
+ # PyCharm
185
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
186
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
187
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
188
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
189
+ # .idea/
190
+
191
+ # Abstra
192
+ # Abstra is an AI-powered process automation framework.
193
+ # Ignore directories containing user credentials, local state, and settings.
194
+ # Learn more at https://abstra.io/docs
195
+ .abstra/
196
+
197
+ # Visual Studio Code
198
+ # Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
199
+ # that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
200
+ # and can be added to the global gitignore or merged into this file. However, if you prefer,
201
+ # you could uncomment the following to ignore the entire vscode folder
202
+ # .vscode/
203
+ # Temporary file for partial code execution
204
+ tempCodeRunnerFile.py
205
+
206
+ # Ruff stuff:
207
+ .ruff_cache/
208
+
209
+ # PyPI configuration file
210
+ .pypirc
211
+
212
+ # Marimo
213
+ marimo/_static/
214
+ marimo/_lsp/
215
+ __marimo__/
216
+
217
+ # Streamlit
218
+ .streamlit/secrets.toml
rwkv_api-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 宁静致远
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,285 @@
1
+ Metadata-Version: 2.4
2
+ Name: rwkv-api
3
+ Version: 0.1.0
4
+ Summary: RWKV-Server Task API 的 Python SDK,支持同步/异步双接口和流式 SSE 响应
5
+ Project-URL: Homepage, https://github.com/RWKV/RWKV-API
6
+ Project-URL: Repository, https://github.com/RWKV/RWKV-API
7
+ Author: RWKV Contributors
8
+ License: MIT License
9
+
10
+ Copyright (c) 2026 宁静致远
11
+
12
+ Permission is hereby granted, free of charge, to any person obtaining a copy
13
+ of this software and associated documentation files (the "Software"), to deal
14
+ in the Software without restriction, including without limitation the rights
15
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
16
+ copies of the Software, and to permit persons to whom the Software is
17
+ furnished to do so, subject to the following conditions:
18
+
19
+ The above copyright notice and this permission notice shall be included in all
20
+ copies or substantial portions of the Software.
21
+
22
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
23
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
24
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
25
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
26
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
27
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
28
+ SOFTWARE.
29
+ License-File: LICENSE
30
+ Keywords: api,llm,rwkv,sdk,sse,streaming
31
+ Classifier: Development Status :: 3 - Alpha
32
+ Classifier: Intended Audience :: Developers
33
+ Classifier: License :: OSI Approved :: MIT License
34
+ Classifier: Programming Language :: Python :: 3
35
+ Classifier: Programming Language :: Python :: 3.10
36
+ Classifier: Programming Language :: Python :: 3.11
37
+ Classifier: Programming Language :: Python :: 3.12
38
+ Classifier: Programming Language :: Python :: 3.13
39
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
40
+ Requires-Python: >=3.10
41
+ Requires-Dist: httpx>=0.24.0
42
+ Requires-Dist: pydantic>=2.0
43
+ Provides-Extra: dev
44
+ Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
45
+ Requires-Dist: pytest>=7.0; extra == 'dev'
46
+ Requires-Dist: respx>=0.20.0; extra == 'dev'
47
+ Description-Content-Type: text/markdown
48
+
49
+ # RWKV-API
50
+
51
+ RWKV-Server 的官方 Python SDK,**以最优雅的方式调用 RWKV 推理**。
52
+
53
+ > 无需关心批量调度、状态管理、GPU 显存分配——只需创建一个 `Task`,RWKV-Server 会在后台自动完成一切。
54
+
55
+ ## 配套服务端
56
+
57
+ 本 SDK 对应 [RWKV-Server](https://github.com/AUXStar/RWKV-Server) 的 Task API。服务端基于 RWKV7 线性注意力架构,支持动态批量推理、2的幂次自动扩缩容、模板持久化等高级特性,单机可达 **11,081 tok/s** 峰值吞吐。
58
+
59
+ ## 安装
60
+
61
+ ```bash
62
+ pip install rwkv-api
63
+ ```
64
+
65
+ ## 快速开始
66
+
67
+ ### 同步调用
68
+
69
+ ```python
70
+ from rwkv_api import Client
71
+
72
+ client = Client("http://localhost:8000")
73
+
74
+ # 创建任务并等待结果 —— 服务端自动批量调度
75
+ task = client.create("Hello world", max_tokens=50)
76
+ result = task.wait()
77
+ print(result.result)
78
+
79
+ # FIM 代码补全
80
+ result = client.fim(
81
+ prefix="def add(a, b):\n ",
82
+ suffix="\n return result",
83
+ max_tokens=30
84
+ )
85
+ print(result.result)
86
+ ```
87
+
88
+ ### 异步调用
89
+
90
+ ```python
91
+ import asyncio
92
+ from rwkv_api import AsyncClient
93
+
94
+ async def main():
95
+ client = AsyncClient("http://localhost:8000")
96
+ task = await client.create("Hello world", max_tokens=50)
97
+ result = await task.wait()
98
+ print(result.result)
99
+
100
+ asyncio.run(main())
101
+ ```
102
+
103
+ ### 实时流式(创建时)
104
+
105
+ ```python
106
+ from rwkv_api import Client
107
+
108
+ client = Client("http://localhost:8000")
109
+
110
+ # 方式1:create_stream 显式流式
111
+ for chunk in client.create_stream("Tell me a story", max_tokens=200):
112
+ print(chunk, end="", flush=True)
113
+
114
+ # 方式2:create(stream=True) 自动分发
115
+ for chunk in client.create("Tell me a story", max_tokens=200, stream=True):
116
+ print(chunk, end="", flush=True)
117
+ ```
118
+
119
+ ### 对已创建任务订阅流式
120
+
121
+ ```python
122
+ from rwkv_api import Client
123
+
124
+ client = Client("http://localhost:8000")
125
+
126
+ # 先创建任务(非流式)
127
+ task = client.create("Hello", max_tokens=200)
128
+
129
+ # 稍后订阅实时流式输出
130
+ for chunk in task.stream():
131
+ print(chunk, end="", flush=True)
132
+
133
+ # 或者直接用 task_id
134
+ for chunk in client.stream_task(task.task_id):
135
+ print(chunk, end="", flush=True)
136
+ ```
137
+
138
+ ### 异步实时流式
139
+
140
+ ```python
141
+ import asyncio
142
+ from rwkv_api import AsyncClient
143
+
144
+ async def main():
145
+ client = AsyncClient("http://localhost:8000")
146
+
147
+ # 创建时流式
148
+ async for chunk in client.create_stream("Tell me a story", max_tokens=200):
149
+ print(chunk, end="", flush=True)
150
+
151
+ # 对已创建任务订阅流式
152
+ task = await client.create("Hello", max_tokens=200)
153
+ async for chunk in task.stream():
154
+ print(chunk, end="", flush=True)
155
+
156
+ asyncio.run(main())
157
+ ```
158
+
159
+ ## Task 对象 —— 优雅调用 RWKV 推理的核心
160
+
161
+ SDK 的核心是 `Task` 对象。你只需关注**要生成什么**,所有底层细节(批量调度、状态生命周期、GPU 显存管理)都由服务端自动处理:
162
+
163
+ ```python
164
+ task = client.create("Hello", max_tokens=50)
165
+
166
+ # 属性
167
+ print(task.task_id) # "TASK_xxx"
168
+ print(task.finished) # False
169
+ print(task.result) # 生成结果(完成后可用)
170
+ print(task.speed) # tok/s
171
+ print(task.prefill_time) # prefill 耗时(秒)
172
+ print(task.gen_time) # 生成耗时(秒)
173
+
174
+ # 操作
175
+ result = task.wait() # 阻塞等待完成,返回 self
176
+ task.stop() # 停止生成
177
+ forked = task.fork() # Fork 新任务
178
+ continued = task.continue_() # 继续生成
179
+ tpl = task.as_template() # 转为模板
180
+ task.delete() # 删除任务
181
+
182
+ # 流式订阅(任务创建后仍可实时获取输出)
183
+ for chunk in task.stream():
184
+ print(chunk, end="", flush=True)
185
+ ```
186
+
187
+ ## 流式说明
188
+
189
+ SDK 提供两种流式方式:
190
+
191
+ | 方式 | 方法 | 返回值 | 适用场景 |
192
+ |------|------|--------|---------|
193
+ | **创建时流式** | `create_stream()` / `create(stream=True)` | `Iterator[str]` / `AsyncIterator[str]` | 创建任务的同时实时输出 |
194
+ | **任务订阅流式** | `task.stream()` / `client.stream_task(id)` | `Iterator[str]` / `AsyncIterator[str]` | 对已创建任务实时订阅 |
195
+
196
+ **注意**:
197
+ - `create(stream=True)` 返回**生成器**,不是 `Task` 对象
198
+ - `task.stream()` 支持多消费者同时连接,服务端广播到所有订阅者
199
+ - 如果任务已完成,`stream()` 一次性返回完整结果
200
+
201
+ ## 异常处理
202
+
203
+ ```python
204
+ from rwkv_api import Client
205
+ from rwkv_api.exceptions import (
206
+ RWKVValidationError, # 422 参数错误
207
+ TaskNotFoundError, # 404 任务不存在
208
+ RWKVServerError, # 500 服务端错误
209
+ TimeoutError, # 等待超时
210
+ ConnectionError, # 连接失败
211
+ )
212
+
213
+ client = Client("http://localhost:8000")
214
+
215
+ try:
216
+ task = client.create("Hello", max_tokens=50)
217
+ result = task.wait(timeout=10)
218
+ except RWKVValidationError as e:
219
+ print(f"参数错误: {e}")
220
+ except TimeoutError:
221
+ print("等待超时")
222
+ except RWKVServerError as e:
223
+ print(f"服务端错误: {e}")
224
+ ```
225
+
226
+ ## API 参考
227
+
228
+ ### Client / AsyncClient
229
+
230
+ | 方法 | 说明 |
231
+ |------|------|
232
+ | `create(prompt, ...)` | 创建任务,`stream=True` 时返回生成器 |
233
+ | `create_stream(prompt, ...)` | **实时流式**,返回生成器 |
234
+ | `create_tmp(prompt, ...)` | 创建临时任务 |
235
+ | `create_persistent(prompt, ...)` | 创建持久化任务 |
236
+ | `fim(prefix, suffix, ...)` | FIM 代码补全,`stream=True` 时返回生成器 |
237
+ | `fim_stream(prefix, suffix, ...)` | **FIM 实时流式** |
238
+ | `stream_task(task_id)` | **订阅已创建任务的流式输出** |
239
+ | `list_tasks()` | 列出所有任务 |
240
+ | `get_task_result(task_id)` | 获取任务结果 |
241
+ | `get_task_status(task_id)` | 获取任务状态 |
242
+ | `stop_task(task_id)` | 停止任务 |
243
+ | `delete_task(task_id, force=False)` | 删除任务 |
244
+ | `fork_task(task_id, ...)` | Fork 任务 |
245
+ | `continue_task(task_id, ...)` | 继续生成 |
246
+ | `as_template(task_id)` | 转为模板 |
247
+
248
+ ### Task / AsyncTask
249
+
250
+ | 方法 | 说明 |
251
+ |------|------|
252
+ | `task.wait()` | 等待任务完成 |
253
+ | `task.stop()` | 停止生成 |
254
+ | `task.fork()` | Fork 新任务 |
255
+ | `task.continue_()` | 继续生成 |
256
+ | `task.as_template()` | 转为模板 |
257
+ | `task.delete()` | 删除任务 |
258
+ | `task.stream()` | **订阅此任务的实时流式输出** |
259
+
260
+ ### 生成参数
261
+
262
+ | 参数 | 类型 | 默认值 | 说明 |
263
+ |------|------|--------|------|
264
+ | `max_tokens` | `int` | 256 | 最大生成 token 数 |
265
+ | `temperature` | `float` | 1.0 | 采样温度 (0~2) |
266
+ | `top_k` | `int` | 0 | Top-K 采样 (0=禁用) |
267
+ | `top_p` | `float` | 0.8 | Top-P 采样 |
268
+ | `presence_penalty` | `float` | 0.0 | 存在惩罚 (0~10) |
269
+ | `repetition_penalty` | `float` | 0.0 | 重复惩罚 (0~10) |
270
+ | `penalty_decay` | `float` | 0.0 | 惩罚衰减 (0~1) |
271
+ | `seed` | `int \| None` | None | 随机种子 |
272
+ | `stream` | `bool` | False | 流式输出 |
273
+
274
+ ## 开发
275
+
276
+ ```bash
277
+ git clone https://github.com/AUXStar/RWKV-API.git
278
+ cd RWKV-API
279
+ pip install -e ".[dev]"
280
+ pytest tests/ -v
281
+ ```
282
+
283
+ ## 许可证
284
+
285
+ MIT