98 lines
2.9 KiB
Python
98 lines
2.9 KiB
Python
from langchain.agents import create_agent
|
|
from langchain.tools import tool
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from pydantic import SecretStr
|
|
|
|
|
|
llm = ChatOpenAI(
|
|
model="local-model",
|
|
base_url="http://localhost:1234/v1",
|
|
api_key=SecretStr("fake"),
|
|
temperature=0.7,
|
|
)
|
|
|
|
|
|
@tool
|
|
def get_price(product: str, city: str) -> str:
|
|
"""Возвращает примерную цену продукта в указанном городе."""
|
|
|
|
price_agent = create_agent(
|
|
model=llm,
|
|
tools=[],
|
|
system_prompt=(
|
|
"Ты субагент оценки цен. Верни одну markdown-таблицу "
|
|
"с колонками: Продукт, Цена (руб.), Магазин."
|
|
),
|
|
)
|
|
answer = price_agent.invoke(
|
|
{
|
|
"messages": [
|
|
{
|
|
"role": "human",
|
|
"content": (
|
|
f"Оцени реалистичную цену продукта '{product}' "
|
|
f"в городе {city}, опираясь на исторические данные о ценах."
|
|
),
|
|
}
|
|
]
|
|
}
|
|
)
|
|
return answer["messages"][-1].content
|
|
|
|
|
|
def format_message(message) -> str:
|
|
if message.content:
|
|
return message.content
|
|
if message.tool_calls:
|
|
call = message.tool_calls[0]
|
|
return f"{call['name']}({call['args']})"
|
|
return ""
|
|
|
|
|
|
def stream_shopping_agent() -> None:
|
|
memory = MemorySaver()
|
|
agent = create_agent(
|
|
model=llm,
|
|
tools=[get_price],
|
|
system_prompt="Ты помощник по планированию покупок",
|
|
checkpointer=memory,
|
|
)
|
|
|
|
stream = agent.stream(
|
|
{
|
|
"messages": [
|
|
{
|
|
"role": "human",
|
|
"content": (
|
|
"Помоги составить список покупок: молоко, хлеб, яблоки. "
|
|
"Я нахожусь в Казани."
|
|
),
|
|
}
|
|
]
|
|
},
|
|
stream_mode=["messages", "updates"],
|
|
config={"configurable": {"thread_id": "shopping-stream-demo"}},
|
|
)
|
|
|
|
step = None
|
|
for chunk_type, chunk_data in stream:
|
|
if chunk_type == "messages":
|
|
message, meta = chunk_data
|
|
current_step = meta.get("langgraph_step")
|
|
if current_step != step:
|
|
step = current_step
|
|
print("\n --- --- --- \n")
|
|
if message.content:
|
|
print(message.content, end="", flush=True)
|
|
|
|
if chunk_type == "updates" and chunk_data.get("model"):
|
|
last_message = chunk_data["model"]["messages"][-1]
|
|
formatted = format_message(last_message)
|
|
if formatted:
|
|
print(formatted)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
stream_shopping_agent()
|