2323 ENV_LLM_MAX_CONCURRENT ,
2424 ENV_LLM_TIMEOUT ,
2525)
26+ from ..metrics import get_metrics_collector
27+ from .response_models import TokenUsage
2628
2729# Seed applied to every Groq request for deterministic behavior.
2830DEFAULT_LLM_SEED = 4242
@@ -174,6 +176,7 @@ async def call(
174176 max_backoff : float = 60.0 ,
175177 skip_validation : bool = False ,
176178 strict_schema : bool = False ,
179+ return_usage : bool = False ,
177180 ) -> Any :
178181 """
179182 Make an LLM API call with retry logic.
@@ -189,9 +192,11 @@ async def call(
189192 max_backoff: Maximum backoff time in seconds.
190193 skip_validation: Return raw JSON without Pydantic validation.
191194 strict_schema: Use strict JSON schema enforcement (OpenAI only). Guarantees all required fields.
195+ return_usage: If True, return tuple (result, TokenUsage) instead of just result.
192196
193197 Returns:
194- Parsed response if response_format is provided, otherwise text content.
198+ If return_usage=False: Parsed response if response_format is provided, otherwise text content.
199+ If return_usage=True: Tuple of (result, TokenUsage) with token counts from the LLM call.
195200
196201 Raises:
197202 OutputTooLongError: If output exceeds token limits.
@@ -203,7 +208,14 @@ async def call(
203208 # Handle Gemini provider separately
204209 if self .provider == "gemini" :
205210 return await self ._call_gemini (
206- messages , response_format , max_retries , initial_backoff , max_backoff , skip_validation , start_time
211+ messages ,
212+ response_format ,
213+ max_retries ,
214+ initial_backoff ,
215+ max_backoff ,
216+ skip_validation ,
217+ start_time ,
218+ return_usage ,
207219 )
208220
209221 # Handle Anthropic provider separately
@@ -217,6 +229,7 @@ async def call(
217229 max_backoff ,
218230 skip_validation ,
219231 start_time ,
232+ return_usage ,
220233 )
221234
222235 # Handle Ollama with native API for structured output (better schema enforcement)
@@ -231,6 +244,7 @@ async def call(
231244 max_backoff ,
232245 skip_validation ,
233246 start_time ,
247+ return_usage ,
234248 )
235249
236250 call_params = {
@@ -379,21 +393,41 @@ async def call(
379393 response = await self ._client .chat .completions .create (** call_params )
380394 result = response .choices [0 ].message .content
381395
382- # Log slow calls
396+ # Record token usage metrics
383397 duration = time .time () - start_time
384398 usage = response .usage
385- if duration > 10.0 :
386- ratio = max (1 , usage .completion_tokens ) / usage .prompt_tokens
399+ input_tokens = usage .prompt_tokens or 0 if usage else 0
400+ output_tokens = usage .completion_tokens or 0 if usage else 0
401+ total_tokens = usage .total_tokens or 0 if usage else 0
402+
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+ )
410+
411+ # Log slow calls
412+ if duration > 10.0 and usage :
413+ ratio = max (1 , output_tokens ) / max (1 , input_tokens )
387414 cached_tokens = 0
388415 if hasattr (usage , "prompt_tokens_details" ) and usage .prompt_tokens_details :
389416 cached_tokens = getattr (usage .prompt_tokens_details , "cached_tokens" , 0 ) or 0
390417 cache_info = f", cached_tokens={ cached_tokens } " if cached_tokens > 0 else ""
391418 logger .info (
392419 f"slow llm call: model={ self .provider } /{ self .model } , "
393- f"input_tokens={ usage . prompt_tokens } , output_tokens={ usage . completion_tokens } , "
394- f"total_tokens={ usage . total_tokens } { cache_info } , time={ duration :.3f} s, ratio out/in={ ratio :.2f} "
420+ f"input_tokens={ input_tokens } , output_tokens={ output_tokens } , "
421+ f"total_tokens={ total_tokens } { cache_info } , time={ duration :.3f} s, ratio out/in={ ratio :.2f} "
395422 )
396423
424+ if return_usage :
425+ token_usage = TokenUsage (
426+ input_tokens = input_tokens ,
427+ output_tokens = output_tokens ,
428+ total_tokens = total_tokens ,
429+ )
430+ return result , token_usage
397431 return result
398432
399433 except LengthFinishReasonError as e :
@@ -452,6 +486,7 @@ async def _call_anthropic(
452486 max_backoff : float ,
453487 skip_validation : bool ,
454488 start_time : float ,
489+ return_usage : bool = False ,
455490 ) -> Any :
456491 """Handle Anthropic-specific API calls."""
457492 from anthropic import APIConnectionError , APIStatusError , RateLimitError
@@ -524,17 +559,35 @@ async def _call_anthropic(
524559 else :
525560 result = content
526561
527- # Log slow calls
562+ # Record token usage metrics
528563 duration = time .time () - start_time
529- if duration > 10.0 :
530- input_tokens = response .usage .input_tokens
531- output_tokens = response .usage .output_tokens
564+ input_tokens = response .usage .input_tokens or 0 if response .usage else 0
565+ output_tokens = response .usage .output_tokens or 0 if response .usage else 0
566+ total_tokens = input_tokens + output_tokens
567+
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+ )
575+
576+ # Log slow calls
577+ if duration > 10.0 and response .usage :
532578 logger .info (
533579 f"slow llm call: model={ self .provider } /{ self .model } , "
534580 f"input_tokens={ input_tokens } , output_tokens={ output_tokens } , "
535581 f"time={ duration :.3f} s"
536582 )
537583
584+ if return_usage :
585+ token_usage = TokenUsage (
586+ input_tokens = input_tokens ,
587+ output_tokens = output_tokens ,
588+ total_tokens = total_tokens ,
589+ )
590+ return result , token_usage
538591 return result
539592
540593 except json .JSONDecodeError as e :
@@ -589,6 +642,7 @@ async def _call_ollama_native(
589642 max_backoff : float ,
590643 skip_validation : bool ,
591644 start_time : float ,
645+ return_usage : bool = False ,
592646 ) -> Any :
593647 """
594648 Call Ollama using native API with JSON schema enforcement.
@@ -663,11 +717,35 @@ async def _call_ollama_native(
663717 else :
664718 raise
665719
720+ # Extract token usage from Ollama response
721+ # Ollama returns prompt_eval_count (input) and eval_count (output)
722+ input_tokens = result .get ("prompt_eval_count" , 0 ) or 0
723+ output_tokens = result .get ("eval_count" , 0 ) or 0
724+ total_tokens = input_tokens + output_tokens
725+
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+ )
734+
666735 # Validate against Pydantic model or return raw JSON
667736 if skip_validation :
668- return json_data
737+ validated_result = json_data
669738 else :
670- return response_format .model_validate (json_data )
739+ validated_result = response_format .model_validate (json_data )
740+
741+ if return_usage :
742+ token_usage = TokenUsage (
743+ input_tokens = input_tokens ,
744+ output_tokens = output_tokens ,
745+ total_tokens = total_tokens ,
746+ )
747+ return validated_result , token_usage
748+ return validated_result
671749
672750 except httpx .HTTPStatusError as e :
673751 last_exception = e
@@ -710,6 +788,7 @@ async def _call_gemini(
710788 max_backoff : float ,
711789 skip_validation : bool ,
712790 start_time : float ,
791+ return_usage : bool = False ,
713792 ) -> Any :
714793 """Handle Gemini-specific API calls."""
715794 # Convert OpenAI-style messages to Gemini format
@@ -786,16 +865,36 @@ async def _call_gemini(
786865 else :
787866 result = content
788867
789- # Log slow calls
868+ # Record token usage metrics
790869 duration = time .time () - start_time
791- if duration > 10.0 and hasattr (response , "usage_metadata" ) and response .usage_metadata :
870+ input_tokens = 0
871+ output_tokens = 0
872+ if hasattr (response , "usage_metadata" ) and response .usage_metadata :
792873 usage = response .usage_metadata
793- logger .info (
794- f"slow llm call: model={ self .provider } /{ self .model } , "
795- f"input_tokens={ usage .prompt_token_count } , output_tokens={ usage .candidates_token_count } , "
796- f"time={ duration :.3f} s"
874+ input_tokens = usage .prompt_token_count or 0
875+ 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 ,
797881 )
798882
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+ )
890+
891+ if return_usage :
892+ token_usage = TokenUsage (
893+ input_tokens = input_tokens ,
894+ output_tokens = output_tokens ,
895+ total_tokens = input_tokens + output_tokens ,
896+ )
897+ return result , token_usage
799898 return result
800899
801900 except json .JSONDecodeError as e :
0 commit comments