add main.py
This commit is contained in:
+62
@@ -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
|
||||
Reference in New Issue
Block a user