diff --git a/main.py b/main.py index 1bcc80e..c6cca96 100644 --- a/main.py +++ b/main.py @@ -1,154 +1,74 @@ """ -Simple hierarchical AI agent for shopping list planning. +Shopping‑list AI assistant. -The script demonstrates: -* Connection to a local LLM via the OpenAI compatible API. -* A tool that internally creates a sub‑agent to estimate product prices. -* A main agent that orchestrates calls to the price tool and aggregates results. +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. -Run with: - python -m venv .venv && source .venv/bin/activate - pip install -r requirements.txt - python main.py +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 -import json -from typing import Dict, Any, List +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 -# LangChain imports – the exact versions are pinned in requirements.txt from langchain_openai import ChatOpenAI from langchain.tools import tool from langchain.agents import create_agent from langchain_core.messages import HumanMessage, SystemMessage -from pydantic import SecretStr -# --------------------------------------------------------------------------- -# 1. LLM configuration – local LM Studio server -# --------------------------------------------------------------------------- -LLM_MODEL = os.getenv("LM_MODEL", "gpt-4o-mini") # default model name in LM Studio -BASE_URL = os.getenv("LM_BASE_URL", "http://localhost:1234/v1") -API_KEY = SecretStr("fake") # LM Studio does not require a real key +# Import the price‑tool from the helper module. +from agent import get_price -llm = ChatOpenAI( - model=LLM_MODEL, - base_url=BASE_URL, - api_key=API_KEY, - temperature=0.7, +# 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, ) -# --------------------------------------------------------------------------- -# 2. Tool that internally creates a sub‑agent to estimate price -# --------------------------------------------------------------------------- -@tool -def get_price(product: str, city: str) -> str: - """ - Estimate the price of *product* in *city*. - - The function builds a tiny sub‑agent that asks the LLM for a realistic - price table. The sub‑agent is created on every call – this keeps the - implementation simple and avoids persisting state between calls. - """ - - # Sub‑agent system prompt – we keep it short to reduce token usage - sub_prompt = ( - f"You are a local market price estimator for {city}. Provide a single table with columns: - | Product | Price (rub.) | Store | - The product is '{product}'. Use realistic Russian prices.") - - # Create the sub‑agent – it only has one tool: none, so it just replies - sub_agent = create_agent( - llm=llm, - tools=[], - system_prompt=sub_prompt, - ) - - # Ask the sub‑agent for a price table - response = sub_agent.invoke({"messages": [HumanMessage(content="Generate the table.")], "configurable": {}}) - # The last message contains the answer - return response["messages"][-1].content.strip() - -# --------------------------------------------------------------------------- -# 3. Main agent – orchestrates calls to get_price and aggregates results -# --------------------------------------------------------------------------- -main_agent = create_agent( - llm=llm, +# 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.", ) -# --------------------------------------------------------------------------- -# 4. Helper to format the final output nicely -# --------------------------------------------------------------------------- +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}." -def aggregate_prices(products: List[str], city: str) -> Dict[str, Any]: - """Call get_price for each product and sum up total cost. + # Invoke the agent. + result = shopping_agent.invoke({"messages": [HumanMessage(content=prompt)]}) - The function returns a dictionary with keys: - - tables: list of price tables (strings) - - total: estimated total in rubles (int or float) - """ - tables = [] - total = 0.0 - for prod in products: - table = get_price(prod, city) - tables.append(table) - # Extract numeric price from the table – naive regex - try: - lines = table.splitlines() - if len(lines) >= 2: - row = lines[1] - parts = [p.strip() for p in row.split('|') if p.strip()] - if len(parts) >= 2: - price_str = parts[1] - # Remove non‑digits - digits = ''.join(ch for ch in price_str if ch.isdigit()) - if digits: - total += float(digits) - except Exception: - pass - return {"tables": tables, "total": total} - -# --------------------------------------------------------------------------- -# 5. Main entry point – parse user input and run the agent -# --------------------------------------------------------------------------- -if __name__ == "__main__": - # Example prompt – in real usage this would come from stdin or a UI - user_prompt = ( - "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани.") - - # Run the main agent - result = main_agent.invoke({"messages": [HumanMessage(content=user_prompt)], "configurable": {}}) - - # Print all messages – tool calls and final answer + # Pretty‑print the chain of messages. + print("\n=== Example: " + prompt + " ===") for msg in result["messages"]: if hasattr(msg, "content") and msg.content: - print(msg.content) + print(f"Assistant: {msg.content}") elif hasattr(msg, "tool_calls") and msg.tool_calls: - for call in msg.tool_calls: - name = call.get("name") - args = json.dumps(call.get("args")) - print(f"{name}({args})") + call = msg.tool_calls[0] + print(f"Tool call – {call['name']}({call['args']})") + print("\n--- End of example ---\n") - # Additionally show aggregated price summary (for demonstration) - # Extract products and city from the user prompt – simple split logic - try: - parts = user_prompt.split(":", 1)[1] - prod_part, city_part = parts.split(". Я нахожусь в ") - products = [p.strip() for p in prod_part.replace("составить список покупок", "").split(",") if p.strip()] - city = city_part.rstrip("") - except Exception: - products, city = [], "" +if __name__ == "__main__": + # Three distinct examples. + run_example(["milk", "bread", "apples"], "Kazan") + run_example(["eggs", "cheese"], "Moscow") + run_example(["coffee"], "Saint‑Petersburg") - if products and city: - agg = aggregate_prices(products, city) - print("\n--- Aggregated price tables ---") - for t in agg["tables"]: - print(t + "\n") - print(f"**Итого:** ~{int(agg['total'])} руб.") - else: - print("Не удалось извлечь список продуктов и город из запроса.") - -# End of file +# End of script.