[NOTICKET] add total token logging

src/agents/chatbot.py

@@ -1,5 +1,6 @@
 """Chatbot agent with RAG capabilities."""
 
+import tiktoken
 from langchain_openai import AzureChatOpenAI
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.output_parsers import StrOutputParser
@@ -9,6 +10,14 @@ from langchain_core.messages import HumanMessage, AIMessage
 
 logger = get_logger("chatbot")
 
+_enc = tiktoken.get_encoding("cl100k_base")
+
+
+def _count_tokens(messages: list, context: str) -> dict:
+    msg_tokens = sum(len(_enc.encode(m.content)) for m in messages)
+    ctx_tokens = len(_enc.encode(context))
+    return {"messages_tokens": msg_tokens, "context_tokens": ctx_tokens, "total": msg_tokens + ctx_tokens}
+
 
 class ChatbotAgent:
     """Chatbot agent with RAG capabilities."""
@@ -64,7 +73,8 @@ class ChatbotAgent:
     async def astream_response(self, messages: list, context: str = ""):
         """Stream response tokens as they are generated."""
         try:
-            logger.info("Streaming chatbot response")
+            token_counts = _count_tokens(messages, context)
+            logger.info("LLM input tokens", **token_counts)
             async for token in self.chain.astream({"messages": messages, "context": context}):
                 yield token
         except Exception as e: