fix bug causing issue with tools

Author: rohit-kadhe

src/examples/langchain_example/__init__.py

@@ -2,6 +2,7 @@
 from langtrace_python_sdk import with_langtrace_root_span
 
 from .groq_example import groq_basic, groq_streaming
+from .langgraph_example_tools import basic_graph_tools
 
 
 class LangChainRunner:
@@ -10,6 +11,7 @@ def run(self):
         basic_app()
         rag()
         load_and_split()
+        basic_graph_tools()
 
 
 class GroqRunner:

src/examples/langchain_example/langgraph_example_tools.py

@@ -0,0 +1,151 @@
+from typing import Annotated
+
+from langchain_anthropic import ChatAnthropic
+from langchain_core.messages import HumanMessage
+from langchain_core.pydantic_v1 import BaseModel
+from typing_extensions import TypedDict
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_core.tools import Tool
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import StateGraph
+from langgraph.graph.message import add_messages
+from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_core.messages import AIMessage, ToolMessage
+
+from langtrace_python_sdk import langtrace
+
+langtrace.init()
+
+primes = {998: 7901, 999: 7907, 1000: 7919}
+
+
+class PrimeInput(BaseModel):
+    n: int = Field()
+
+
+def is_prime(n: int) -> bool:
+    if n <= 1 or (n % 2 == 0 and n > 2):
+        return False
+    for i in range(3, int(n**0.5) + 1, 2):
+        if n % i == 0:
+            return False
+    return True
+
+
+def get_prime(n: int, primes: dict = primes) -> str:
+    return str(primes.get(int(n)))
+
+
+async def aget_prime(n: int, primes: dict = primes) -> str:
+    return str(primes.get(int(n)))
+
+
+class State(TypedDict):
+    messages: Annotated[list, add_messages]
+    # This flag is new
+    ask_human: bool
+
+
+class RequestAssistance(BaseModel):
+    """Escalate the conversation to an expert. Use this if you are unable to assist directly or if the user requires support beyond your permissions.
+
+    To use this function, relay the user's 'request' so the expert can provide the right guidance.
+    """
+
+    request: str
+
+
+llm = ChatAnthropic(model="claude-3-haiku-20240307")
+# We can bind the llm to a tool definition, a pydantic model, or a json schema
+llm_with_tools = llm.bind_tools([RequestAssistance])
+tools = [
+    Tool(
+        name="GetPrime",
+        func=get_prime,
+        description="A tool that returns the `n`th prime number",
+        args_schema=PrimeInput,
+        coroutine=aget_prime,
+    ),
+]
+
+
+def chatbot(state: State):
+    response = llm_with_tools.invoke(state["messages"])
+    ask_human = False
+    if (
+        response.tool_calls
+        and response.tool_calls[0]["name"] == RequestAssistance.__name__
+    ):
+        ask_human = True
+    return {"messages": [response], "ask_human": ask_human}
+
+
+graph_builder = StateGraph(State)
+
+graph_builder.add_node("chatbot", chatbot)
+graph_builder.add_node("tools", ToolNode(tools=tools))
+
+
+def create_response(response: str, ai_message: AIMessage):
+    return ToolMessage(
+        content=response,
+        tool_call_id=ai_message.tool_calls[0]["id"],
+    )
+
+
+def human_node(state: State):
+    new_messages = []
+    if not isinstance(state["messages"][-1], ToolMessage):
+        # Typically, the user will have updated the state during the interrupt.
+        # If they choose not to, we will include a placeholder ToolMessage to
+        # let the LLM continue.
+        new_messages.append(
+            create_response("No response from human.", state["messages"][-1])
+        )
+    return {
+        # Append the new messages
+        "messages": new_messages,
+        # Unset the flag
+        "ask_human": False,
+    }
+
+
+def select_next_node(state: State):
+    if state["ask_human"]:
+        return "human"
+    # Otherwise, we can route as before
+    return tools_condition(state)
+
+
+def basic_graph_tools():
+    graph_builder.add_node("human", human_node)
+    graph_builder.add_conditional_edges(
+        "chatbot",
+        select_next_node,
+        {"human": "human", "tools": "tools", "__end__": "__end__"},
+    )
+    graph_builder.add_edge("tools", "chatbot")
+    graph_builder.add_edge("human", "chatbot")
+    graph_builder.set_entry_point("chatbot")
+    memory = MemorySaver()
+    graph = graph_builder.compile(
+        checkpointer=memory,
+        interrupt_before=["human"],
+    )
+
+    config = {"configurable": {"thread_id": "1"}}
+    events = graph.stream(
+        {
+            "messages": [
+                (
+                    "user",
+                    "I'm learning LangGraph. Could you do some research on it for me?",
+                )
+            ]
+        },
+        config,
+        stream_mode="values",
+    )
+    for event in events:
+        if "messages" in event:
+            event["messages"][-1]

src/langtrace_python_sdk/instrumentation/anthropic/patch.py

@@ -61,7 +61,7 @@ def traced_method(wrapped, instance, args, kwargs):
             **get_llm_request_attributes(kwargs, prompts=prompts),
             **get_llm_url(instance),
             SpanAttributes.LLM_PATH: APIS["MESSAGES_CREATE"]["ENDPOINT"],
-            **get_extra_attributes(),
+            **get_extra_attributes(),  # type: ignore
         }
 
         attributes = LLMSpanAttributes(**span_attributes)
@@ -88,22 +88,26 @@ def traced_method(wrapped, instance, args, kwargs):
     @silently_fail
     def set_response_attributes(result, span, kwargs):
         if not is_streaming(kwargs):
+
             if hasattr(result, "content") and result.content is not None:
                 set_span_attribute(
                     span, SpanAttributes.LLM_RESPONSE_MODEL, result.model
                 )
-                completion = [
-                    {
-                        "role": result.role if result.role else "assistant",
-                        "content": result.content[0].text,
-                        "type": result.content[0].type,
-                    }
-                ]
-                set_event_completion(span, completion)
-
-            else:
-                responses = []
-                set_event_completion(span, responses)
+                if hasattr(result, "content") and result.content[0] is not None:
+                    content = result.content[0]
+                    typ = content.type
+                    role = result.role if result.role else "assistant"
+                    if typ == "tool_result" or typ == "tool_use":
+                        content = content.json()
+                        set_span_attribute(
+                            span, SpanAttributes.LLM_TOOL_RESULTS, json.dumps(content)
+                        )
+                    if typ == "text":
+                        content = result.content[0].text
+                        content = content.text
+                        set_event_completion(
+                            span, [{type: typ, role: role, content: content}]
+                        )
 
             if (
                 hasattr(result, "system_fingerprint")