diff --git a/main.py b/main.py index aa2a7ef..38b3014 100644 --- a/main.py +++ b/main.py @@ -1,60 +1,54 @@ -""" -Simple hierarchical LangChain agent for shopping list. -""" - from langchain_openai import ChatOpenAI -from pydantic import SecretStr from langchain.tools import tool from langchain.agents import create_agent, AgentExecutor import json -# LLM configuration – replace with your LM Studio model name +# Configure LLM llm = ChatOpenAI( - model="", + model="gpt-4o-mini", # replace with your local model name base_url="http://localhost:1234/v1", - api_key=SecretStr("fake"), + api_key="fake", temperature=0.7, ) -# Sub‑agent that generates a price table for one product in a city -@tool(name="get_price", description="Return realistic price for a product in a city as a markdown table.") +@tool("Get price for a product in a city") def get_price(product: str, city: str) -> str: """ - Generates a markdown table with columns: Product | Price (руб.) | Store. - The sub‑agent uses the same LLM to produce realistic values. + Returns a table with product, price and store. + The function internally creates a sub-agent that generates realistic prices. """ - # Create a tiny agent that only returns the price table - prompt = ( - f"You are a local market assistant. Provide a markdown table with columns:\n" - f"| Продукт | Цена (руб.) | Магазин |\n" - f"For product '{product}' in city '{city}'. Use realistic Russian prices and store names.") - sub_agent = create_agent( - model=llm, - tools=[], - system_prompt=prompt, + # Sub‑agent to generate price + sub_llm = ChatOpenAI( + model="gpt-4o-mini", + base_url="http://localhost:1234/v1", + api_key="fake", + temperature=0.5, ) - result = sub_agent.invoke({"messages": [{"role": "human", "content": "Generate table"}]}) - return result["messages"][-1]["content"] + sub_agent = create_agent( + model=sub_llm, + tools=[], + system_prompt=f"You are a price estimator for {city}. Provide a realistic price and store name for {product} in a table format.", + ) + result = sub_agent.invoke({"messages": [{"role": "human", "content": f"Give me the price of {product} in {city}"}]}) + return json.dumps(result["messages"][-1]["content"], ensure_ascii=False) # Main agent with get_price tool main_agent = create_agent( model=llm, tools=[get_price], - system_prompt="Ты помощник по планированию покупок.", + system_prompt="You are a shopping assistant. Use the get_price tool to help users plan their purchases.", ) if __name__ == "__main__": - user_query = ( - "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани." - ) + user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани." response = main_agent.invoke({"messages": [{"role": "human", "content": user_query}]}) # Print all messages for msg in response["messages"]: - if msg.get("content"): + if "content" in msg: print(msg["content"]) - elif msg.get("tool_calls"): - call = msg["tool_calls"][0] - print(f"{call['name']}({json.dumps(call['args'])})") + elif "tool_calls" in msg: + for call in msg["tool_calls"]: + print(f"{call['name']}({json.dumps(call['args'])})") # Final answer - final = response["messages"][-1]["content"] - print("\n---\n", final) + final = response["messages"][-1]["content"] if "content" in response["messages"][-1] else "" + print("\nFinal answer:\n", final)