123 lines
4.9 KiB
Python
123 lines
4.9 KiB
Python
from typing import TypedDict, List, Optional, Annotated
|
||
import operator
|
||
from langchain_community.llms import Ollama
|
||
from langchain_community.vectorstores import Qdrant
|
||
from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader
|
||
from langgraph.graph import StateGraph, END
|
||
from langgraph.graph.message import add_messages
|
||
from langchain_core.messages import AIMessage, SystemMessage
|
||
from langchain.tools import tool
|
||
# from langchain.agents import AgentExecutor, create_react_agent
|
||
import numpy as np
|
||
|
||
# Определяем состояние агента
|
||
class AgentState(TypedDict):
|
||
messages: Annotated[List, add_messages] # История сообщений
|
||
knowledge_base: Optional[str] # Релевантные документы
|
||
current_step: str # Текущий шаг выполнения
|
||
needs_search: bool # Нужен ли поиск в RAG
|
||
final_answer: Optional[str] # Финальный ответ
|
||
|
||
# Инициализация LLM через Ollama
|
||
llm = Ollama(
|
||
model="qwen2.5-coder-16k:14b",
|
||
temperature=0.1, # Низкая температура для консистентности
|
||
num_predict=1024, # Максимальная длина ответа
|
||
)
|
||
|
||
# Создаем инструменты для агента
|
||
@tool
|
||
def search_knowledge_base(query: str) -> str:
|
||
"""Поиск информации в локальной базе знаний"""
|
||
# Здесь будет подключение к Qdrant
|
||
return "Найденная информация из базы знаний"
|
||
|
||
@tool
|
||
def calculate(expression: str) -> str:
|
||
"""Выполнение математических вычислений"""
|
||
try:
|
||
result = eval(expression)
|
||
return f"Результат: {result}"
|
||
except:
|
||
return "Ошибка в выражении"
|
||
|
||
@tool
|
||
def web_search(query: str) -> str:
|
||
"""Поиск в интернете (если нужно)"""
|
||
# Можно подключить локальный поиск через DuckDuckGo
|
||
return "Результаты поиска из интернета"
|
||
|
||
# Создаем граф агента
|
||
def create_agent_graph():
|
||
workflow = StateGraph(AgentState)
|
||
|
||
# Узел: анализ запроса
|
||
def analyze_query(state: AgentState):
|
||
messages = state["messages"]
|
||
last_message = messages[-1].content if messages else ""
|
||
|
||
# Простой анализ: проверяем, нужен ли поиск
|
||
search_keywords = ["информация", "документ", "найди", "ищи", "база знаний"]
|
||
needs_search = any(keyword in last_message.lower() for keyword in search_keywords)
|
||
|
||
return {
|
||
"needs_search": needs_search,
|
||
"current_step": "analyzing_query"
|
||
}
|
||
|
||
# Узел: поиск в RAG
|
||
def rag_search(state: AgentState):
|
||
if not state["needs_search"]:
|
||
return {"knowledge_base": None, "current_step": "generating_answer"}
|
||
|
||
# Получаем последний запрос
|
||
query = state["messages"][-1].content
|
||
|
||
# Здесь должен быть реальный поиск в Qdrant
|
||
# Пока заглушка
|
||
results = ["Документ 1: Информация о...", "Документ 2: Данные по..."]
|
||
|
||
return {
|
||
"knowledge_base": "\n".join(results),
|
||
"current_step": "generating_answer"
|
||
}
|
||
|
||
# Узел: генерация ответа
|
||
def generate_answer(state: AgentState):
|
||
messages = state["messages"]
|
||
knowledge = state.get("knowledge_base", "")
|
||
|
||
# Формируем промпт с контекстом
|
||
prompt = f"""Ты — интеллектуальный ассистент. Используй следующую информацию если она релевантна:
|
||
|
||
Контекст из базы знаний:
|
||
{knowledge}
|
||
|
||
История разговора:
|
||
{messages[-5:] if len(messages) > 5 else messages}
|
||
|
||
Текущий запрос: {messages[-1].content if messages else ''}
|
||
|
||
Ответь максимально полезно и точно:"""
|
||
|
||
response = llm.invoke(prompt)
|
||
|
||
return {
|
||
"final_answer": response,
|
||
"current_step": "completed",
|
||
"messages": messages + [AIMessage(content=response)]
|
||
}
|
||
|
||
# Добавляем узлы в граф
|
||
workflow.add_node("analyze", analyze_query)
|
||
workflow.add_node("search", rag_search)
|
||
workflow.add_node("generate", generate_answer)
|
||
|
||
# Определяем edges (переходы)
|
||
workflow.set_entry_point("analyze")
|
||
workflow.add_edge("analyze", "search")
|
||
workflow.add_edge("search", "generate")
|
||
workflow.add_edge("generate", END)
|
||
|
||
return workflow.compile()
|