1 line
1.3 KiB
Python
1 line
1.3 KiB
Python
import os\nfrom dotenv import load_dotenv\nfrom langgraph import create_agent\nfrom langchain.chat_models import ChatOpenAI\nfrom langchain.tools import CalculatorTool\n\nload_dotenv()\n\ndef format_message(message) -> str:\n if message.content:\n return message.content\n return f\"{message.tool_calls[0]['name']}({message.tool_calls[0]['args']})\"\n\n\ndef main():\n llm = ChatOpenAI(temperature=0)\n tools = [CalculatorTool()]\n agent = create_agent(llm=llm, tools=tools)\n\n stream = agent.stream(\n {\n \"messages\": [{\"role\": \"human\", \"content\": \"What is 12 * 34?\"}]\n },\n stream_mode=['messages', 'updates']\n )\n\n step = 1\n for chunk in stream:\n chunk_type, chunk_data = chunk\n if chunk_type == \"messages\":\n message, meta = chunk_data\n if meta.get('langgraph_step') != step:\n step = meta.get('langgraph_step')\n print('\\n --- --- --- \\n')\n if message.content:\n print(message.content, end='', flush=True)\n elif chunk_type == \"updates\":\n if chunk_data.get('model'):\n last_message = chunk_data['model']['messages'][-1]\n print(format_message(last_message))\n\n print() # Final newline\n\n\nif __name__ == \"__main__\":\n main()\n |