diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..386ac3a --- /dev/null +++ b/src/main.py @@ -0,0 +1,62 @@ +""" +Simple hierarchical LangChain agent for shopping list. + +The script demonstrates: +* Connection to a local LLM via OpenAI‑compatible API. +* A tool `get_price` that internally creates a sub‑agent to generate a price table. +* A main agent that uses the tool and prints all intermediate calls and final answer. +""" + +from langchain_openai import ChatOpenAI +from langchain.tools import tool, BaseTool +from langchain.agents import create_agent +from pydantic import SecretStr +import json + +# --- LLM setup ----------------------------------------------------------- +llm = ChatOpenAI( + model="gpt-4o-mini", # replace with your LM Studio model name + base_url="http://localhost:1234/v1", + api_key=SecretStr("fake"), + temperature=0.7, +) + +# --- Sub‑agent that generates a price table ------------------------------ +def create_price_agent(product: str, city: str): + """Return a simple string with a fake price table for the product.""" + # In a real scenario you would query a database or API. + prices = { + ("молоко", "Казань"): ("89", "Магнит"), + ("хлеб", "Казань"): ("45", "Пятёрочка"), + ("яблоки", "Казань"): ("120/кг", "Перекрёсток"), + } + price, store = prices.get((product.lower(), city), ("?", "?")) + table = f"| Продукт | Цена (руб.) | Магазин |\n|---------|-------------|---------|\n| {product} | {price} | {store} | +" + return table + +# --- Tool that calls the sub‑agent --------------------------------------- +@tool(name="get_price", description="Get price for a product in a city") +def get_price(product: str, city: str) -> str: + return create_price_agent(product, city) + +# --- Main agent ---------------------------------------------------------- +main_agent = create_agent( + model=llm, + tools=[get_price], + system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цен.", +) + +# --- Run the agent with a sample query ----------------------------------- +query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани." +result = main_agent.invoke({"messages": [{"role": "human", "content": query}]}) + +# Print all messages (intermediate tool calls and final answer) +for msg in result["messages"]: + if msg.get("tool_calls"): + for call in msg["tool_calls"]: + print(f"{call['name']}({json.dumps(call['args'])})") + else: + print(msg["content"]) + +# End of script