diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..d46fbe4 --- /dev/null +++ b/agent.py @@ -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()