Add agent.py

This commit is contained in:
2026-05-28 11:28:31 +00:00
parent 840fb21754
commit b92422e01a
+73
View File
@@ -0,0 +1,73 @@
from langchain_openai import ChatOpenAI
from langchain.schema import format_message
from langchain.tools import tool
from langchain.agents import create_agent
@tool
def get_price(product: str, city: str) -> str:
"""Return a price table for the specified product in the given city."""
llm = ChatOpenAI(
base_url="http://localhost:1234/v1",
api_key="ollama",
model="<название модели в LM Studio>",
temperature=0.7,
)
system_prompt = f"Generate a price table for {product} in {city}."
sub_agent = create_agent(
llm=llm,
system_prompt=system_prompt,
tools=[],
)
response = sub_agent.invoke(
messages=[{"role": "user", "content": "Please provide the price table."}]
)
return response.content
def main() -> None:
llm = ChatOpenAI(
base_url="http://localhost:1234/v1",
api_key="ollama",
model="<название модели в LM Studio>",
temperature=0.7,
)
agent = create_agent(
llm=llm,
system_prompt="Ты помощник по планированию покупок",
tools=[get_price],
)
user_query = input("Введите запрос: ")
attempts = 2
for attempt in range(attempts):
try:
stream = agent.stream(
messages=[{"role": "user", "content": user_query}],
stream_mode=["messages", "updates"],
)
for chunk in stream:
if chunk.type == "messages":
message = chunk.data[0]
# Print incremental content as it arrives character by character
if hasattr(message, "tool_calls") and message.tool_calls:
formatted = format_message(message)
for ch in formatted:
print(ch, end="")
else:
for ch in message.content:
print(ch, end="")
# Separator after each message chunk
print("\n--- --- ---\n")
elif chunk.type == "updates":
update = chunk.data[0]
if update.get("model") == "tool" and "content" in update:
for ch in update["content"]:
print(ch, end="")
print("\n--- --- ---\n")
break
except Exception as e:
print(f"\nError during streaming (attempt {attempt + 1}): {e}")
if attempt == attempts - 1:
raise
if __name__ == "__main__":
main()