55 lines
2.0 KiB
Python
55 lines
2.0 KiB
Python
from langchain_openai import ChatOpenAI
|
||
from langchain.tools import tool
|
||
from langchain.agents import create_agent, AgentExecutor
|
||
import json
|
||
|
||
# Configure LLM
|
||
llm = ChatOpenAI(
|
||
model="gpt-4o-mini", # replace with your local model name
|
||
base_url="http://localhost:1234/v1",
|
||
api_key="fake",
|
||
temperature=0.7,
|
||
)
|
||
|
||
@tool("Get price for a product in a city")
|
||
def get_price(product: str, city: str) -> str:
|
||
"""
|
||
Returns a table with product, price and store.
|
||
The function internally creates a sub-agent that generates realistic prices.
|
||
"""
|
||
# Sub‑agent to generate price
|
||
sub_llm = ChatOpenAI(
|
||
model="gpt-4o-mini",
|
||
base_url="http://localhost:1234/v1",
|
||
api_key="fake",
|
||
temperature=0.5,
|
||
)
|
||
sub_agent = create_agent(
|
||
model=sub_llm,
|
||
tools=[],
|
||
system_prompt=f"You are a price estimator for {city}. Provide a realistic price and store name for {product} in a table format.",
|
||
)
|
||
result = sub_agent.invoke({"messages": [{"role": "human", "content": f"Give me the price of {product} in {city}"}]})
|
||
return json.dumps(result["messages"][-1]["content"], ensure_ascii=False)
|
||
|
||
# Main agent with get_price tool
|
||
main_agent = create_agent(
|
||
model=llm,
|
||
tools=[get_price],
|
||
system_prompt="You are a shopping assistant. Use the get_price tool to help users plan their purchases.",
|
||
)
|
||
|
||
if __name__ == "__main__":
|
||
user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
|
||
response = main_agent.invoke({"messages": [{"role": "human", "content": user_query}]})
|
||
# Print all messages
|
||
for msg in response["messages"]:
|
||
if "content" in msg:
|
||
print(msg["content"])
|
||
elif "tool_calls" in msg:
|
||
for call in msg["tool_calls"]:
|
||
print(f"{call['name']}({json.dumps(call['args'])})")
|
||
# Final answer
|
||
final = response["messages"][-1]["content"] if "content" in response["messages"][-1] else ""
|
||
print("\nFinal answer:\n", final)
|