feat: solution for 'Untitled Task'

src/main.py

@@ -1,52 +1 @@
-import os
-from dotenv import load_dotenv
-from langgraph import AgentBuilder
-from langchain_openai import ChatOpenAI
-
-load_dotenv()
-
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-if not OPENAI_API_KEY:
-    raise ValueError("OPENAI_API_KEY not set in environment")
-
-llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo")
-
-builder = AgentBuilder(llm=llm)
-agent = builder.build()
-
-def format_message(message) -> str:
-    if message.content:
-        return message.content
-    if message.tool_calls:
-        tool = message.tool_calls[0]
-        return f"{tool['name']}({tool['args']})"
-    return ""
-
-step = 1
-
-def format_chunk_message(chunk):
-    global step
-    message, meta = chunk
-    if meta.get("langgraph_step") != step:
-        step = meta.get("langgraph_step")
-        print("\n --- --- --- \n")
-    if message.content:
-        print(message.content, end="", flush=True)
-
-def main():
-    user_input = input("Enter your question: ")
-    stream = agent.stream(
-        {"messages": [{"role": "human", "content": user_input}]},
-        stream_mode=["messages", "updates"]
-    )
-    for chunk_type, chunk_data in stream:
-        if chunk_type == "messages":
-            format_chunk_message(chunk_data)
-        elif chunk_type == "updates":
-            if chunk_data.get("model"):
-                last_message = chunk_data["model"]["messages"][-1]
-                print(format_message(last_message))
-    print("\n\nDone.")
-
-if __name__ == "__main__":
-    main()
+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