diff --git a/src/crud/representation.py b/src/crud/representation.py index 09eaeae7..13aefce7 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -14,7 +14,7 @@ from src.dependencies import tracked_db from src.dreamer.dream_scheduler import check_and_schedule_dream from src.embedding_client import embedding_client from src.utils.formatting import format_datetime_utc -from src.utils.logging import accumulate_metric, conditional_observe +from src.utils.logging import accumulate_metric from src.utils.representation import ( DeductiveObservation, ExplicitObservation, @@ -41,7 +41,6 @@ class RepresentationManager: self.observer: str = observer self.observed: str = observed - @conditional_observe async def save_representation( self, representation: Representation, diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index 20e0b1b0..6cafc0fc 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -7,13 +7,11 @@ from rich.console import Console from sqlalchemy import select from src import models -from src.config import settings from src.dependencies import tracked_db from src.deriver.deriver import process_representation_tasks_batch from src.dreamer.dreamer import process_dream from src.models import Message from src.utils import summarizer -from src.utils.langfuse_client import get_langfuse_client from src.utils.logging import log_performance_metrics from src.utils.queue_payload import ( DreamPayload, @@ -27,8 +25,6 @@ logging.getLogger("sqlalchemy.engine.Engine").disabled = True console = Console(markup=True) -lf = get_langfuse_client() if settings.LANGFUSE_PUBLIC_KEY else None - async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None: """Process a single item from the queue.""" @@ -80,39 +76,16 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None: message_public_id = message.public_id with sentry_sdk.start_transaction(name="process_summary_task", op="deriver"): - if lf: - with lf.start_as_current_span( - name="summary_processing", - input={ - "workspace_name": validated.workspace_name, - "session_name": validated.session_name, - "message_id": validated.message_id, - }, - metadata={ - "summary_model": settings.SUMMARY.MODEL, - }, - ): - await summarizer.summarize_if_needed( - validated.workspace_name, - validated.session_name, - validated.message_id, - validated.message_seq_in_session, - message_public_id, - ) - log_performance_metrics( - "summary", f"{validated.workspace_name}_{validated.message_id}" - ) - else: - await summarizer.summarize_if_needed( - validated.workspace_name, - validated.session_name, - validated.message_id, - validated.message_seq_in_session, - message_public_id, - ) - log_performance_metrics( - "summary", f"{validated.workspace_name}_{validated.message_id}" - ) + await summarizer.summarize_if_needed( + validated.workspace_name, + validated.session_name, + validated.message_id, + validated.message_seq_in_session, + message_public_id, + ) + log_performance_metrics( + "summary", f"{validated.workspace_name}_{validated.message_id}" + ) elif task_type == "dream": with sentry_sdk.start_transaction(name="process_dream_task", op="deriver"): @@ -162,28 +135,6 @@ async def process_representation_batch( len(messages), ) - if lf: - with lf.start_as_current_span( - name="representation_processing", - input={ - "payloads": [ - { - "message_id": msg.id, - "observer": observer, - "observed": observed, - "session_name": msg.session_name, - } - for msg in messages - ] - }, - metadata={ - "critical_analysis_model": settings.DERIVER.MODEL, - }, - ): - await process_representation_tasks_batch( - messages, observer=observer, observed=observed - ) - else: - await process_representation_tasks_batch( - messages, observer=observer, observed=observed - ) + await process_representation_tasks_batch( + messages, observer=observer, observed=observed + ) diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index 15900e86..5990a29b 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -12,7 +12,6 @@ from src.models import Message from src.utils import summarizer from src.utils.clients import honcho_llm_call from src.utils.formatting import format_new_turn_with_timestamp -from src.utils.langfuse_client import get_langfuse_client from src.utils.logging import ( accumulate_metric, conditional_observe, @@ -33,9 +32,8 @@ from .prompts import ( logger = logging.getLogger(__name__) logging.getLogger("sqlalchemy.engine.Engine").disabled = True -lf = get_langfuse_client() if settings.LANGFUSE_PUBLIC_KEY else None - +@conditional_observe(name="Critical Analysis Call") async def critical_analysis_call( peer_id: str, peer_card: list[str] | None, @@ -78,6 +76,7 @@ async def critical_analysis_call( return response.content +@conditional_observe(name="Peer Card Call") async def peer_card_call( old_peer_card: list[str] | None, new_observations: Representation, @@ -281,9 +280,6 @@ async def process_representation_tasks_batch( log_performance_metrics("deriver", f"{latest_message.id}_{observer}") - if lf: - lf.update_current_trace(output=final_observations.format_as_markdown()) - class CertaintyReasoner: """Certainty reasoner for analyzing and deriving insights.""" @@ -308,7 +304,7 @@ class CertaintyReasoner: self.observer = observer self.estimated_input_tokens: int = estimated_input_tokens - @conditional_observe + @conditional_observe(name="Deriver") @sentry_sdk.trace async def reason( self, @@ -364,11 +360,6 @@ class CertaintyReasoner: latest_message.created_at, ) - if lf: - lf.update_current_generation( - output=reasoning_response.format_as_markdown(), - ) - analysis_duration_ms = (time.perf_counter() - analysis_start) * 1000 accumulate_metric( f"deriver_{latest_message.id}_{self.observer}", @@ -413,7 +404,6 @@ class CertaintyReasoner: return reasoning_response - @conditional_observe @sentry_sdk.trace async def _update_peer_card( self, diff --git a/src/deriver/enqueue.py b/src/deriver/enqueue.py index 9405f6fa..9731b91c 100644 --- a/src/deriver/enqueue.py +++ b/src/deriver/enqueue.py @@ -32,7 +32,7 @@ async def enqueue(payload: list[dict[str, Any]]) -> None: # Generate work unit keys for dreams that might be affected by this message dream_keys: list[str] = get_affected_dream_keys(message) for dream_key in dream_keys: - if dream_scheduler.cancel_dream(dream_key): + if await dream_scheduler.cancel_dream(dream_key): cancelled_dreams.add(dream_key) if cancelled_dreams: diff --git a/src/dialectic/chat.py b/src/dialectic/chat.py index b134b043..d90977e9 100644 --- a/src/dialectic/chat.py +++ b/src/dialectic/chat.py @@ -18,9 +18,9 @@ from src.config import settings from src.dependencies import tracked_db from src.utils import summarizer from src.utils.clients import HonchoLLMCallStreamChunk, honcho_llm_call -from src.utils.langfuse_client import get_langfuse_client from src.utils.logging import ( accumulate_metric, + conditional_observe, log_performance_metrics, ) from src.utils.representation import Representation @@ -34,9 +34,6 @@ logger = logging.getLogger(__name__) # Load environment variables load_dotenv() -# Create langfuse client -lf = get_langfuse_client() if settings.LANGFUSE_PUBLIC_KEY else None - async def dialectic_call( query: str, @@ -151,6 +148,7 @@ async def dialectic_stream( return response +@conditional_observe(name="Dialectic") async def chat( workspace_name: str, session_name: str | None, @@ -189,14 +187,6 @@ async def chat( context_window_size -= estimate_tokens(query) - if lf: - lf.update_current_trace( - metadata={ - "query_generation_model": settings.DIALECTIC.QUERY_GENERATION_MODEL, - "query_generation_provider": settings.DIALECTIC.QUERY_GENERATION_PROVIDER, - "dialectic_model": settings.DIALECTIC.MODEL, - } - ) accumulate_metric( f"dialectic_chat_{dialectic_chat_uuid}", "query", diff --git a/src/dreamer/dream_scheduler.py b/src/dreamer/dream_scheduler.py index 8c572e1c..a56edd3f 100644 --- a/src/dreamer/dream_scheduler.py +++ b/src/dreamer/dream_scheduler.py @@ -1,4 +1,5 @@ import asyncio +import contextlib from datetime import datetime, timezone from logging import getLogger from typing import Any @@ -80,7 +81,7 @@ class DreamScheduler: cls._instance = None cls._initialized = False - def schedule_dream( + async def schedule_dream( self, work_unit_key: str, workspace_name: str, @@ -95,7 +96,7 @@ class DreamScheduler: return # Cancel any existing dream for this collection - self.cancel_dream(work_unit_key) + await self.cancel_dream(work_unit_key) task = asyncio.create_task( self._delayed_dream( @@ -110,12 +111,14 @@ class DreamScheduler: self.pending_dreams[work_unit_key] = task task.add_done_callback(lambda t: self.pending_dreams.pop(work_unit_key, None)) - def cancel_dream(self, work_unit_key: str) -> bool: + async def cancel_dream(self, work_unit_key: str) -> bool: """Cancel a pending dream. Returns True if a dream was cancelled.""" if work_unit_key in self.pending_dreams: task = self.pending_dreams.pop(work_unit_key) task.cancel() - logger.debug(f"Cancelled pending dream for {work_unit_key}") + # Wait for the task to actually finish (including its done callback) + with contextlib.suppress(asyncio.CancelledError): + await task return True return False @@ -331,7 +334,7 @@ async def check_and_schedule_dream( } ) - dream_scheduler.schedule_dream( + await dream_scheduler.schedule_dream( collection_work_unit_key, collection.workspace_name, current_document_count, diff --git a/src/dreamer/dreamer.py b/src/dreamer/dreamer.py index bac32db9..9ee43bb8 100644 --- a/src/dreamer/dreamer.py +++ b/src/dreamer/dreamer.py @@ -11,6 +11,7 @@ from src.dreamer.prompts import consolidation_prompt from src.embedding_client import embedding_client from src.utils.clients import honcho_llm_call from src.utils.formatting import format_datetime_utc +from src.utils.logging import conditional_observe from src.utils.queue_payload import DreamPayload from src.utils.representation import ( ExplicitObservation, @@ -176,6 +177,7 @@ async def _consolidate_cluster( ) +@conditional_observe(name="[Dream] Consolidate Call") async def consolidate_call( representation: Representation, ) -> Representation: diff --git a/src/utils/clients.py b/src/utils/clients.py index 920af15a..9a0628ef 100644 --- a/src/utils/clients.py +++ b/src/utils/clients.py @@ -1,7 +1,6 @@ import json import logging -from collections.abc import AsyncIterator, Callable -from functools import wraps +from collections.abc import AsyncIterator from typing import Any, Generic, Literal, TypeVar, cast, overload from anthropic import AsyncAnthropic @@ -18,7 +17,7 @@ from tenacity import retry, stop_after_attempt, wait_exponential from src.config import settings from src.utils.json_parser import validate_and_repair_json -from src.utils.langfuse_client import get_langfuse_client +from src.utils.logging import conditional_observe from src.utils.representation import PromptRepresentation from src.utils.types import SupportedProviders @@ -27,8 +26,6 @@ logger = logging.getLogger(__name__) T = TypeVar("T") M = TypeVar("M", bound=BaseModel) -lf = get_langfuse_client() if settings.LANGFUSE_PUBLIC_KEY else None - CLIENTS: dict[ SupportedProviders, AsyncAnthropic | AsyncOpenAI | genai.Client | AsyncGroq, @@ -168,6 +165,7 @@ async def honcho_llm_call( ) -> AsyncIterator[HonchoLLMCallStreamChunk]: ... +@conditional_observe(name="LLM Call") async def honcho_llm_call( provider: SupportedProviders, model: str, @@ -191,10 +189,6 @@ async def honcho_llm_call( decorated = honcho_llm_call_inner - # apply langfuse if enabled - if settings.LANGFUSE_PUBLIC_KEY: - decorated = with_langfuse(decorated) - # apply tracking if track_name: decorated = ai_track(track_name)(decorated) @@ -781,15 +775,3 @@ async def handle_streaming_response( is_done=True, finish_reasons=[chunk.choices[0].finish_reason], ) - - -def with_langfuse(func: Callable[..., Any]) -> Callable[..., Any]: - @wraps(func) - async def wrapper(*args: Any, **kwargs: Any) -> Any: - if lf: - with lf.start_as_current_generation(name="LLM Call"): - return await func(*args, **kwargs) - else: - return await func(*args, **kwargs) - - return wrapper diff --git a/src/utils/langfuse_client.py b/src/utils/langfuse_client.py deleted file mode 100644 index c6e68c0b..00000000 --- a/src/utils/langfuse_client.py +++ /dev/null @@ -1,40 +0,0 @@ -""" -Centralized Langfuse client management. - -This module provides a singleton Langfuse client to avoid multiple initialization -errors when modules import get_client() at the module level. -""" - -from typing import Any - -from langfuse import get_client - -_langfuse_client: Any = None - - -def get_langfuse_client() -> Any: - """ - Get the singleton Langfuse client instance. - - This function ensures that get_client() is only called once, regardless of - how many modules import this function. This prevents multiple authentication - error messages when LANGFUSE_PUBLIC_KEY is not configured. - - Returns: - Any: The singleton Langfuse client instance - """ - global _langfuse_client - - if _langfuse_client is None: - _langfuse_client = get_client() - - return _langfuse_client - - -# For backward compatibility, provide the client as a module-level variable -# but only initialize it when first accessed -def __getattr__(name: str): - """Lazy initialization of module-level 'lf' attribute.""" - if name == "lf": - return get_langfuse_client() - raise AttributeError(f"module '{__name__}' has no attribute '{name}'") diff --git a/src/utils/logging.py b/src/utils/logging.py index ca615899..2002b7bd 100644 --- a/src/utils/logging.py +++ b/src/utils/logging.py @@ -6,9 +6,10 @@ and a conditional observe decorator that only applies when Langfuse is configure import datetime from collections.abc import Callable -from typing import Any +from typing import ParamSpec, TypeVar, overload from fastapi import Request +from langfuse import observe # pyright: ignore from rich import box from rich.console import Console, Group, RenderableType from rich.panel import Panel @@ -27,25 +28,56 @@ console = Console(markup=True) COLLECT_METRICS_LOCAL = settings.COLLECT_METRICS_LOCAL +P = ParamSpec("P") +R = TypeVar("R") -def conditional_observe(func: Callable[..., Any]) -> Callable[..., Any]: + +@overload +def conditional_observe( + func: Callable[P, R], +) -> Callable[P, R]: ... + + +@overload +def conditional_observe( + *, + name: str, +) -> Callable[[Callable[P, R]], Callable[P, R]]: ... + + +def conditional_observe( + func: Callable[P, R] | None = None, + *, + name: str | None = None, +) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]: """ Conditionally apply the @observe decorator only when LANGFUSE_PUBLIC_KEY is present. + Can be used in two ways: + 1. As a decorator: @conditional_observe + 2. As a decorator factory: @conditional_observe(name="...") + Args: - func: The function to potentially decorate + func: The function to potentially decorate (when used as @conditional_observe) + name: Optional name for the observation (when used as @conditional_observe(name="...")) Returns: The decorated function if Langfuse is configured, otherwise the original function """ - if settings.LANGFUSE_PUBLIC_KEY: - # Import here to avoid circular imports and only import when needed - from langfuse import observe # pyright: ignore - return observe()(func) + def decorator(f: Callable[P, R]) -> Callable[P, R]: + if settings.LANGFUSE_PUBLIC_KEY: + observe_name = name if name is not None else f.__name__ + return observe(name=observe_name)(f) + else: + return f + + if func is not None: + # Used as @conditional_observe (without parentheses) + return decorator(func) else: - # Return the function unchanged if Langfuse is not configured - return func + # Used as @conditional_observe(name="...") (with parentheses and keyword args) + return decorator # dict[task_name, list[tuple[metric_name, metric_value, metric_unit]]] diff --git a/src/utils/summarizer.py b/src/utils/summarizer.py index be21009d..065e3f99 100644 --- a/src/utils/summarizer.py +++ b/src/utils/summarizer.py @@ -14,7 +14,7 @@ from src.dependencies import tracked_db from src.exceptions import ResourceNotFoundException from src.utils.clients import HonchoLLMCallResponse, honcho_llm_call from src.utils.formatting import utc_now_iso -from src.utils.logging import accumulate_metric +from src.utils.logging import accumulate_metric, conditional_observe from .. import crud, models @@ -79,6 +79,7 @@ class SummaryType(Enum): LONG = "honcho_chat_summary_long" +@conditional_observe(name="Create Short Summary") async def create_short_summary( messages: list[models.Message], input_tokens: int, @@ -129,6 +130,7 @@ Produce as thorough a summary as possible in {output_words} words or less. ) +@conditional_observe(name="Create Long Summary") async def create_long_summary( messages: list[models.Message], previous_summary: str | None = None, diff --git a/tests/utils/test_clients.py b/tests/utils/test_clients.py index c6552926..54909d8d 100644 --- a/tests/utils/test_clients.py +++ b/tests/utils/test_clients.py @@ -32,7 +32,6 @@ from src.utils.clients import ( handle_streaming_response, honcho_llm_call, honcho_llm_call_inner, - with_langfuse, ) @@ -89,24 +88,6 @@ class TestLLMCallResponse: assert chunk.finish_reasons == [] -class TestLangfuseIntegration: - """Tests for Langfuse integration""" - - @pytest.mark.asyncio - async def test_with_langfuse_decorator(self): - """Test Langfuse decorator functionality""" - - @with_langfuse - async def test_func(): - return "decorated" - - # Mock the langfuse client - with patch("src.utils.clients.lf") as mock_lf: - result = await test_func() - assert result == "decorated" - mock_lf.start_as_current_generation.assert_called_once_with(name="LLM Call") - - @pytest.mark.asyncio class TestAnthropicClient: """Tests for Anthropic client functionality"""