91 lines
3.8 KiB
Python
91 lines
3.8 KiB
Python
from langchain_openai import ChatOpenAI
|
|
from langchain.agents import create_agent
|
|
from langchain.tools import tool
|
|
import re
|
|
import json
|
|
|
|
# Initialize main LLM
|
|
llm = ChatOpenAI(
|
|
base_url="http://localhost:11434/v1",
|
|
api_key="ollama",
|
|
model="<название модели в LM Studio>",
|
|
temperature=0.7,
|
|
)
|
|
|
|
# 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:
|
|
# For demonstration, use a subagent that simply returns a static price
|
|
sub_llm = ChatOpenAI(
|
|
base_url="http://localhost:11434/v1",
|
|
api_key="ollama",
|
|
model="<название модели в LM Studio>",
|
|
temperature=0.2,
|
|
)
|
|
sub_agent = create_agent(
|
|
llm=sub_llm,
|
|
tools=[],
|
|
system_prompt="You are a simple price lookup tool. Return price as 'price: <amount>, store: <store>'.",
|
|
verbose=False,
|
|
)
|
|
response = sub_agent.invoke({"input": f"Give price for {product} in {city}."})
|
|
# The sub_agent will return a dict with messages; extract content
|
|
if isinstance(response, dict) and "messages" in response:
|
|
# Take the first assistant message content
|
|
for msg in response["messages"]:
|
|
if msg.get("role") == "assistant" and msg.get("content"):
|
|
return msg["content"].strip()
|
|
# Fallback to static response if sub_agent fails
|
|
return "price: 100, store: SuperMarket"
|
|
|
|
|
|
def main():
|
|
# Create main agent
|
|
agent = create_agent(
|
|
llm=llm,
|
|
tools=[get_price],
|
|
system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цен.",
|
|
verbose=False,
|
|
)
|
|
user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
|
|
result = agent.invoke({"input": user_query})
|
|
messages = result.get("messages", [])
|
|
# Print each message and tool calls
|
|
for msg in messages:
|
|
if msg.get("role") == "assistant":
|
|
if msg.get("content"):
|
|
print(msg["content"])
|
|
if 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}")
|
|
# Build table from tool calls
|
|
price_entries = []
|
|
total = 0.0
|
|
for msg in messages:
|
|
if msg.get("role") == "assistant" and msg.get("tool_calls"):
|
|
for call in msg["tool_calls"]:
|
|
if call.get("name") == "get_price":
|
|
arg_str = call.get("arguments") or call.get("function", {}).get("arguments")
|
|
try:
|
|
args = json.loads(arg_str)
|
|
product = args.get("product")
|
|
except Exception:
|
|
product = "unknown"
|
|
result_text = call.get("function", {}).get("arguments", "")
|
|
m = re.search(r"price:\s*([0-9]+)\s*,\s*store:\s*([A-Za-z0-9 ]+)", result_text, re.IGNORECASE)
|
|
if m:
|
|
price = float(m.group(1))
|
|
store = m.group(2).strip()
|
|
price_entries.append((product, price, store))
|
|
total += price
|
|
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()
|