Initial curated RAGcore source from c16667f0c9397e2bfc8dd0255b114fba23c8d05e

tests/integration/models/generation/test_generation_profile_store.py

@@ -0,0 +1,112 @@
+"""Real Postgres coverage for the generation profile registry store."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncEngine
+
+from ragcore.config import Settings
+from ragcore.domain.control.entities import DataPolicy
+from ragcore.domain.models import GenerationProfile, GenerationProvider
+from ragcore.infrastructure.database import create_database_engine, create_session_factory
+from ragcore.infrastructure.database.models.generation import (
+    GenerationProfileRepositoryError,
+    PostgresGenerationProfileStore,
+)
+
+pytestmark = pytest.mark.integration
+
+
+def _profile(**overrides: object) -> GenerationProfile:
+    values: dict[str, object] = {
+        "id": uuid4(),
+        "profile": "answer-hybrid",
+        "version": uuid4().hex[:12],
+        "provider": GenerationProvider.OLLAMA,
+        "model_identifier": "llama3.1:8b-instruct-q4_0",
+        "model_digest": "a" * 64,
+        "context_limit": 8192,
+        "max_output_tokens": 1024,
+        "temperature": 0.4,
+        "data_policy": DataPolicy.LOCAL_ONLY,
+    }
+    values.update(overrides)
+    return GenerationProfile(**values)
+
+
+@pytest.fixture
+async def db_engine() -> AsyncIterator[AsyncEngine]:
+    engine = create_database_engine(Settings(env="test"))
+    try:
+        yield engine
+    finally:
+        await engine.dispose()
+
+
+@pytest.fixture
+async def store(db_engine: AsyncEngine) -> AsyncIterator[PostgresGenerationProfileStore]:
+    async with db_engine.connect() as connection:
+        await connection.execute(text("TRUNCATE models.generation_profiles"))
+        await connection.commit()
+    yield PostgresGenerationProfileStore(create_session_factory(db_engine))
+
+
+async def test_add_and_get_round_trip(store: PostgresGenerationProfileStore) -> None:
+    profile = _profile()
+    await store.add(profile)
+    fetched = await store.get(profile.id)
+    assert fetched == profile
+
+
+async def test_get_returns_none_for_an_unknown_id(store: PostgresGenerationProfileStore) -> None:
+    assert await store.get(uuid4()) is None
+
+
+async def test_adding_the_same_id_twice_is_rejected(store: PostgresGenerationProfileStore) -> None:
+    profile = _profile()
+    await store.add(profile)
+    with pytest.raises(GenerationProfileRepositoryError):
+        await store.add(profile)
+
+
+async def test_adding_the_same_profile_version_twice_is_rejected(
+    store: PostgresGenerationProfileStore,
+) -> None:
+    profile = _profile()
+    await store.add(profile)
+    with pytest.raises(GenerationProfileRepositoryError):
+        await store.add(_profile(id=uuid4(), profile=profile.profile, version=profile.version))
+
+
+async def test_activate_requires_a_registered_profile(
+    store: PostgresGenerationProfileStore,
+) -> None:
+    with pytest.raises(GenerationProfileRepositoryError):
+        await store.activate(uuid4())
+
+
+async def test_activating_a_second_profile_deactivates_the_first(
+    store: PostgresGenerationProfileStore,
+) -> None:
+    first = _profile()
+    second = _profile(id=uuid4(), version=uuid4().hex[:12])
+    await store.add(first)
+    await store.add(second)
+
+    await store.activate(first.id)
+    assert (await store.get_active()).id == first.id  # type: ignore[union-attr]
+
+    await store.activate(second.id)
+    active = await store.get_active()
+    assert active is not None
+    assert active.id == second.id
+
+
+async def test_get_active_returns_none_when_nothing_is_activated(
+    store: PostgresGenerationProfileStore,
+) -> None:
+    assert await store.get_active() is None

tests/integration/models/generation/test_ollama_generation_adapter.py

@@ -0,0 +1,256 @@
+"""OllamaGenerationAdapter: local-host validation, streaming, retries, breaker."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from uuid import uuid4
+
+import httpx
+import pytest
+
+from ragcore.application.limits import ConcurrencyBudget
+from ragcore.domain.control.entities import DataPolicy
+from ragcore.domain.models import GenerationProfile, GenerationProvider
+from ragcore.infrastructure.models.ollama_generation import (
+    CIRCUIT_FAILURE_THRESHOLD,
+    CircuitOpenError,
+    GenerationProviderError,
+    OllamaGenerationAdapter,
+    assemble,
+)
+
+pytestmark = pytest.mark.integration
+
+
+def profile(**overrides: object) -> GenerationProfile:
+    values: dict[str, object] = {
+        "id": uuid4(),
+        "profile": "answer-hybrid",
+        "version": "v1",
+        "provider": GenerationProvider.OLLAMA,
+        "model_identifier": "llama3.1:8b-instruct-q4_0",
+        "model_digest": "a" * 64,
+        "context_limit": 4096,
+        "max_output_tokens": 512,
+        "temperature": 0.2,
+        "data_policy": DataPolicy.LOCAL_ONLY,
+        "timeout_seconds": 5.0,
+        "max_retries": 2,
+    }
+    values.update(overrides)
+    return GenerationProfile(**values)
+
+
+async def no_sleep(_: float) -> None:
+    return None
+
+
+def _ndjson(*lines: dict[str, object]) -> bytes:
+    return b"\n".join(json.dumps(line).encode() for line in lines)
+
+
+def _client(handler) -> httpx.AsyncClient:
+    return httpx.AsyncClient(transport=httpx.MockTransport(handler), follow_redirects=False)
+
+
+def _adapter(handler, **overrides: object) -> OllamaGenerationAdapter:
+    clock = overrides.pop("clock", None)
+    kwargs: dict[str, object] = {"base_url": "http://ollama", "sleep": no_sleep}
+    if clock is not None:
+        kwargs["clock"] = clock
+    return OllamaGenerationAdapter(profile(**overrides), _client(handler), **kwargs)
+
+
+def test_adapter_rejects_nonlocal_provider_hosts() -> None:
+    with pytest.raises(ValueError, match="allowlisted local"):
+        OllamaGenerationAdapter(
+            profile(), _client(lambda request: httpx.Response(200)), base_url="http://169.254.169.254"
+        )
+
+
+async def test_generate_stream_yields_each_chunk_in_order() -> None:
+    body = _ndjson(
+        {"response": "Hel", "done": False},
+        {"response": "lo", "done": False},
+        {"response": "", "done": True},
+    )
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        assert json.loads(request.content)["stream"] is True
+        return httpx.Response(200, content=body)
+
+    adapter = _adapter(handler)
+    chunks = [chunk async for chunk in adapter.generate_stream("hi")]
+    assert [chunk.text for chunk in chunks] == ["Hel", "lo", ""]
+    assert chunks[-1].done is True
+
+
+async def test_generate_assembles_the_same_content_a_stream_would() -> None:
+    body = _ndjson(
+        {"response": "The answer is ", "done": False},
+        {"response": "42.", "done": True},
+    )
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(200, content=body)
+
+    adapter = _adapter(handler)
+    result = await adapter.generate("what is the answer?")
+    streamed = [chunk async for chunk in adapter.generate_stream("what is the answer?")]
+
+    assert result.text == "The answer is 42."
+    assert result == assemble(streamed)
+
+
+async def test_breaking_out_of_the_stream_early_does_not_raise() -> None:
+    body = _ndjson(
+        {"response": "one", "done": False},
+        {"response": "two", "done": False},
+        {"response": "three", "done": True},
+    )
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(200, content=body)
+
+    adapter = _adapter(handler)
+    seen = []
+    async for chunk in adapter.generate_stream("hi"):
+        seen.append(chunk.text)
+        if len(seen) == 1:
+            break
+    assert seen == ["one"]
+
+
+async def test_retryable_status_retries_before_any_streaming_then_succeeds() -> None:
+    attempts = {"count": 0}
+    body = _ndjson({"response": "ok", "done": True})
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        attempts["count"] += 1
+        if attempts["count"] < 3:
+            return httpx.Response(503)
+        return httpx.Response(200, content=body)
+
+    adapter = _adapter(handler)
+    result = await adapter.generate("hi")
+    assert result.text == "ok"
+    assert attempts["count"] == 3
+
+
+async def test_exhausting_retries_raises_generation_unavailable() -> None:
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(503)
+
+    adapter = _adapter(handler, max_retries=1)
+    with pytest.raises(GenerationProviderError, match="generation_unavailable"):
+        await adapter.generate("hi")
+
+
+async def test_non_retryable_status_raises_generation_rejected() -> None:
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(400)
+
+    adapter = _adapter(handler)
+    with pytest.raises(GenerationProviderError, match="generation_rejected"):
+        await adapter.generate("hi")
+
+
+async def test_invalid_response_line_is_a_stable_error() -> None:
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(200, content=b"not json at all")
+
+    adapter = _adapter(handler)
+    with pytest.raises(GenerationProviderError, match="generation_response_invalid"):
+        await adapter.generate("hi")
+
+
+async def test_stream_ending_without_a_done_line_is_a_stable_error() -> None:
+    body = _ndjson({"response": "partial", "done": False})
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(200, content=body)
+
+    adapter = _adapter(handler)
+    with pytest.raises(GenerationProviderError, match="generation_response_invalid"):
+        await adapter.generate("hi")
+
+
+async def test_empty_and_oversized_prompts_are_rejected() -> None:
+    adapter = _adapter(lambda request: httpx.Response(200))
+    with pytest.raises(ValueError, match="empty"):
+        await adapter.generate("")
+    with pytest.raises(ValueError, match="bounded"):
+        await adapter.generate("x" * 200_001)
+
+
+async def test_circuit_opens_after_consecutive_failures_and_blocks_further_calls() -> None:
+    attempts = {"count": 0}
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        attempts["count"] += 1
+        return httpx.Response(503)
+
+    adapter = _adapter(handler, max_retries=0)
+    for _ in range(CIRCUIT_FAILURE_THRESHOLD):
+        with pytest.raises(GenerationProviderError, match="generation_unavailable"):
+            await adapter.generate("hi")
+
+    calls_before_circuit_check = attempts["count"]
+    with pytest.raises(CircuitOpenError):
+        await adapter.generate("hi")
+    assert attempts["count"] == calls_before_circuit_check  # the breaker skipped the network call
+
+
+async def test_circuit_half_opens_after_cooldown_and_closes_on_success() -> None:
+    outcomes = iter(
+        [httpx.Response(503)] * CIRCUIT_FAILURE_THRESHOLD
+        + [httpx.Response(200, content=_ndjson({"response": "ok", "done": True}))]
+    )
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        return next(outcomes)
+
+    clock_time = {"value": 0.0}
+
+    def clock() -> float:
+        return clock_time["value"]
+
+    adapter = _adapter(handler, max_retries=0, clock=clock)
+    for _ in range(CIRCUIT_FAILURE_THRESHOLD):
+        with pytest.raises(GenerationProviderError, match="generation_unavailable"):
+            await adapter.generate("hi")
+
+    # Still within the cooldown window: the breaker rejects without a network call.
+    with pytest.raises(CircuitOpenError):
+        await adapter.generate("hi")
+
+    # Advance past the cooldown: the next call is a real, half-open attempt.
+    clock_time["value"] += 3600.0
+    result = await adapter.generate("hi")
+    assert result.text == "ok"
+
+
+async def test_adapter_surfaces_concurrency_exhaustion_as_a_provider_error() -> None:
+    budget = ConcurrencyBudget(name="generation-calls", max_concurrent=1, max_queued=0)
+    release = asyncio.Event()
+
+    async def occupy() -> None:
+        async with budget.acquire():
+            await release.wait()
+
+    holder = asyncio.create_task(occupy())
+    await asyncio.sleep(0)
+
+    adapter = OllamaGenerationAdapter(
+        profile(),
+        _client(lambda request: httpx.Response(200)),
+        base_url="http://ollama",
+        sleep=no_sleep,
+        concurrency_budget=budget,
+    )
+    with pytest.raises(GenerationProviderError, match="generation_concurrency_exhausted"):
+        await adapter.generate("hi")
+
+    release.set()
+    await holder