53 lines
2.1 KiB
Python
53 lines
2.1 KiB
Python
from langchain_openai import ChatOpenAI
|
||
from langchain.tools import tool, BaseTool
|
||
from langchain.agents import create_agent
|
||
from pydantic import SecretStr
|
||
import json
|
||
|
||
# Connect to local LLM via OpenAI-compatible API
|
||
llm = ChatOpenAI(
|
||
model="gpt-4o-mini", # replace with your LM Studio model name
|
||
base_url="http://localhost:1234/v1",
|
||
api_key=SecretStr("fake"),
|
||
temperature=0.7,
|
||
)
|
||
|
||
# Sub-agent that generates a price table for a product in a city
|
||
@tool
|
||
def get_price(product: str, city: str) -> str:
|
||
"""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=llm,
|
||
tools=[],
|
||
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: Продукт, Цена (руб.), Магазин.",
|
||
)
|
||
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 that uses get_price to build shopping list
|
||
main_agent = create_agent(
|
||
model=llm,
|
||
tools=[get_price],
|
||
system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цены каждого продукта.",
|
||
)
|
||
|
||
# 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 ---")
|