figureout 0.2.2__tar.gz → 0.2.4__tar.gz

{figureout-0.2.2/src/figureout.egg-info → figureout-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: figureout
-Version: 0.2.2
+Version: 0.2.4
 Summary: FigureOut - a lightweight and resilient multi-LLM orchestrator
 Author: balajeekalyan
 License: MIT
@@ -8,6 +8,7 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: fastmcp
+Requires-Dist: mcp<2,>=1.25
 Requires-Dist: python-dotenv
 Provides-Extra: openai
 Requires-Dist: openai; extra == "openai"
@@ -222,9 +223,12 @@ With `verbose=True`, the response is wrapped under a `"response"` key and a `"de
 ```python
 concierge = FigureOut(..., verbose=True)
 result = await concierge.run("Find flights to Tokyo")
-# {"response": {"flights": [...]}, "debug": {"input_tokens": 512, "output_tokens": 256, ...}}
+# {"response": {"flights": [...]}, "debug": {"input_tokens": 512, "output_tokens": 256,
+#   "classifier_input_tokens": 180, "classifier_output_tokens": 8, ...}}
 ```
 
+`input_tokens` and `output_tokens` are totals for the whole call, including the role-classification call. `classifier_input_tokens` and `classifier_output_tokens` give the classification share on its own; both are 0 when `role=` is passed explicitly.
+
 ## Examples
 
 The `examples/` directory contains complete domain-specific backends with React+Vite frontends:

{figureout-0.2.2 → figureout-0.2.4}/README.md

@@ -186,9 +186,12 @@ With `verbose=True`, the response is wrapped under a `"response"` key and a `"de
 ```python
 concierge = FigureOut(..., verbose=True)
 result = await concierge.run("Find flights to Tokyo")
-# {"response": {"flights": [...]}, "debug": {"input_tokens": 512, "output_tokens": 256, ...}}
+# {"response": {"flights": [...]}, "debug": {"input_tokens": 512, "output_tokens": 256,
+#   "classifier_input_tokens": 180, "classifier_output_tokens": 8, ...}}
 ```
 
+`input_tokens` and `output_tokens` are totals for the whole call, including the role-classification call. `classifier_input_tokens` and `classifier_output_tokens` give the classification share on its own; both are 0 when `role=` is passed explicitly.
+
 ## Examples
 
 The `examples/` directory contains complete domain-specific backends with React+Vite frontends:

{figureout-0.2.2 → figureout-0.2.4}/pyproject.toml

@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "figureout"
-version = "0.2.2"
+version = "0.2.4"
 description = "FigureOut - a lightweight and resilient multi-LLM orchestrator"
 readme = "README.md"
 requires-python = ">=3.8"
 license = {text = "MIT"}
-dependencies = ["fastmcp", "python-dotenv"]
+dependencies = ["fastmcp", "mcp>=1.25,<2", "python-dotenv"]  # mcp 2.x renamed FastMCP, whose internals FigureOut reads
 authors = [
     {name = "balajeekalyan"},
 ]

{figureout-0.2.2 → figureout-0.2.4}/src/figureout/engine.py

@@ -23,13 +23,20 @@ def _cache_key(query: str, context: str | None = None) -> str:
 
 async def classify_role(chat, user_query: str, classification_prompt: str, valid_roles: set[str]) -> list[str]:
     """Use the LLM to determine the appropriate roles from a user query."""
+    roles, _, _ = await _classify_role_with_usage(chat, user_query, classification_prompt, valid_roles)
+    return roles
+
+
+async def _classify_role_with_usage(chat, user_query: str, classification_prompt: str, valid_roles: set[str]) -> tuple[list[str], int, int]:
+    """Classify the query and also return the classifier call's input and output token counts."""
     messages = [
         {"role": "system", "content": classification_prompt},
         {"role": "user", "content": user_query},
     ]
     result = await chat(messages, json_mode=True)
     parsed = _extract_json(result["text"])
-    return [r for r in parsed["roles"] if r in valid_roles]
+    roles = [r for r in parsed["roles"] if r in valid_roles]
+    return roles, result.get("input_tokens", 0), result.get("output_tokens", 0)
 
 
 class FigureOut:
@@ -241,6 +248,8 @@ class FigureOut:
         output_schema: str | None = None,
     ) -> dict:
         """Send query to LLM with role-based system prompt, return result."""
+        classifier_input_tokens = 0
+        classifier_output_tokens = 0
         if role is not None:
             roles = [role]
             cache_key = None
@@ -258,7 +267,9 @@ class FigureOut:
                 return cached_response
 
             try:
-                roles = await classify_role(self.chat_lite, user_query, self._classification_prompt, self._valid_roles)
+                roles, classifier_input_tokens, classifier_output_tokens = await _classify_role_with_usage(
+                    self.chat_lite, user_query, self._classification_prompt, self._valid_roles
+                )
             except Exception as exc:
                 raise RuntimeError(f"Role classification failed: {exc}") from exc
 
@@ -275,8 +286,10 @@ class FigureOut:
                         "context": context,
                         "roles_matched": roles,
                         "roles_selected": [],
-                        "input_tokens": 0,
-                        "output_tokens": 0,
+                        "input_tokens": classifier_input_tokens,
+                        "output_tokens": classifier_output_tokens,
+                        "classifier_input_tokens": classifier_input_tokens,
+                        "classifier_output_tokens": classifier_output_tokens,
                         "tools_used": [],
                         "assistant_messages": [],
                     },
@@ -289,8 +302,9 @@ class FigureOut:
             user_content = user_query
 
         assistant_messages = []
-        total_input_tokens = 0
-        total_output_tokens = 0
+        # Totals include the classifier call so reported cost covers the whole pipeline
+        total_input_tokens = classifier_input_tokens
+        total_output_tokens = classifier_output_tokens
 
         # LLM calls — one per selected role, in parallel, with native tool calling loop
         async def _run_role(role_to_run: str) -> tuple[str, dict]:
@@ -422,6 +436,8 @@ class FigureOut:
                         "roles_selected": selected_roles,
                         "input_tokens": total_input_tokens,
                         "output_tokens": total_output_tokens,
+                        "classifier_input_tokens": classifier_input_tokens,
+                        "classifier_output_tokens": classifier_output_tokens,
                         "tools_used": all_tools_used,
                         "assistant_messages": assistant_messages,
                     },
@@ -442,6 +458,8 @@ class FigureOut:
                     "roles_selected": selected_roles,
                     "input_tokens": total_input_tokens,
                     "output_tokens": total_output_tokens,
+                    "classifier_input_tokens": classifier_input_tokens,
+                    "classifier_output_tokens": classifier_output_tokens,
                     "tools_used": all_tools_used,
                     "assistant_messages": assistant_messages,
                 },

{figureout-0.2.2 → figureout-0.2.4/src/figureout.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: figureout
-Version: 0.2.2
+Version: 0.2.4
 Summary: FigureOut - a lightweight and resilient multi-LLM orchestrator
 Author: balajeekalyan
 License: MIT
@@ -8,6 +8,7 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: fastmcp
+Requires-Dist: mcp<2,>=1.25
 Requires-Dist: python-dotenv
 Provides-Extra: openai
 Requires-Dist: openai; extra == "openai"
@@ -222,9 +223,12 @@ With `verbose=True`, the response is wrapped under a `"response"` key and a `"de
 ```python
 concierge = FigureOut(..., verbose=True)
 result = await concierge.run("Find flights to Tokyo")
-# {"response": {"flights": [...]}, "debug": {"input_tokens": 512, "output_tokens": 256, ...}}
+# {"response": {"flights": [...]}, "debug": {"input_tokens": 512, "output_tokens": 256,
+#   "classifier_input_tokens": 180, "classifier_output_tokens": 8, ...}}
 ```
 
+`input_tokens` and `output_tokens` are totals for the whole call, including the role-classification call. `classifier_input_tokens` and `classifier_output_tokens` give the classification share on its own; both are 0 when `role=` is passed explicitly.
+
 ## Examples
 
 The `examples/` directory contains complete domain-specific backends with React+Vite frontends:

{figureout-0.2.2 → figureout-0.2.4}/src/figureout.egg-info/requires.txt

@@ -1,4 +1,5 @@
 fastmcp
+mcp<2,>=1.25
 python-dotenv
 
 [all]

{figureout-0.2.2 → figureout-0.2.4}/tests/test_figureout.py

@@ -388,3 +388,66 @@ def test_interpret_tool_response_false_returns_raw_output(mock_get_llm):
     mock_tool.fn.assert_called_once_with(query="test")
     assert isinstance(result, list)
     assert result[0]["id"] == 1
+
+
+@patch("figureout.engine.get_llm_client")
+def test_verbose_tokens_include_classifier(mock_get_llm):
+    mock_classify_chat = AsyncMock(return_value={
+        "text": json.dumps({"roles": ["discovery"]}),
+        "input_tokens": 10,
+        "output_tokens": 5,
+        "tools_used": [],
+    })
+    mock_main_chat = AsyncMock(return_value={
+        "text": json.dumps({"results": [], "summary": "No results"}),
+        "input_tokens": 20,
+        "output_tokens": 10,
+        "tools_used": [],
+        "tool_calls": [],
+        "assistant_message": {"role": "assistant", "content": "{}"},
+    })
+    mock_get_llm.side_effect = [mock_main_chat, mock_classify_chat]
+    pkg = FigureOut(roles=ROLES, verbose=True, cache_enabled=False)
+
+    debug = asyncio.run(pkg.run("find something"))["debug"]
+    assert debug["classifier_input_tokens"] == 10
+    assert debug["classifier_output_tokens"] == 5
+    assert debug["input_tokens"] == 30
+    assert debug["output_tokens"] == 15
+
+
+@patch("figureout.engine.get_llm_client")
+def test_verbose_off_topic_reports_classifier_tokens(mock_get_llm):
+    mock_classify_chat = AsyncMock(return_value={
+        "text": json.dumps({"roles": []}),
+        "input_tokens": 10,
+        "output_tokens": 5,
+        "tools_used": [],
+    })
+    mock_main_chat = AsyncMock()
+    mock_get_llm.side_effect = [mock_main_chat, mock_classify_chat]
+    pkg = FigureOut(roles=ROLES, verbose=True, cache_enabled=False)
+
+    debug = asyncio.run(pkg.run("what's the weather"))["debug"]
+    assert mock_main_chat.call_count == 0
+    assert debug["input_tokens"] == 10
+    assert debug["output_tokens"] == 5
+    assert debug["classifier_input_tokens"] == 10
+
+
+@patch("figureout.engine.get_llm_client")
+def test_verbose_explicit_role_has_no_classifier_tokens(mock_get_llm):
+    mock_main_chat = AsyncMock(return_value={
+        "text": json.dumps({"results": [], "summary": "No results"}),
+        "input_tokens": 20,
+        "output_tokens": 10,
+        "tools_used": [],
+        "tool_calls": [],
+        "assistant_message": {"role": "assistant", "content": "{}"},
+    })
+    mock_get_llm.return_value = mock_main_chat
+    pkg = FigureOut(roles=ROLES, verbose=True)
+
+    debug = asyncio.run(pkg.run("find something", role="discovery"))["debug"]
+    assert debug["classifier_input_tokens"] == 0
+    assert debug["input_tokens"] == 20