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.
@@ -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,12 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ .venv/
4
+ venv/
5
+ dist/
6
+ build/
7
+ *.egg-info/
8
+ .env
9
+ .mypy_cache/
10
+ .ruff_cache/
11
+ .pytest_cache/
12
+ .DS_Store
@@ -0,0 +1,7 @@
1
+ # Changelog
2
+ All notable changes follow [Keep a Changelog] and [SemVer].
3
+
4
+ ## [Unreleased]
5
+ ### Added
6
+ - Public interfaces: ModelClient, EmbeddingService, JobQueue, MetricsSink,
7
+ BaseServiceSettings, DB base (stub).
@@ -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,6 @@
1
+ .PHONY: install lint type test all
2
+ install: ; pip install -e ".[dev]"
3
+ lint: ; ruff check .
4
+ type: ; mypy src
5
+ test: ; pytest -q
6
+ all: lint type test
@@ -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,6 @@
1
+ """aikit — shared platform layer for AI-backend systems.
2
+
3
+ Public interfaces (SOLID contracts) live here. Concrete implementations are
4
+ optional and gated behind extras so consumers pull only what they use.
5
+ """
6
+ __version__ = "0.3.0"
@@ -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