84 lines
3.1 KiB
Python
84 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 ToolExecutor, AgentExecutor, create_openai_tools_agent
|
|
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]
|
|
tools_dict = {t.name: t for t in tools}
|
|
|
|
# ---------- Создание агента ----------
|
|
prompt_template = ChatPromptTemplate.from_messages(
|
|
[
|
|
SystemMessage(content="Ты помощник. Используй инструменты при необходимости."),
|
|
("human", "{input}"),
|
|
]
|
|
)
|
|
|
|
agent = create_openai_tools_agent(llm, tools, prompt=prompt_template)
|
|
executor = AgentExecutor(agent=agent, tools=tools, verbose=False)
|
|
|
|
# ---------- Функции форматирования ----------
|
|
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 = {"messages": [HumanMessage(content=user_input)]}
|
|
stream = 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) |