flowstate-sdk 0.1.11__py3-none-any.whl → 0.1.12__py3-none-any.whl

flowstate_sdk/langchain/callback_handler.py

@@ -252,6 +252,31 @@ def _is_tool_call_message(message: Any, gen_info: Optional[Dict[str, Any]]) -> b
     return False
 
 
+def _estimate_token_count(text: Optional[str], model: Optional[str]) -> Optional[int]:
+    if not text:
+        return 0
+    try:
+        import tiktoken  # type: ignore
+    except Exception:
+        tiktoken = None
+
+    if tiktoken is not None:
+        try:
+            try:
+                enc = tiktoken.encoding_for_model(model or "")
+            except Exception:
+                enc = tiktoken.get_encoding("cl100k_base")
+            return len(enc.encode(text))
+        except Exception:
+            return None
+
+    if os.getenv("FLOWSTATE_APPROX_TOKENS"):
+        approx = (len(text) + 3) // 4
+        return approx if approx > 0 else 0
+
+    return None
+
+
 class FlowstateCallbackHandler(BaseCallbackHandler):
     def __init__(self, provider: str, model: str) -> None:
         self.provider = provider
@@ -341,6 +366,25 @@ class FlowstateCallbackHandler(BaseCallbackHandler):
         if finish_reason is None and _is_tool_call_message(ai_message_chunk, gen_info):
             return
 
+        llm_output = getattr(response, "llm_output", None)
+        resolved_model = self.model
+        if isinstance(llm_output, dict):
+            for key in ("model_name", "model", "model_id"):
+                val = llm_output.get(key)
+                if isinstance(val, str) and val.strip():
+                    resolved_model = val
+                    break
+        resolved_model = _normalize_model_name(resolved_model)
+
+        # Extract output text robustly (chat message content or plain .text)
+        output_text = ""
+        if response.generations and response.generations[0]:
+            generation = response.generations[0][0]
+            if getattr(generation, "message", None) is not None:
+                output_text = _content_to_text(generation.message.content)
+            else:
+                output_text = generation.text or ""
+
         usage_candidates: List[Dict[str, Any]] = []
         gen_response_metadata = getattr(generation_chunk, "response_metadata", None)
         if isinstance(gen_response_metadata, dict):
@@ -360,20 +404,15 @@ class FlowstateCallbackHandler(BaseCallbackHandler):
                 usage_candidates.append(additional_kwargs)
         if isinstance(gen_info, dict) and gen_info:
             usage_candidates.append(gen_info)
-        llm_output = getattr(response, "llm_output", None)
         if isinstance(llm_output, dict):
             usage_candidates.append(llm_output)
 
         input_tokens, output_tokens = _extract_token_usage(usage_candidates)
+        if input_tokens is None:
+            input_tokens = _estimate_token_count(self._input_str, resolved_model)
+        if output_tokens is None:
+            output_tokens = _estimate_token_count(output_text, resolved_model)
 
-        resolved_model = self.model
-        if isinstance(llm_output, dict):
-            for key in ("model_name", "model", "model_id"):
-                val = llm_output.get(key)
-                if isinstance(val, str) and val.strip():
-                    resolved_model = val
-                    break
-        resolved_model = _normalize_model_name(resolved_model)
         if os.getenv("FLOWSTATE_DEBUG_LLM_USAGE"):
             try:
                 import json
@@ -410,15 +449,6 @@ class FlowstateCallbackHandler(BaseCallbackHandler):
             if cost_value != 0.0:
                 cost_usd = cost_value
 
-        # Extract output text robustly (chat message content or plain .text)
-        output_text = ""
-        if response.generations and response.generations[0]:
-            generation = response.generations[0][0]
-            if getattr(generation, "message", None) is not None:
-                output_text = _content_to_text(generation.message.content)
-            else:
-                output_text = generation.text or ""
-
         provider_metrics = ProviderMetrics(
             run_id=context.current_run.get(),
             provider=self.provider,

{flowstate_sdk-0.1.11.dist-info → flowstate_sdk-0.1.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flowstate_sdk
-Version: 0.1.11
+Version: 0.1.12
 Summary: SDK for Agentic AI workflow management with Flowstate
 Author: Flowstate
 License: Nobody can use this

{flowstate_sdk-0.1.11.dist-info → flowstate_sdk-0.1.12.dist-info}/RECORD

@@ -9,9 +9,9 @@ flowstate_sdk/shared_dataclasses.py,sha256=-hq4NUdBfK02NI1m8cEScPYBtURUcC7ll7q_q
 flowstate_sdk/task.py,sha256=3quyOV30H2YY9j7mNNLNJo3IKBvglfqyCfOvLLcksas,702
 flowstate_sdk/task_context.py,sha256=DtM3o4dsJzMNHwFygRq2474wSwJMPiZOMJDrZDInAS4,5944
 flowstate_sdk/workflow.py,sha256=CqKLNJ_1hY93NKXIvanNSYX-fALYbzZs5m0VnRsaNzI,1988
-flowstate_sdk/langchain/callback_handler.py,sha256=eZev8ezE27AWVWwe7dYNWoiTknci0Ulmfduwr8uUOWs,15998
-flowstate_sdk-0.1.11.dist-info/licenses/LICENSE,sha256=gBfhEg0GV1iXXcGHVELsXUX-y_ewaTaFvE2MWfJoNUI,19
-flowstate_sdk-0.1.11.dist-info/METADATA,sha256=xOGJCB2EgbxLshUUrBT6oAOvWMhLbHr2ngKzJ_A_Fxg,376
-flowstate_sdk-0.1.11.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-flowstate_sdk-0.1.11.dist-info/top_level.txt,sha256=dZ-q9vpo98TaUsZA3d8NNe1XxK2RaPmoIbDXuidZ1uk,14
-flowstate_sdk-0.1.11.dist-info/RECORD,,
+flowstate_sdk/langchain/callback_handler.py,sha256=5f_zOAWquP8XBknPmFxM9Vpat_LrhTW1GGhnKpQzpiQ,16892
+flowstate_sdk-0.1.12.dist-info/licenses/LICENSE,sha256=gBfhEg0GV1iXXcGHVELsXUX-y_ewaTaFvE2MWfJoNUI,19
+flowstate_sdk-0.1.12.dist-info/METADATA,sha256=A1giTx1fkBdZLPM2KhGhMRof_xww7Wvg10P8Ejs2cPU,376
+flowstate_sdk-0.1.12.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+flowstate_sdk-0.1.12.dist-info/top_level.txt,sha256=dZ-q9vpo98TaUsZA3d8NNe1XxK2RaPmoIbDXuidZ1uk,14
+flowstate_sdk-0.1.12.dist-info/RECORD,,