89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
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
|
||
|
||
# ---------- Настройки LLM и эмбеддингов ----------
|
||
LLM_MODEL = os.getenv("OLLAMA_MODEL", "llama3")
|
||
EMBEDDING_MODEL = os.getenv("NOMIC_EMBEDDING_MODEL", "nomic-embed-text")
|
||
|
||
llm = ChatOllama(model=LLM_MODEL, temperature=0.2)
|
||
embedding = NomicEmbedText()
|
||
|
||
# ---------- Пример инструмента ----------
|
||
class DummyTool(BaseTool):
|
||
name: str = "get_price"
|
||
description: str = (
|
||
"Получает цену товара в указанном городе. "
|
||
"Аргументы: {'product': 'название', 'city': 'город'}."
|
||
)
|
||
|
||
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}"),
|
||
]
|
||
)
|
||
|
||
agent_executor = initialize_agent(
|
||
tools,
|
||
llm,
|
||
agent="openai-tools",
|
||
verbose=False,
|
||
prompt=prompt_template,
|
||
)
|
||
|
||
# ---------- Функции форматирования ----------
|
||
def format_message(message) -> str:
|
||
if message.content:
|
||
return message.content
|
||
# Если нет content, выводим вызов инструмента
|
||
call = message.tool_calls[0]
|
||
return f"{call['name']}({call['args']})"
|
||
|
||
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)
|
||
|
||
# ---------- Запуск агента в режиме 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"])
|
||
|
||
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")
|
||
|
||
if __name__ == "__main__":
|
||
user_query = input("Введите запрос: ")
|
||
run_agent(user_query) |