From 9ed3708ab39babb075b7035d7a51d79325acf9ca Mon Sep 17 00:00:00 2001 From: Zac Li Date: Thu, 15 Jun 2023 11:20:27 +0800 Subject: [PATCH] fix: add tokens number --- lcserve/backend/langchain_helper.py | 7 ++++--- requirements.txt | 4 ---- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/lcserve/backend/langchain_helper.py b/lcserve/backend/langchain_helper.py index ae997b19..20f05eeb 100644 --- a/lcserve/backend/langchain_helper.py +++ b/lcserve/backend/langchain_helper.py @@ -39,6 +39,7 @@ class TraceInfo: action: str prompts: Optional[List[str]] = None outputs: str = "" + tokens: Optional[int] = None cost: Optional[float] = None @@ -48,6 +49,7 @@ def __init__(self, tracer: Tracer, parent_span: Span): self.tracer = tracer self.parent_span = parent_span self.logger = get_tracing_logger() + self.total_tokens = 0 self.total_cost = 0 def _register_span(self, run_id, span): @@ -118,6 +120,7 @@ def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any) -> Non span=span_context.span_id, action="on_llm_end", outputs=texts, + tokens=round(self.total_tokens, 3) if self.total_tokens else None, cost=round(self.total_cost, 3) if self.total_cost else None, ) self.logger.info(json.dumps(trace_info.__dict__)) @@ -182,8 +185,6 @@ def on_agent_action( if not self.tracer: return - operation = "langchain.agent" - try: span = self._current_span(run_id) span.add_event( @@ -238,7 +239,7 @@ class TracingCallbackHandler(TracingCallbackHandlerMixin): class OpenAITracingCallbackHandler(TracingCallbackHandlerMixin, OpenAICallbackHandler): def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any) -> None: - # Set the computed total cost first with OpenAICallbackHandler and then handle the tracing + # Set the computed total token used and total cost first with OpenAICallbackHandler and then handle the tracing OpenAICallbackHandler.on_llm_end(self, response, run_id=run_id, **kwargs) TracingCallbackHandlerMixin.on_llm_end(self, response, run_id=run_id, **kwargs) diff --git a/requirements.txt b/requirements.txt index f546f709..e6a4bd3a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,7 +7,3 @@ jina-hubble-sdk nest-asyncio textual toml -# Below libs caused issue https://github.com/hwchase17/langchain/issues/5113, -# can unpin once it's resolved -typing-inspect==0.8.0 -typing_extensions==4.5.0