Skip to content

Commit 5e1f13e

Browse files
nicoloboschinicoloboschi
authored
feat: add metrics for llm call latency (#120)
* feat: add metrics for llm call latency * feat: add metrics for llm call latency * fix
1 parent 67c1a42 commit 5e1f13e

6 files changed

Lines changed: 724 additions & 122 deletions

File tree

‎hindsight-api/hindsight_api/engine/llm_wrapper.py‎

Lines changed: 57 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -400,13 +400,17 @@ async def call(
400400
output_tokens = usage.completion_tokens or 0 if usage else 0
401401
total_tokens = usage.total_tokens or 0 if usage else 0
402402

403-
if usage:
404-
get_metrics_collector().record_tokens(
405-
operation=scope,
406-
bank_id="llm",
407-
input_tokens=input_tokens,
408-
output_tokens=output_tokens,
409-
)
403+
# Record LLM metrics
404+
metrics = get_metrics_collector()
405+
metrics.record_llm_call(
406+
provider=self.provider,
407+
model=self.model,
408+
scope=scope,
409+
duration=duration,
410+
input_tokens=input_tokens,
411+
output_tokens=output_tokens,
412+
success=True,
413+
)
410414

411415
# Log slow calls
412416
if duration > 10.0 and usage:
@@ -559,24 +563,28 @@ async def _call_anthropic(
559563
else:
560564
result = content
561565

562-
# Record token usage metrics
566+
# Record metrics and log slow calls
563567
duration = time.time() - start_time
564568
input_tokens = response.usage.input_tokens or 0 if response.usage else 0
565569
output_tokens = response.usage.output_tokens or 0 if response.usage else 0
566570
total_tokens = input_tokens + output_tokens
567571

568-
if response.usage:
569-
get_metrics_collector().record_tokens(
570-
operation="memory",
571-
bank_id="llm",
572-
input_tokens=input_tokens,
573-
output_tokens=output_tokens,
574-
)
572+
# Record LLM metrics
573+
metrics = get_metrics_collector()
574+
metrics.record_llm_call(
575+
provider=self.provider,
576+
model=self.model,
577+
scope="memory",
578+
duration=duration,
579+
input_tokens=input_tokens,
580+
output_tokens=output_tokens,
581+
success=True,
582+
)
575583

576584
# Log slow calls
577-
if duration > 10.0 and response.usage:
585+
if duration > 10.0:
578586
logger.info(
579-
f"slow llm call: model={self.provider}/{self.model}, "
587+
f"slow llm call: scope=memory, model={self.provider}/{self.model}, "
580588
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
581589
f"time={duration:.3f}s"
582590
)
@@ -719,18 +727,22 @@ async def _call_ollama_native(
719727

720728
# Extract token usage from Ollama response
721729
# Ollama returns prompt_eval_count (input) and eval_count (output)
730+
duration = time.time() - start_time
722731
input_tokens = result.get("prompt_eval_count", 0) or 0
723732
output_tokens = result.get("eval_count", 0) or 0
724733
total_tokens = input_tokens + output_tokens
725734

726-
# Record to metrics
727-
if input_tokens > 0 or output_tokens > 0:
728-
get_metrics_collector().record_tokens(
729-
operation="memory",
730-
bank_id="llm",
731-
input_tokens=input_tokens,
732-
output_tokens=output_tokens,
733-
)
735+
# Record LLM metrics
736+
metrics = get_metrics_collector()
737+
metrics.record_llm_call(
738+
provider=self.provider,
739+
model=self.model,
740+
scope="memory",
741+
duration=duration,
742+
input_tokens=input_tokens,
743+
output_tokens=output_tokens,
744+
success=True,
745+
)
734746

735747
# Validate against Pydantic model or return raw JSON
736748
if skip_validation:
@@ -865,28 +877,34 @@ async def _call_gemini(
865877
else:
866878
result = content
867879

868-
# Record token usage metrics
880+
# Record metrics and log slow calls
869881
duration = time.time() - start_time
870882
input_tokens = 0
871883
output_tokens = 0
872884
if hasattr(response, "usage_metadata") and response.usage_metadata:
873885
usage = response.usage_metadata
874886
input_tokens = usage.prompt_token_count or 0
875887
output_tokens = usage.candidates_token_count or 0
876-
get_metrics_collector().record_tokens(
877-
operation="memory",
878-
bank_id="llm",
879-
input_tokens=input_tokens,
880-
output_tokens=output_tokens,
881-
)
882888

883-
# Log slow calls
884-
if duration > 10.0:
885-
logger.info(
886-
f"slow llm call: model={self.provider}/{self.model}, "
887-
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
888-
f"time={duration:.3f}s"
889-
)
889+
# Record LLM metrics
890+
metrics = get_metrics_collector()
891+
metrics.record_llm_call(
892+
provider=self.provider,
893+
model=self.model,
894+
scope="memory",
895+
duration=duration,
896+
input_tokens=input_tokens,
897+
output_tokens=output_tokens,
898+
success=True,
899+
)
900+
901+
# Log slow calls
902+
if duration > 10.0 and input_tokens > 0:
903+
logger.info(
904+
f"slow llm call: scope=memory, model={self.provider}/{self.model}, "
905+
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
906+
f"time={duration:.3f}s"
907+
)
890908

891909
if return_usage:
892910
token_usage = TokenUsage(

‎hindsight-api/hindsight_api/engine/memory_engine.py‎

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from typing import TYPE_CHECKING, Any
1919

2020
from ..config import get_config
21+
from ..metrics import get_metrics_collector
2122

2223
# Context variable for current schema (async-safe, per-task isolation)
2324
_current_schema: contextvars.ContextVar[str] = contextvars.ContextVar("current_schema", default="public")
@@ -3162,16 +3163,20 @@ async def reflect_async(
31623163

31633164
# Steps 1-3: Run multi-fact-type search (12-way retrieval: 4 methods × 3 fact types)
31643165
recall_start = time.time()
3165-
search_result = await self.recall_async(
3166-
bank_id=bank_id,
3167-
query=query,
3168-
budget=budget,
3169-
max_tokens=4096,
3170-
enable_trace=False,
3171-
fact_type=["experience", "world", "opinion"],
3172-
include_entities=True,
3173-
request_context=request_context,
3174-
)
3166+
metrics = get_metrics_collector()
3167+
with metrics.record_operation(
3168+
"recall", bank_id=bank_id, source="reflect", budget=budget.value if budget else None
3169+
):
3170+
search_result = await self.recall_async(
3171+
bank_id=bank_id,
3172+
query=query,
3173+
budget=budget,
3174+
max_tokens=4096,
3175+
enable_trace=False,
3176+
fact_type=["experience", "world", "opinion"],
3177+
include_entities=True,
3178+
request_context=request_context,
3179+
)
31753180
recall_time = time.time() - recall_start
31763181

31773182
all_results = search_result.results

0 commit comments

Comments
 (0)