116 lines
4.5 KiB
Python
116 lines
4.5 KiB
Python
# agent.py
|
|
|
|
"""
|
|
Простой LangChain‑агент с поддержкой потокового вывода через `agent.stream()`.
|
|
Внутри используется один инструмент – get_price, который вызывает под‑агента,
|
|
обрабатывающего запросы к «базе цен».
|
|
"""
|
|
|
|
from langchain.agents import create_agent
|
|
from langchain.tools import tool
|
|
from rich import print as rprint
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 1. Определяем LLM (Ollama) – это единственный доступный LLM в задании.
|
|
# --------------------------------------------------------------------------- #
|
|
try:
|
|
from langchain_ollama import ChatOllama
|
|
except ImportError: # если Ollama недоступен, используем простой Mock‑LLM
|
|
class DummyLLM:
|
|
def __call__(self, *args, **kwargs):
|
|
return "Mock response"
|
|
|
|
llm = DummyLLM()
|
|
else:
|
|
llm = ChatOllama(model="llama3.1") # пример модели
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 2. Создаём инструмент get_price.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
@tool
|
|
def get_price(product: str, city: str) -> str:
|
|
"""
|
|
Возвращает таблицу цен на указанный продукт в заданном городе.
|
|
Для демонстрации используется под‑агент, который формирует ответ.
|
|
"""
|
|
# Под‑агент – простая функция, которая возвращает строку
|
|
sub_agent = create_agent(
|
|
llm=llm,
|
|
tools=[],
|
|
verbose=False,
|
|
agent_type="openai-tools",
|
|
)
|
|
|
|
# Запускаем под‑агента в режиме invoke (т.к. он короткий)
|
|
result = sub_agent.invoke({
|
|
"messages": [
|
|
{"role": "human", "content": f"Сколько стоит {product} в городе {city}?"}
|
|
]
|
|
})
|
|
return result["output"]
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 3. Создаём главный агент с инструментом get_price.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
agent = create_agent(
|
|
llm=llm,
|
|
tools=[get_price],
|
|
verbose=False,
|
|
agent_type="openai-tools",
|
|
)
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 4. Функции для форматирования и вывода чанков.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
step = 1
|
|
|
|
def format_chunk_message(chunk):
|
|
"""
|
|
Выводит токен текста без перевода строки.
|
|
При смене шага печатает разделитель.
|
|
"""
|
|
global step
|
|
message, meta = chunk
|
|
if meta["langgraph_step"] != step:
|
|
step = meta["langgraph_step"]
|
|
rprint("\n--- --- ---\n")
|
|
if message.content:
|
|
print(message.content, end="", flush=True)
|
|
|
|
def format_message(message):
|
|
"""
|
|
Форматирует сообщение для вывода в режиме invoke.
|
|
"""
|
|
if message.content:
|
|
return message.content
|
|
# Если нет content – это вызов инструмента
|
|
tool_call = message.tool_calls[0]
|
|
return f"{tool_call['name']}({tool_call['args']})"
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 5. Запускаем потоковый вывод.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
if __name__ == "__main__":
|
|
# Пример запроса к агенту
|
|
stream = agent.stream(
|
|
{
|
|
"messages": [
|
|
{"role": "human", "content": "Покажи мне цены на молоко и хлеб в Казани."}
|
|
]
|
|
},
|
|
stream_mode=["messages", "updates"],
|
|
)
|
|
|
|
for chunk in stream:
|
|
chunk_type, chunk_data = chunk
|
|
if chunk_type == "messages":
|
|
format_chunk_message(chunk_data)
|
|
elif chunk_type == "updates":
|
|
# При завершении шага выводим итоговое сообщение
|
|
if chunk_data.get("model"):
|
|
last_msg = chunk_data["model"]["messages"][-1]
|
|
rprint("\n" + format_message(last_msg) + "\n") |