108 lines
4.0 KiB
Python
108 lines
4.0 KiB
Python
from langchain_openai import ChatOpenAI
|
|
from langchain.agents import create_agent
|
|
from langchain.tools import tool
|
|
from pydantic import SecretStr
|
|
import re
|
|
import sys
|
|
|
|
# Utility to initialize LLM with error handling
|
|
|
|
def init_llm(base_url: str, model: str, api_key: str = "ollama") -> ChatOpenAI:
|
|
try:
|
|
return ChatOpenAI(
|
|
base_url=base_url,
|
|
api_key=SecretStr(api_key),
|
|
model=model,
|
|
temperature=0.7,
|
|
)
|
|
except Exception as e:
|
|
print(f"Ошибка при подключении к LLM: {e}")
|
|
sys.exit(1)
|
|
|
|
# Main LLM
|
|
LLM_URL = "http://localhost:1234/v1"
|
|
LLM_MODEL = "llama3"
|
|
llm = init_llm(LLM_URL, LLM_MODEL)
|
|
|
|
# Define get_price tool with subagent
|
|
@tool(name="get_price", description="Get price for a product in a city.")
|
|
def get_price(product: str, city: str) -> str:
|
|
"""Subagent that looks up price for a product in a given city.
|
|
For demonstration it returns a static price string, but the structure
|
|
mirrors a real subagent that could call another LLM.
|
|
"""
|
|
sub_llm = init_llm(LLM_URL, LLM_MODEL, api_key="ollama")
|
|
sub_agent = create_agent(
|
|
llm=sub_llm,
|
|
tools=[],
|
|
system_prompt="You are a simple price lookup tool. Return price as 'price: <amount>, store: <store>'.",
|
|
)
|
|
try:
|
|
response = sub_agent.invoke({"input": f"Give price for {product} in {city}."})
|
|
if isinstance(response, dict) and "messages" in response:
|
|
for msg in response["messages"]:
|
|
if msg.get("role") == "assistant" and msg.get("content"):
|
|
return msg["content"].strip()
|
|
except Exception as e:
|
|
print(f"Ошибка при вызове subagent: {e}")
|
|
return "price: 100, store: SuperMarket"
|
|
|
|
|
|
def format_message(message) -> str:
|
|
if message.get("content"):
|
|
return message["content"]
|
|
if message.get("tool_calls"):
|
|
parts = []
|
|
for call in message["tool_calls"]:
|
|
name = call.get("name")
|
|
args = call.get("arguments") or call.get("function", {}).get("arguments")
|
|
parts.append(f"{name}({args})")
|
|
return ", ".join(parts)
|
|
return ""
|
|
|
|
|
|
def main():
|
|
agent = create_agent(
|
|
llm=llm,
|
|
tools=[get_price],
|
|
system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цен.",
|
|
)
|
|
user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
|
|
result = agent.invoke({"input": user_query})
|
|
messages = result.get("messages", [])
|
|
|
|
# Print tool calls in order
|
|
for msg in messages:
|
|
if msg.get("role") == "assistant" and msg.get("tool_calls"):
|
|
for call in msg["tool_calls"]:
|
|
name = call.get("name")
|
|
args = call.get("arguments") or call.get("function", {}).get("arguments")
|
|
print(f"Tool call: {name} with args {args}")
|
|
|
|
# Aggregate prices by calling the tool directly for each product
|
|
products = ["молоко", "хлеб", "яблоки"]
|
|
city = "Казань"
|
|
price_entries = []
|
|
total = 0.0
|
|
for prod in products:
|
|
try:
|
|
price_str = get_price(prod, city)
|
|
m = re.search(r"price:\s*([0-9]+)\s*,\s*store:\s*([A-Za-z0-9 ]+)", price_str, re.IGNORECASE)
|
|
if m:
|
|
price = float(m.group(1))
|
|
store = m.group(2).strip()
|
|
price_entries.append((prod, price, store))
|
|
total += price
|
|
except Exception:
|
|
continue
|
|
|
|
# Print final table
|
|
print("\nТаблица цен:")
|
|
print("{:<15} {:<10} {:<15}".format("Товар", "Цена", "Магазин"))
|
|
for prod, price, store in price_entries:
|
|
print("{:<15} {:<10} {:<15}".format(prod, f"{price} руб", store))
|
|
print(f"\nОбщая стоимость: {total} руб")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|