75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
"""
|
||
Shopping‑list AI assistant.
|
||
|
||
The script demonstrates a hierarchical LangChain agent that can estimate prices for
|
||
products in a city. The top‑level agent uses the ``get_price`` tool, which itself
|
||
creates a short‑lived sub‑agent to generate a markdown table row with the price.
|
||
|
||
Three example calls are executed when the module is run as a script:
|
||
1. Milk, bread, apples – Kazan
|
||
2. Eggs, cheese – Moscow
|
||
3. Coffee – Saint‑Petersburg
|
||
|
||
Each example prints the full message chain (tool calls and final answer).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from typing import List
|
||
|
||
# Local LLM configuration – environment variables allow CI to override.
|
||
LOCAL_LLM_MODEL = os.getenv("LOCAL_LLM_MODEL", "gpt-3.5-turbo")
|
||
LOCAL_LLM_BASE_URL = os.getenv("LOCAL_LLM_BASE_URL", "http://localhost:1234/v1")
|
||
LOCAL_LLM_API_KEY = os.getenv("LOCAL_LLM_API_KEY", "fake") # LM Studio dummy key
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain.tools import tool
|
||
from langchain.agents import create_agent
|
||
from langchain_core.messages import HumanMessage, SystemMessage
|
||
|
||
# Import the price‑tool from the helper module.
|
||
from agent import get_price
|
||
|
||
# Base LLM used by all agents.
|
||
_base_llm = ChatOpenAI(
|
||
model=LOCAL_LLM_MODEL,
|
||
base_url=LOCAL_LLM_BASE_URL,
|
||
api_key=LOCAL_LLM_API_KEY,
|
||
temperature=0.2,
|
||
)
|
||
|
||
# Main agent – it only has the ``get_price`` tool.
|
||
shopping_agent = create_agent(
|
||
llm=_base_llm,
|
||
tools=[get_price],
|
||
system_prompt="You are a helpful assistant for planning shopping lists.",
|
||
)
|
||
|
||
def run_example(products: List[str], city: str) -> None:
|
||
"""Run the agent on *products* in *city* and print the full chain."""
|
||
# Build the user message.
|
||
product_list = ", ".join(products)
|
||
prompt = f"Help me plan a shopping list: {product_list}. I am in {city}."
|
||
|
||
# Invoke the agent.
|
||
result = shopping_agent.invoke({"messages": [HumanMessage(content=prompt)]})
|
||
|
||
# Pretty‑print the chain of messages.
|
||
print("\n=== Example: " + prompt + " ===")
|
||
for msg in result["messages"]:
|
||
if hasattr(msg, "content") and msg.content:
|
||
print(f"Assistant: {msg.content}")
|
||
elif hasattr(msg, "tool_calls") and msg.tool_calls:
|
||
call = msg.tool_calls[0]
|
||
print(f"Tool call – {call['name']}({call['args']})")
|
||
print("\n--- End of example ---\n")
|
||
|
||
if __name__ == "__main__":
|
||
# Three distinct examples.
|
||
run_example(["milk", "bread", "apples"], "Kazan")
|
||
run_example(["eggs", "cheese"], "Moscow")
|
||
run_example(["coffee"], "Saint‑Petersburg")
|
||
|
||
# End of script.
|