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()