gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: Simplify health check methods by using utility functions for LLM and embedding connections

Browse source
This commit is contained in:
Daulet Amirkhanov 2025-10-07 21:26:22 +01:00
parent c1e467209f
commit 405b65f857
1 changed files with 4 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
cognee/api/health.py
View file

@ -189,21 +189,12 @@ class HealthChecker:
start_time = time.time() start_time = time.time()
try: try:
from cognee.infrastructure.llm.config import get_llm_config from cognee.infrastructure.llm.config import get_llm_config
from cognee.infrastructure.llm import LLMGateway
config = get_llm_config() config = get_llm_config()
from pydantic import BaseModel from cognee.infrastructure.llm.utils import test_llm_connection
class TestResponse(BaseModel): await test_llm_connection()
status: str
# Use a simple test to verify LLM connectivity
await LLMGateway.acreate_structured_output(
text_input="Respond with status: ok",
system_prompt="You are a test assistant. Respond with a simple status.",
response_model=TestResponse,
)
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
return ComponentHealth( return ComponentHealth(
@ -226,13 +217,9 @@ class HealthChecker:
"""Check embedding service health (non-critical).""" """Check embedding service health (non-critical)."""
start_time = time.time() start_time = time.time()
try: try:
from cognee.infrastructure.databases.vector.embeddings.get_embedding_engine import ( from cognee.infrastructure.llm.utils import test_embedding_connection
get_embedding_engine,
)
# Test actual embedding generation with minimal text await test_embedding_connection()
engine = get_embedding_engine()
await engine.embed_text(["test"])
response_time = int((time.time() - start_time) * 1000) response_time = int((time.time() - start_time) * 1000)
return ComponentHealth( return ComponentHealth(