Stream-режим AI-агента: agent.py

solutions/699cc158d6d3a5544a3ed35b_Stream-режим_AI-агента/agent.py

@@ -1,32 +1,23 @@
-
 from langchain.agents import create_agent
-from langchain_community.chat_models import GigaChat
 from langchain.tools import tool
-import os
+from langchain_ollama import OllamaLLM
 
-GPT2GIGA_PORT=8090
-GIGACHAT_CREDENTIALS="MDE5YzBlMzUtNTJlYi03ODFiLTg1ZWUtZTc2MDFiZGUxYmM2OmNjZmIwODI2LTQ0ZTQtNDQwNC04NGE3LTgzNmM5ZDJhYmMzMg=="
-GIGACHAT_SCOPE="GIGACHAT_API_B2B"
-GIGACHAT_MODEL="GigaChat-MAX"
-GIGACHAT_VERIFY_SSL_CERTS=False
-
-llm = GigaChat(
-    credentials=GIGACHAT_CREDENTIALS,
-    scope=GIGACHAT_SCOPE,
-    model=GIGACHAT_MODEL,
-    verify_ssl_certs=False,  
+# --- LLM --------------------------------------------------------------------
+llm = OllamaLLM(
+    model="llama3",
     temperature=0.7,
     timeout=60,
 )
 
+# --- Инструмент -------------------------------------------------------------
 @tool
 def check_wish(wish: str) -> str:
     """Инструмент для проверки желания на наличие подвоха"""
-    
     genie_agent = create_agent(
         model=llm,
         tools=[],
-        system_prompt="""Ты - коварный джинн, который ищет подвох в любом желании.
+        system_prompt="""
+Ты - коварный джинн, который ищет подвох в любом желании.
 Проанализируй желание человека и найди скрытую опасность, буквальное толкование, неожиданные последствия.
 Если подвох найден - коротко предупреди о нем, например:
 "хочу много денег - деньги будут фальшивыми."
@@ -36,84 +27,102 @@ def check_wish(wish: str) -> str:
 Ты должен говорить только на русском языке."""
     )
 
-    result = genie_agent.invoke({
-        "messages": [
-            {"role": "human", "content": f"Проверь желание: {wish}"}
-        ]
-    })
+    result = genie_agent.invoke(
+        {
+            "messages": [
+                {"role": "human", "content": f"Проверь желание: {wish}"}
+            ]
+        }
+    )
+    return result["messages"][-1].content
 
-    return result['messages'][-1].content
 
+# --- Агент человека ---------------------------------------------------------
 human_agent = create_agent(
     model=llm,
     tools=[check_wish],
-    system_prompt="""Ты - человек, который загадывает желания джинну.
+    system_prompt="""
+Ты - человек, который загадывает желания джинну.
 Твоя задача - передать желание джинну через инструмент check_wish и сообщить результат.
 Говори только на русском языке."""
 )
 
+# --- Ввод -------------------------------------------------------------------
 current_wish = "Хочу читать мысли"
 
 print("ДЖИНН ГОТОВ ИСПОЛНЯТЬ ЖЕЛАНИЯ!")
 print(f"Человек: {current_wish}\n")
 print("Джинн:", end=" ", flush=True)
 
+
+# --- Форматирование ---------------------------------------------------------
 def format_message(message) -> str:
-    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']})"
+    """Возвращает строку, которую нужно вывести в консоль."""
+    if message.content:
+        return message.content
+    if message.tool_calls:
+        tool_call = message.tool_calls[0]
+        return f"{tool_call.name}({tool_call.args})"
     return ""
 
+
 step = 1
 
 def format_chunk_message(chunk):
+    """Обрабатывает чанк типа 'messages'."""
     global step
     message, meta = chunk
 
-    if meta['langgraph_step'] != step:
-        step = meta['langgraph_step']
-        print('\n --- --- --- \n')
+    if meta["langgraph_step"] != step:
+        step = meta["langgraph_step"]
+        print("\n --- --- --- \n")
 
-    if message.get('content'):
-        print(message['content'], end='', flush=False)
+    if message.content:
+        # выводим токен без перевода строки, чтобы текст «тёк»
+        print(message.content, end="", flush=True)
 
+
+# --- Потоковый вызов --------------------------------------------------------
 stream = human_agent.stream(
     {
         "messages": [
-            {"role": "human", "content": f"Вот мое желание: '{current_wish}'. Проверь его у джинна через инструмент check_wish и скажи мне результат."}
+            {
+                "role": "human",
+                "content": (
+                    f"Вот мое желание: '{current_wish}'. Проверь его у джинна "
+                    "через инструмент check_wish и скажи мне результат."
+                ),
+            }
         ]
     },
-    stream_mode=['messages', 'updates']
+    stream_mode=["messages", "updates"],
 )
 
 full_response = ""
-for chunk in stream:
-    chunk_type, chunk_data = chunk
 
-    if chunk_type == 'messages':
+for chunk_type, chunk_data in stream:
+    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)
+        full_response += message.content or ""
+    elif chunk_type == "updates":
+        # событие завершения шага (например, вызов инструмента)
+        if chunk_data.get("model"):
+            last_msg = chunk_data["model"]["messages"][-1]
+            formatted = format_message(last_msg)
             if formatted:
                 print(f"\n[Вызов инструмента: {formatted}]")
                 full_response += f"\n[Вызов инструмента: {formatted}]\n"
 
 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()
+        clean_response = full_response.replace("\n", " ").strip()
         print(f" Причина: {clean_response}")
     else:
         print(" Джинн устал...")