59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
import os
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain.tools import tool
|
|
from langchain.agents import create_agent
|
|
from pydantic import SecretStr
|
|
|
|
# 1. Connect to local LLM
|
|
llm = ChatOpenAI(
|
|
model="<название модели в LM Studio>",
|
|
base_url="http://localhost:1234/v1",
|
|
api_key=SecretStr("fake"),
|
|
temperature=0.7,
|
|
)
|
|
|
|
# 2. Sub-agent tool to get price
|
|
@tool
|
|
def get_price(product: str, city: str) -> str:
|
|
"""Return a realistic price table for a product in a city."""
|
|
# Sub-agent that generates a price table
|
|
sub_llm = ChatOpenAI(
|
|
model="<название модели в LM Studio>",
|
|
base_url="http://localhost:1234/v1",
|
|
api_key=SecretStr("fake"),
|
|
temperature=0.7,
|
|
)
|
|
sub_agent = create_agent(
|
|
model=sub_llm,
|
|
tools=[],
|
|
system_prompt=f"You are a price estimator for {city}. Provide a table with columns: Продукт, Цена (руб.), Магазин. Use realistic Russian prices.",
|
|
)
|
|
prompt = f"Generate a price table for product '{product}' in city '{city}'."
|
|
result = sub_agent.invoke({"messages": [{"role": "human", "content": prompt}]} )
|
|
# Extract the last message content
|
|
return result["messages"][-1]["content"]
|
|
|
|
# 3. Main agent
|
|
main_agent = create_agent(
|
|
model=llm,
|
|
tools=[get_price],
|
|
system_prompt="Ты помощник по планированию покупок.",
|
|
)
|
|
|
|
# 4. Query
|
|
question = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
|
|
response = main_agent.invoke({"messages": [{"role": "human", "content": question}]})
|
|
|
|
# 5. 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']}({call['args']})")
|
|
else:
|
|
print(msg)
|
|
|
|
if __name__ == "__main__":
|
|
pass
|