74 lines
2.7 KiB
Python
74 lines
2.7 KiB
Python
import json
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain.tools import tool
|
|
from langchain.agents import create_agent
|
|
from pydantic import SecretStr
|
|
|
|
# Конфигурация локальной модели LM Studio
|
|
llm = ChatOpenAI(
|
|
model="gpt-4o-mini",
|
|
base_url="http://localhost:1234/v1",
|
|
api_key=SecretStr("fake"),
|
|
temperature=0.7,
|
|
)
|
|
|
|
@tool
|
|
def get_price(product: str, city: str) -> str:
|
|
"""Получить примерную цену продукта в указанном городе."""
|
|
# Создаём субагента, который генерирует таблицу цен
|
|
sub_agent = create_agent(
|
|
model=llm,
|
|
tools=[],
|
|
system_prompt=(
|
|
f"Ты генератор цены для продукта {product} в городе {city}. "
|
|
"Сгенерируй таблицу в формате:\n"
|
|
"| Продукт | Цена (руб.) | Магазин |\n"
|
|
"|---------|-------------|---------|\n"
|
|
f"Например:\n"
|
|
f"| {product} | 100 | Магнит |\n"
|
|
"Ответ должен содержать только таблицу без лишних слов."
|
|
),
|
|
)
|
|
# Запускаем субагента
|
|
result = sub_agent.invoke(
|
|
{"messages": [{"role": "user", "content": f"Сгенерируй цену для {product} в {city}"}]}
|
|
)
|
|
# Извлекаем текст ответа
|
|
content = result["messages"][-1]["content"]
|
|
return content
|
|
|
|
# Главный агент
|
|
main_agent = create_agent(
|
|
model=llm,
|
|
tools=[get_price],
|
|
system_prompt="Ты помощник по планированию покупок.",
|
|
)
|
|
|
|
def format_message(message: dict) -> str:
|
|
"""Форматируем сообщение для вывода."""
|
|
if "content" in message and message["content"]:
|
|
return message["content"]
|
|
if "tool_calls" in message and message["tool_calls"]:
|
|
calls = []
|
|
for call in message["tool_calls"]:
|
|
name = call["name"]
|
|
args_str = call.get("arguments", "{}")
|
|
try:
|
|
args = json.loads(args_str)
|
|
except json.JSONDecodeError:
|
|
args = args_str
|
|
calls.append(f"{name}({args})")
|
|
return " | ".join(calls)
|
|
return ""
|
|
|
|
def main():
|
|
user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
|
|
result = main_agent.invoke(
|
|
{"messages": [{"role": "human", "content": user_query}]}
|
|
)
|
|
for msg in result["messages"]:
|
|
print(format_message(msg))
|
|
print("---")
|
|
|
|
if __name__ == "__main__":
|
|
main() |