Spaces:
Runtime error
Runtime error
File size: 2,650 Bytes
90f65f7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | # file: 02_basic_agent.py
import os
from dotenv import load_dotenv
load_dotenv()
from typing import TypedDict, Annotated, List
from langchain_core.messages import BaseMessage, ToolMessage
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langchain_openai import ChatOpenAI
from tools.word_counter import count_words
# Define the constant for the OpenRouter API URL
OPENROUTER_API_URL = "https://openrouter.ai/api/v1"
# 1. Define the Agent State
class AgentState(TypedDict):
messages: Annotated[List[BaseMessage], lambda x, y: x + y]
def main():
"""Main function to set up and run a LangGraph agent."""
if not os.getenv("OPENROUTER_API_KEY"):
print("Error: OPENROUTER_API_KEY is not set in .env file.")
return
print("--- Modern Agent with LangGraph ---")
tools = [count_words]
llm = ChatOpenAI(
model="anthropic/claude-3.5-haiku",
temperature=0,
openai_api_key=os.getenv("OPENROUTER_API_KEY"),
openai_api_base=OPENROUTER_API_URL,
)
llm_with_tools = llm.bind_tools(tools)
# 2. Define the Nodes
def call_model(state):
"""The node that calls the LLM."""
messages = state["messages"]
response = llm_with_tools.invoke(messages)
return {"messages": [response]}
tool_node = ToolNode(tools)
# 3. Define the Conditional Edge
def should_continue(state):
"""Decides whether to continue or end the loop."""
last_message = state["messages"][-1]
if last_message.tool_calls:
return "continue"
return "end"
# 4. Build the Graph
workflow = StateGraph(AgentState)
workflow.add_node("agent", call_model)
workflow.add_node("action", tool_node)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
"agent",
should_continue,
{"continue": "action", "end": END}
)
workflow.add_edge("action", "agent")
app = workflow.compile()
# 5. Invoke the Agent
query = "How many words are in the text 'I am a skilled AI engineer'?"
inputs = {"messages": [("user", query)]}
print(f"Human: {query}")
for chunk in app.stream(inputs, stream_mode="values"):
last_message = chunk["messages"][-1]
if isinstance(last_message, ToolMessage):
print(f"Tool Output: {last_message.content}")
final_response = app.invoke(inputs)
final_answer = final_response["messages"][-1].content
print("\n--- Agent Response ---")
print(final_answer)
print("----------------------")
if __name__ == "__main__":
main() |