diff --git a/docker-compose.yml b/docker-compose.yml index 7c9f4ec4..28ce60c8 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -16,7 +16,7 @@ services: timeout: 5s retries: 10 db: - image: postgres:17 + image: pgvector/pgvector:pg18 environment: POSTGRES_DB: database_app_test POSTGRES_USER: app diff --git a/docs/adr/063-ai-embeddings.md b/docs/adr/063-ai-embeddings.md new file mode 100644 index 00000000..3ad247ae --- /dev/null +++ b/docs/adr/063-ai-embeddings.md @@ -0,0 +1,61 @@ +# 063 — AI embeddings with a Laravel-style fake + +- Date: 2026-11-01 +- Status: Accepted + +## Problem + +The `ai` package covers agents, images, audio and transcription, but has no way to turn text into vectors. Apps that need semantic search or retrieval call a provider SDK by hand, so they cannot switch providers through `AIConfig` and have nothing to fake in tests. Laravel's AI SDK offers `Embeddings::for([...])->dimensions(n)->generate($provider, $model)` and an `Embeddings::fake()` with `assertGenerated`, `assertNotGenerated`, `assertNothingGenerated`, `preventStrayEmbeddings` and `fakeEmbedding`. + +## Alternatives + +1. **An `EmbeddingsFactory` ABC with one class per vendor SDK**, like `ImageFactory` and `AudioFactory`. Rejected: it rewrites what LangChain already ships, and chat agents already go through LangChain's `init_chat_model`. +2. **Have apps build LangChain `Embeddings` objects themselves.** Rejected: no shared config, no fake, no assertions. +3. **A context-manager fake, like `AgentFake`.** Rejected for this API: Laravel's call shape is `Embeddings::fake()` followed by static assertions, and nothing has to be entered or exited to embed. +4. **Port embedding caching.** Deferred: it needs a cache store, which the package does not have, and it is not needed to generate or fake embeddings. +5. **Accept text in the vector query methods and embed it automatically, as Laravel does.** Rejected: the query builder is synchronous and an embedding call is async. Callers run `await Str.of(text).to_embeddings()` first, and passing a `str` raises `TypeError` with that hint. +6. **Bind vectors through a pgvector-specific asyncpg codec.** Rejected: pgvector accepts the JSON text form `[0.1, 0.2]` for an untyped parameter, which is also how Laravel binds them. No driver registration is needed. +7. **A fluent builder over LangChain's `init_embeddings`, with a class-level fake.** Chosen. + +## Decision + +- `Embeddings.of(inputs).dimensions(n).generate(provider=None, model=None)` returns an `EmbeddingsResponse(embeddings, provider, model)` with `first()`. `of` replaces Laravel's `for`, which is a Python keyword; `Image.of` and `Audio.of` use the same name. +- The provider comes from `AIConfig.default_embeddings` (env `AI_DEFAULT_EMBEDDINGS_PROVIDER`, default `openai`) through a new `ModelType.EMBEDDINGS`, whose key is `default_embeddings`. Default models: OpenAI `text-embedding-3-small`, Google `gemini-embedding-001`, Bedrock `amazon.titan-embed-text-v2:0`. +- The real call is `init_embeddings(model, provider=...)` and then `aembed_documents(inputs)`. A table maps each `Lab` to its LangChain embeddings provider key and dimensions option: OpenAI `dimensions`, Google `output_dimensionality`, Bedrock `model_kwargs={"dimensions": n}`. The key from config is passed as `api_key` for OpenAI and Google. Bedrock gets `Lab.get_client_options()` (region, credentials, profile). A provider without embeddings (Anthropic, ElevenLabs) raises `ValueError`. +- Every call builds an `EmbeddingsPrompt(inputs, dimensions, provider, model)` with `contains(text)`, the same object the fake records and passes to callbacks. +- `Embeddings.fake(responses=None)` installs a class-level `EmbeddingsFake` and returns it: + - with no responses, each input gets a random unit vector of the prompt's dimensions (1536 when unset), from `Embeddings.fake_embedding(dimensions)`; + - a list gives one response per call, each a list of vectors; once the list runs out, calls fall back to random vectors; + - a callable receives the `EmbeddingsPrompt` and returns the vectors. + - `prevent_stray_embeddings()` makes a call with no queued or callable response raise `RuntimeError`. +- `Embeddings.assert_generated`, `assert_not_generated` and `assert_nothing_generated` check the recorded prompts. They accept a callable predicate or a string, which matches through `contains`. They raise `RuntimeError` when embeddings are not faked. `Embeddings.reset_fake()` removes the fake, and `Ai.reset_fakes()` removes it too, so a test that only resets agent fakes does not leak the embeddings fake. +- `Str.of(text).to_embeddings(provider=None, model=None, dimensions=None)` is an async `Stringable` method that returns one vector through `Embeddings`, so it is faked by `Embeddings.fake()`. +- pgvector columns: `table.vector(column, dimensions)` adds a `VECTOR(n)` column on PostgreSQL and raises `ValueError` on any other platform. `.index()` picks the index type from the column type: when any indexed column was added as `vector` in the same blueprint, whether named or taken from the last column, it creates an HNSW index with `vector_cosine_ops`. A btree index on `vector(1536)` fails every insert with `index row size exceeds btree maximum`. `table.vector_index(column=None, name=None)` creates the HNSW index explicitly, for a vector column that already exists, and raises `ValueError` on other platforms. `schema.ensure_vector_extension_exists()` runs `CREATE EXTENSION IF NOT EXISTS vector`. +- Vector queries, PostgreSQL only (other grammars raise `ValueError`), using cosine distance `<=>`: + - `select_vector_distance(column, vector, as_=None)` adds `(col <=> ?) AS col_distance`; + - `where_vector_distance_less_than(column, vector, max_distance)` and its `or_` variant; + - `order_by_vector_distance(column, vector)`, closest first; + - `where_vector_similar_to(column, vector, min_similarity=0.6, order=True)` filters on `distance <= 1 - min_similarity` and orders by distance, like Laravel's `whereVectorSimilarTo`. + The vector is bound as a JSON array of floats. The same methods are available on `Model`. + Column and alias names go through `PostgresGrammar.wrap_identifier`, which doubles any embedded `"`, so a name cannot break out of its identifier. An unqualified column is prefixed with the query alias or table (`"documents"."embedding"`), like other where clauses, so it stays unambiguous in joins. +- Model fields annotated with parameterized `list[...]` or `dict[...]` (for example `embedding: list[float]`) now get `JsonCast`, as bare `list` and `dict` already did. Before this they fell back to the string cast. `JsonCast.set(None)` now returns `None` (SQL `NULL`) instead of the JSON text `'null'`, which pgvector rejects and which other casts never produced. +- A raw `ORDER BY` that follows another order term is now separated by a comma. Before this, `order_by("a").order_by_raw("b")` produced invalid SQL. +- Laravel's queued assertions are not ported, because the package has no queued generation for any modality. + +## Validation + +`fastapi_startkit/tests/ai/test_embeddings.py`: + +- `generate` passes the resolved provider key, model, dimension option, API key and Bedrock client options to `init_embeddings` and returns the vectors from `aembed_documents` (LangChain patched); +- the default provider and model come from `AIConfig`, and explicit arguments override them; +- an unsupported provider raises `ValueError`; +- `fake()` returns unit vectors of the requested or default dimensions, one per input; +- list responses are used in order, then fall back to random vectors; +- callable responses receive the prompt; +- `prevent_stray_embeddings` raises once responses run out; +- `assert_generated` / `assert_not_generated` / `assert_nothing_generated` pass and fail as expected with callables and strings; +- assertions without a fake raise. + +`tests/support/test_string.py` checks that `to_embeddings` returns the faked vector and forwards provider, model and dimensions. `tests/masoniteorm/query/grammars/test_vector_grammar.py` checks the SQL and binding order of each vector query method, table and alias prefixes, the escaping of an injected `"` in columns and aliases, the order-by comma fix, the `TypeError` for text and the `ValueError` on non-Postgres grammars. `tests/masoniteorm/schema/test_vector_schema.py` checks `VECTOR(n)`, the HNSW index on create and alter for bare and named `.index()`, `vector_index()` on an existing column, btree for other columns, the extension SQL and the errors on other platforms. `tests/masoniteorm/models/test_caster.py` checks the JSON cast for generic annotations and that `None` is stored as `NULL`. `test_embeddings.py` checks that `Ai.reset_fakes()` clears the embeddings fake. `tests/masoniteorm/postgres/test_postgres_vector.py` runs against a live `pgvector/pgvector` database: it round-trips `list[float]` and checks similarity filtering, distance selection and ordering. + +`uv run pytest --ignore=tests/masoniteorm/postgres --cov` passes with coverage at or above 80%. diff --git a/docs/adr/AGENTS.md b/docs/adr/AGENTS.md index 22271811..ef2fc3e9 100644 --- a/docs/adr/AGENTS.md +++ b/docs/adr/AGENTS.md @@ -51,3 +51,4 @@ Read this index first, then open the records relevant to the change. Before impl | [059](059-subquery-where-in-and-select-exists.md) | 2026-10-31 | Subqueries in `where_in`/`where_not_in`, `where_not_exists` and `add_select_exists` | Membership, anti-join and per-row exists flags needed `DB.select`: `where_in(col, builder)` raised `'QueryBuilder' object is not iterable` and no method emitted `NOT EXISTS` or an `EXISTS` column. `where_in`/`where_not_in` now take a builder or a callable (like `select_sub`) and render `col IN (subquery)` via `SubSelectExpression`; iterables are unchanged. `where_exists`/`or_where_exists` also take a callable; new `where_not_exists`/`or_where_not_exists`. `add_select_exists(alias, query)` appends `EXISTS(subquery) AS alias` (selecting `{table}.*` first when nothing is selected) through a `SubGroupExpression(exists=True)` and the grammar's `select_exists_string()`; MSSQL renders `CASE WHEN EXISTS(...) THEN 1 ELSE 0 END`. A shared `_resolve_subquery` raises one `TypeError`. Subquery bindings stay bound and land in SQL order. `Model` proxies all. Limits: Postgres returns `bool`, others `1`/`0`; `NOT IN` over a subquery with `NULL` matches nothing. Chosen over raw SQL or separate `*_sub` method names. | | [061](061-builder-row-locks.md) | 2026-10-31 | Row locks on the query builder: `lock_for_update` and `shared_lock` | The grammars already rendered a per-dialect lock clause, but no builder method set it and `get_grammar()` never passed `lock`, so apps wrote `SELECT ... FOR UPDATE` by hand. `QueryBuilder.lock_for_update()`/`shared_lock()` set `_lock` to `"update"`/`"share"` (last call wins) and `get_grammar()` passes it: Postgres `FOR UPDATE`/`FOR SHARE`, MySQL `FOR UPDATE`/`LOCK IN SHARE MODE`, SQLite a no-op, and MSSQL now `WITH(ROWLOCK, UPDLOCK, HOLDLOCK)`/`WITH(ROWLOCK, HOLDLOCK)` instead of a bare `ROWLOCK` hint. `Model` proxies both. Only selects lock; `new()` groups and the `paginate()` count start unlocked. Chosen over update-first, whose `FOR NO KEY UPDATE` does not wait for an in-flight foreign-key insert, and over a stringly `lock(mode)`. Known limit: Postgres refuses a lock with aggregates, `DISTINCT` or `GROUP BY`. | | [062](062-update-uses-the-updated-table.md) | 2026-11-01 | `update()` compiles against the updated table | `update()` built its grammar with `self.grammar()`, keeping the default table `users`, so MySQL and MSSQL prefixed every `SET`/`WHERE` column with `users` whatever table was updated (`UPDATE [queues] SET [users].[status] = ?`). It now passes `table=self._table`. Postgres and SQLite do not prefix update columns. `update()` still returns the row count; a queue claim reads locked rows, iterates and updates each by key, as Laravel's database queue does, so no `UPDATE ... RETURNING` is added. | +| [063](063-ai-embeddings.md) | 2026-11-01 | AI embeddings with a Laravel-style fake | The `ai` package had no way to embed text, so apps called vendor SDKs by hand, with no shared config and nothing to fake. `Embeddings.of(inputs).dimensions(n).generate(provider, model)` returns an `EmbeddingsResponse`, using LangChain `init_embeddings` + `aembed_documents`. A per-`Lab` table maps the provider key and dimension option (OpenAI `dimensions`, Google `output_dimensionality`, Bedrock `model_kwargs`). The default comes from `AIConfig.default_embeddings` / `ModelType.EMBEDDINGS`. `Embeddings.fake()` (random unit vectors, a list of responses, or a callable taking `EmbeddingsPrompt`), `prevent_stray_embeddings`, `assert_generated`/`assert_not_generated`/`assert_nothing_generated` and `fake_embedding` mirror Laravel. Chosen over a per-vendor `EmbeddingsFactory` ABC. `Str.of(text).to_embeddings()` returns one vector. `table.vector(col, n)` adds a pgvector `VECTOR(n)` column, `.index()` on a vector column (named or not) and `vector_index()` add an HNSW cosine index, and `schema.ensure_vector_extension_exists()` enables the extension. `select_vector_distance`, `where_vector_distance_less_than`, `order_by_vector_distance` and `where_vector_similar_to` use `<=>` on Postgres, with escaped, table-prefixed identifiers. Generic `list[...]`/`dict[...]` fields cast as JSON. Raw order-by now gets a comma. Caching and queued assertions are deferred. | diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index e2802e4f..511f7cc7 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py @@ -6,6 +6,7 @@ from .config.ai import AIConfig from .decorators import max_steps, max_tokens, model, provider, timeout, top_p from .document import Document +from .embeddings import Embeddings, EmbeddingsFake, EmbeddingsPrompt, EmbeddingsResponse from .image import Image, ImageResponse from .image_factory import ImageFactory from .ai import Ai @@ -31,6 +32,10 @@ "AudioFactory", "BedrockConfig", "Document", + "Embeddings", + "EmbeddingsFake", + "EmbeddingsPrompt", + "EmbeddingsResponse", "ElevenLabsConfig", "GoogleConfig", "Image", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/ai.py b/fastapi_startkit/src/fastapi_startkit/ai/ai.py index 8c662486..1486abaf 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/ai.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/ai.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any +from .embeddings import Embeddings from .lab import Lab if TYPE_CHECKING: @@ -45,6 +46,7 @@ def forget(cls, agent: "Agent | str") -> None: @classmethod def reset_fakes(cls) -> None: cls._fakes.clear() + Embeddings.reset_fake() def get_model_for( self, diff --git a/fastapi_startkit/src/fastapi_startkit/ai/config/ai.py b/fastapi_startkit/src/fastapi_startkit/ai/config/ai.py index 5a5e6ed3..1a6220d3 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/config/ai.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/config/ai.py @@ -14,6 +14,7 @@ class AIConfig(Config): default_image: str = env("AI_DEFAULT_IMAGE_PROVIDER", "openai") default_audio: str = env("AI_DEFAULT_AUDIO_PROVIDER", "openai") default_transcribe: str = env("AI_DEFAULT_TRANSCRIBE_PROVIDER", "openai") + default_embeddings: str = env("AI_DEFAULT_EMBEDDINGS_PROVIDER", "openai") providers: dict[str, Config] = field( default_factory=lambda: { diff --git a/fastapi_startkit/src/fastapi_startkit/ai/config/config.py b/fastapi_startkit/src/fastapi_startkit/ai/config/config.py index 321cba5b..630ff5b7 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/config/config.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/config/config.py @@ -36,6 +36,7 @@ class OpenAIConfig(Config): default_factory=lambda: { "default": "gpt-4o", "default_image": "dall-e-3", + "default_embeddings": "text-embedding-3-small", } ) @@ -51,6 +52,7 @@ class GoogleConfig(Config): default_factory=lambda: { "default": "gemini-2.5-flash-lite", "default_image": "imagen-3.0-generate-002", + "default_embeddings": "gemini-embedding-001", } ) @@ -83,5 +85,6 @@ class BedrockConfig(Config): models: dict[str, str] = field( default_factory=lambda: { "default": os.getenv("AWS_BEDROCK_MODEL") or "us.anthropic.claude-sonnet-4-6", + "default_embeddings": "amazon.titan-embed-text-v2:0", } ) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/embeddings.py b/fastapi_startkit/src/fastapi_startkit/ai/embeddings.py new file mode 100644 index 00000000..3806f049 --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/embeddings.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import math +import random +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, ClassVar, NamedTuple + +from .lab import Lab, ModelType + +Vectors = list[list[float]] +EmbeddingsResponder = Callable[["EmbeddingsPrompt"], Vectors] +PromptMatcher = str | Callable[["EmbeddingsPrompt"], bool] + +DEFAULT_FAKE_DIMENSIONS = 1536 + + +class EmbeddingsProvider(NamedTuple): + key: str + dimensions_option: Callable[[int], dict[str, Any]] + accepts_api_key: bool + + +EMBEDDINGS_PROVIDERS: dict[Lab, EmbeddingsProvider] = { + Lab.OPENAI: EmbeddingsProvider("openai", lambda size: {"dimensions": size}, True), + Lab.GOOGLE: EmbeddingsProvider("google_genai", lambda size: {"output_dimensionality": size}, True), + Lab.BEDROCK: EmbeddingsProvider("bedrock", lambda size: {"model_kwargs": {"dimensions": size}}, False), +} + + +@dataclass(frozen=True) +class EmbeddingsPrompt: + inputs: list[str] + dimensions: int | None + provider: str + model: str + + def contains(self, text: str) -> bool: + return any(text in value for value in self.inputs) + + +@dataclass(frozen=True) +class EmbeddingsResponse: + embeddings: Vectors + provider: str + model: str + + def first(self) -> list[float]: + return self.embeddings[0] + + +class EmbeddingsFake: + def __init__(self, responses: list[Vectors] | EmbeddingsResponder | None = None) -> None: + self._responder = responses if callable(responses) else None + self._queue = [] if callable(responses) or responses is None else list(responses) + self._prevent_stray = False + self.prompts: list[EmbeddingsPrompt] = [] + + def prevent_stray_embeddings(self) -> EmbeddingsFake: + self._prevent_stray = True + return self + + def respond(self, prompt: EmbeddingsPrompt) -> Vectors: + self.prompts.append(prompt) + if self._responder is not None: + return self._responder(prompt) + if self._queue: + return self._queue.pop(0) + if self._prevent_stray: + raise RuntimeError(f"No fake embeddings response defined for inputs {prompt.inputs!r}.") + size = prompt.dimensions or DEFAULT_FAKE_DIMENSIONS + return [Embeddings.fake_embedding(size) for _ in prompt.inputs] + + +class Embeddings: + _faked: ClassVar[EmbeddingsFake | None] = None + + def __init__(self, inputs: list[str]) -> None: + self._inputs = list(inputs) + self._dimensions: int | None = None + + @classmethod + def of(cls, inputs: list[str]) -> Embeddings: + return cls(inputs) + + def dimensions(self, size: int) -> Embeddings: + self._dimensions = size + return self + + async def generate(self, provider: str | None = None, model: str | None = None) -> EmbeddingsResponse: + lab = Lab.get_provider(provider, ModelType.EMBEDDINGS) + prompt = EmbeddingsPrompt( + inputs=self._inputs, + dimensions=self._dimensions, + provider=lab.value, + model=lab.get_model(model, ModelType.EMBEDDINGS), + ) + fake = Embeddings._faked + embeddings = fake.respond(prompt) if fake is not None else await self._embed(lab, prompt) + return EmbeddingsResponse(embeddings=embeddings, provider=prompt.provider, model=prompt.model) + + async def _embed(self, lab: Lab, prompt: EmbeddingsPrompt) -> Vectors: + from langchain.embeddings import init_embeddings # noqa: PLC0415 # pyright: ignore[reportMissingImports] + + return await init_embeddings(prompt.model, **self._options(lab, prompt)).aembed_documents(prompt.inputs) + + @staticmethod + def _options(lab: Lab, prompt: EmbeddingsPrompt) -> dict[str, Any]: + if lab not in EMBEDDINGS_PROVIDERS: + supported = ", ".join(sorted(EMBEDDINGS_PROVIDERS)) + raise ValueError(f"Provider {lab.value!r} does not support embeddings. Use one of: {supported}.") + settings = EMBEDDINGS_PROVIDERS[lab] + options: dict[str, Any] = {"provider": settings.key, **lab.get_client_options()} + api_key = lab.get_api_key() + if settings.accepts_api_key and api_key: + options["api_key"] = api_key + if prompt.dimensions: + options.update(settings.dimensions_option(prompt.dimensions)) + return options + + @classmethod + def fake(cls, responses: list[Vectors] | EmbeddingsResponder | None = None) -> EmbeddingsFake: + Embeddings._faked = EmbeddingsFake(responses) + return Embeddings._faked + + @classmethod + def reset_fake(cls) -> None: + Embeddings._faked = None + + @staticmethod + def fake_embedding(dimensions: int = DEFAULT_FAKE_DIMENSIONS) -> list[float]: + vector = [random.gauss(0.0, 1.0) for _ in range(dimensions)] + norm = math.sqrt(sum(value * value for value in vector)) or 1.0 + return [value / norm for value in vector] + + @classmethod + def assert_generated(cls, matcher: PromptMatcher) -> None: + if not any(_matches(matcher, prompt) for prompt in cls._recorded()): + raise AssertionError(f"Expected embeddings matching {matcher!r} to be generated.") + + @classmethod + def assert_not_generated(cls, matcher: PromptMatcher) -> None: + if any(_matches(matcher, prompt) for prompt in cls._recorded()): + raise AssertionError(f"Unexpected embeddings matching {matcher!r} were generated.") + + @classmethod + def assert_nothing_generated(cls) -> None: + prompts = cls._recorded() + if prompts: + raise AssertionError(f"Expected no embeddings, but {len(prompts)} were generated.") + + @classmethod + def _recorded(cls) -> list[EmbeddingsPrompt]: + if Embeddings._faked is None: + raise RuntimeError("Embeddings are not faked. Call Embeddings.fake() first.") + return Embeddings._faked.prompts + + +def _matches(matcher: PromptMatcher, prompt: EmbeddingsPrompt) -> bool: + return prompt.contains(matcher) if isinstance(matcher, str) else bool(matcher(prompt)) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/lab.py b/fastapi_startkit/src/fastapi_startkit/ai/lab.py index f1c7c260..93378786 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/lab.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/lab.py @@ -18,6 +18,7 @@ class ModelType(StrEnum): IMAGE = "image" AUDIO = "audio" TRANSCRIBE = "transcribe" + EMBEDDINGS = "embeddings" @property def key(self) -> str: @@ -26,6 +27,7 @@ def key(self) -> str: ModelType.IMAGE: "default_image", ModelType.AUDIO: "default_audio", ModelType.TRANSCRIBE: "default_transcribe", + ModelType.EMBEDDINGS: "default_embeddings", }[self] diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py index c48d86f2..2f0e52b6 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py @@ -1,5 +1,6 @@ import inspect -from collections.abc import Callable, Iterable +import json +from collections.abc import Callable, Iterable, Sequence from typing import TYPE_CHECKING, Any, Generic, Literal, Self, TypeVar, overload from fastapi_startkit.masoniteorm.expressions.expressions import ( @@ -671,6 +672,51 @@ def or_where_raw(self, expression: str, bindings: tuple[Any, ...] = ()) -> "Self self._wheres += (QueryExpression(expression, "=", None, raw=True, keyword="or", bindings=bindings),) return self + def select_vector_distance(self, column: str, vector: Sequence[float], as_: str | None = None) -> "Self": + alias = as_ or f"{column.split('.')[-1]}_distance" + return self.add_select_raw( + f"({self._vector_column(column)} <=> ?) AS {self._vector_identifier(alias)}", [self._vector_binding(vector)] + ) + + def where_vector_distance_less_than(self, column: str, vector: Sequence[float], max_distance: float) -> "Self": + return self.where_raw( + f"({self._vector_column(column)} <=> ?) <= ?", (self._vector_binding(vector), max_distance) + ) + + def or_where_vector_distance_less_than(self, column: str, vector: Sequence[float], max_distance: float) -> "Self": + return self.or_where_raw( + f"({self._vector_column(column)} <=> ?) <= ?", (self._vector_binding(vector), max_distance) + ) + + def order_by_vector_distance(self, column: str, vector: Sequence[float]) -> "Self": + expression = f"({self._vector_column(column)} <=> ?) ASC" + self._order_by += (OrderByExpression(expression, raw=True, bindings=(self._vector_binding(vector),)),) + return self + + def where_vector_similar_to( + self, column: str, vector: Sequence[float], min_similarity: float = 0.6, order: bool = True + ) -> "Self": + self.where_vector_distance_less_than(column, vector, 1 - min_similarity) + return self.order_by_vector_distance(column, vector) if order else self + + def _vector_column(self, column: str) -> str: + table = self._alias or self._table + qualified = column if "." in column or not table else f"{table}.{column}" + return ".".join(self._vector_identifier(part) for part in qualified.split(".")) + + def _vector_identifier(self, name: str) -> str: + if not self.grammar.supports_vectors: + raise ValueError(f"{self.grammar.__name__} does not support vector queries; use PostgreSQL with pgvector.") + return self.grammar.wrap_identifier(name) + + @staticmethod + def _vector_binding(vector: Sequence[float]) -> str: + if isinstance(vector, str): + raise TypeError( + "Vector queries need a list of floats; embed text first with await Str.of(text).to_embeddings()." + ) + return json.dumps([float(value) for value in vector]) + def join(self, table: str, column1: str, equality: str, column2: str, clause: str = "join") -> "Self": join_clause = JoinClause(table, clause=clause) join_clause.on(column1, equality, column2) diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/caster.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/caster.py index 6d8309d4..91b7336f 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/caster.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/caster.py @@ -56,6 +56,9 @@ def get(self, value): return value def set(self, value): + if value is None: + return None + if isinstance(value, str): # make sure the string is valid JSON json.loads(value) @@ -315,7 +318,7 @@ def normalize_type(t): return "float" if t is bool: return "bool" - if t is dict or t is list: + if t in (dict, list) or get_origin(t) in (dict, list): return "json" if t is pendulum.DateTime or t is datetime.datetime or t is datetime.date or t is Carbon: return "date" diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/model.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/model.py index d2770037..9fa27eda 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/model.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/model.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Self, dataclass_transform, overload import inflection @@ -243,6 +244,28 @@ def not_between(cls, column: str, low: Any, high: Any) -> "QueryBuilder[Self]": def distinct(cls) -> "QueryBuilder[Self]": return cls.query().distinct() + @classmethod + def select_vector_distance( + cls, column: str, vector: Sequence[float], as_: str | None = None + ) -> "QueryBuilder[Self]": + return cls.query().select_vector_distance(column, vector, as_) + + @classmethod + def where_vector_distance_less_than( + cls, column: str, vector: Sequence[float], max_distance: float + ) -> "QueryBuilder[Self]": + return cls.query().where_vector_distance_less_than(column, vector, max_distance) + + @classmethod + def order_by_vector_distance(cls, column: str, vector: Sequence[float]) -> "QueryBuilder[Self]": + return cls.query().order_by_vector_distance(column, vector) + + @classmethod + def where_vector_similar_to( + cls, column: str, vector: Sequence[float], min_similarity: float = 0.6, order: bool = True + ) -> "QueryBuilder[Self]": + return cls.query().where_vector_similar_to(column, vector, min_similarity, order) + @classmethod def lock_for_update(cls) -> "QueryBuilder[Self]": return cls.query().lock_for_update() diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/query/grammars/BaseGrammar.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/query/grammars/BaseGrammar.py index cab3e71c..e2d53083 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/query/grammars/BaseGrammar.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/query/grammars/BaseGrammar.py @@ -25,6 +25,7 @@ class BaseGrammar: """ table = "users" + supports_vectors = False if TYPE_CHECKING: # Dialect contract: every concrete grammar (MySQL, Postgres, SQLite, MSSQL) @@ -433,6 +434,8 @@ def process_order_by(self, qmark=False): order_crit = "" for order_bys in self._order_by: if order_bys.raw: + if len(order_crit): + order_crit += ", " order_crit += order_bys.column if not isinstance(order_bys.bindings, (list, tuple)): raise ValueError(f"Bindings must be tuple or list. Received {type(order_bys.bindings)}") diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/query/grammars/PostgresGrammar.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/query/grammars/PostgresGrammar.py index 9440695d..b276e8ba 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/query/grammars/PostgresGrammar.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/query/grammars/PostgresGrammar.py @@ -10,6 +10,8 @@ class PostgresGrammar(BaseGrammar): """Postgres grammar class.""" + supports_vectors = True + types_without_lengths: list[str] _creates: Any @@ -151,6 +153,11 @@ def to_sql(self): def table_string(self): return '"{table}"' + @staticmethod + def wrap_identifier(name: str) -> str: + escaped = name.replace('"', '""') + return f'"{escaped}"' + def order_by_format(self): return "{column} {direction}" diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/Blueprint.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/Blueprint.py index fc8732ff..842338a9 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/Blueprint.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/Blueprint.py @@ -636,6 +636,12 @@ def geometry(self, column: str, nullable: bool = False): self._last_column = self.table.add_column(column, "geometry", nullable=nullable) return self + def vector(self, column: str, dimensions: int, nullable: bool = False): + if not self.platform.supports_vectors: + raise ValueError(f"{self.platform.__name__} does not support vector columns; use PostgreSQL with pgvector.") + self._last_column = self.table.add_column(column, "vector", length=dimensions, nullable=nullable) + return self + def year(self, column: str, length: int = 4, default: object = None, nullable: bool = False): """Sets a column to be the year representation for the table. @@ -807,18 +813,25 @@ def index(self, column: str | list[str] | None = None, name: str | None = None): Returns: self """ - if not column: - column = self._column().name + columns = self._index_columns(column) + index_type = "vector" if any(self._is_vector_column(name) for name in columns) else "index" + return self._add_index(columns, name, index_type) - if not isinstance(column, list): - column = [column] + def vector_index(self, column: str | None = None, name: str | None = None): + if not self.platform.supports_vectors: + raise ValueError(f"{self.platform.__name__} does not support vector indexes; use PostgreSQL with pgvector.") + return self._add_index(self._index_columns(column), name, "vector") - self.table.add_index( - column, - name or f"{self.table.name}_{'_'.join(column)}_index", - "index", - ) + def _index_columns(self, column: str | list[str] | None) -> list[str]: + column = column or self._column().name + return column if isinstance(column, list) else [column] + + def _is_vector_column(self, name: str) -> bool: + added = self.table.added_columns.get(name) + return added is not None and added.column_type == "vector" + def _add_index(self, columns: list[str], name: str | None, index_type: str): + self.table.add_index(columns, name or f"{self.table.name}_{'_'.join(columns)}_index", index_type) return self def fulltext(self, column: str | list[str] | None = None, name: str | None = None): diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/platforms/Platform.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/platforms/Platform.py index 119981f7..5695533b 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/platforms/Platform.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/platforms/Platform.py @@ -75,6 +75,11 @@ def get_column_string(self) -> str: def get_foreign_key_constraint_string(self) -> str: raise NotImplementedError + supports_vectors = False + + def compile_ensure_vector_extension(self) -> str: + raise ValueError(f"{type(self).__name__} does not support vector columns; use PostgreSQL with pgvector.") + def create_column_length(self, column_type): if column_type in self.types_without_lengths: return "" diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/platforms/PostgresPlatform.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/platforms/PostgresPlatform.py index ca59ed85..fe3b51c4 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/platforms/PostgresPlatform.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/platforms/PostgresPlatform.py @@ -57,8 +57,11 @@ class PostgresPlatform(Platform): "datetime": "TIMESTAMPTZ", "tiny_increments": "TINYINT AUTO_INCREMENT", "unsigned": "INT", + "vector": "VECTOR", } + supports_vectors = True + table_info_map = { "CHARACTER VARYING": "string", "TIMESTAMP WITH TIME ZONE": "datetime", @@ -73,6 +76,17 @@ class PostgresPlatform(Platform): premapped_nulls = {True: "NULL", False: "NOT NULL"} + def compile_ensure_vector_extension(self) -> str: + return "CREATE EXTENSION IF NOT EXISTS vector" + + def compile_create_index(self, table_name, index): + if index.index_type == "vector": + return ( + f"CREATE INDEX {index.name} ON {self.wrap_table(table_name)} " + f"USING hnsw ({','.join(index.column)} vector_cosine_ops)" + ) + return f"CREATE INDEX {index.name} ON {self.wrap_table(table_name)}({','.join(index.column)})" + def compile_create_sql(self, table, if_not_exists=False): sql = [] table_create_format = self.create_if_not_exists_format() if if_not_exists else self.create_format() @@ -94,14 +108,8 @@ def compile_create_sql(self, table, if_not_exists=False): ) if table.added_indexes: - for name, index in table.added_indexes.items(): - sql.append( - "CREATE INDEX {name} ON {table}({column})".format( - name=index.name, - table=self.wrap_table(table.name), - column=",".join(index.column), - ) - ) + for index in table.added_indexes.values(): + sql.append(self.compile_create_index(table.name, index)) for name, column in table.get_added_columns().items(): if column.comment: @@ -306,14 +314,8 @@ def compile_alter_sql(self, table): sql.append(f"ALTER TABLE {self.wrap_table(table.name)} DROP CONSTRAINT {constraint}") if table.added_indexes: - for name, index in table.added_indexes.items(): - sql.append( - "CREATE INDEX {name} ON {table}({column})".format( - name=index.name, - table=self.wrap_table(table.name), - column=",".join(index.column), - ) - ) + for index in table.added_indexes.values(): + sql.append(self.compile_create_index(table.name, index)) if table.added_constraints: for name, constraint in table.added_constraints.items(): diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/schema.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/schema.py index 29c168d0..550bdc98 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/schema.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/schema/schema.py @@ -47,6 +47,7 @@ class Schema: "tiny_increments": int, "unsigned": int, "unsigned_integer": int, + "vector": list, } def __init__(self, manager) -> None: @@ -158,6 +159,9 @@ async def enable_foreign_key_constraints(self) -> None: sql = connection.get_default_platform()().enable_foreign_key_constraints() await connection.statement(sql, ()) + async def ensure_vector_extension_exists(self) -> None: + await self.get_connection().statement(self.platform().compile_ensure_vector_extension(), ()) + async def get_all_tables(self) -> list[Any]: """Gets all tables in the database.""" connection = self.get_connection() diff --git a/fastapi_startkit/src/fastapi_startkit/support/string.py b/fastapi_startkit/src/fastapi_startkit/support/string.py index 091c5b7c..9ca9f2c3 100644 --- a/fastapi_startkit/src/fastapi_startkit/support/string.py +++ b/fastapi_startkit/src/fastapi_startkit/support/string.py @@ -20,6 +20,17 @@ def camel_case(self) -> "Stringable": def snake_case(self) -> "Stringable": return Stringable(Str.snake_case(self.text)) + async def to_embeddings( + self, provider: str | None = None, model: str | None = None, dimensions: int | None = None + ) -> list[float]: + from fastapi_startkit.ai.embeddings import Embeddings + + embeddings = Embeddings.of([self.text]) + if dimensions: + embeddings.dimensions(dimensions) + response = await embeddings.generate(provider, model) + return response.first() + class Str: @classmethod diff --git a/fastapi_startkit/tests/ai/test_config.py b/fastapi_startkit/tests/ai/test_config.py index 396279d6..20a41356 100644 --- a/fastapi_startkit/tests/ai/test_config.py +++ b/fastapi_startkit/tests/ai/test_config.py @@ -23,6 +23,14 @@ def test_aiconfig_default_reads_ai_provider_env(self): self._patch_env({"AI_PROVIDER": "anthropic"}) self.assertEqual(AIConfig().default, "anthropic") + def test_aiconfig_default_embeddings_provider_is_openai(self): + self._patch_env(unset=["AI_DEFAULT_EMBEDDINGS_PROVIDER"]) + self.assertEqual(AIConfig().default_embeddings, "openai") + + def test_aiconfig_default_embeddings_reads_env(self): + self._patch_env({"AI_DEFAULT_EMBEDDINGS_PROVIDER": "google"}) + self.assertEqual(AIConfig().default_embeddings, "google") + def test_aiconfig_providers_has_anthropic_key(self): self.assertIn("anthropic", AIConfig().providers) diff --git a/fastapi_startkit/tests/ai/test_embeddings.py b/fastapi_startkit/tests/ai/test_embeddings.py new file mode 100644 index 00000000..f5b3efed --- /dev/null +++ b/fastapi_startkit/tests/ai/test_embeddings.py @@ -0,0 +1,310 @@ +import math +import sys +from types import ModuleType +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from fastapi_startkit.ai import Ai, AIConfig, Embeddings, EmbeddingsFake, EmbeddingsPrompt, EmbeddingsResponse +from fastapi_startkit.application import app + + +@pytest.fixture(autouse=True) +def ai_config(): + config = AIConfig() + config.default_embeddings = "openai" + config.providers["openai"].key = "sk-test" + config.providers["google"].key = "g-test" + config.providers["bedrock"].region = "us-east-1" + app().make("config").set("ai", config) + yield config + Embeddings.reset_fake() + + +@pytest.fixture +def init_embeddings(): + model = MagicMock() + model.aembed_documents = AsyncMock(return_value=[[0.1, 0.2], [0.3, 0.4]]) + factory = MagicMock(return_value=model) + module = ModuleType("langchain.embeddings") + module.init_embeddings = factory # type: ignore[attr-defined] + with patch.dict(sys.modules, {"langchain.embeddings": module}): + yield factory + + +async def test_generate_embeds_inputs_with_default_provider_and_model(init_embeddings): + response = await Embeddings.of(["Napa Valley has great wine.", "Laravel is a PHP framework."]).generate() + + assert response.embeddings == [[0.1, 0.2], [0.3, 0.4]] + assert response.first() == [0.1, 0.2] + assert (response.provider, response.model) == ("openai", "text-embedding-3-small") + init_embeddings.assert_called_once_with("text-embedding-3-small", provider="openai", api_key="sk-test") + init_embeddings.return_value.aembed_documents.assert_awaited_once_with( + ["Napa Valley has great wine.", "Laravel is a PHP framework."] + ) + + +@pytest.mark.parametrize( + ("provider", "model", "expected"), + [ + ("openai", "text-embedding-3-large", {"provider": "openai", "api_key": "sk-test", "dimensions": 256}), + ( + "google", + "gemini-embedding-001", + {"provider": "google_genai", "api_key": "g-test", "output_dimensionality": 256}, + ), + ( + "bedrock", + "amazon.titan-embed-text-v2:0", + {"provider": "bedrock", "region_name": "us-east-1", "model_kwargs": {"dimensions": 256}}, + ), + ], +) +async def test_generate_passes_provider_options_and_dimensions(init_embeddings, provider, model, expected): + response = await Embeddings.of(["hello"]).dimensions(256).generate(provider, model) + + assert (response.provider, response.model) == (provider, model) + init_embeddings.assert_called_once_with(model, **expected) + + +async def test_generate_omits_empty_api_key(init_embeddings, ai_config): + ai_config.providers["openai"].key = "" + app().make("config").set("ai", ai_config) + + await Embeddings.of(["hello"]).generate() + + init_embeddings.assert_called_once_with("text-embedding-3-small", provider="openai") + + +async def test_provider_without_embeddings_raises(init_embeddings): + with pytest.raises(ValueError, match="'anthropic' does not support embeddings"): + await Embeddings.of(["hello"]).generate("anthropic", "claude-sonnet-4-6") + init_embeddings.assert_not_called() + + +async def test_fake_generates_unit_vectors_of_requested_dimensions(init_embeddings): + Embeddings.fake() + + response = await Embeddings.of(["a", "b"]).dimensions(8).generate() + default = await Embeddings.of(["c"]).generate() + + assert [len(vector) for vector in response.embeddings] == [8, 8] + assert all(math.isclose(sum(value * value for value in vector), 1.0) for vector in response.embeddings) + assert len(default.first()) == 1536 + init_embeddings.assert_not_called() + + +async def test_fake_uses_listed_responses_in_order_then_random_vectors(): + Embeddings.fake([[[1.0, 0.0]], [[0.0, 1.0]]]) + + first = await Embeddings.of(["a"]).generate() + second = await Embeddings.of(["b"]).generate() + third = await Embeddings.of(["c"]).dimensions(3).generate() + + assert first.embeddings == [[1.0, 0.0]] + assert second.embeddings == [[0.0, 1.0]] + assert len(third.first()) == 3 + + +async def test_fake_callable_receives_the_prompt(): + received: list[EmbeddingsPrompt] = [] + + def respond(prompt: EmbeddingsPrompt): + received.append(prompt) + return [Embeddings.fake_embedding(prompt.dimensions or 4) for _ in prompt.inputs] + + Embeddings.fake(respond) + + response = await Embeddings.of(["a", "b"]).dimensions(4).generate("google") + + assert received == [EmbeddingsPrompt(["a", "b"], 4, "google", "gemini-embedding-001")] + assert [len(vector) for vector in response.embeddings] == [4, 4] + + +async def test_prevent_stray_embeddings_raises_once_responses_run_out(): + Embeddings.fake([[[1.0]]]).prevent_stray_embeddings() + + await Embeddings.of(["a"]).generate() + with pytest.raises(RuntimeError, match="No fake embeddings response"): + await Embeddings.of(["b"]).generate() + + +async def test_generated_assertions(): + Embeddings.fake() + Embeddings.assert_nothing_generated() + + await Embeddings.of(["Laravel is a PHP framework."]).dimensions(1536).generate() + + Embeddings.assert_generated(lambda prompt: prompt.contains("Laravel") and prompt.dimensions == 1536) + Embeddings.assert_generated("PHP") + Embeddings.assert_not_generated(lambda prompt: prompt.contains("Other")) + Embeddings.assert_not_generated("Other") + with pytest.raises(AssertionError): + Embeddings.assert_generated("Other") + with pytest.raises(AssertionError): + Embeddings.assert_not_generated("Laravel") + with pytest.raises(AssertionError): + Embeddings.assert_nothing_generated() + + +def test_assertions_require_a_fake(): + with pytest.raises(RuntimeError, match="not faked"): + Embeddings.assert_nothing_generated() + + +def test_prompt_contains_matches_substrings_of_any_input(): + prompt = EmbeddingsPrompt(["Napa Valley has great wine.", "Laravel is a PHP framework."], None, "openai", "m") + + assert prompt.contains("Valley") + assert prompt.contains("PHP framework") + assert not prompt.contains("Django") + + +def test_prompt_contains_is_false_without_inputs(): + assert not EmbeddingsPrompt([], None, "openai", "m").contains("anything") + + +def test_response_first_returns_the_first_vector(): + response = EmbeddingsResponse([[1.0, 0.0], [0.0, 1.0]], "openai", "m") + + assert response.first() == [1.0, 0.0] + + +def test_fake_embedding_is_a_unit_vector_of_the_given_size(): + vector = Embeddings.fake_embedding(64) + + assert len(vector) == 64 + assert math.isclose(math.sqrt(sum(value * value for value in vector)), 1.0) + + +def test_fake_embedding_defaults_to_1536_dimensions(): + assert len(Embeddings.fake_embedding()) == 1536 + + +def test_fake_embedding_vectors_differ(): + assert Embeddings.fake_embedding(16) != Embeddings.fake_embedding(16) + + +def test_embeddings_fake_records_prompts_in_order(): + fake = EmbeddingsFake() + first = EmbeddingsPrompt(["a"], 2, "openai", "m") + second = EmbeddingsPrompt(["b", "c"], 3, "openai", "m") + + assert len(fake.respond(first)[0]) == 2 + assert [len(vector) for vector in fake.respond(second)] == [3, 3] + assert fake.prompts == [first, second] + + +def test_embeddings_fake_returns_no_vectors_for_no_inputs(): + assert EmbeddingsFake().respond(EmbeddingsPrompt([], None, "openai", "m")) == [] + + +def test_embeddings_fake_callable_wins_over_prevent_stray(): + fake = EmbeddingsFake(lambda prompt: [[0.5]] * len(prompt.inputs)).prevent_stray_embeddings() + + assert fake.respond(EmbeddingsPrompt(["a", "b"], None, "openai", "m")) == [[0.5], [0.5]] + + +def test_prevent_stray_without_responses_raises_on_first_call(): + fake = EmbeddingsFake().prevent_stray_embeddings() + + with pytest.raises(RuntimeError, match=r"\['a'\]"): + fake.respond(EmbeddingsPrompt(["a"], None, "openai", "m")) + assert len(fake.prompts) == 1 + + +def test_prevent_stray_embeddings_returns_the_fake(): + fake = EmbeddingsFake() + + assert fake.prevent_stray_embeddings() is fake + + +def test_listed_responses_are_copied(): + responses = [[[1.0]]] + fake = EmbeddingsFake(responses) + + fake.respond(EmbeddingsPrompt(["a"], None, "openai", "m")) + + assert responses == [[[1.0]]] + + +async def test_fake_replaces_the_previous_fake(): + Embeddings.fake() + await Embeddings.of(["old"]).generate() + + Embeddings.fake() + + Embeddings.assert_nothing_generated() + + +async def test_reset_fake_sends_calls_to_the_provider_again(init_embeddings): + Embeddings.fake() + Embeddings.reset_fake() + + response = await Embeddings.of(["a", "b"]).generate() + + assert response.embeddings == [[0.1, 0.2], [0.3, 0.4]] + init_embeddings.assert_called_once() + with pytest.raises(RuntimeError, match="not faked"): + Embeddings.assert_nothing_generated() + + +async def test_fake_records_resolved_provider_model_and_dimensions(): + fake = Embeddings.fake() + + await Embeddings.of(["a"]).dimensions(32).generate("bedrock") + await Embeddings.of(["b"]).generate(model="text-embedding-3-large") + + assert fake.prompts == [ + EmbeddingsPrompt(["a"], 32, "bedrock", "amazon.titan-embed-text-v2:0"), + EmbeddingsPrompt(["b"], None, "openai", "text-embedding-3-large"), + ] + + +async def test_default_provider_follows_config(ai_config): + ai_config.default_embeddings = "google" + app().make("config").set("ai", ai_config) + Embeddings.fake() + + response = await Embeddings.of(["a"]).generate() + + assert (response.provider, response.model) == ("google", "gemini-embedding-001") + + +async def test_builder_copies_inputs(): + inputs = ["a"] + builder = Embeddings.of(inputs) + inputs.append("b") + fake = Embeddings.fake() + + await builder.generate() + + assert fake.prompts[0].inputs == ["a"] + + +def test_dimensions_returns_the_builder(): + builder = Embeddings.of(["a"]) + + assert builder.dimensions(8) is builder + + +async def test_unknown_provider_raises(): + with pytest.raises(ValueError): + await Embeddings.of(["a"]).generate("nope") + + +async def test_assert_not_generated_passes_with_no_prompts(): + Embeddings.fake() + + Embeddings.assert_not_generated("anything") + with pytest.raises(AssertionError, match="Expected embeddings matching 'anything'"): + Embeddings.assert_generated("anything") + + +async def test_ai_reset_fakes_clears_the_embeddings_fake(init_embeddings): + Embeddings.fake() + Ai.reset_fakes() + + response = await Embeddings.of(["hello"]).generate() + + assert response.embeddings == [[0.1, 0.2], [0.3, 0.4]] diff --git a/fastapi_startkit/tests/ai/test_lab.py b/fastapi_startkit/tests/ai/test_lab.py index 644c0bac..88be0759 100644 --- a/fastapi_startkit/tests/ai/test_lab.py +++ b/fastapi_startkit/tests/ai/test_lab.py @@ -31,6 +31,17 @@ def test_get_model_image_default(self): self.assertEqual(Lab.OPENAI.get_model(model_type=ModelType.IMAGE), "dall-e-3") self.assertEqual(Lab.GOOGLE.get_model(model_type=ModelType.IMAGE), "imagen-3.0-generate-002") + def test_embeddings_model_type_reads_default_embeddings(self): + self.assertEqual(ModelType.EMBEDDINGS.key, "default_embeddings") + + def test_get_model_embeddings_default(self): + self.assertEqual(Lab.OPENAI.get_model(model_type=ModelType.EMBEDDINGS), "text-embedding-3-small") + self.assertEqual(Lab.GOOGLE.get_model(model_type=ModelType.EMBEDDINGS), "gemini-embedding-001") + self.assertEqual(Lab.BEDROCK.get_model(model_type=ModelType.EMBEDDINGS), "amazon.titan-embed-text-v2:0") + + def test_get_provider_embeddings_uses_default_embeddings(self): + self.assertIs(Lab.get_provider(model_type=ModelType.EMBEDDINGS), Lab.OPENAI) + def test_get_provider_uses_config_default(self): self.assertEqual(Lab.get_provider().value, "google") diff --git a/fastapi_startkit/tests/masoniteorm/models/test_caster.py b/fastapi_startkit/tests/masoniteorm/models/test_caster.py index 28b36397..50bc2f60 100644 --- a/fastapi_startkit/tests/masoniteorm/models/test_caster.py +++ b/fastapi_startkit/tests/masoniteorm/models/test_caster.py @@ -176,7 +176,7 @@ def test_bare_annotations_keep_their_none_behaviour(): assert "plain_bool" not in caster.nullable assert caster.get("plain_bool", None) is False - assert caster.set("plain_dict", None) == "null" + assert caster.set("plain_dict", None) is None def test_annotated_nullable_type_infers_the_inner_cast(): @@ -224,3 +224,21 @@ class DefaultedProfile(Model): assert isinstance(caster.casts["level"], IntCast) assert caster.get("level", None) == 1 assert caster.get("level", "5") == 5 + + +@pytest.mark.parametrize("annotation", [list, dict, list[float], dict[str, int], list[float] | None]) +def test_list_and_dict_annotations_cast_as_json(annotation): + class CasterJsonDocument(Model): + payload: annotation # pyright: ignore[reportInvalidTypeForm] + + cast = Caster(CasterJsonDocument).casts["payload"] + + assert isinstance(cast, JsonCast) + assert cast.get("[0.1, 0.2]") == [0.1, 0.2] + + +def test_json_cast_stores_none_as_sql_null(): + class CasterVectorDocument(Model): + embedding: list[float] + + assert Caster(CasterVectorDocument).set("embedding", None) is None diff --git a/fastapi_startkit/tests/masoniteorm/postgres/test_postgres_vector.py b/fastapi_startkit/tests/masoniteorm/postgres/test_postgres_vector.py new file mode 100644 index 00000000..c835c039 --- /dev/null +++ b/fastapi_startkit/tests/masoniteorm/postgres/test_postgres_vector.py @@ -0,0 +1,58 @@ +from fastapi_startkit.masoniteorm import Model + +from .fixtures.db import schema +from .test_case import TestCase + + +class VecDocument(Model): + __table__ = "vec_documents" + __timestamps__ = None + __fillable__ = ["title", "embedding"] + id: int + title: str + embedding: list[float] + + +class TestPostgresVector(TestCase): + async def asyncSetUp(self): + await super().asyncSetUp() + await schema.ensure_vector_extension_exists() + await schema.drop_table_if_exists("vec_documents") + async with await schema.create("vec_documents") as table: + table.increments("id") + table.string("title") + table.vector("embedding", dimensions=3).index() + await VecDocument.create({"title": "wine", "embedding": [1.0, 0.0, 0.0]}) + await VecDocument.create({"title": "vineyard", "embedding": [0.9, 0.1, 0.0]}) + await VecDocument.create({"title": "framework", "embedding": [0.0, 0.0, 1.0]}) + + async def asyncTearDown(self): + await schema.drop_table_if_exists("vec_documents") + await super().asyncTearDown() + + async def test_vectors_round_trip_through_the_model(self): + document = await VecDocument.where("title", "wine").first() + + self.assertEqual(document.embedding, [1.0, 0.0, 0.0]) + + async def test_where_vector_similar_to_returns_closest_first(self): + documents = await VecDocument.where_vector_similar_to("embedding", [1.0, 0.05, 0.0], min_similarity=0.5).get() + + self.assertEqual([document.title for document in documents], ["wine", "vineyard"]) + + async def test_select_vector_distance(self): + documents = ( + await VecDocument.select("title") + .select_vector_distance("embedding", [0.0, 0.0, 1.0]) + .order_by_vector_distance("embedding", [0.0, 0.0, 1.0]) + .get() + ) + + self.assertEqual(documents[0].title, "framework") + self.assertAlmostEqual(documents[0].embedding_distance, 0.0) + self.assertAlmostEqual(documents[2].embedding_distance, 1.0) + + async def test_where_vector_distance_less_than(self): + documents = await VecDocument.where_vector_distance_less_than("embedding", [0.0, 0.0, 1.0], 0.1).get() + + self.assertEqual([document.title for document in documents], ["framework"]) diff --git a/fastapi_startkit/tests/masoniteorm/query/grammars/test_vector_grammar.py b/fastapi_startkit/tests/masoniteorm/query/grammars/test_vector_grammar.py new file mode 100644 index 00000000..85dbc3b3 --- /dev/null +++ b/fastapi_startkit/tests/masoniteorm/query/grammars/test_vector_grammar.py @@ -0,0 +1,139 @@ +import pytest + +from fastapi_startkit.masoniteorm.models.builder import QueryBuilder +from fastapi_startkit.masoniteorm.query.grammars.MSSQLGrammar import MSSQLGrammar +from fastapi_startkit.masoniteorm.query.grammars.MySQLGrammar import MySQLGrammar +from fastapi_startkit.masoniteorm.query.grammars.PostgresGrammar import PostgresGrammar +from fastapi_startkit.masoniteorm.query.grammars.SQLiteGrammar import SQLiteGrammar + +VECTOR = [0.1, 0.2, 0.3] +LITERAL = "[0.1, 0.2, 0.3]" + + +def qb(grammar=PostgresGrammar): + q = QueryBuilder(connection=None, grammar=grammar, processor=None) + q._table = "documents" + return q + + +def test_select_vector_distance_adds_the_distance_after_all_columns(): + q = qb().select_vector_distance("embedding", VECTOR) + + assert ( + q.to_qmark() == 'SELECT "documents".*, ("documents"."embedding" <=> ?) AS "embedding_distance" FROM "documents"' + ) + assert list(q.get_bindings()) == [LITERAL] + + +def test_select_vector_distance_accepts_an_alias_and_a_qualified_column(): + q = qb().select("id").select_vector_distance("documents.embedding", VECTOR, as_="distance") + + assert q.to_qmark() == 'SELECT "documents"."id", ("documents"."embedding" <=> ?) AS "distance" FROM "documents"' + + +def test_where_vector_distance_less_than_binds_vector_and_distance(): + q = qb().where("team_id", 7).where_vector_distance_less_than("embedding", VECTOR, 0.3) + + assert ( + q.to_qmark() + == 'SELECT * FROM "documents" WHERE "documents"."team_id" = ? AND ("documents"."embedding" <=> ?) <= ?' + ) + assert list(q.get_bindings()) == [7, LITERAL, 0.3] + + +def test_or_where_vector_distance_less_than(): + q = qb().where("id", 1).or_where_vector_distance_less_than("embedding", VECTOR, 0.2) + + assert ( + q.to_qmark() == 'SELECT * FROM "documents" WHERE "documents"."id" = ? OR ("documents"."embedding" <=> ?) <= ?' + ) + assert list(q.get_bindings()) == [1, LITERAL, 0.2] + + +def test_order_by_vector_distance_follows_other_orders(): + q = qb().order_by("title").order_by_vector_distance("embedding", VECTOR) + + assert q.to_qmark() == 'SELECT * FROM "documents" ORDER BY "title" ASC, ("documents"."embedding" <=> ?) ASC' + assert list(q.get_bindings()) == [LITERAL] + + +def test_raw_order_by_after_another_order_is_comma_separated(): + q = qb().order_by("title").order_by_raw("id DESC") + + assert q.to_qmark() == 'SELECT * FROM "documents" ORDER BY "title" ASC, id DESC' + + +def test_where_vector_similar_to_filters_by_distance_and_orders(): + q = qb().where_vector_similar_to("embedding", VECTOR, min_similarity=0.4).limit(10) + + assert q.to_qmark() == ( + 'SELECT * FROM "documents" WHERE ("documents"."embedding" <=> ?) <= ? ORDER BY ("documents"."embedding" <=> ?) ASC LIMIT 10' + ) + bindings = list(q.get_bindings()) + assert bindings[0] == bindings[2] == LITERAL + assert bindings[1] == pytest.approx(0.6) + + +def test_where_vector_similar_to_can_skip_ordering(): + q = qb().where_vector_similar_to("embedding", VECTOR, order=False) + + assert q.to_qmark() == 'SELECT * FROM "documents" WHERE ("documents"."embedding" <=> ?) <= ?' + assert list(q.get_bindings())[1] == pytest.approx(0.4) + + +def test_bindings_follow_sql_order(): + q = ( + qb() + .select("id") + .select_vector_distance("embedding", [1, 0]) + .where_vector_distance_less_than("embedding", [0, 1], 0.5) + .order_by_vector_distance("embedding", [1, 1]) + ) + q.to_qmark() + + assert list(q.get_bindings()) == ["[1.0, 0.0]", "[0.0, 1.0]", 0.5, "[1.0, 1.0]"] + + +def test_vectors_accept_any_numeric_sequence(): + q = qb().order_by_vector_distance("embedding", (1, 2.5)) + q.to_qmark() + + assert list(q.get_bindings()) == ["[1.0, 2.5]"] + + +def test_text_vectors_are_rejected(): + with pytest.raises(TypeError, match="to_embeddings"): + qb().where_vector_similar_to("embedding", "best wineries") + + +def test_vector_columns_use_the_query_alias(): + q = qb().table("documents", "d").order_by_vector_distance("embedding", VECTOR) + + assert q.to_qmark() == 'SELECT * FROM "documents" AS "d" ORDER BY ("d"."embedding" <=> ?) ASC' + + +def test_vector_columns_keep_a_schema_qualified_table(): + q = qb().table("search.documents").order_by_vector_distance("embedding", VECTOR) + + assert '("search"."documents"."embedding" <=> ?) ASC' in q.to_qmark() + + +def test_vector_identifiers_escape_embedded_quotes(): + injected = 'emb" <=> ?) AS x FROM pg_shadow --' + q = ( + qb() + .select_vector_distance(injected, VECTOR, as_=injected) + .where_vector_distance_less_than(injected, VECTOR, 0.5) + ) + + assert q.to_qmark() == ( + 'SELECT "documents".*, ("documents"."emb"" <=> ?) AS x FROM pg_shadow --" <=> ?) ' + 'AS "emb"" <=> ?) AS x FROM pg_shadow --" FROM "documents" ' + 'WHERE ("documents"."emb"" <=> ?) AS x FROM pg_shadow --" <=> ?) <= ?' + ) + + +@pytest.mark.parametrize("grammar", [SQLiteGrammar, MySQLGrammar, MSSQLGrammar]) +def test_vector_queries_require_postgres(grammar): + with pytest.raises(ValueError, match="does not support vector queries"): + qb(grammar).where_vector_similar_to("embedding", VECTOR) diff --git a/fastapi_startkit/tests/masoniteorm/schema/test_vector_schema.py b/fastapi_startkit/tests/masoniteorm/schema/test_vector_schema.py new file mode 100644 index 00000000..0ccf5794 --- /dev/null +++ b/fastapi_startkit/tests/masoniteorm/schema/test_vector_schema.py @@ -0,0 +1,123 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from fastapi_startkit.masoniteorm.schema.Blueprint import Blueprint +from fastapi_startkit.masoniteorm.schema.platforms import MSSQLPlatform, MySQLPlatform, PostgresPlatform, SQLitePlatform +from fastapi_startkit.masoniteorm.schema.schema import Schema +from fastapi_startkit.masoniteorm.schema.Table import Table +from fastapi_startkit.masoniteorm.schema.TableDiff import TableDiff + + +def blueprint(table, platform=PostgresPlatform): + return Blueprint(grammar=None, table=table, platform=platform, dry=True) + + +def test_vector_column_renders_with_dimensions(): + table = Table("documents") + blueprint(table).vector("embedding", dimensions=1536) + + assert PostgresPlatform().compile_create_sql(table) == [ + 'CREATE TABLE "documents" ("embedding" VECTOR(1536) NOT NULL)' + ] + + +def test_nullable_vector_column(): + table = Table("documents") + blueprint(table).vector("embedding", dimensions=3, nullable=True) + + assert PostgresPlatform().compile_create_sql(table) == ['CREATE TABLE "documents" ("embedding" VECTOR(3) NULL)'] + + +def test_index_on_a_vector_column_creates_an_hnsw_cosine_index(): + table = Table("documents") + blueprint(table).vector("embedding", dimensions=1536).index() + + assert PostgresPlatform().compile_create_sql(table)[1] == ( + 'CREATE INDEX documents_embedding_index ON "documents" USING hnsw (embedding vector_cosine_ops)' + ) + + +def test_index_on_other_columns_stays_a_btree_index(): + table = Table("documents") + blueprint(table).string("title").index() + + assert PostgresPlatform().compile_create_sql(table)[1] == ( + 'CREATE INDEX documents_title_index ON "documents"(title)' + ) + + +def test_adding_an_indexed_vector_column_to_an_existing_table(): + table = TableDiff("documents") + blueprint(table).vector("embedding", dimensions=8).index("embedding") + blueprint(table).vector("summary_embedding", dimensions=8).index() + + assert PostgresPlatform().compile_alter_sql(table) == [ + 'ALTER TABLE "documents" ADD COLUMN "embedding" VECTOR(8) NOT NULL, ' + 'ADD COLUMN "summary_embedding" VECTOR(8) NOT NULL', + 'CREATE INDEX documents_embedding_index ON "documents" USING hnsw (embedding vector_cosine_ops)', + 'CREATE INDEX documents_summary_embedding_index ON "documents" USING hnsw (summary_embedding vector_cosine_ops)', + ] + + +def test_named_index_on_a_vector_column_creates_an_hnsw_index(): + table = Table("documents") + blueprint(table).vector("embedding", dimensions=1536).string("title") + blueprint(table).index("embedding", name="documents_embedding_hnsw") + + assert PostgresPlatform().compile_create_sql(table)[1] == ( + 'CREATE INDEX documents_embedding_hnsw ON "documents" USING hnsw (embedding vector_cosine_ops)' + ) + + +def test_named_index_on_a_regular_column_stays_a_btree_index(): + table = Table("documents") + blueprint(table).vector("embedding", dimensions=3).string("title") + blueprint(table).index("title") + + assert PostgresPlatform().compile_create_sql(table)[1] == ( + 'CREATE INDEX documents_title_index ON "documents"(title)' + ) + + +def test_vector_index_on_an_existing_column(): + table = TableDiff("documents") + blueprint(table).vector_index("embedding") + + assert PostgresPlatform().compile_alter_sql(table) == [ + 'CREATE INDEX documents_embedding_index ON "documents" USING hnsw (embedding vector_cosine_ops)' + ] + + +@pytest.mark.parametrize("platform", [SQLitePlatform, MySQLPlatform, MSSQLPlatform]) +def test_vector_indexes_require_postgres(platform): + with pytest.raises(ValueError, match="does not support vector indexes"): + blueprint(Table("documents"), platform).vector_index("embedding") + + +@pytest.mark.parametrize("platform", [SQLitePlatform, MySQLPlatform, MSSQLPlatform]) +def test_vector_columns_require_postgres(platform): + with pytest.raises(ValueError, match="does not support vector columns"): + blueprint(Table("documents"), platform).vector("embedding", dimensions=3) + + +def test_postgres_compiles_the_vector_extension(): + assert PostgresPlatform().compile_ensure_vector_extension() == "CREATE EXTENSION IF NOT EXISTS vector" + + +@pytest.mark.parametrize("platform", [SQLitePlatform, MySQLPlatform, MSSQLPlatform]) +def test_other_platforms_reject_the_vector_extension(platform): + with pytest.raises(ValueError, match="does not support vector columns"): + platform().compile_ensure_vector_extension() + + +async def test_schema_ensures_the_vector_extension(): + connection = MagicMock() + connection.get_default_platform.return_value = PostgresPlatform + connection.statement = AsyncMock() + manager = MagicMock() + manager.connection.return_value = connection + + await Schema(manager).ensure_vector_extension_exists() + + connection.statement.assert_awaited_once_with("CREATE EXTENSION IF NOT EXISTS vector", ()) diff --git a/fastapi_startkit/tests/support/test_string.py b/fastapi_startkit/tests/support/test_string.py index 96f9af31..b17c89eb 100644 --- a/fastapi_startkit/tests/support/test_string.py +++ b/fastapi_startkit/tests/support/test_string.py @@ -1,5 +1,9 @@ """Unit tests for the Str / Stringable string helpers (task #1214).""" +import pytest + +from fastapi_startkit.ai import AIConfig, Embeddings, EmbeddingsPrompt +from fastapi_startkit.application import app from fastapi_startkit.support.string import Str, Stringable @@ -61,3 +65,31 @@ def test_slugify_returns_stringable(self): def test_camel_case_returns_stringable(self): assert str(Str.of("make_model").camel_case()) == "makeModel" + + +class TestStringableToEmbeddings: + @pytest.fixture(autouse=True) + def fake_embeddings(self): + config = AIConfig() + config.default_embeddings = "openai" + app().make("config").set("ai", config) + fake = Embeddings.fake([[[0.1, 0.2, 0.3]]]) + yield fake + Embeddings.reset_fake() + + async def test_returns_the_vector_for_the_text(self, fake_embeddings): + vector = await Str.of("Napa Valley has great wine.").to_embeddings() + + assert vector == [0.1, 0.2, 0.3] + assert fake_embeddings.prompts == [ + EmbeddingsPrompt(["Napa Valley has great wine."], None, "openai", "text-embedding-3-small") + ] + + async def test_passes_provider_model_and_dimensions(self, fake_embeddings): + await Str.of("Laravel").to_embeddings(provider="google", model="gemini-embedding-001", dimensions=3) + + Embeddings.assert_generated( + lambda prompt: ( + prompt.provider == "google" and prompt.model == "gemini-embedding-001" and prompt.dimensions == 3 + ) + )