langgraph-demo / src /doc_analysis_graph.py

hibam1

add dynamic_tool_agent

1870dee 7 days ago

4.21 kB

import base64
from typing import List, TypedDict, Annotated, Optional
from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
from langgraph.graph.message import add_messages
from langgraph.graph import START, END, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from IPython.display import Image, display
from dotenv import load_dotenv

load_dotenv()  # reads OPENAI_API_KEY from .env

# How state got updated? (by Node return partial State update)
# In LangGraph, a Node typically does not mutate the shared State directly.
# It returns a partial State update, and LangGraph applies that update
# Node receives State → performs work → returns partial updates → LangGraph merges updates into State.
# lang chain llm message returned like this -> its pre defined
# AIMessage(
#     content="Hello! How can I help you?",
#     tool_calls=[],
#     id="run-123",
#     response_metadata={...},
#     usage_metadata={...}
# )

class AgentState(TypedDict):
    input_file: Optional[str]
    messages: Annotated[list[AnyMessage], add_messages]

vision_llm = ChatOpenAI(model="gpt-4o") # model supports reading iages

def extract_text(img_path: str) -> str:
    """Extract text from an image using a vision model."""
    with open(img_path, "rb") as image_file:
        image_bytes = image_file.read()
    image_base64 = base64.b64encode(image_bytes).decode("utf-8") # base64 encode the image
    message = [
        HumanMessage(
            content=[
                {
                    "type": "text",
                    "text": "Extract all the text from this image. Return only the extracted text."
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{image_base64}"
                    }
                }
            ]
        )
    ]

    response = vision_llm.invoke(message)
    return response.content.strip()


def divide(a: int, b: int) -> float:
    """Divide a and b."""
    return a / b

tools = [divide, extract_text]

llm = ChatOpenAI(model="gpt-4o")

llm_with_tools = llm.bind_tools(
    tools,
    parallel_tool_calls=False
)

#  بتستدعي الـ LLM وبترجع جوابه أو طلب استخدام أداة.
def assistant(state: AgentState):
    image = state["input_file"]
    sys_msg = SystemMessage(
        content=f"""
        You are a helpful butler named Alfred
        that serves Mr. Wayne and Batman.
        You can analyse documents and run computations
        with provided tools:
        extract_text(img_path: str) -> str:
            Extract text from an image file.
        divide(a: int, b: int) -> float:
            Divide a and b.
        Currently the loaded image is: {image}
        """
    )

    return {
        "messages": [
            llm_with_tools.invoke( # return AIMessage
                [sys_msg] + state["messages"]
            )
        ],
        "input_file": state["input_file"]
    }

# connect the graph:
# 1. Create the graph builder
builder = StateGraph(AgentState)
# 2. Add nodes
# The assistant node is responsible for both calling tools and generating the final response.
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
# 3. Set the starting point
builder.add_edge(START, "assistant")
# 4. Route based on tool calls work the same as tools_condition
def route_assistant(state: AgentState):
    last_message = state["messages"][-1]
    if last_message.tool_calls:
        return "tools"
    return END

builder.add_conditional_edges(
    "assistant",
    route_assistant,
    {
        "tools": "tools",
        END: END
    }
)
# 5. Return from tools to assistant
builder.add_edge("tools", "assistant")
# 6. Compile
react_graph = builder.compile()

############ call example: ############
initial_state = {
    "input_file": None,
    "messages": [
        HumanMessage(content="Divide 100 by 5")
    ]
}

#invoke() — executes the Graph and returns the final State.
#stream() — lets us observe what's happening during execution.

result = react_graph.invoke(initial_state)
print(result)

	import base64
	from typing import List, TypedDict, Annotated, Optional
	from langchain_openai import ChatOpenAI
	from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
	from langgraph.graph.message import add_messages
	from langgraph.graph import START, END, StateGraph
	from langgraph.prebuilt import ToolNode, tools_condition
	from IPython.display import Image, display
	from dotenv import load_dotenv

	load_dotenv() # reads OPENAI_API_KEY from .env

	# How state got updated? (by Node return partial State update)
	# In LangGraph, a Node typically does not mutate the shared State directly.
	# It returns a partial State update, and LangGraph applies that update
	# Node receives State → performs work → returns partial updates → LangGraph merges updates into State.
	# lang chain llm message returned like this -> its pre defined
	# AIMessage(
	# content="Hello! How can I help you?",
	# tool_calls=[],
	# id="run-123",
	# response_metadata={...},
	# usage_metadata={...}
	# )

	class AgentState(TypedDict):
	input_file: Optional[str]
	messages: Annotated[list[AnyMessage], add_messages]

	vision_llm = ChatOpenAI(model="gpt-4o") # model supports reading iages

	def extract_text(img_path: str) -> str:
	"""Extract text from an image using a vision model."""
	with open(img_path, "rb") as image_file:
	image_bytes = image_file.read()
	image_base64 = base64.b64encode(image_bytes).decode("utf-8") # base64 encode the image
	message = [
	HumanMessage(
	content=[
	{
	"type": "text",
	"text": "Extract all the text from this image. Return only the extracted text."
	},
	{
	"type": "image_url",
	"image_url": {
	"url": f"data:image/png;base64,{image_base64}"
	}
	}
	]
	)
	]

	response = vision_llm.invoke(message)
	return response.content.strip()


	def divide(a: int, b: int) -> float:
	"""Divide a and b."""
	return a / b

	tools = [divide, extract_text]

	llm = ChatOpenAI(model="gpt-4o")

	llm_with_tools = llm.bind_tools(
	tools,
	parallel_tool_calls=False
	)

	# بتستدعي الـ LLM وبترجع جوابه أو طلب استخدام أداة.
	def assistant(state: AgentState):
	image = state["input_file"]
	sys_msg = SystemMessage(
	content=f"""
	You are a helpful butler named Alfred
	that serves Mr. Wayne and Batman.

	You can analyse documents and run computations
	with provided tools:

	extract_text(img_path: str) -> str:
	Extract text from an image file.

	divide(a: int, b: int) -> float:
	Divide a and b.

	Currently the loaded image is: {image}
	"""
	)

	return {
	"messages": [
	llm_with_tools.invoke( # return AIMessage
	[sys_msg] + state["messages"]
	)
	],
	"input_file": state["input_file"]
	}

	# connect the graph:
	# 1. Create the graph builder
	builder = StateGraph(AgentState)
	# 2. Add nodes
	# The assistant node is responsible for both calling tools and generating the final response.
	builder.add_node("assistant", assistant)
	builder.add_node("tools", ToolNode(tools))
	# 3. Set the starting point
	builder.add_edge(START, "assistant")
	# 4. Route based on tool calls work the same as tools_condition
	def route_assistant(state: AgentState):
	last_message = state["messages"][-1]
	if last_message.tool_calls:
	return "tools"
	return END

	builder.add_conditional_edges(
	"assistant",
	route_assistant,
	{
	"tools": "tools",
	END: END
	}
	)
	# 5. Return from tools to assistant
	builder.add_edge("tools", "assistant")
	# 6. Compile
	react_graph = builder.compile()

	############ call example: ############
	initial_state = {
	"input_file": None,
	"messages": [
	HumanMessage(content="Divide 100 by 5")
	]
	}

	#invoke() — executes the Graph and returns the final State.
	#stream() — lets us observe what's happening during execution.

	result = react_graph.invoke(initial_state)
	print(result)