diff --git a/main.py b/main.py new file mode 100644 index 0000000..1f28d53 --- /dev/null +++ b/main.py @@ -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