RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat: rework langfuse setup to work more cleanly; fix bug in dream scheduling

This commit is contained in:
Benjamin McCormick 2025-10-29 16:10:11 -04:00
parent 6df41265ed
commit a90156e113
12 changed files with 77 additions and 185 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/crud/representation.py
View File

@ -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,

75
src/deriver/consumer.py
View File

@ -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
)

16
src/deriver/deriver.py
View File

@ -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,

2
src/deriver/enqueue.py
View File

@ -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:

14
src/dialectic/chat.py
View File

@ -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",

13
src/dreamer/dream_scheduler.py
View File

@ -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,

2
src/dreamer/dreamer.py
View File

@ -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:

24
src/utils/clients.py
View File

@ -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

40
src/utils/langfuse_client.py
View File

@ -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}'")

50
src/utils/logging.py
View File

@ -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]]]

4
src/utils/summarizer.py
View File

@ -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,

19
tests/utils/test_clients.py
View File

@ -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"""