96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
"""
|
||
Simple hierarchical AI agent for shopping list using LangChain.
|
||
|
||
The script demonstrates:
|
||
* Connection to a local LLM via OpenAI-compatible API.
|
||
* A tool that internally creates a sub‑agent to generate realistic prices.
|
||
* A top‑level agent that orchestrates the price queries and aggregates results.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from typing import Dict, Any
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain.tools import tool
|
||
from langchain.agents import create_agent
|
||
from pydantic import SecretStr
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. LLM configuration – adjust model name to your LM Studio instance.
|
||
# ---------------------------------------------------------------------------
|
||
llm = ChatOpenAI(
|
||
model="gpt-4o-mini", # replace with the actual model name in LM Studio
|
||
base_url="http://localhost:1234/v1",
|
||
api_key=SecretStr("fake"),
|
||
temperature=0.7,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2. Sub‑agent that generates a price table for a single product.
|
||
# ---------------------------------------------------------------------------
|
||
@tool
|
||
def get_price(product: str, city: str) -> str:
|
||
"""Return a realistic price table for *product* in *city*.
|
||
|
||
The function internally creates a small LangChain agent that asks the LLM to
|
||
produce a markdown table with columns Product | Price (руб.) | Store.
|
||
"""
|
||
# Create a tiny agent that only has the task of generating a price row.
|
||
sub_agent = create_agent(
|
||
model=llm,
|
||
tools=[],
|
||
system_prompt=(
|
||
f"You are an assistant that provides realistic prices for products in {city}. "
|
||
"Respond with a markdown table containing columns: Product, Price (руб.), Store."
|
||
),
|
||
)
|
||
|
||
# Ask the sub‑agent to produce the price.
|
||
response = sub_agent.invoke(
|
||
{
|
||
"messages": [
|
||
{"role": "human", "content": f"Provide a price for {product} in {city}."}
|
||
]
|
||
}
|
||
)
|
||
|
||
# Extract the final message content.
|
||
messages = response.get("messages", [])
|
||
if not messages:
|
||
return f"| {product} | N/A | Unknown |
|
||
"
|
||
last_msg = messages[-1]
|
||
content = last_msg.get("content") or ""
|
||
# Ensure the table has a header row.
|
||
if "|" in content and "Product" not in content:
|
||
content = f"| Product | Price (руб.) | Store |\n{content}"
|
||
return content
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3. Main agent that orchestrates price queries for a shopping list.
|
||
# ---------------------------------------------------------------------------
|
||
shopping_agent = create_agent(
|
||
model=llm,
|
||
tools=[get_price],
|
||
system_prompt="You are an assistant that helps plan a shopping list and calculates total cost.",
|
||
)
|
||
|
||
def main() -> None:
|
||
user_query = (
|
||
"Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
|
||
)
|
||
result = shopping_agent.invoke({"messages": [{"role": "human", "content": user_query}]})
|
||
|
||
# Pretty‑print all messages.
|
||
for msg in result.get("messages", []):
|
||
if msg.get("content"):
|
||
print(msg["content"])
|
||
elif msg.get("tool_calls"):
|
||
call = msg["tool_calls"][0]
|
||
print(f"{call['name']}({json.dumps(call['args'], ensure_ascii=False)})")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|