tests/test_cognee_bootstrap.py

import os
import sys
import types

import pytest
from pydantic import BaseModel


def test_cognee_bootstrap_forces_cerebras_gemma_strict_json_schema(monkeypatch):
    from app.services.cognee_bootstrap import configure_cognee_llm

    calls = []

    class FakeCogneeConfig:
        @staticmethod
        def set_llm_config(config):
            calls.append(config)

    def clear_cache():
        calls.append("clear_cache")

    monkeypatch.setenv("CEREBRAS_API_KEY", "test-key")

    config = configure_cognee_llm(FakeCogneeConfig, clear_llm_client_cache=clear_cache)

    assert config["llm_provider"] == "openai"
    assert config["llm_model"] == "openai/gemma-4-31b"
    assert config["llm_endpoint"] == "https://api.cerebras.ai/v1"
    assert config["llm_api_key"] == "test-key"
    assert config["llm_instructor_mode"] == "json_schema_mode"
    assert config["llm_temperature"] == 0.0
    assert config["llm_args"]["temperature"] == 0
    assert "extra_body" not in config["llm_args"]
    assert calls == [config, "clear_cache"]
    assert os.environ["LLM_INSTRUCTOR_MODE"] == "json_schema_mode"
    assert os.environ["COGNEE_SKIP_CONNECTION_TEST"] == "true"
    assert os.environ["LLM_ARGS"]


@pytest.mark.asyncio
async def test_cognee_openai_adapter_uses_native_cerebras_structured_output(monkeypatch):
    from app.services import cognee_bootstrap

    class SummaryModel(BaseModel):
        summary: str
        description: str = ""

    class FakeOpenAIAdapter:
        def __init__(self):
            self.model = "openai/gemma-4-31b"
            self.endpoint = "https://api.cerebras.ai/v1"

        async def acreate_str_output(self, text_input, system_prompt, **kwargs):
            return "raw text"

    module_name = "fake_cognee_openai_adapter"
    fake_module = types.ModuleType(module_name)
    fake_module.OpenAIAdapter = FakeOpenAIAdapter
    monkeypatch.setitem(sys.modules, module_name, fake_module)
    monkeypatch.setattr(cognee_bootstrap, "COGNEE_OPENAI_ADAPTER_MODULE", module_name)

    calls = []

    class FakeCerebrasClient:
        def structured_complete(self, messages, output_model, model=None, **kwargs):
            calls.append((messages, output_model, model, kwargs))
            return output_model(summary="Recovered strict summary", description="")

    monkeypatch.setattr(cognee_bootstrap, "CerebrasClient", FakeCerebrasClient)

    assert cognee_bootstrap.patch_cognee_cerebras_structured_output() is True
    adapter = FakeOpenAIAdapter()
    result = await adapter.acreate_structured_output("chunk text", "summarize exactly", SummaryModel)

    assert result == SummaryModel(summary="Recovered strict summary", description="")
    assert calls[0][1] is SummaryModel
    assert calls[0][2] == "gemma-4-31b"
    assert calls[0][0] == [
        {"role": "system", "content": "summarize exactly"},
        {"role": "user", "content": "chunk text"},
    ]