Add main.py
This commit is contained in:
@@ -0,0 +1,58 @@
|
||||
import os
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain.tools import tool
|
||||
from langchain.agents import create_agent
|
||||
from pydantic import SecretStr
|
||||
|
||||
# 1. Connect to local LLM
|
||||
llm = ChatOpenAI(
|
||||
model="<название модели в LM Studio>",
|
||||
base_url="http://localhost:1234/v1",
|
||||
api_key=SecretStr("fake"),
|
||||
temperature=0.7,
|
||||
)
|
||||
|
||||
# 2. Sub-agent tool to get price
|
||||
@tool
|
||||
def get_price(product: str, city: str) -> str:
|
||||
"""Return a realistic price table for a product in a city."""
|
||||
# Sub-agent that generates a price table
|
||||
sub_llm = ChatOpenAI(
|
||||
model="<название модели в LM Studio>",
|
||||
base_url="http://localhost:1234/v1",
|
||||
api_key=SecretStr("fake"),
|
||||
temperature=0.7,
|
||||
)
|
||||
sub_agent = create_agent(
|
||||
model=sub_llm,
|
||||
tools=[],
|
||||
system_prompt=f"You are a price estimator for {city}. Provide a table with columns: Продукт, Цена (руб.), Магазин. Use realistic Russian prices.",
|
||||
)
|
||||
prompt = f"Generate a price table for product '{product}' in city '{city}'."
|
||||
result = sub_agent.invoke({"messages": [{"role": "human", "content": prompt}]} )
|
||||
# Extract the last message content
|
||||
return result["messages"][-1]["content"]
|
||||
|
||||
# 3. Main agent
|
||||
main_agent = create_agent(
|
||||
model=llm,
|
||||
tools=[get_price],
|
||||
system_prompt="Ты помощник по планированию покупок.",
|
||||
)
|
||||
|
||||
# 4. Query
|
||||
question = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
|
||||
response = main_agent.invoke({"messages": [{"role": "human", "content": question}]})
|
||||
|
||||
# 5. 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']}({call['args']})")
|
||||
else:
|
||||
print(msg)
|
||||
|
||||
if __name__ == "__main__":
|
||||
pass
|
||||
Reference in New Issue
Block a user