diff --git a/dz/lab01/simpleAI.py b/dz/lab01/simpleAI.py new file mode 100644 index 0000000..e8f4776 --- /dev/null +++ b/dz/lab01/simpleAI.py @@ -0,0 +1,106 @@ +import uuid + +from langchain.agents import create_agent +from langchain_openai import ChatOpenAI +from langchain.tools import tool +from langgraph.checkpoint.memory import MemorySaver +from rich.console import Console + +console = Console() +llm = ChatOpenAI( + model = 'openai/gpt-oss-20b', + base_url= 'http://localhost:1234/v1', + temperature=.7, + api_key=('fake') +) + +@tool +def get_price(product: str, city:str) -> str: + """Это инструмент, который по городу и списку продуктов считает итоговую стоимость корзины""" + product_agent = create_agent( + model=llm, + system_prompt=""" + Требуется подсчитать итоговую стоимость корзины пользователя для указанного города в виде таблицы + | Продукт | Цена (руб.) | Магазин | + + Если данных нет, то сформируй реалистичный ответ, заполни все ячейки таблицы. + +""" + ) + answer = product_agent.invoke({ + "messages":[ + { + "role": "human", + "content": f"Какая итоговая стоимость корзины {product} в городе {city}?" + } + ] + }) + return answer['messages'][-1].content + +memory = MemorySaver() + +agent = create_agent( + model=llm, + tools=[get_price], + system_prompt="Ты помощник по планированию покупок", + checkpointer=memory, + interrupt_before=['tools'] +) + +config = {"configurable": {"thread_id": "user_thread"}} + +user_input = input('\nВы: ') + +step = 1 + +def format_chunk_message(chunk): #вынести потом в отдельный класс бы, а то и так кода многовато как-то :) + message, meta = chunk + global step + + if meta['langgraph_step'] != step: + step = meta['langgraph_step'] + console.print('\n --- --- --- \n') + + if message.content: + console.print(message.content, end='') + +def format_message(message) -> str: + """Формат сообщения для вывода""" + if message.content: + return message.content + return f"{message.tool_calls[0]['name']}({message.tool_calls[0]['args']})" + +def ask_and_run(user_input,config): + for chunk in agent.stream(user_input, config=config, stream_mode=['messages', 'updates']): + state = agent.get_state(config) + chunk_type, chunk_data = chunk + + if chunk_type == 'messages': + format_chunk_message(chunk_data) + + if chunk_type == 'updates': + if chunk_data.get('model',None): + console.print(format_message(chunk_data['model']['messages'][-1]), sep='\n---\n') + + if '__interrupt__' in chunk_data and state.next == ('tools',): + tool_call = state.values['messages'][-1].tool_calls[0] + console.print(f'Агент хочет вызвать тулзу {tool_call['name']}({tool_call['args']})') + + answer = input('\nРазрешить? (Y/N): ') + + if answer.lower().strip() == 'y': + ask_and_run(None, config) + else: + console.print('Отменено') + break + +while True: + if user_input == 'exit': + break + + ask_and_run( + {"messages": [{"role": "human", "content": user_input}]}, + config, + ) + +