diff --git a/cognee/api/v1/cognify/code_graph_pipeline.py b/cognee/api/v1/cognify/code_graph_pipeline.py index fd0b723ef9..3777b405cc 100644 --- a/cognee/api/v1/cognify/code_graph_pipeline.py +++ b/cognee/api/v1/cognify/code_graph_pipeline.py @@ -1,18 +1,17 @@ import os import pathlib import asyncio -from cognee.shared.logging_utils import get_logger from uuid import NAMESPACE_OID, uuid5 +from cognee.shared.logging_utils import get_logger +from cognee.modules.observability.get_observe import get_observe from cognee.api.v1.search import SearchType, search from cognee.api.v1.visualize.visualize import visualize_graph -from cognee.base_config import get_base_config from cognee.modules.cognify.config import get_cognify_config from cognee.modules.pipelines import run_tasks from cognee.modules.pipelines.tasks.task import Task from cognee.modules.users.methods import get_default_user -from cognee.shared.data_models import KnowledgeGraph, MonitoringTool -from cognee.shared.utils import render_graph +from cognee.shared.data_models import KnowledgeGraph from cognee.tasks.documents import classify_documents, extract_chunks_from_documents from cognee.tasks.graph import extract_graph_from_data from cognee.tasks.ingestion import ingest_data @@ -22,11 +21,7 @@ from cognee.tasks.summarization import summarize_text from cognee.infrastructure.llm import get_max_chunk_tokens -monitoring = get_base_config().monitoring_tool - -if monitoring == MonitoringTool.LANGFUSE: - from langfuse.decorators import observe - +observe = get_observe() logger = get_logger("code_graph_pipeline") diff --git a/cognee/base_config.py b/cognee/base_config.py index c5693d18f8..8378342736 100644 --- a/cognee/base_config.py +++ b/cognee/base_config.py @@ -1,14 +1,14 @@ import os from typing import Optional from functools import lru_cache -from pydantic_settings import BaseSettings, SettingsConfigDict from cognee.root_dir import get_absolute_path -from cognee.shared.data_models import MonitoringTool +from cognee.modules.observability.observers import Observer +from pydantic_settings import BaseSettings, SettingsConfigDict class BaseConfig(BaseSettings): data_root_directory: str = get_absolute_path(".data_storage") - monitoring_tool: object = MonitoringTool.LANGFUSE + monitoring_tool: object = Observer.LANGFUSE graphistry_username: Optional[str] = os.getenv("GRAPHISTRY_USERNAME") graphistry_password: Optional[str] = os.getenv("GRAPHISTRY_PASSWORD") langfuse_public_key: Optional[str] = os.getenv("LANGFUSE_PUBLIC_KEY") diff --git a/cognee/infrastructure/llm/gemini/adapter.py b/cognee/infrastructure/llm/gemini/adapter.py index c3aa30a441..e62755f6a9 100644 --- a/cognee/infrastructure/llm/gemini/adapter.py +++ b/cognee/infrastructure/llm/gemini/adapter.py @@ -1,9 +1,10 @@ -from typing import Type, Optional -from pydantic import BaseModel -from cognee.shared.logging_utils import get_logger import litellm +from pydantic import BaseModel +from typing import Type, Optional from litellm import acompletion, JSONSchemaValidationError -from cognee.shared.data_models import MonitoringTool + +from cognee.shared.logging_utils import get_logger +from cognee.modules.observability.get_observe import get_observe from cognee.exceptions import InvalidValueError from cognee.infrastructure.llm.llm_interface import LLMInterface from cognee.infrastructure.llm.prompts import read_query_prompt @@ -11,14 +12,9 @@ rate_limit_async, sleep_and_retry_async, ) -from cognee.base_config import get_base_config logger = get_logger() - -monitoring = get_base_config().monitoring_tool - -if monitoring == MonitoringTool.LANGFUSE: - from langfuse.decorators import observe +observe = get_observe() class GeminiAdapter(LLMInterface): diff --git a/cognee/infrastructure/llm/openai/adapter.py b/cognee/infrastructure/llm/openai/adapter.py index 417af85dff..4a135c1ccf 100644 --- a/cognee/infrastructure/llm/openai/adapter.py +++ b/cognee/infrastructure/llm/openai/adapter.py @@ -1,14 +1,11 @@ import os import base64 -from pathlib import Path -from typing import Type - import litellm import instructor +from typing import Type from pydantic import BaseModel from cognee.modules.data.processing.document_types.open_data_file import open_data_file -from cognee.shared.data_models import MonitoringTool from cognee.exceptions import InvalidValueError from cognee.infrastructure.llm.llm_interface import LLMInterface from cognee.infrastructure.llm.prompts import read_query_prompt @@ -18,12 +15,9 @@ sleep_and_retry_async, sleep_and_retry_sync, ) -from cognee.base_config import get_base_config - -monitoring = get_base_config().monitoring_tool +from cognee.modules.observability.get_observe import get_observe -if monitoring == MonitoringTool.LANGFUSE: - from langfuse.decorators import observe +observe = get_observe() class OpenAIAdapter(LLMInterface): diff --git a/cognee/modules/observability/get_observe.py b/cognee/modules/observability/get_observe.py new file mode 100644 index 0000000000..db3655482d --- /dev/null +++ b/cognee/modules/observability/get_observe.py @@ -0,0 +1,11 @@ +from cognee.base_config import get_base_config +from .observers import Observer + + +def get_observe(): + monitoring = get_base_config().monitoring_tool + + if monitoring == Observer.LANGFUSE: + from langfuse.decorators import observe + + return observe diff --git a/cognee/modules/observability/observers.py b/cognee/modules/observability/observers.py new file mode 100644 index 0000000000..7bd0380ec6 --- /dev/null +++ b/cognee/modules/observability/observers.py @@ -0,0 +1,9 @@ +from enum import Enum + + +class Observer(str, Enum): + """Monitoring tools""" + + LANGFUSE = "langfuse" + LLMLITE = "llmlite" + LANGSMITH = "langsmith" diff --git a/cognee/shared/data_models.py b/cognee/shared/data_models.py index f6d66e53db..e0e90e3221 100644 --- a/cognee/shared/data_models.py +++ b/cognee/shared/data_models.py @@ -350,11 +350,3 @@ class ChunkSummaries(BaseModel): """Relevant summary and chunk id""" summaries: List[ChunkSummary] - - -class MonitoringTool(str, Enum): - """Monitoring tools""" - - LANGFUSE = "langfuse" - LLMLITE = "llmlite" - LANGSMITH = "langsmith"