cleanup code

Author: jsonbailey

packages/ai-providers/server-ai-langchain/src/ldai_langchain/langgraph_agent_graph_runner.py

@@ -17,13 +17,6 @@
 from ldai_langchain.langgraph_callback_handler import LDMetricsCallbackHandler
 
 
-def _tool_call_id_from_entry(tc: Any) -> Any:
-    """Return tool_call id from a dict or ToolCall-like object."""
-    if isinstance(tc, dict):
-        return tc.get('id')
-    return getattr(tc, 'id', None)
-
-
 def _make_handoff_tool(child_key: str, description: str) -> Any:
     """
     Create a tool that transfers control to ``child_key``.
@@ -58,114 +51,6 @@ def handoff(
     return handoff
 
 
-def _coalesce_tool_messages_for_openai(msgs: List[Any]) -> List[Any]:
-    """
-    Rewind shared LangGraph message state into OpenAI's required shape.
-
-    Multi-agent graphs append multiple AIMessages before tool outputs are written, so
-    ToolMessages can appear *after* later assistant turns. OpenAI requires each
-    assistant ``tool_calls`` block to be followed immediately by matching ToolMessages.
-
-    This walks non-tool messages in order; after each AIMessage with ``tool_calls``,
-    appends the corresponding ToolMessages (looked up by ``tool_call_id``). Any
-    ToolMessage not referenced by some assistant ``tool_calls`` in the list is dropped.
-
-    An AIMessage whose tool_calls have *no* matching ToolMessages in the state is also
-    dropped.  This handles the parallel fan-out case where a sibling branch's tool
-    execution hasn't run yet (or routes to END) by the time a downstream node reads the
-    accumulated state — sending such an AIMessage to OpenAI would cause a 400 error.
-    """
-    from langchain_core.messages import AIMessage, ToolMessage
-
-    pending: Dict[str, Any] = {}
-    for m in msgs:
-        if isinstance(m, ToolMessage):
-            tid = getattr(m, 'tool_call_id', None)
-            if tid:
-                if tid in pending:
-                    log.warning(
-                        'LangGraphAgentGraphRunner: duplicate ToolMessage for tool_call_id=%r; keeping last',
-                        tid,
-                    )
-                pending[tid] = m
-
-    output: List[Any] = []
-    for m in msgs:
-        if isinstance(m, ToolMessage):
-            continue
-        if isinstance(m, AIMessage) and getattr(m, 'tool_calls', None):
-            call_ids = [_tool_call_id_from_entry(tc) for tc in m.tool_calls]
-            resolved = [tid for tid in call_ids if tid and tid in pending]
-            if not resolved:
-                # None of this AIMessage's tool_calls have responses in the current
-                # state — it belongs to a sibling branch whose ToolMessages aren't
-                # available here.  Drop it to avoid an OpenAI 400.
-                log.warning(
-                    'LangGraphAgentGraphRunner: dropping AIMessage with unresolvable '
-                    'tool_calls %s (no matching ToolMessages in state — likely from a '
-                    'sibling branch in a parallel fan-out)',
-                    call_ids,
-                )
-                continue
-            output.append(m)
-            for tid in resolved:
-                output.append(pending.pop(tid))
-        else:
-            output.append(m)
-
-    if pending:
-        log.warning(
-            'LangGraphAgentGraphRunner: dropping %s orphan ToolMessage(s) (no assistant '
-            'tool_calls in history): %s',
-            len(pending),
-            list(pending.keys())[:32],
-        )
-
-    return output
-
-
-def _message_content_len(msg: Any) -> int:
-    c = getattr(msg, 'content', None)
-    if c is None:
-        return 0
-    if isinstance(c, str):
-        return len(c)
-    return len(str(c))
-
-
-def _format_chat_messages_for_log(msgs: List[Any]) -> str:
-    """
-    One line per message index — matches OpenAI error indices (e.g. messages.[5]).
-    Logged at DEBUG before each model ainvoke.
-    """
-    from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
-
-    lines: List[str] = []
-    for idx, m in enumerate(msgs):
-        cls = type(m).__name__
-        if isinstance(m, AIMessage):
-            tc = getattr(m, 'tool_calls', None) or []
-            pairs = []
-            for x in tc:
-                if isinstance(x, dict):
-                    pairs.append((x.get('name'), x.get('id')))
-                else:
-                    pairs.append((getattr(x, 'name', None), getattr(x, 'id', None)))
-            lines.append(
-                f'[{idx}] {cls} tool_calls={pairs!r} content_len={_message_content_len(m)}'
-            )
-        elif isinstance(m, ToolMessage):
-            lines.append(
-                f'[{idx}] {cls} tool_call_id={getattr(m, "tool_call_id", None)!r} '
-                f'name={getattr(m, "name", None)!r} content_len={_message_content_len(m)}'
-            )
-        elif isinstance(m, (HumanMessage, SystemMessage)):
-            lines.append(f'[{idx}] {cls} content_len={_message_content_len(m)}')
-        else:
-            lines.append(f'[{idx}] {cls}')
-    return '\n'.join(lines)
-
-
 class LangGraphAgentGraphRunner(AgentGraphRunner):
     """
     CAUTION:
@@ -268,15 +153,9 @@ def make_node_fn(bound_model: Any, node_instructions: Any, nk: str):
                 async def invoke(state: WorkflowState) -> dict:
                     if not bound_model:
                         return {'messages': []}
-                    msgs = _coalesce_tool_messages_for_openai(list(state['messages']))
+                    msgs = list(state['messages'])
                     if node_instructions:
                         msgs = [SystemMessage(content=node_instructions)] + msgs
-                    log.debug(
-                        'LangGraphAgentGraphRunner node=%s: CHAT_INPUT (%s messages)\n%s',
-                        nk,
-                        len(msgs),
-                        _format_chat_messages_for_log(msgs),
-                    )
                     response = await bound_model.ainvoke(msgs)
                     return {'messages': [response]}
 

packages/ai-providers/server-ai-langchain/src/ldai_langchain/langgraph_callback_handler.py

@@ -7,6 +7,8 @@
 from ldai.agent_graph import AgentGraphDefinition
 from ldai.tracker import TokenUsage
 
+from ldai_langchain.langchain_helper import get_ai_usage_from_response
+
 
 class LDMetricsCallbackHandler(BaseCallbackHandler):
     """
@@ -137,7 +139,11 @@ def on_llm_end(
         if node_key is None:
             return
 
-        usage = self._extract_token_usage(response)
+        try:
+            message = response.generations[0][0].message
+        except (IndexError, AttributeError, TypeError):
+            return
+        usage = get_ai_usage_from_response(message)
         if usage is None:
             return
 
@@ -210,42 +216,3 @@ def flush(self, graph: AgentGraphDefinition, graph_tracker: Any) -> None:
             for tool_key in self._node_tool_calls.get(node_key, []):
                 config_tracker.track_tool_call(tool_key, graph_key=gk)
 
-    # ------------------------------------------------------------------
-    # Internal helpers
-    # ------------------------------------------------------------------
-
-    @staticmethod
-    def _extract_token_usage(response: LLMResult) -> Optional[TokenUsage]:
-        """Extract token usage from an LLMResult, trying multiple locations."""
-        llm_output = response.llm_output or {}
-
-        # Primary: llm_output['token_usage'] or llm_output['tokenUsage']
-        tu = llm_output.get('token_usage') or llm_output.get('tokenUsage')
-        if tu and isinstance(tu, dict):
-            total = tu.get('total_tokens') or tu.get('totalTokens') or 0
-            inp = tu.get('prompt_tokens') or tu.get('promptTokens') or 0
-            out = tu.get('completion_tokens') or tu.get('completionTokens') or 0
-            if total or inp or out:
-                return TokenUsage(total=total, input=inp, output=out)
-
-        # Fallback: first generation's generation_info
-        try:
-            gen_info = response.generations[0][0].generation_info or {}
-        except (IndexError, AttributeError, TypeError):
-            gen_info = {}
-
-        for key in ('usage_metadata', 'token_usage'):
-            meta = gen_info.get(key)
-            if meta and isinstance(meta, dict):
-                total = (
-                    meta.get('total_tokens')
-                    or meta.get('totalTokens')
-                    or (meta.get('input_tokens', 0) + meta.get('output_tokens', 0))
-                    or 0
-                )
-                inp = meta.get('input_tokens') or meta.get('prompt_tokens') or meta.get('promptTokens') or 0
-                out = meta.get('output_tokens') or meta.get('completion_tokens') or meta.get('completionTokens') or 0
-                if total or inp or out:
-                    return TokenUsage(total=total, input=inp, output=out)
-
-        return None

packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py

@@ -518,11 +518,11 @@ async def test_returns_failure_when_exception_thrown(self):
 
 
 class TestBuildTools:
-    """Tests for langchain_helper.build_tools (sync vs async registry callables)."""
+    """Tests for build_structured_tools (sync vs async registry callables)."""
 
     def test_registers_sync_callable_as_structured_tool_func(self):
         from ldai.models import AIAgentConfig, ModelConfig, ProviderConfig
-        from ldai_langchain.langchain_helper import build_tools
+        from ldai_langchain.langchain_helper import build_structured_tools
 
         def sync_tool(x: str = '') -> str:
             return 'ok'
@@ -538,14 +538,14 @@ def sync_tool(x: str = '') -> str:
             instructions='',
             tracker=MagicMock(),
         )
-        tools = build_tools(cfg, {'my_tool': sync_tool})
+        tools = build_structured_tools(cfg, {'my_tool': sync_tool})
         assert len(tools) == 1
         assert tools[0].func is sync_tool
         assert getattr(tools[0], 'coroutine', None) is None
 
     def test_registers_async_callable_as_structured_tool_coroutine(self):
         from ldai.models import AIAgentConfig, ModelConfig, ProviderConfig
-        from ldai_langchain.langchain_helper import build_tools
+        from ldai_langchain.langchain_helper import build_structured_tools
 
         async def async_tool(x: str = '') -> str:
             return 'ok'
@@ -561,6 +561,6 @@ async def async_tool(x: str = '') -> str:
             instructions='',
             tracker=MagicMock(),
         )
-        tools = build_tools(cfg, {'my_tool': async_tool})
+        tools = build_structured_tools(cfg, {'my_tool': async_tool})
         assert len(tools) == 1
         assert tools[0].coroutine is async_tool

packages/ai-providers/server-ai-langchain/tests/test_langgraph_callback_handler.py

@@ -69,14 +69,14 @@ def _make_graph(mock_ld_client: MagicMock, node_key: str = 'root-agent', graph_k
 
 def _llm_result(total: int, prompt: int, completion: int) -> LLMResult:
     return LLMResult(
-        generations=[[ChatGeneration(message=AIMessage(content='ok'), text='ok')]],
-        llm_output={
-            'token_usage': {
-                'total_tokens': total,
-                'prompt_tokens': prompt,
-                'completion_tokens': completion,
-            }
-        },
+        generations=[[ChatGeneration(
+            message=AIMessage(
+                content='ok',
+                usage_metadata={'total_tokens': total, 'input_tokens': prompt, 'output_tokens': completion},
+            ),
+            text='ok',
+        )]],
+        llm_output={},
     )
 
 
@@ -242,14 +242,17 @@ def test_on_llm_end_unknown_parent_run_id_ignored():
 
 
 def test_on_llm_end_camel_case_token_keys():
-    """camelCase token_usage keys (tokenUsage, totalTokens, etc.) are parsed."""
+    """camelCase token keys in response_metadata (e.g. some AWS Bedrock models) are parsed."""
     handler = LDMetricsCallbackHandler({'root-agent'}, {})
     node_run_id = uuid4()
     handler.on_chain_start({}, {}, run_id=node_run_id, name='root-agent')
 
+    msg = AIMessage(content='ok', response_metadata={
+        'tokenUsage': {'totalTokens': 20, 'promptTokens': 12, 'completionTokens': 8}
+    })
     result = LLMResult(
-        generations=[[ChatGeneration(message=AIMessage(content='ok'), text='ok')]],
-        llm_output={'tokenUsage': {'totalTokens': 20, 'promptTokens': 12, 'completionTokens': 8}},
+        generations=[[ChatGeneration(message=msg, text='ok')]],
+        llm_output={},
     )
     handler.on_llm_end(result, run_id=uuid4(), parent_run_id=node_run_id)
 
@@ -260,22 +263,6 @@ def test_on_llm_end_camel_case_token_keys():
     assert tokens.output == 8
 
 
-def test_on_llm_end_fallback_generation_info():
-    """Token usage in generation_info is used as fallback when llm_output is absent."""
-    handler = LDMetricsCallbackHandler({'root-agent'}, {})
-    node_run_id = uuid4()
-    handler.on_chain_start({}, {}, run_id=node_run_id, name='root-agent')
-
-    gen = ChatGeneration(message=AIMessage(content='ok'), text='ok')
-    gen.generation_info = {'usage_metadata': {'input_tokens': 4, 'output_tokens': 2, 'total_tokens': 6}}
-    result = LLMResult(generations=[[gen]], llm_output={})
-    handler.on_llm_end(result, run_id=uuid4(), parent_run_id=node_run_id)
-
-    tokens = handler.node_tokens.get('root-agent')
-    assert tokens is not None
-    assert tokens.total == 6
-
-
 # ---------------------------------------------------------------------------
 # on_tool_end tests
 # ---------------------------------------------------------------------------