Spaces:
Running
Running
File size: 2,905 Bytes
2e818da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | 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"},
]
|