aikit-platform 0.3.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.
- aikit_platform-0.3.0/.github/workflows/ci.yml +15 -0
- aikit_platform-0.3.0/.github/workflows/release.yml +16 -0
- aikit_platform-0.3.0/.gitignore +12 -0
- aikit_platform-0.3.0/CHANGELOG.md +7 -0
- aikit_platform-0.3.0/LICENSE +21 -0
- aikit_platform-0.3.0/Makefile +6 -0
- aikit_platform-0.3.0/PKG-INFO +147 -0
- aikit_platform-0.3.0/README.md +81 -0
- aikit_platform-0.3.0/pyproject.toml +60 -0
- aikit_platform-0.3.0/src/aikit/__init__.py +6 -0
- aikit_platform-0.3.0/src/aikit/config/__init__.py +13 -0
- aikit_platform-0.3.0/src/aikit/db/__init__.py +65 -0
- aikit_platform-0.3.0/src/aikit/embeddings/__init__.py +20 -0
- aikit_platform-0.3.0/src/aikit/embeddings/sentence_transformer.py +46 -0
- aikit_platform-0.3.0/src/aikit/jobqueue/__init__.py +15 -0
- aikit_platform-0.3.0/src/aikit/jobqueue/arq_queue.py +56 -0
- aikit_platform-0.3.0/src/aikit/model_client/__init__.py +64 -0
- aikit_platform-0.3.0/src/aikit/model_client/hosted.py +108 -0
- aikit_platform-0.3.0/src/aikit/model_client/ollama.py +87 -0
- aikit_platform-0.3.0/src/aikit/observability/__init__.py +77 -0
- aikit_platform-0.3.0/src/aikit/py.typed +0 -0
- aikit_platform-0.3.0/tests/test_db.py +46 -0
- aikit_platform-0.3.0/tests/test_jobqueue_arq.py +37 -0
- aikit_platform-0.3.0/tests/test_model_client_hosted.py +73 -0
- aikit_platform-0.3.0/tests/test_model_client_ollama.py +75 -0
- aikit_platform-0.3.0/tests/test_observability.py +15 -0
- aikit_platform-0.3.0/tests/test_smoke.py +17 -0
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
name: ci
|
|
2
|
+
on:
|
|
3
|
+
push: { branches: [main] }
|
|
4
|
+
pull_request:
|
|
5
|
+
jobs:
|
|
6
|
+
test:
|
|
7
|
+
runs-on: ubuntu-latest
|
|
8
|
+
steps:
|
|
9
|
+
- uses: actions/checkout@v4
|
|
10
|
+
- uses: actions/setup-python@v5
|
|
11
|
+
with: { python-version: "3.11" }
|
|
12
|
+
- run: pip install -e ".[eval,dev]"
|
|
13
|
+
- run: ruff check .
|
|
14
|
+
- run: mypy src
|
|
15
|
+
- run: pytest -q
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
name: release
|
|
2
|
+
on:
|
|
3
|
+
push:
|
|
4
|
+
tags: ["v*"]
|
|
5
|
+
permissions:
|
|
6
|
+
id-token: write # PyPI trusted publishing (OIDC) — no API token needed
|
|
7
|
+
jobs:
|
|
8
|
+
publish:
|
|
9
|
+
runs-on: ubuntu-latest
|
|
10
|
+
steps:
|
|
11
|
+
- uses: actions/checkout@v4
|
|
12
|
+
- uses: actions/setup-python@v5
|
|
13
|
+
with: { python-version: "3.12" }
|
|
14
|
+
- run: pip install build hatchling
|
|
15
|
+
- run: python -m build
|
|
16
|
+
- uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 YOUR NAME
|
|
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,147 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: aikit-platform
|
|
3
|
+
Version: 0.3.0
|
|
4
|
+
Summary: Shared platform layer for AI-backend systems: model client, embeddings, observability, config, db, queue.
|
|
5
|
+
Project-URL: Homepage, https://github.com/Aaryan123456679/aikit
|
|
6
|
+
Project-URL: Repository, https://github.com/Aaryan123456679/aikit
|
|
7
|
+
Author: Aaryan Mahajan
|
|
8
|
+
License: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Keywords: embeddings,infrastructure,llm,observability
|
|
11
|
+
Requires-Python: >=3.11
|
|
12
|
+
Requires-Dist: pydantic-settings>=2.3
|
|
13
|
+
Requires-Dist: pydantic>=2.7
|
|
14
|
+
Provides-Extra: db
|
|
15
|
+
Requires-Dist: asyncpg>=0.29; extra == 'db'
|
|
16
|
+
Requires-Dist: pgvector>=0.3; extra == 'db'
|
|
17
|
+
Requires-Dist: sqlalchemy[asyncio]>=2.0; extra == 'db'
|
|
18
|
+
Provides-Extra: dev
|
|
19
|
+
Requires-Dist: aiosqlite>=0.20; extra == 'dev'
|
|
20
|
+
Requires-Dist: mypy>=1.10; extra == 'dev'
|
|
21
|
+
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
|
|
22
|
+
Requires-Dist: pytest>=8.2; extra == 'dev'
|
|
23
|
+
Requires-Dist: ruff>=0.5; extra == 'dev'
|
|
24
|
+
Provides-Extra: embeddings
|
|
25
|
+
Requires-Dist: sentence-transformers>=3.0; extra == 'embeddings'
|
|
26
|
+
Provides-Extra: eval
|
|
27
|
+
Requires-Dist: arq>=0.26; extra == 'eval'
|
|
28
|
+
Requires-Dist: asyncpg>=0.29; extra == 'eval'
|
|
29
|
+
Requires-Dist: httpx>=0.27; extra == 'eval'
|
|
30
|
+
Requires-Dist: pgvector>=0.3; extra == 'eval'
|
|
31
|
+
Requires-Dist: prometheus-client>=0.20; extra == 'eval'
|
|
32
|
+
Requires-Dist: redis>=5.0; extra == 'eval'
|
|
33
|
+
Requires-Dist: sentence-transformers>=3.0; extra == 'eval'
|
|
34
|
+
Requires-Dist: sqlalchemy[asyncio]>=2.0; extra == 'eval'
|
|
35
|
+
Requires-Dist: structlog>=24.1; extra == 'eval'
|
|
36
|
+
Provides-Extra: gateway
|
|
37
|
+
Requires-Dist: asyncpg>=0.29; extra == 'gateway'
|
|
38
|
+
Requires-Dist: httpx>=0.27; extra == 'gateway'
|
|
39
|
+
Requires-Dist: pgvector>=0.3; extra == 'gateway'
|
|
40
|
+
Requires-Dist: prometheus-client>=0.20; extra == 'gateway'
|
|
41
|
+
Requires-Dist: redis>=5.0; extra == 'gateway'
|
|
42
|
+
Requires-Dist: sentence-transformers>=3.0; extra == 'gateway'
|
|
43
|
+
Requires-Dist: sqlalchemy[asyncio]>=2.0; extra == 'gateway'
|
|
44
|
+
Requires-Dist: structlog>=24.1; extra == 'gateway'
|
|
45
|
+
Provides-Extra: http
|
|
46
|
+
Requires-Dist: httpx>=0.27; extra == 'http'
|
|
47
|
+
Provides-Extra: memory
|
|
48
|
+
Requires-Dist: arq>=0.26; extra == 'memory'
|
|
49
|
+
Requires-Dist: asyncpg>=0.29; extra == 'memory'
|
|
50
|
+
Requires-Dist: httpx>=0.27; extra == 'memory'
|
|
51
|
+
Requires-Dist: pgvector>=0.3; extra == 'memory'
|
|
52
|
+
Requires-Dist: prometheus-client>=0.20; extra == 'memory'
|
|
53
|
+
Requires-Dist: redis>=5.0; extra == 'memory'
|
|
54
|
+
Requires-Dist: sentence-transformers>=3.0; extra == 'memory'
|
|
55
|
+
Requires-Dist: sqlalchemy[asyncio]>=2.0; extra == 'memory'
|
|
56
|
+
Requires-Dist: structlog>=24.1; extra == 'memory'
|
|
57
|
+
Provides-Extra: observability
|
|
58
|
+
Requires-Dist: prometheus-client>=0.20; extra == 'observability'
|
|
59
|
+
Requires-Dist: structlog>=24.1; extra == 'observability'
|
|
60
|
+
Provides-Extra: queue
|
|
61
|
+
Requires-Dist: arq>=0.26; extra == 'queue'
|
|
62
|
+
Requires-Dist: redis>=5.0; extra == 'queue'
|
|
63
|
+
Provides-Extra: redis
|
|
64
|
+
Requires-Dist: redis>=5.0; extra == 'redis'
|
|
65
|
+
Description-Content-Type: text/markdown
|
|
66
|
+
|
|
67
|
+
# aikit
|
|
68
|
+
|
|
69
|
+
*The shared library behind three AI-backend services — see the [AI Infrastructure Suite overview](https://github.com/Aaryan123456679/ai-infrastructure-suite).*
|
|
70
|
+
|
|
71
|
+
Shared platform layer for a suite of AI-backend systems (LLM eval platform,
|
|
72
|
+
inference gateway, agent-memory service). Publishes clean, SOLID interfaces
|
|
73
|
+
and reusable infrastructure so each downstream service depends on contracts,
|
|
74
|
+
not implementations.
|
|
75
|
+
|
|
76
|
+
## Modules
|
|
77
|
+
| Module | Contract | Concrete implementation |
|
|
78
|
+
|---|---|---|
|
|
79
|
+
| `model_client` | `ModelClient` | `OllamaClient`, `HostedClient` (OpenAI-compatible) — `complete` / `stream` (text + terminal `Usage`) / `health` / `aclose` |
|
|
80
|
+
| `embeddings` | `EmbeddingService` | `SentenceTransformerEmbeddingService` — local batched embeddings + cosine |
|
|
81
|
+
| `db` | async SQLAlchemy `Base` | `make_engine` / `make_sessionmaker` / `session_scope` |
|
|
82
|
+
| `jobqueue` | `JobQueue` | `ArqJobQueue` (Redis/arq) |
|
|
83
|
+
| `observability` | `MetricsSink`, logging | `PrometheusMetricsSink`, structlog JSON |
|
|
84
|
+
| `config` | `BaseServiceSettings` | env-driven |
|
|
85
|
+
|
|
86
|
+
## Status
|
|
87
|
+
|
|
88
|
+
**v0.1.0** validated end-to-end by `eval-platform`'s live Docker smoke test:
|
|
89
|
+
a real Postgres migration + advisory-lock-guarded prompt versioning +
|
|
90
|
+
CAS-guarded run finalization, a real Redis-backed arq worker actually
|
|
91
|
+
dispatching jobs, and a real Ollama model producing a scored, completed run
|
|
92
|
+
through the API. That process caught and fixed real bugs — not just code
|
|
93
|
+
review findings:
|
|
94
|
+
|
|
95
|
+
- `sentence-transformers` version drift breaking `EmbeddingService`'s
|
|
96
|
+
structural match (`dim` typed wider than `int` in a newer release).
|
|
97
|
+
- `ArqJobQueue` never dispatched a single job — arq names a job function by
|
|
98
|
+
`coroutine.__qualname__`, not `__name__`, so every registration collided
|
|
99
|
+
under the same closure qualname regardless of task name.
|
|
100
|
+
- CI installed only `dev` extras, so it silently only ever type-checked and
|
|
101
|
+
tested the core Protocol stubs, never the concrete implementations.
|
|
102
|
+
|
|
103
|
+
**v0.2.0** (breaking: `stream()` now yields `str | Usage`, terminal `Usage`
|
|
104
|
+
carries `tokens_in`/`tokens_out`/`finish_reason` — additive to `complete()`,
|
|
105
|
+
`eval-platform` stays pinned at v0.1.0 and is unaffected since it never
|
|
106
|
+
calls `stream()`). Adds `HostedClient` (OpenAI-compatible: OpenAI, Groq,
|
|
107
|
+
Together, Fireworks, OpenRouter, ...) for `inference-gateway`. Both closed
|
|
108
|
+
gaps `eval-platform`'s status page called out:
|
|
109
|
+
|
|
110
|
+
- `OllamaClient.stream()` is now **live-validated** against a real running
|
|
111
|
+
Ollama server, including a production-sized model (`llama3.1:8b`, not
|
|
112
|
+
just the small model used for eval-platform's fast smoke run) — text
|
|
113
|
+
deltas followed by exactly one terminal `Usage`, asserted live and
|
|
114
|
+
pinned down as a regression test replaying the real captured response.
|
|
115
|
+
- `HostedClient` is new and mock-tested against the OpenAI-compatible wire
|
|
116
|
+
format (including the `stream_options.include_usage` SSE shape); it has
|
|
117
|
+
**not** been live-validated against a real hosted provider (needs a real
|
|
118
|
+
API key) — that's `inference-gateway`'s job before it goes live there.
|
|
119
|
+
|
|
120
|
+
**Still open**: `JudgeScorer` (in eval-platform) has no live-model test.
|
|
121
|
+
|
|
122
|
+
## Install
|
|
123
|
+
Published on PyPI as `aikit-platform` (the import name is still `aikit` -
|
|
124
|
+
only the distribution name differs, because the name `aikit` itself is
|
|
125
|
+
already taken by an unrelated package). Core is light; heavy deps are
|
|
126
|
+
extras:
|
|
127
|
+
```bash
|
|
128
|
+
pip install "aikit-platform[eval]" # or [gateway], [memory]
|
|
129
|
+
pip install "aikit-platform[embeddings,db,queue,observability,http]" # à la carte
|
|
130
|
+
```
|
|
131
|
+
```python
|
|
132
|
+
import aikit # same import either way
|
|
133
|
+
```
|
|
134
|
+
Pin an exact tag from GitHub instead, if you want the git history alongside
|
|
135
|
+
the code (this is what eval-platform/gateway/agent-memory all do):
|
|
136
|
+
```bash
|
|
137
|
+
pip install "aikit-platform[gateway] @ git+https://github.com/Aaryan123456679/aikit@v0.2.0"
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
## Develop
|
|
141
|
+
```bash
|
|
142
|
+
make install && make all
|
|
143
|
+
```
|
|
144
|
+
|
|
145
|
+
## Release
|
|
146
|
+
Tag-driven: `git tag vX.Y.Z && git push origin vX.Y.Z` → CI builds and
|
|
147
|
+
publishes to PyPI via trusted publishing (OIDC, no secrets in the repo).
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
# aikit
|
|
2
|
+
|
|
3
|
+
*The shared library behind three AI-backend services — see the [AI Infrastructure Suite overview](https://github.com/Aaryan123456679/ai-infrastructure-suite).*
|
|
4
|
+
|
|
5
|
+
Shared platform layer for a suite of AI-backend systems (LLM eval platform,
|
|
6
|
+
inference gateway, agent-memory service). Publishes clean, SOLID interfaces
|
|
7
|
+
and reusable infrastructure so each downstream service depends on contracts,
|
|
8
|
+
not implementations.
|
|
9
|
+
|
|
10
|
+
## Modules
|
|
11
|
+
| Module | Contract | Concrete implementation |
|
|
12
|
+
|---|---|---|
|
|
13
|
+
| `model_client` | `ModelClient` | `OllamaClient`, `HostedClient` (OpenAI-compatible) — `complete` / `stream` (text + terminal `Usage`) / `health` / `aclose` |
|
|
14
|
+
| `embeddings` | `EmbeddingService` | `SentenceTransformerEmbeddingService` — local batched embeddings + cosine |
|
|
15
|
+
| `db` | async SQLAlchemy `Base` | `make_engine` / `make_sessionmaker` / `session_scope` |
|
|
16
|
+
| `jobqueue` | `JobQueue` | `ArqJobQueue` (Redis/arq) |
|
|
17
|
+
| `observability` | `MetricsSink`, logging | `PrometheusMetricsSink`, structlog JSON |
|
|
18
|
+
| `config` | `BaseServiceSettings` | env-driven |
|
|
19
|
+
|
|
20
|
+
## Status
|
|
21
|
+
|
|
22
|
+
**v0.1.0** validated end-to-end by `eval-platform`'s live Docker smoke test:
|
|
23
|
+
a real Postgres migration + advisory-lock-guarded prompt versioning +
|
|
24
|
+
CAS-guarded run finalization, a real Redis-backed arq worker actually
|
|
25
|
+
dispatching jobs, and a real Ollama model producing a scored, completed run
|
|
26
|
+
through the API. That process caught and fixed real bugs — not just code
|
|
27
|
+
review findings:
|
|
28
|
+
|
|
29
|
+
- `sentence-transformers` version drift breaking `EmbeddingService`'s
|
|
30
|
+
structural match (`dim` typed wider than `int` in a newer release).
|
|
31
|
+
- `ArqJobQueue` never dispatched a single job — arq names a job function by
|
|
32
|
+
`coroutine.__qualname__`, not `__name__`, so every registration collided
|
|
33
|
+
under the same closure qualname regardless of task name.
|
|
34
|
+
- CI installed only `dev` extras, so it silently only ever type-checked and
|
|
35
|
+
tested the core Protocol stubs, never the concrete implementations.
|
|
36
|
+
|
|
37
|
+
**v0.2.0** (breaking: `stream()` now yields `str | Usage`, terminal `Usage`
|
|
38
|
+
carries `tokens_in`/`tokens_out`/`finish_reason` — additive to `complete()`,
|
|
39
|
+
`eval-platform` stays pinned at v0.1.0 and is unaffected since it never
|
|
40
|
+
calls `stream()`). Adds `HostedClient` (OpenAI-compatible: OpenAI, Groq,
|
|
41
|
+
Together, Fireworks, OpenRouter, ...) for `inference-gateway`. Both closed
|
|
42
|
+
gaps `eval-platform`'s status page called out:
|
|
43
|
+
|
|
44
|
+
- `OllamaClient.stream()` is now **live-validated** against a real running
|
|
45
|
+
Ollama server, including a production-sized model (`llama3.1:8b`, not
|
|
46
|
+
just the small model used for eval-platform's fast smoke run) — text
|
|
47
|
+
deltas followed by exactly one terminal `Usage`, asserted live and
|
|
48
|
+
pinned down as a regression test replaying the real captured response.
|
|
49
|
+
- `HostedClient` is new and mock-tested against the OpenAI-compatible wire
|
|
50
|
+
format (including the `stream_options.include_usage` SSE shape); it has
|
|
51
|
+
**not** been live-validated against a real hosted provider (needs a real
|
|
52
|
+
API key) — that's `inference-gateway`'s job before it goes live there.
|
|
53
|
+
|
|
54
|
+
**Still open**: `JudgeScorer` (in eval-platform) has no live-model test.
|
|
55
|
+
|
|
56
|
+
## Install
|
|
57
|
+
Published on PyPI as `aikit-platform` (the import name is still `aikit` -
|
|
58
|
+
only the distribution name differs, because the name `aikit` itself is
|
|
59
|
+
already taken by an unrelated package). Core is light; heavy deps are
|
|
60
|
+
extras:
|
|
61
|
+
```bash
|
|
62
|
+
pip install "aikit-platform[eval]" # or [gateway], [memory]
|
|
63
|
+
pip install "aikit-platform[embeddings,db,queue,observability,http]" # à la carte
|
|
64
|
+
```
|
|
65
|
+
```python
|
|
66
|
+
import aikit # same import either way
|
|
67
|
+
```
|
|
68
|
+
Pin an exact tag from GitHub instead, if you want the git history alongside
|
|
69
|
+
the code (this is what eval-platform/gateway/agent-memory all do):
|
|
70
|
+
```bash
|
|
71
|
+
pip install "aikit-platform[gateway] @ git+https://github.com/Aaryan123456679/aikit@v0.2.0"
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
## Develop
|
|
75
|
+
```bash
|
|
76
|
+
make install && make all
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
## Release
|
|
80
|
+
Tag-driven: `git tag vX.Y.Z && git push origin vX.Y.Z` → CI builds and
|
|
81
|
+
publishes to PyPI via trusted publishing (OIDC, no secrets in the repo).
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "aikit-platform"
|
|
7
|
+
dynamic = ["version"] # sourced from src/aikit/__init__.py; bump there before tagging
|
|
8
|
+
description = "Shared platform layer for AI-backend systems: model client, embeddings, observability, config, db, queue."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.11"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [{ name = "Aaryan Mahajan" }]
|
|
13
|
+
keywords = ["llm", "embeddings", "observability", "infrastructure"]
|
|
14
|
+
|
|
15
|
+
# Core stays dependency-light on purpose. Heavy deps live in extras.
|
|
16
|
+
dependencies = [
|
|
17
|
+
"pydantic>=2.7",
|
|
18
|
+
"pydantic-settings>=2.3",
|
|
19
|
+
]
|
|
20
|
+
|
|
21
|
+
[project.optional-dependencies]
|
|
22
|
+
embeddings = ["sentence-transformers>=3.0"]
|
|
23
|
+
db = ["sqlalchemy[asyncio]>=2.0", "asyncpg>=0.29", "pgvector>=0.3"]
|
|
24
|
+
redis = ["redis>=5.0"]
|
|
25
|
+
queue = ["aikit-platform[redis]", "arq>=0.26"]
|
|
26
|
+
observability = ["prometheus-client>=0.20", "structlog>=24.1"]
|
|
27
|
+
http = ["httpx>=0.27"]
|
|
28
|
+
|
|
29
|
+
# Per-project bundles so downstream repos install exactly what they need.
|
|
30
|
+
eval = ["aikit-platform[db,queue,observability,http,embeddings]"]
|
|
31
|
+
# No `queue`/arq: the gateway is a synchronous request path, not a fan-out
|
|
32
|
+
# worker system - it uses raw `redis` (cache, circuit breaker, rate limit)
|
|
33
|
+
# directly, never arq.
|
|
34
|
+
gateway = ["aikit-platform[db,redis,observability,http,embeddings]"]
|
|
35
|
+
memory = ["aikit-platform[db,queue,observability,http,embeddings]"]
|
|
36
|
+
|
|
37
|
+
dev = ["pytest>=8.2", "pytest-asyncio>=0.23", "aiosqlite>=0.20", "ruff>=0.5", "mypy>=1.10"]
|
|
38
|
+
|
|
39
|
+
[project.urls]
|
|
40
|
+
Homepage = "https://github.com/Aaryan123456679/aikit"
|
|
41
|
+
Repository = "https://github.com/Aaryan123456679/aikit"
|
|
42
|
+
|
|
43
|
+
[tool.hatch.version]
|
|
44
|
+
path = "src/aikit/__init__.py"
|
|
45
|
+
|
|
46
|
+
[tool.hatch.build.targets.wheel]
|
|
47
|
+
packages = ["src/aikit"]
|
|
48
|
+
|
|
49
|
+
[tool.ruff]
|
|
50
|
+
line-length = 100
|
|
51
|
+
target-version = "py311"
|
|
52
|
+
[tool.ruff.lint]
|
|
53
|
+
select = ["E", "F", "I", "UP", "B"]
|
|
54
|
+
|
|
55
|
+
[tool.pytest.ini_options]
|
|
56
|
+
asyncio_mode = "auto"
|
|
57
|
+
|
|
58
|
+
[tool.mypy]
|
|
59
|
+
python_version = "3.11"
|
|
60
|
+
strict = true
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
"""Config — env-driven settings base. Each service subclasses this."""
|
|
2
|
+
from __future__ import annotations
|
|
3
|
+
|
|
4
|
+
from pydantic_settings import BaseSettings, SettingsConfigDict
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class BaseServiceSettings(BaseSettings):
|
|
8
|
+
model_config = SettingsConfigDict(env_file=".env", extra="ignore")
|
|
9
|
+
|
|
10
|
+
env: str = "local"
|
|
11
|
+
log_level: str = "INFO"
|
|
12
|
+
database_url: str = "postgresql+asyncpg://localhost/app"
|
|
13
|
+
redis_url: str = "redis://localhost:6379/0"
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
"""DB — async SQLAlchemy engine/session base + pgvector helpers.
|
|
2
|
+
|
|
3
|
+
Impl requires the `db` extra. Provides:
|
|
4
|
+
- async engine factory
|
|
5
|
+
- transaction-scoped session context manager
|
|
6
|
+
- DeclarativeBase subclass every service builds models on
|
|
7
|
+
"""
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
from collections.abc import AsyncIterator
|
|
11
|
+
from contextlib import asynccontextmanager
|
|
12
|
+
from typing import Any
|
|
13
|
+
|
|
14
|
+
from sqlalchemy.ext.asyncio import (
|
|
15
|
+
AsyncEngine,
|
|
16
|
+
AsyncSession,
|
|
17
|
+
async_sessionmaker,
|
|
18
|
+
create_async_engine,
|
|
19
|
+
)
|
|
20
|
+
from sqlalchemy.orm import DeclarativeBase
|
|
21
|
+
|
|
22
|
+
__all__ = ["Base", "make_engine", "make_sessionmaker", "session_scope"]
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class Base(DeclarativeBase):
|
|
26
|
+
"""Declarative base every downstream service builds its ORM models on."""
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def make_engine(
|
|
30
|
+
database_url: str,
|
|
31
|
+
*,
|
|
32
|
+
echo: bool = False,
|
|
33
|
+
pool_size: int | None = None,
|
|
34
|
+
max_overflow: int | None = None,
|
|
35
|
+
) -> AsyncEngine:
|
|
36
|
+
# pool_size/max_overflow default to None (omitted entirely) rather
|
|
37
|
+
# than SQLAlchemy's own QueuePool defaults (5/10), because those two
|
|
38
|
+
# kwargs are QueuePool-specific and blow up with a TypeError against
|
|
39
|
+
# any other pool class - including SQLite's StaticPool, which this
|
|
40
|
+
# project's own test suite uses. Omitting them when unset preserves
|
|
41
|
+
# exact prior behavior for every caller that doesn't pass them.
|
|
42
|
+
# Callers on Postgres that want a bigger pool (found necessary load-
|
|
43
|
+
# testing the gateway: every request that logs to Postgres holds a
|
|
44
|
+
# connection for that write, so 50 concurrent requests need more than
|
|
45
|
+
# QueuePool's own default 15-connection ceiling) pass them explicitly.
|
|
46
|
+
kwargs: dict[str, Any] = {"echo": echo, "pool_pre_ping": True}
|
|
47
|
+
if pool_size is not None:
|
|
48
|
+
kwargs["pool_size"] = pool_size
|
|
49
|
+
if max_overflow is not None:
|
|
50
|
+
kwargs["max_overflow"] = max_overflow
|
|
51
|
+
return create_async_engine(database_url, **kwargs)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def make_sessionmaker(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
|
|
55
|
+
return async_sessionmaker(engine, expire_on_commit=False)
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
@asynccontextmanager
|
|
59
|
+
async def session_scope(
|
|
60
|
+
sessionmaker: async_sessionmaker[AsyncSession],
|
|
61
|
+
) -> AsyncIterator[AsyncSession]:
|
|
62
|
+
"""Open a session bound to a single transaction: commits on success,
|
|
63
|
+
rolls back on any exception."""
|
|
64
|
+
async with sessionmaker() as session, session.begin():
|
|
65
|
+
yield session
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
"""EmbeddingService — local, batched text embeddings.
|
|
2
|
+
|
|
3
|
+
Reused by: eval (similarity scorer), gateway (semantic cache), memory (retrieval).
|
|
4
|
+
Concrete impl requires the `embeddings` extra (sentence-transformers).
|
|
5
|
+
"""
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
from typing import Protocol, runtime_checkable
|
|
9
|
+
|
|
10
|
+
Vector = list[float]
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@runtime_checkable
|
|
14
|
+
class EmbeddingService(Protocol):
|
|
15
|
+
dim: int
|
|
16
|
+
|
|
17
|
+
async def embed(self, texts: list[str]) -> list[Vector]: ...
|
|
18
|
+
|
|
19
|
+
@staticmethod
|
|
20
|
+
def cosine(a: Vector, b: Vector) -> float: ...
|
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
"""SentenceTransformerEmbeddingService — local, batched embeddings.
|
|
2
|
+
|
|
3
|
+
Requires the `embeddings` extra (sentence-transformers). Runs the (blocking,
|
|
4
|
+
CPU-bound) model in a thread executor so it never blocks the event loop.
|
|
5
|
+
"""
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
import asyncio
|
|
9
|
+
|
|
10
|
+
import numpy as np
|
|
11
|
+
from sentence_transformers import SentenceTransformer
|
|
12
|
+
|
|
13
|
+
from . import Vector
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class SentenceTransformerEmbeddingService:
|
|
17
|
+
"""LSP: substitutable anywhere an `EmbeddingService` is expected."""
|
|
18
|
+
|
|
19
|
+
def __init__(self, model_name: str = "all-MiniLM-L6-v2") -> None:
|
|
20
|
+
# device="cpu" is deliberate, not a missing feature: auto-detection
|
|
21
|
+
# picks MPS/CUDA when available, but a server process calls encode()
|
|
22
|
+
# from multiple concurrent executor threads (one per in-flight
|
|
23
|
+
# request), and MPS does not tolerate concurrent access from
|
|
24
|
+
# multiple threads reliably - it can crash the whole process rather
|
|
25
|
+
# than raise a catchable exception. CPU is also the right choice
|
|
26
|
+
# for this workload's shape (many small, single-item, latency-
|
|
27
|
+
# sensitive calls), where GPU dispatch overhead would dominate
|
|
28
|
+
# anyway.
|
|
29
|
+
self._model = SentenceTransformer(model_name, device="cpu")
|
|
30
|
+
dim = self._model.get_sentence_embedding_dimension()
|
|
31
|
+
if dim is None:
|
|
32
|
+
raise ValueError(f"model {model_name!r} did not report an embedding dimension")
|
|
33
|
+
self.dim: int = dim
|
|
34
|
+
|
|
35
|
+
async def embed(self, texts: list[str]) -> list[Vector]:
|
|
36
|
+
loop = asyncio.get_running_loop()
|
|
37
|
+
vectors = await loop.run_in_executor(None, self._model.encode, texts)
|
|
38
|
+
return [vector.tolist() for vector in vectors]
|
|
39
|
+
|
|
40
|
+
@staticmethod
|
|
41
|
+
def cosine(a: Vector, b: Vector) -> float:
|
|
42
|
+
va, vb = np.asarray(a, dtype=float), np.asarray(b, dtype=float)
|
|
43
|
+
denom = float(np.linalg.norm(va) * np.linalg.norm(vb))
|
|
44
|
+
if denom == 0.0:
|
|
45
|
+
return 0.0
|
|
46
|
+
return float(np.dot(va, vb) / denom)
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
"""JobQueue — async task fan-out abstraction.
|
|
2
|
+
|
|
3
|
+
OCP: swap Redis/arq for a Postgres-backed queue without touching producers.
|
|
4
|
+
"""
|
|
5
|
+
from __future__ import annotations
|
|
6
|
+
|
|
7
|
+
from collections.abc import Awaitable, Callable
|
|
8
|
+
from typing import Any, Protocol
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class JobQueue(Protocol):
|
|
12
|
+
async def enqueue(self, task: str, payload: dict[str, Any]) -> str: ...
|
|
13
|
+
def register(
|
|
14
|
+
self, task: str, handler: Callable[[dict[str, Any]], Awaitable[None]]
|
|
15
|
+
) -> None: ...
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
"""ArqJobQueue — JobQueue implementation over Redis via arq.
|
|
2
|
+
|
|
3
|
+
Requires the `queue` extra (redis, arq). OCP: producers depend on the
|
|
4
|
+
`JobQueue` Protocol, so swapping this for a Postgres-backed queue later
|
|
5
|
+
touches no call sites.
|
|
6
|
+
"""
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
from collections.abc import Awaitable, Callable
|
|
10
|
+
from typing import Any
|
|
11
|
+
|
|
12
|
+
from arq import create_pool
|
|
13
|
+
from arq.connections import ArqRedis, RedisSettings
|
|
14
|
+
from arq.worker import Function, func
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class ArqJobQueue:
|
|
18
|
+
def __init__(self, redis_url: str) -> None:
|
|
19
|
+
self._redis_url = redis_url
|
|
20
|
+
self._pool: ArqRedis | None = None
|
|
21
|
+
self._handlers: dict[str, Function] = {}
|
|
22
|
+
|
|
23
|
+
async def connect(self) -> None:
|
|
24
|
+
if self._pool is None:
|
|
25
|
+
self._pool = await create_pool(RedisSettings.from_dsn(self._redis_url))
|
|
26
|
+
|
|
27
|
+
async def enqueue(self, task: str, payload: dict[str, Any]) -> str:
|
|
28
|
+
await self.connect()
|
|
29
|
+
assert self._pool is not None
|
|
30
|
+
job = await self._pool.enqueue_job(task, payload)
|
|
31
|
+
if job is None:
|
|
32
|
+
raise RuntimeError(f"job deduped or rejected by arq: {task}")
|
|
33
|
+
return job.job_id
|
|
34
|
+
|
|
35
|
+
def register(
|
|
36
|
+
self, task: str, handler: Callable[[dict[str, Any]], Awaitable[None]]
|
|
37
|
+
) -> None:
|
|
38
|
+
async def wrapper(ctx: Any, payload: dict[str, Any]) -> None:
|
|
39
|
+
await handler(payload)
|
|
40
|
+
|
|
41
|
+
# arq's `func()` names a job by `name`, defaulting to
|
|
42
|
+
# `coroutine.__qualname__` (NOT `__name__`) if omitted - a closure's
|
|
43
|
+
# qualname is always its def-site path (e.g.
|
|
44
|
+
# "ArqJobQueue.register.<locals>.wrapper"), so the name must be set
|
|
45
|
+
# explicitly here or every registration collides under that one name.
|
|
46
|
+
self._handlers[task] = func(wrapper, name=task)
|
|
47
|
+
|
|
48
|
+
@property
|
|
49
|
+
def functions(self) -> list[Function]:
|
|
50
|
+
"""Pass to `arq.worker.WorkerSettings.functions`."""
|
|
51
|
+
return list(self._handlers.values())
|
|
52
|
+
|
|
53
|
+
async def close(self) -> None:
|
|
54
|
+
if self._pool is not None:
|
|
55
|
+
await self._pool.close()
|
|
56
|
+
self._pool = None
|
|
@@ -0,0 +1,64 @@
|
|
|
1
|
+
"""ModelClient — unified async interface over heterogeneous LLM backends.
|
|
2
|
+
|
|
3
|
+
DIP: consumers depend on this Protocol, never on Ollama/hosted specifics.
|
|
4
|
+
LSP: every implementation is substitutable.
|
|
5
|
+
ISP: streaming and non-streaming are separate methods.
|
|
6
|
+
"""
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
from collections.abc import AsyncIterator
|
|
10
|
+
from dataclasses import dataclass
|
|
11
|
+
from typing import Literal, Protocol, runtime_checkable
|
|
12
|
+
|
|
13
|
+
Role = Literal["system", "user", "assistant"]
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@dataclass(frozen=True)
|
|
17
|
+
class ChatMessage:
|
|
18
|
+
role: Role
|
|
19
|
+
content: str
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
@dataclass(frozen=True)
|
|
23
|
+
class Completion:
|
|
24
|
+
content: str
|
|
25
|
+
model: str
|
|
26
|
+
tokens_in: int
|
|
27
|
+
tokens_out: int
|
|
28
|
+
latency_ms: float
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
@dataclass(frozen=True)
|
|
32
|
+
class Usage:
|
|
33
|
+
"""Terminal event on a `stream()` iterator: exactly one, always last."""
|
|
34
|
+
|
|
35
|
+
tokens_in: int
|
|
36
|
+
tokens_out: int
|
|
37
|
+
finish_reason: str
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
StreamEvent = str | Usage
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
@runtime_checkable
|
|
44
|
+
class ModelClient(Protocol):
|
|
45
|
+
"""A single LLM backend. Implementations: OllamaClient, HostedClient."""
|
|
46
|
+
|
|
47
|
+
name: str
|
|
48
|
+
|
|
49
|
+
async def complete(
|
|
50
|
+
self, messages: list[ChatMessage], *, model: str | None = None
|
|
51
|
+
) -> Completion: ...
|
|
52
|
+
|
|
53
|
+
def stream(
|
|
54
|
+
self, messages: list[ChatMessage], *, model: str | None = None
|
|
55
|
+
) -> AsyncIterator[StreamEvent]:
|
|
56
|
+
"""An async-generator method: called directly (no `await`), then
|
|
57
|
+
iterated with `async for`. Yields text deltas (`str`) followed by
|
|
58
|
+
exactly one terminal `Usage` — callers can accumulate text and read
|
|
59
|
+
cost/finish_reason off the last event without a second round trip."""
|
|
60
|
+
...
|
|
61
|
+
|
|
62
|
+
async def health(self) -> bool: ...
|
|
63
|
+
|
|
64
|
+
async def aclose(self) -> None: ...
|
|
@@ -0,0 +1,108 @@
|
|
|
1
|
+
"""HostedClient — ModelClient implementation over any OpenAI-compatible
|
|
2
|
+
chat-completions API (OpenAI itself, and the many providers that mirror
|
|
3
|
+
its wire format: Groq, Together, Fireworks, OpenRouter, ...).
|
|
4
|
+
|
|
5
|
+
Requires the `http` extra (httpx).
|
|
6
|
+
"""
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import json
|
|
10
|
+
import time
|
|
11
|
+
from collections.abc import AsyncIterator
|
|
12
|
+
|
|
13
|
+
import httpx
|
|
14
|
+
|
|
15
|
+
from . import ChatMessage, Completion, StreamEvent, Usage
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class HostedClient:
|
|
19
|
+
"""LSP: substitutable anywhere a `ModelClient` is expected."""
|
|
20
|
+
|
|
21
|
+
def __init__(
|
|
22
|
+
self,
|
|
23
|
+
base_url: str,
|
|
24
|
+
api_key: str,
|
|
25
|
+
*,
|
|
26
|
+
default_model: str,
|
|
27
|
+
timeout: float = 60.0,
|
|
28
|
+
name: str = "hosted",
|
|
29
|
+
) -> None:
|
|
30
|
+
self.name = name
|
|
31
|
+
self._default_model = default_model
|
|
32
|
+
self._client = httpx.AsyncClient(
|
|
33
|
+
base_url=base_url.rstrip("/"),
|
|
34
|
+
timeout=timeout,
|
|
35
|
+
headers={"Authorization": f"Bearer {api_key}"},
|
|
36
|
+
)
|
|
37
|
+
|
|
38
|
+
async def complete(
|
|
39
|
+
self, messages: list[ChatMessage], *, model: str | None = None
|
|
40
|
+
) -> Completion:
|
|
41
|
+
target_model = model or self._default_model
|
|
42
|
+
payload = {
|
|
43
|
+
"model": target_model,
|
|
44
|
+
"messages": [{"role": m.role, "content": m.content} for m in messages],
|
|
45
|
+
"stream": False,
|
|
46
|
+
}
|
|
47
|
+
start = time.perf_counter()
|
|
48
|
+
resp = await self._client.post("/chat/completions", json=payload)
|
|
49
|
+
resp.raise_for_status()
|
|
50
|
+
data = resp.json()
|
|
51
|
+
latency_ms = (time.perf_counter() - start) * 1000
|
|
52
|
+
choice = data["choices"][0]
|
|
53
|
+
usage = data.get("usage") or {}
|
|
54
|
+
return Completion(
|
|
55
|
+
content=choice.get("message", {}).get("content", ""),
|
|
56
|
+
model=data.get("model", target_model),
|
|
57
|
+
tokens_in=usage.get("prompt_tokens", 0),
|
|
58
|
+
tokens_out=usage.get("completion_tokens", 0),
|
|
59
|
+
latency_ms=latency_ms,
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
async def stream(
|
|
63
|
+
self, messages: list[ChatMessage], *, model: str | None = None
|
|
64
|
+
) -> AsyncIterator[StreamEvent]:
|
|
65
|
+
payload = {
|
|
66
|
+
"model": model or self._default_model,
|
|
67
|
+
"messages": [{"role": m.role, "content": m.content} for m in messages],
|
|
68
|
+
"stream": True,
|
|
69
|
+
# Ask for a final usage-only chunk (OpenAI and most compatible
|
|
70
|
+
# providers support this) so the terminal Usage is exact, not
|
|
71
|
+
# estimated client-side.
|
|
72
|
+
"stream_options": {"include_usage": True},
|
|
73
|
+
}
|
|
74
|
+
finish_reason = "stop"
|
|
75
|
+
async with self._client.stream("POST", "/chat/completions", json=payload) as resp:
|
|
76
|
+
resp.raise_for_status()
|
|
77
|
+
async for line in resp.aiter_lines():
|
|
78
|
+
if not line or not line.startswith("data:"):
|
|
79
|
+
continue
|
|
80
|
+
data = line[len("data:") :].strip()
|
|
81
|
+
if data == "[DONE]":
|
|
82
|
+
break
|
|
83
|
+
chunk = json.loads(data)
|
|
84
|
+
choices = chunk.get("choices") or []
|
|
85
|
+
if choices:
|
|
86
|
+
delta = choices[0].get("delta", {})
|
|
87
|
+
piece = delta.get("content")
|
|
88
|
+
if piece:
|
|
89
|
+
yield piece
|
|
90
|
+
if choices[0].get("finish_reason"):
|
|
91
|
+
finish_reason = choices[0]["finish_reason"]
|
|
92
|
+
usage = chunk.get("usage")
|
|
93
|
+
if usage:
|
|
94
|
+
yield Usage(
|
|
95
|
+
tokens_in=usage.get("prompt_tokens", 0),
|
|
96
|
+
tokens_out=usage.get("completion_tokens", 0),
|
|
97
|
+
finish_reason=finish_reason,
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
async def health(self) -> bool:
|
|
101
|
+
try:
|
|
102
|
+
resp = await self._client.get("/models")
|
|
103
|
+
return resp.status_code == 200
|
|
104
|
+
except httpx.HTTPError:
|
|
105
|
+
return False
|
|
106
|
+
|
|
107
|
+
async def aclose(self) -> None:
|
|
108
|
+
await self._client.aclose()
|
|
@@ -0,0 +1,87 @@
|
|
|
1
|
+
"""OllamaClient — ModelClient implementation over a local Ollama server.
|
|
2
|
+
|
|
3
|
+
Requires the `http` extra (httpx). This is the only concrete ModelClient
|
|
4
|
+
in v1: it doubles as target and judge model per the eval-platform HLD ($0
|
|
5
|
+
cost, no hosted API keys needed).
|
|
6
|
+
"""
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import json
|
|
10
|
+
import time
|
|
11
|
+
from collections.abc import AsyncIterator
|
|
12
|
+
|
|
13
|
+
import httpx
|
|
14
|
+
|
|
15
|
+
from . import ChatMessage, Completion, StreamEvent, Usage
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class OllamaClient:
|
|
19
|
+
"""LSP: substitutable anywhere a `ModelClient` is expected."""
|
|
20
|
+
|
|
21
|
+
def __init__(
|
|
22
|
+
self,
|
|
23
|
+
host: str = "http://localhost:11434",
|
|
24
|
+
*,
|
|
25
|
+
default_model: str = "llama3.1",
|
|
26
|
+
timeout: float = 60.0,
|
|
27
|
+
name: str = "ollama",
|
|
28
|
+
) -> None:
|
|
29
|
+
self.name = name
|
|
30
|
+
self._default_model = default_model
|
|
31
|
+
self._client = httpx.AsyncClient(base_url=host.rstrip("/"), timeout=timeout)
|
|
32
|
+
|
|
33
|
+
async def complete(
|
|
34
|
+
self, messages: list[ChatMessage], *, model: str | None = None
|
|
35
|
+
) -> Completion:
|
|
36
|
+
target_model = model or self._default_model
|
|
37
|
+
payload = {
|
|
38
|
+
"model": target_model,
|
|
39
|
+
"messages": [{"role": m.role, "content": m.content} for m in messages],
|
|
40
|
+
"stream": False,
|
|
41
|
+
}
|
|
42
|
+
start = time.perf_counter()
|
|
43
|
+
resp = await self._client.post("/api/chat", json=payload)
|
|
44
|
+
resp.raise_for_status()
|
|
45
|
+
data = resp.json()
|
|
46
|
+
latency_ms = (time.perf_counter() - start) * 1000
|
|
47
|
+
return Completion(
|
|
48
|
+
content=data.get("message", {}).get("content", ""),
|
|
49
|
+
model=target_model,
|
|
50
|
+
tokens_in=data.get("prompt_eval_count", 0),
|
|
51
|
+
tokens_out=data.get("eval_count", 0),
|
|
52
|
+
latency_ms=latency_ms,
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
async def stream(
|
|
56
|
+
self, messages: list[ChatMessage], *, model: str | None = None
|
|
57
|
+
) -> AsyncIterator[StreamEvent]:
|
|
58
|
+
payload = {
|
|
59
|
+
"model": model or self._default_model,
|
|
60
|
+
"messages": [{"role": m.role, "content": m.content} for m in messages],
|
|
61
|
+
"stream": True,
|
|
62
|
+
}
|
|
63
|
+
async with self._client.stream("POST", "/api/chat", json=payload) as resp:
|
|
64
|
+
resp.raise_for_status()
|
|
65
|
+
async for line in resp.aiter_lines():
|
|
66
|
+
if not line:
|
|
67
|
+
continue
|
|
68
|
+
chunk = json.loads(line)
|
|
69
|
+
piece = chunk.get("message", {}).get("content")
|
|
70
|
+
if piece:
|
|
71
|
+
yield piece
|
|
72
|
+
if chunk.get("done"):
|
|
73
|
+
yield Usage(
|
|
74
|
+
tokens_in=chunk.get("prompt_eval_count", 0),
|
|
75
|
+
tokens_out=chunk.get("eval_count", 0),
|
|
76
|
+
finish_reason=chunk.get("done_reason", "stop"),
|
|
77
|
+
)
|
|
78
|
+
|
|
79
|
+
async def health(self) -> bool:
|
|
80
|
+
try:
|
|
81
|
+
resp = await self._client.get("/api/tags")
|
|
82
|
+
return resp.status_code == 200
|
|
83
|
+
except httpx.HTTPError:
|
|
84
|
+
return False
|
|
85
|
+
|
|
86
|
+
async def aclose(self) -> None:
|
|
87
|
+
await self._client.aclose()
|
|
@@ -0,0 +1,77 @@
|
|
|
1
|
+
"""Observability — structured logging + metrics + health.
|
|
2
|
+
|
|
3
|
+
Impl requires the `observability` extra. Framework-agnostic helpers so any
|
|
4
|
+
FastAPI service wires them identically.
|
|
5
|
+
"""
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
import logging
|
|
9
|
+
import sys
|
|
10
|
+
from typing import Protocol
|
|
11
|
+
|
|
12
|
+
import structlog
|
|
13
|
+
from prometheus_client import CONTENT_TYPE_LATEST as PROMETHEUS_CONTENT_TYPE
|
|
14
|
+
from prometheus_client import CollectorRegistry, Counter, Histogram, generate_latest
|
|
15
|
+
|
|
16
|
+
__all__ = [
|
|
17
|
+
"MetricsSink",
|
|
18
|
+
"PrometheusMetricsSink",
|
|
19
|
+
"PROMETHEUS_CONTENT_TYPE",
|
|
20
|
+
"configure_logging",
|
|
21
|
+
"get_logger",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class MetricsSink(Protocol):
|
|
26
|
+
def incr(self, name: str, **labels: str) -> None: ...
|
|
27
|
+
def observe(self, name: str, value: float, **labels: str) -> None: ...
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def configure_logging(level: str = "INFO") -> None:
|
|
31
|
+
"""Call once at process start. Renders JSON lines to stdout."""
|
|
32
|
+
logging.basicConfig(format="%(message)s", stream=sys.stdout, level=level)
|
|
33
|
+
structlog.configure(
|
|
34
|
+
processors=[
|
|
35
|
+
structlog.contextvars.merge_contextvars,
|
|
36
|
+
structlog.processors.add_log_level,
|
|
37
|
+
structlog.processors.TimeStamper(fmt="iso"),
|
|
38
|
+
structlog.processors.JSONRenderer(),
|
|
39
|
+
],
|
|
40
|
+
wrapper_class=structlog.make_filtering_bound_logger(
|
|
41
|
+
logging.getLevelName(level) if isinstance(level, str) else level
|
|
42
|
+
),
|
|
43
|
+
logger_factory=structlog.PrintLoggerFactory(),
|
|
44
|
+
cache_logger_on_first_use=True,
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def get_logger(name: str) -> structlog.stdlib.BoundLogger:
|
|
49
|
+
return structlog.get_logger(name) # type: ignore[no-any-return]
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
class PrometheusMetricsSink:
|
|
53
|
+
"""Counter/Histogram-backed MetricsSink. Series are created lazily on
|
|
54
|
+
first use, keyed by name; the label set of the first call fixes that
|
|
55
|
+
series' label names for its lifetime."""
|
|
56
|
+
|
|
57
|
+
def __init__(self, registry: CollectorRegistry | None = None):
|
|
58
|
+
self._registry = registry or CollectorRegistry(auto_describe=True)
|
|
59
|
+
self._counters: dict[str, Counter] = {}
|
|
60
|
+
self._histograms: dict[str, Histogram] = {}
|
|
61
|
+
|
|
62
|
+
def incr(self, name: str, **labels: str) -> None:
|
|
63
|
+
counter = self._counters.get(name)
|
|
64
|
+
if counter is None:
|
|
65
|
+
counter = Counter(name, name, labelnames=sorted(labels), registry=self._registry)
|
|
66
|
+
self._counters[name] = counter
|
|
67
|
+
(counter.labels(**labels) if labels else counter).inc()
|
|
68
|
+
|
|
69
|
+
def observe(self, name: str, value: float, **labels: str) -> None:
|
|
70
|
+
hist = self._histograms.get(name)
|
|
71
|
+
if hist is None:
|
|
72
|
+
hist = Histogram(name, name, labelnames=sorted(labels), registry=self._registry)
|
|
73
|
+
self._histograms[name] = hist
|
|
74
|
+
(hist.labels(**labels) if labels else hist).observe(value)
|
|
75
|
+
|
|
76
|
+
def render(self) -> bytes:
|
|
77
|
+
return generate_latest(self._registry)
|
|
File without changes
|
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from sqlalchemy import Column, Integer, String, select
|
|
4
|
+
|
|
5
|
+
from aikit.db import Base, make_engine, make_sessionmaker, session_scope
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class Widget(Base):
|
|
9
|
+
__tablename__ = "widget"
|
|
10
|
+
id = Column(Integer, primary_key=True)
|
|
11
|
+
name = Column(String, nullable=False)
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
async def test_session_scope_commits_on_success():
|
|
15
|
+
engine = make_engine("sqlite+aiosqlite://")
|
|
16
|
+
async with engine.begin() as conn:
|
|
17
|
+
await conn.run_sync(Base.metadata.create_all)
|
|
18
|
+
sessionmaker = make_sessionmaker(engine)
|
|
19
|
+
|
|
20
|
+
async with session_scope(sessionmaker) as session:
|
|
21
|
+
session.add(Widget(id=1, name="a"))
|
|
22
|
+
|
|
23
|
+
async with sessionmaker() as session:
|
|
24
|
+
rows = (await session.execute(select(Widget))).scalars().all()
|
|
25
|
+
assert [r.name for r in rows] == ["a"]
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
async def test_session_scope_rolls_back_on_error():
|
|
29
|
+
engine = make_engine("sqlite+aiosqlite://")
|
|
30
|
+
async with engine.begin() as conn:
|
|
31
|
+
await conn.run_sync(Base.metadata.create_all)
|
|
32
|
+
sessionmaker = make_sessionmaker(engine)
|
|
33
|
+
|
|
34
|
+
class Boom(Exception):
|
|
35
|
+
pass
|
|
36
|
+
|
|
37
|
+
try:
|
|
38
|
+
async with session_scope(sessionmaker) as session:
|
|
39
|
+
session.add(Widget(id=1, name="a"))
|
|
40
|
+
raise Boom
|
|
41
|
+
except Boom:
|
|
42
|
+
pass
|
|
43
|
+
|
|
44
|
+
async with sessionmaker() as session:
|
|
45
|
+
rows = (await session.execute(select(Widget))).scalars().all()
|
|
46
|
+
assert rows == []
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from aikit.jobqueue.arq_queue import ArqJobQueue
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def test_registered_function_is_named_by_task_not_qualname():
|
|
7
|
+
"""Regression test: arq's `func()` names a job from `coroutine.__qualname__`
|
|
8
|
+
when `name` is omitted, and a closure's qualname is always its def-site
|
|
9
|
+
path (e.g. "ArqJobQueue.register.<locals>.wrapper") - never the task
|
|
10
|
+
name. Caught in a live Docker run where the worker registered the
|
|
11
|
+
wrapper under its qualname and every job then failed with
|
|
12
|
+
"function 'execute_case' not found"."""
|
|
13
|
+
queue = ArqJobQueue("redis://localhost:6379/0")
|
|
14
|
+
|
|
15
|
+
async def handler(payload: dict) -> None:
|
|
16
|
+
pass
|
|
17
|
+
|
|
18
|
+
queue.register("execute_case", handler)
|
|
19
|
+
|
|
20
|
+
[registered] = queue.functions
|
|
21
|
+
assert registered.name == "execute_case"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def test_multiple_registrations_keep_distinct_names():
|
|
25
|
+
queue = ArqJobQueue("redis://localhost:6379/0")
|
|
26
|
+
|
|
27
|
+
async def handler_a(payload: dict) -> None:
|
|
28
|
+
pass
|
|
29
|
+
|
|
30
|
+
async def handler_b(payload: dict) -> None:
|
|
31
|
+
pass
|
|
32
|
+
|
|
33
|
+
queue.register("task_a", handler_a)
|
|
34
|
+
queue.register("task_b", handler_b)
|
|
35
|
+
|
|
36
|
+
names = {f.name for f in queue.functions}
|
|
37
|
+
assert names == {"task_a", "task_b"}
|
|
@@ -0,0 +1,73 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import httpx
|
|
4
|
+
|
|
5
|
+
from aikit.model_client import ChatMessage, Usage
|
|
6
|
+
from aikit.model_client.hosted import HostedClient
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def _client_with_transport(handler) -> HostedClient:
|
|
10
|
+
client = HostedClient("http://hosted.local", "sk-test", default_model="gpt-4o-mini")
|
|
11
|
+
client._client = httpx.AsyncClient(
|
|
12
|
+
base_url="http://hosted.local",
|
|
13
|
+
headers={"Authorization": "Bearer sk-test"},
|
|
14
|
+
transport=httpx.MockTransport(handler),
|
|
15
|
+
)
|
|
16
|
+
return client
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
async def test_complete_parses_openai_response():
|
|
20
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
21
|
+
assert request.headers["authorization"] == "Bearer sk-test"
|
|
22
|
+
assert request.url.path == "/chat/completions"
|
|
23
|
+
return httpx.Response(
|
|
24
|
+
200,
|
|
25
|
+
json={
|
|
26
|
+
"model": "gpt-4o-mini",
|
|
27
|
+
"choices": [{"message": {"content": "hi there"}, "finish_reason": "stop"}],
|
|
28
|
+
"usage": {"prompt_tokens": 10, "completion_tokens": 3},
|
|
29
|
+
},
|
|
30
|
+
)
|
|
31
|
+
|
|
32
|
+
client = _client_with_transport(handler)
|
|
33
|
+
result = await client.complete([ChatMessage(role="user", content="hello")])
|
|
34
|
+
assert result.content == "hi there"
|
|
35
|
+
assert result.model == "gpt-4o-mini"
|
|
36
|
+
assert result.tokens_in == 10
|
|
37
|
+
assert result.tokens_out == 3
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
async def test_stream_yields_text_then_exactly_one_terminal_usage():
|
|
41
|
+
"""SSE shape per the OpenAI-compatible spec with stream_options.include_usage."""
|
|
42
|
+
sse_lines = [
|
|
43
|
+
'data: {"choices":[{"delta":{"content":"hi"},"finish_reason":null}]}',
|
|
44
|
+
'data: {"choices":[{"delta":{"content":" there"},"finish_reason":"stop"}]}',
|
|
45
|
+
'data: {"choices":[],"usage":{"prompt_tokens":10,"completion_tokens":2}}',
|
|
46
|
+
"data: [DONE]",
|
|
47
|
+
]
|
|
48
|
+
|
|
49
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
50
|
+
return httpx.Response(200, text="\n\n".join(sse_lines) + "\n\n")
|
|
51
|
+
|
|
52
|
+
client = _client_with_transport(handler)
|
|
53
|
+
events = [event async for event in client.stream([ChatMessage(role="user", content="hi")])]
|
|
54
|
+
|
|
55
|
+
assert events[:-1] == ["hi", " there"]
|
|
56
|
+
assert events[-1] == Usage(tokens_in=10, tokens_out=2, finish_reason="stop")
|
|
57
|
+
assert sum(isinstance(e, Usage) for e in events) == 1
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
async def test_health_true_on_200():
|
|
61
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
62
|
+
return httpx.Response(200, json={"data": []})
|
|
63
|
+
|
|
64
|
+
client = _client_with_transport(handler)
|
|
65
|
+
assert await client.health() is True
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
async def test_health_false_on_error():
|
|
69
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
70
|
+
raise httpx.ConnectError("refused", request=request)
|
|
71
|
+
|
|
72
|
+
client = _client_with_transport(handler)
|
|
73
|
+
assert await client.health() is False
|
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import httpx
|
|
4
|
+
|
|
5
|
+
from aikit.model_client import ChatMessage, Usage
|
|
6
|
+
from aikit.model_client.ollama import OllamaClient
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def _client_with_transport(handler) -> OllamaClient:
|
|
10
|
+
client = OllamaClient(default_model="llama3.1")
|
|
11
|
+
client._client = httpx.AsyncClient(
|
|
12
|
+
base_url="http://ollama.local", transport=httpx.MockTransport(handler)
|
|
13
|
+
)
|
|
14
|
+
return client
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
async def test_complete_parses_ollama_response():
|
|
18
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
19
|
+
assert request.url.path == "/api/chat"
|
|
20
|
+
return httpx.Response(
|
|
21
|
+
200,
|
|
22
|
+
json={
|
|
23
|
+
"message": {"role": "assistant", "content": "hi there"},
|
|
24
|
+
"prompt_eval_count": 5,
|
|
25
|
+
"eval_count": 3,
|
|
26
|
+
},
|
|
27
|
+
)
|
|
28
|
+
|
|
29
|
+
client = _client_with_transport(handler)
|
|
30
|
+
result = await client.complete([ChatMessage(role="user", content="hello")])
|
|
31
|
+
assert result.content == "hi there"
|
|
32
|
+
assert result.model == "llama3.1"
|
|
33
|
+
assert result.tokens_in == 5
|
|
34
|
+
assert result.tokens_out == 3
|
|
35
|
+
assert result.latency_ms >= 0
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
async def test_health_false_on_error():
|
|
39
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
40
|
+
raise httpx.ConnectError("refused", request=request)
|
|
41
|
+
|
|
42
|
+
client = _client_with_transport(handler)
|
|
43
|
+
assert await client.health() is False
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
async def test_stream_yields_text_then_exactly_one_terminal_usage():
|
|
47
|
+
"""Response lines captured verbatim from a real Ollama server (see
|
|
48
|
+
aikit v0.2.0's HLD/LLD B.0 - live-validated separately against a real
|
|
49
|
+
llama3.1:8b before this mock was written)."""
|
|
50
|
+
lines = [
|
|
51
|
+
'{"model":"llama3.2:latest","message":{"role":"assistant","content":"Hi"},"done":false}',
|
|
52
|
+
'{"model":"llama3.2:latest","message":{"role":"assistant","content":" there"},'
|
|
53
|
+
'"done":false}',
|
|
54
|
+
'{"model":"llama3.2:latest","message":{"role":"assistant","content":"!"},"done":false}',
|
|
55
|
+
'{"model":"llama3.2:latest","message":{"role":"assistant","content":""},'
|
|
56
|
+
'"done":true,"done_reason":"stop","prompt_eval_count":32,"eval_count":4}',
|
|
57
|
+
]
|
|
58
|
+
|
|
59
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
60
|
+
return httpx.Response(200, text="\n".join(lines))
|
|
61
|
+
|
|
62
|
+
client = _client_with_transport(handler)
|
|
63
|
+
events = [event async for event in client.stream([ChatMessage(role="user", content="hi")])]
|
|
64
|
+
|
|
65
|
+
assert events[:-1] == ["Hi", " there", "!"]
|
|
66
|
+
assert events[-1] == Usage(tokens_in=32, tokens_out=4, finish_reason="stop")
|
|
67
|
+
assert sum(isinstance(e, Usage) for e in events) == 1
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
async def test_health_true_on_200():
|
|
71
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
72
|
+
return httpx.Response(200, json={"models": []})
|
|
73
|
+
|
|
74
|
+
client = _client_with_transport(handler)
|
|
75
|
+
assert await client.health() is True
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from aikit.observability import PrometheusMetricsSink
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def test_incr_and_observe_render_as_prometheus_text():
|
|
7
|
+
sink = PrometheusMetricsSink()
|
|
8
|
+
sink.incr("jobs_total", status="ok")
|
|
9
|
+
sink.incr("jobs_total", status="ok")
|
|
10
|
+
sink.observe("latency_ms", 12.5, backend="ollama")
|
|
11
|
+
|
|
12
|
+
body = sink.render().decode()
|
|
13
|
+
assert 'jobs_total{status="ok"} 2.0' in body
|
|
14
|
+
assert "latency_ms_sum" in body
|
|
15
|
+
assert 'backend="ollama"' in body
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
from aikit import __version__
|
|
2
|
+
from aikit.embeddings import EmbeddingService
|
|
3
|
+
from aikit.model_client import ChatMessage, ModelClient
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def test_version():
|
|
7
|
+
assert isinstance(__version__, str)
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def test_chat_message_immutable():
|
|
11
|
+
m = ChatMessage(role="user", content="hi")
|
|
12
|
+
assert m.role == "user"
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def test_protocols_importable():
|
|
16
|
+
assert ModelClient is not None
|
|
17
|
+
assert EmbeddingService is not None
|