8. Самописный поисковый агент на основе deep agents from scratch: client.py
This commit is contained in:
+96
-62
@@ -1,73 +1,98 @@
|
||||
import os
|
||||
import sys
|
||||
from typing import List
|
||||
|
||||
from langchain_ollama import OllamaEmbeddings, OllamaLLM
|
||||
from langchain_qdrant import QdrantVectorStore
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http.models import Distance, VectorParams
|
||||
# Импортируем необходимые компоненты из LangChain и связанных библиотек
|
||||
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
||||
from langchain_qdrant import QdrantStore
|
||||
from langchain_core.documents import Document
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
from langchain.chains import RetrievalQA
|
||||
|
||||
# Конфигурация подключения к Qdrant
|
||||
QDRANT_HOST = "localhost"
|
||||
QDRANT_PORT = 6333
|
||||
# Конфигурация подключения к Qdrant (используем локальный сервер по умолчанию)
|
||||
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
|
||||
COLLECTION_NAME = "documents"
|
||||
|
||||
# Инициализация клиента Qdrant
|
||||
qdrant_client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
|
||||
|
||||
# Создание коллекции при необходимости
|
||||
if COLLECTION_NAME not in qdrant_client.get_collections().collections:
|
||||
qdrant_client.create_collection(
|
||||
collection_name=COLLECTION_NAME,
|
||||
vectors_config=VectorParams(size=768, distance=Distance.COSINE),
|
||||
)
|
||||
|
||||
# Инициализация эмбеддинговой модели (nomic-embed-text)
|
||||
# Инициализируем Ollama embeddings и LLM
|
||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
llm = ChatOllama(model="llama3")
|
||||
|
||||
# Создание векторного хранилища
|
||||
vector_store = QdrantVectorStore(
|
||||
client=qdrant_client,
|
||||
# Создаём или подключаемся к коллекции в Qdrant
|
||||
vector_store = QdrantStore(
|
||||
url=QDRANT_URL,
|
||||
collection_name=COLLECTION_NAME,
|
||||
embeddings=embeddings,
|
||||
)
|
||||
|
||||
# Инициализация LLM (llama3)
|
||||
llm = OllamaLLM(model="llama3")
|
||||
# Шаблон запроса для RetrievalQA
|
||||
prompt_template = PromptTemplate.from_template(
|
||||
"""
|
||||
Ты – интеллектуальный помощник. На основе предоставленного контекста ответь на вопрос пользователя.
|
||||
|
||||
Контекст:
|
||||
{context}
|
||||
|
||||
Вопрос: {question}
|
||||
Ответ:
|
||||
"""
|
||||
)
|
||||
|
||||
qa_chain = RetrievalQA.from_chain_type(
|
||||
llm=llm,
|
||||
chain_type="stuff",
|
||||
retriever=vector_store.as_retriever(),
|
||||
return_source_documents=True,
|
||||
question_generator=None,
|
||||
combine_docs_chain_kwargs={"prompt": prompt_template},
|
||||
)
|
||||
|
||||
|
||||
def add_document(text: str) -> None:
|
||||
"""Добавляет документ в хранилище."""
|
||||
vector_store.add_texts([text])
|
||||
print("✅ Документ добавлен.")
|
||||
"""
|
||||
Добавляет новый документ в Qdrant.
|
||||
"""
|
||||
doc = Document(page_content=text)
|
||||
vector_store.add_documents([doc])
|
||||
print("Документ добавлен.")
|
||||
|
||||
def search_documents(query: str, k: int = 3) -> List[str]:
|
||||
"""Ищет похожие документы и возвращает их тексты."""
|
||||
results = vector_store.similarity_search_with_score(query, k=k)
|
||||
docs = [doc.page_content for doc, _ in results]
|
||||
return docs
|
||||
|
||||
def answer_query(query: str) -> str:
|
||||
"""Получает ответ LLM на основе найденных документов."""
|
||||
docs = search_documents(query)
|
||||
context = "\n\n".join(docs)
|
||||
prompt = f"Контекст:\n{context}\n\nВопрос: {query}\nОтвет:"
|
||||
return llm.invoke(prompt)
|
||||
def search(query: str) -> None:
|
||||
"""
|
||||
Выполняет поиск по запросу и выводит ответ LLM.
|
||||
"""
|
||||
result = qa_chain({"question": query})
|
||||
answer = result["answer"]
|
||||
sources = result.get("source_documents", [])
|
||||
print("\nОтвет:\n" + answer)
|
||||
if sources:
|
||||
print("\nИсточники:")
|
||||
for i, src in enumerate(sources, 1):
|
||||
print(f"{i}. {src.page_content[:200]}...")
|
||||
|
||||
|
||||
def print_help() -> None:
|
||||
"""Показывает список команд."""
|
||||
"""
|
||||
Выводит справку по доступным командам.
|
||||
"""
|
||||
help_text = """
|
||||
Команды:
|
||||
/add <текст> - добавить документ
|
||||
/search <запрос> - поиск по запросу и вывод ответа LLM
|
||||
/quit - выйти из программы
|
||||
Доступные команды:
|
||||
/add <текст> - Добавить новый документ
|
||||
/search <запрос> - Выполнить поиск и получить ответ
|
||||
/quit - Выйти из программы
|
||||
/help - Показать эту справку
|
||||
"""
|
||||
print(help_text)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""
|
||||
Основной цикл CLI.
|
||||
"""
|
||||
print("=== Самописный поисковый агент ===")
|
||||
print_help()
|
||||
while True:
|
||||
try:
|
||||
user_input = input("> ").strip()
|
||||
user_input = input("\n> ").strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\nВыход.")
|
||||
break
|
||||
@@ -75,27 +100,36 @@ def main() -> None:
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
if user_input.startswith("/add "):
|
||||
text = user_input[5:].strip()
|
||||
if text:
|
||||
add_document(text)
|
||||
if user_input.startswith("/"):
|
||||
parts = user_input.split(maxsplit=1)
|
||||
cmd = parts[0].lower()
|
||||
arg = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
if cmd == "/add":
|
||||
if not arg:
|
||||
print("Ошибка: необходимо указать текст для добавления.")
|
||||
else:
|
||||
add_document(arg)
|
||||
|
||||
elif cmd == "/search":
|
||||
if not arg:
|
||||
print("Ошибка: необходимо указать поисковый запрос.")
|
||||
else:
|
||||
search(arg)
|
||||
|
||||
elif cmd == "/quit":
|
||||
print("Выход.")
|
||||
break
|
||||
|
||||
elif cmd == "/help":
|
||||
print_help()
|
||||
|
||||
else:
|
||||
print("⚠️ Текст не указан.")
|
||||
elif user_input.startswith("/search "):
|
||||
query = user_input[8:].strip()
|
||||
if query:
|
||||
answer = answer_query(query)
|
||||
print("\nОтвет LLM:")
|
||||
print(answer)
|
||||
else:
|
||||
print("⚠️ Запрос не указан.")
|
||||
elif user_input == "/quit":
|
||||
print("Выход.")
|
||||
break
|
||||
elif user_input == "/help":
|
||||
print_help()
|
||||
print(f"Неизвестная команда '{cmd}'. Введите /help для справки.")
|
||||
|
||||
else:
|
||||
print("❓ Неизвестная команда. Введите /help для списка команд.")
|
||||
print("Команды начинаются с '/'. Введите /help для списка команд.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user