diff --git a/main.py b/main.py index 38b3014..1b396dc 100644 --- a/main.py +++ b/main.py @@ -1,54 +1,52 @@ from langchain_openai import ChatOpenAI -from langchain.tools import tool -from langchain.agents import create_agent, AgentExecutor +from langchain.tools import tool, BaseTool +from langchain.agents import create_agent +from pydantic import SecretStr import json -# Configure LLM +# Connect to local LLM via OpenAI-compatible API llm = ChatOpenAI( - model="gpt-4o-mini", # replace with your local model name + model="gpt-4o-mini", # replace with your LM Studio model name base_url="http://localhost:1234/v1", - api_key="fake", + api_key=SecretStr("fake"), temperature=0.7, ) -@tool("Get price for a product in a city") +# Sub-agent that generates a price table for a product in a city +@tool def get_price(product: str, city: str) -> str: - """ - Returns a table with product, price and store. - The function internally creates a sub-agent that generates realistic prices. - """ - # Sub‑agent to generate price - sub_llm = ChatOpenAI( - model="gpt-4o-mini", - base_url="http://localhost:1234/v1", - api_key="fake", - temperature=0.5, - ) + """Return a realistic price table for the given product and city.""" + # Create a sub‑agent with a simple prompt to generate a table sub_agent = create_agent( - model=sub_llm, + model=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.", + system_prompt=f"You are a market analyst. Provide a realistic price for {product} in {city}. Return the result as a markdown table with columns: Продукт, Цена (руб.), Магазин.", ) - 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) + response = sub_agent.invoke({"messages": [{"role": "human", "content": f"Generate price for {product} in {city}"}]}) + # The agent returns a dict with messages; the last message contains the table + return response["messages"][-1]["content"] -# Main agent with get_price tool +# Main agent that uses get_price to build shopping list main_agent = create_agent( model=llm, tools=[get_price], - system_prompt="You are a shopping assistant. Use the get_price tool to help users plan their purchases.", + system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цены каждого продукта.", ) -if __name__ == "__main__": - user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани." - response = main_agent.invoke({"messages": [{"role": "human", "content": user_query}]}) - # Print all messages - for msg in response["messages"]: - if "content" in msg: - print(msg["content"]) - 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"] if "content" in response["messages"][-1] else "" - print("\nFinal answer:\n", final) +# Example query +query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани." +result = main_agent.invoke({"messages": [{"role": "human", "content": query}]}) + +# Pretty‑print all messages +for msg in result["messages"]: + if "content" in msg and msg["content"]: + print(msg["content"]) + elif "tool_calls" in msg and msg["tool_calls"]: + for call in msg["tool_calls"]: + name = call["name"] + args = json.dumps(call["args"], ensure_ascii=False) + print(f"{name}({args})") + else: + print(msg) + +print("\n--- End of conversation ---")