Обновить solutions/699cc158d6d3a5544a3ed35b_Stream-режим_AI-агента/agent.py
This commit is contained in:
@@ -1,89 +1,119 @@
|
||||
|
||||
from langchain.agents import create_agent
|
||||
from langchain_community.chat_models import GigaChat
|
||||
from langchain.tools import tool
|
||||
import os
|
||||
from typing import Dict, Any
|
||||
|
||||
from langchain_ollama import ChatOllama
|
||||
from nomic_embed_text import NomicEmbedText
|
||||
from langchain.tools import BaseTool
|
||||
from langchain.schema import HumanMessage, AIMessage, SystemMessage
|
||||
from langchain.agents import AgentExecutor, initialize_agent, load_tools
|
||||
from langchain.prompts import ChatPromptTemplate
|
||||
GPT2GIGA_PORT=8090
|
||||
GIGACHAT_CREDENTIALS="MDE5YzBlMzUtNTJlYi03ODFiLTg1ZWUtZTc2MDFiZGUxYmM2OmNjZmIwODI2LTQ0ZTQtNDQwNC04NGE3LTgzNmM5ZDJhYmMzMg=="
|
||||
GIGACHAT_SCOPE="GIGACHAT_API_B2B"
|
||||
GIGACHAT_MODEL="GigaChat-MAX"
|
||||
GIGACHAT_VERIFY_SSL_CERTS=False
|
||||
|
||||
# ---------- Настройки LLM и эмбеддингов ----------
|
||||
LLM_MODEL = os.getenv("OLLAMA_MODEL", "llama3")
|
||||
EMBEDDING_MODEL = os.getenv("NOMIC_EMBEDDING_MODEL", "nomic-embed-text")
|
||||
llm = GigaChat(
|
||||
credentials=GIGACHAT_CREDENTIALS,
|
||||
scope=GIGACHAT_SCOPE,
|
||||
model=GIGACHAT_MODEL,
|
||||
verify_ssl_certs=False,
|
||||
temperature=0.7,
|
||||
timeout=60,
|
||||
)
|
||||
|
||||
llm = ChatOllama(model=LLM_MODEL, temperature=0.2)
|
||||
embedding = NomicEmbedText()
|
||||
@tool
|
||||
def check_wish(wish: str) -> str:
|
||||
"""Инструмент для проверки желания на наличие подвоха"""
|
||||
|
||||
genie_agent = create_agent(
|
||||
model=llm,
|
||||
tools=[],
|
||||
system_prompt="""Ты - коварный джинн, который ищет подвох в любом желании.
|
||||
Проанализируй желание человека и найди скрытую опасность, буквальное толкование, неожиданные последствия.
|
||||
Если подвох найден - коротко предупреди о нем, например:
|
||||
"хочу много денег - деньги будут фальшивыми."
|
||||
"хочу деньги на счет в банке - хорошо, но банк завтра обанкротится"
|
||||
|
||||
# ---------- Пример инструмента ----------
|
||||
class DummyTool(BaseTool):
|
||||
name: str = "get_price"
|
||||
description: str = (
|
||||
"Получает цену товара в указанном городе. "
|
||||
"Аргументы: {'product': 'название', 'city': 'город'}."
|
||||
Если желание безопасно и не имеет подвоха - ответь "Желание безопасно! Исполняю."
|
||||
Ты должен говорить только на русском языке."""
|
||||
)
|
||||
|
||||
result = genie_agent.invoke({
|
||||
"messages": [
|
||||
{"role": "human", "content": f"Проверь желание: {wish}"}
|
||||
]
|
||||
})
|
||||
|
||||
return result['messages'][-1].content
|
||||
|
||||
def _run(self, product: str, city: str) -> str:
|
||||
# В реальном коде здесь будет запрос к API
|
||||
return f"Цена {product} в {city}: 100₽"
|
||||
|
||||
tool = DummyTool()
|
||||
tools = [tool]
|
||||
|
||||
# ---------- Создание агента ----------
|
||||
prompt_template = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
SystemMessage(content="Ты помощник. Используй инструменты при необходимости."),
|
||||
("human", "{input}"),
|
||||
]
|
||||
human_agent = create_agent(
|
||||
model=llm,
|
||||
tools=[check_wish],
|
||||
system_prompt="""Ты - человек, который загадывает желания джинну.
|
||||
Твоя задача - передать желание джинну через инструмент check_wish и сообщить результат.
|
||||
Говори только на русском языке."""
|
||||
)
|
||||
|
||||
agent_executor = initialize_agent(
|
||||
tools,
|
||||
llm,
|
||||
agent="openai-tools",
|
||||
verbose=False,
|
||||
prompt=prompt_template,
|
||||
)
|
||||
current_wish = "Хочу читать мысли"
|
||||
|
||||
print("ДЖИНН ГОТОВ ИСПОЛНЯТЬ ЖЕЛАНИЯ!")
|
||||
print(f"Человек: {current_wish}\n")
|
||||
print("Джинн:", end=" ", flush=True)
|
||||
|
||||
# ---------- Функции форматирования ----------
|
||||
def format_message(message) -> str:
|
||||
if message.content:
|
||||
return message.content
|
||||
# Если нет content, выводим вызов инструмента
|
||||
call = message.tool_calls[0]
|
||||
return f"{call['name']}({call['args']})"
|
||||
if message.get('content'):
|
||||
return message['content']
|
||||
elif message.get('tool_calls'):
|
||||
tool_call = message['tool_calls'][0]
|
||||
return f"{tool_call['name']}({tool_call['args']})"
|
||||
return ""
|
||||
|
||||
step = 1
|
||||
|
||||
def format_chunk_message(chunk):
|
||||
global step
|
||||
message, meta = chunk
|
||||
if meta.get("langgraph_step") != step:
|
||||
step = meta["langgraph_step"]
|
||||
print("\n --- --- --- \n")
|
||||
if message.content:
|
||||
print(message.content, end="", flush=True)
|
||||
|
||||
if meta['langgraph_step'] != step:
|
||||
step = meta['langgraph_step']
|
||||
print('\n --- --- --- \n')
|
||||
|
||||
if message.get('content'):
|
||||
print(message['content'], end='', flush=False)
|
||||
|
||||
# ---------- Запуск агента в режиме stream ----------
|
||||
def run_agent(user_input: str):
|
||||
global step
|
||||
step = 1
|
||||
# Инициализируем состояние с пользовательским сообщением
|
||||
init_state = {"input": user_input}
|
||||
# В LangChain 0.2 AgentExecutor имеет метод stream()
|
||||
stream = agent_executor.stream(init_state, stream_mode=["messages", "updates"])
|
||||
stream = human_agent.stream(
|
||||
{
|
||||
"messages": [
|
||||
{"role": "human", "content": f"Вот мое желание: '{current_wish}'. Проверь его у джинна через инструмент check_wish и скажи мне результат."}
|
||||
]
|
||||
},
|
||||
stream_mode=['messages', 'updates']
|
||||
)
|
||||
|
||||
for chunk_type, chunk_data in stream:
|
||||
if chunk_type == "messages":
|
||||
format_chunk_message(chunk_data)
|
||||
elif chunk_type == "updates":
|
||||
# При завершении шага выводим итоговое сообщение
|
||||
model_info = chunk_data.get("model")
|
||||
if model_info:
|
||||
last_msg = model_info["messages"][-1]
|
||||
print("\n" + format_message(last_msg) + "\n")
|
||||
full_response = ""
|
||||
for chunk in stream:
|
||||
chunk_type, chunk_data = chunk
|
||||
|
||||
if chunk_type == 'messages':
|
||||
format_chunk_message(chunk_data)
|
||||
message, _ = chunk_data
|
||||
if message.get('content'):
|
||||
full_response += message['content']
|
||||
|
||||
elif chunk_type == 'updates':
|
||||
if chunk_data.get('model', None):
|
||||
last_message = chunk_data['model']['messages'][-1]
|
||||
formatted = format_message(last_message)
|
||||
if formatted:
|
||||
print(f"\n[Вызов инструмента: {formatted}]")
|
||||
full_response += f"\n[Вызов инструмента: {formatted}]\n"
|
||||
|
||||
if __name__ == "__main__":
|
||||
user_query = input("Введите запрос: ")
|
||||
run_agent(user_query)
|
||||
print("\n")
|
||||
|
||||
if "исполняю" in full_response.lower() or "безопасно" in full_response.lower():
|
||||
print(f" Желание исполнено! Финальная версия: {current_wish}")
|
||||
else:
|
||||
print(f" Джинн отказался исполнять желание: {current_wish}")
|
||||
if full_response:
|
||||
clean_response = full_response.replace('\n', ' ').strip()
|
||||
print(f" Причина: {clean_response}")
|
||||
else:
|
||||
print(" Джинн устал...")
|
||||
|
||||
Reference in New Issue
Block a user