Add main.py

This commit is contained in:
2026-05-28 10:11:23 +00:00
parent dd2ce7ca22
commit 1f44b4a822
+32 -34
View File
@@ -1,54 +1,52 @@
from langchain_openai import ChatOpenAI from langchain_openai import ChatOpenAI
from langchain.tools import tool from langchain.tools import tool, BaseTool
from langchain.agents import create_agent, AgentExecutor from langchain.agents import create_agent
from pydantic import SecretStr
import json import json
# Configure LLM # Connect to local LLM via OpenAI-compatible API
llm = ChatOpenAI( 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", base_url="http://localhost:1234/v1",
api_key="fake", api_key=SecretStr("fake"),
temperature=0.7, 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: def get_price(product: str, city: str) -> str:
""" """Return a realistic price table for the given product and city."""
Returns a table with product, price and store. # Create a subagent with a simple prompt to generate a table
The function internally creates a sub-agent that generates realistic prices.
"""
# Subagent to generate price
sub_llm = ChatOpenAI(
model="gpt-4o-mini",
base_url="http://localhost:1234/v1",
api_key="fake",
temperature=0.5,
)
sub_agent = create_agent( sub_agent = create_agent(
model=sub_llm, model=llm,
tools=[], 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}"}]}) response = sub_agent.invoke({"messages": [{"role": "human", "content": f"Generate price for {product} in {city}"}]})
return json.dumps(result["messages"][-1]["content"], ensure_ascii=False) # 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( main_agent = create_agent(
model=llm, model=llm,
tools=[get_price], 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__": # Example query
user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани." query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
response = main_agent.invoke({"messages": [{"role": "human", "content": user_query}]}) result = main_agent.invoke({"messages": [{"role": "human", "content": query}]})
# Print all messages
for msg in response["messages"]: # Prettyprint all messages
if "content" in msg: for msg in result["messages"]:
if "content" in msg and msg["content"]:
print(msg["content"]) print(msg["content"])
elif "tool_calls" in msg: elif "tool_calls" in msg and msg["tool_calls"]:
for call in msg["tool_calls"]: for call in msg["tool_calls"]:
print(f"{call['name']}({json.dumps(call['args'])})") name = call["name"]
# Final answer args = json.dumps(call["args"], ensure_ascii=False)
final = response["messages"][-1]["content"] if "content" in response["messages"][-1] else "" print(f"{name}({args})")
print("\nFinal answer:\n", final) else:
print(msg)
print("\n--- End of conversation ---")