feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'

2026-06-29 11:49:00 +03:00
parent c0aac04442
commit d6805973d6
6 changed files with 371 additions and 9 deletions
@@ -7,11 +7,15 @@
 EN
 Экзамен: RAG-агент с ChromaDB и веб-поиском
 Зачёт
-Версия 1
+Версия 2
 Дедлайн сдачи: 31.08.2026

 В работе

+Требуется доработка
+
+Переделайте решение: используйте QDrant вместо текущего векторного хранилища.
+
 Редактирование ответа

 Заполните ответ и отправьте работу на проверку преподавателю.
@@ -20,7 +24,7 @@ EN
 Текст
 Ссылка
 Файлы
-Текст ответа
+Ссылка (URL)
 Прикреплённые файлы
 Загрузить файл
 Отправить на проверку
@@ -28,7 +32,4 @@ EN

 Задание

-Практическое задание: RAG-агент с ChromaDB и веб-поиском
-Цель
-
-Построить AI-агента с локальным RAG-хранилищем на ChromaDB и веб-п
+Практическое задание: RAG-аге
@@ -0,0 +1,99 @@
+import os
+from typing import Any, Dict
+
+from langchain_ollama import Ollama
+from langchain.agents import Tool, AgentExecutor, initialize_agent, AgentType
+from langchain.schema import AgentAction, AgentFinish
+from langchain.tools import tool
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.output_parsers import StrOutputParser
+
+from vectorstore import QdrantVectorStore
+from tools import search_local_kb, web_search
+
+
+def create_agent(
+    vectorstore: QdrantVectorStore,
+    tavily_api_key: str,
+    llm_model: str = "llama3",
+) -> AgentExecutor:
+    """
+    Create a LangChain agent that routes queries to either the local KB or the web.
+
+    Parameters
+    ----------
+    vectorstore : QdrantVectorStore
+        The vector store for local knowledge base.
+    tavily_api_key : str
+        Tavily API key.
+    llm_model : str
+        Ollama model name for LLM.
+
+    Returns
+    -------
+    AgentExecutor
+        Configured agent executor.
+    """
+    # LLM
+    llm = Ollama(model=llm_model)
+
+    # Define tools with partial application of required arguments
+    tools = [
+        Tool(
+            name="search_local_kb",
+            func=lambda query, top_k=3: search_local_kb(
+                query=query, top_k=top_k, vectorstore=vectorstore
+            ),
+            description=(
+                "Use this tool to search the local knowledge base. "
+                "Return the most relevant snippets."
+            ),
+        ),
+        Tool(
+            name="web_search",
+            func=lambda query, max_results=3: web_search(
+                query=query, tavily_api_key=tavily_api_key, max_results=max_results
+            ),
+            description=(
+                "Use this tool to search the web via Tavily. "
+                "Return the most relevant snippets."
+            ),
+        ),
+    ]
+
+    # Prompt template for the agent
+    system_prompt = (
+        "You are an AI assistant that can answer questions using either "
+        "the local knowledge base or the web. If the question is about "
+        "recent events, news, or requires up-to-date information, "
+        "use the web_search tool. If the question is about "
+        "information contained in the local documents, use the "
+        "search_local_kb tool. After retrieving the information, "
+        "provide a concise answer and state the source (chromadb or tavily)."
+    )
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt),
+            MessagesPlaceholder("history"),
+            ("human", "{input}"),
+            MessagesPlaceholder("agent_scratchpad"),
+        ]
+    )
+
+    # Output parser
+    output_parser = StrOutputParser()
+
+    # Agent
+    agent = initialize_agent(
+        tools=tools,
+        llm=llm,
+        agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+        verbose=False,
+        handle_parsing_errors=True,
+        max_iterations=5,
+        early_stopping_method="generate",
+        system_message=system_prompt,
+    )
+
+    return agent
@@ -0,0 +1,65 @@
+import os
+import sys
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+from vectorstore import create_vectorstore, load_documents, collection_exists
+from agent import create_agent
+
+def main():
+    # Load environment variables
+    load_dotenv()
+
+    # Configuration
+    qdrant_path = os.getenv("QDRANT_PATH", "./qdrant_db")
+    embedding_model = os.getenv("EMBEDDING_MODEL", "nomic-embed-text")
+    llm_model = os.getenv("LLM_MODEL", "llama3")
+    tavily_api_key = os.getenv("TAVILY_API_KEY")
+    if not tavily_api_key:
+        print("Error: TAVILY_API_KEY not set in .env")
+        sys.exit(1)
+
+    # Create vector store
+    vectorstore = create_vectorstore(
+        persist_directory=qdrant_path,
+        collection_name="documents",
+        embedding_model=embedding_model,
+    )
+
+    # Load documents if collection is empty
+    if not collection_exists(vectorstore):
+        print("Loading documents into Qdrant...")
+        docs_dir = Path("documents")
+        if not docs_dir.exists():
+            print(f"Documents directory '{docs_dir}' not found.")
+            sys.exit(1)
+        load_documents(str(docs_dir), vectorstore)
+        print("Documents loaded.")
+    else:
+        print("Qdrant collection already exists. Skipping document load.")
+
+    # Create agent
+    agent = create_agent(vectorstore, tavily_api_key, llm_model=llm_model)
+
+    print("Chat agent ready. Type 'exit' to quit.")
+    while True:
+        try:
+            user_input = input("\nYou: ")
+        except (KeyboardInterrupt, EOFError):
+            print("\nExiting.")
+            break
+
+        if user_input.strip().lower() in {"exit", "quit"}:
+            print("Goodbye!")
+            break
+
+        # Run agent
+        try:
+            response = agent.run(user_input)
+            print(f"\nAssistant: {response}")
+        except Exception as e:
+            print(f"Error: {e}")
+
+if __name__ == "__main__":
+    main()
@@ -1,7 +1,8 @@
 langchain
-langchain-chroma
-langchain-tavily
 langchain-ollama
+langchain-qdrant
+langchain-tavily
 tavily-python
 chromadb
 python-dotenv
+qdrant-client
@@ -0,0 +1,71 @@
+from typing import List
+
+from langchain.tools import tool
+from langchain_ollama import Ollama
+from langchain_qdrant import QdrantVectorStore
+from langchain_tavily import TavilySearchResults
+
+
+@tool
+def search_local_kb(
+    query: str,
+    top_k: int,
+    vectorstore: QdrantVectorStore,
+) -> List[str]:
+    """
+    Perform a semantic search in the local knowledge base stored in Qdrant.
+
+    Parameters
+    ----------
+    query : str
+        The user's query.
+    top_k : int
+        Number of top results to return.
+    vectorstore : QdrantVectorStore
+        The vector store to search.
+
+    Returns
+    -------
+    List[str]
+        List of relevant document snippets.
+    """
+    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
+    docs = retriever.get_relevant_documents(query)
+    return [doc.page_content for doc in docs]
+
+
+@tool
+def web_search(
+    query: str,
+    tavily_api_key: str,
+    max_results: int = 3,
+) -> List[str]:
+    """
+    Perform a web search using Tavily.
+
+    Parameters
+    ----------
+    query : str
+        The user's query.
+    tavily_api_key : str
+        Tavily API key.
+    max_results : int
+        Number of search results to return.
+
+    Returns
+    -------
+    List[str]
+        List of search result snippets.
+    """
+    tavily = TavilySearchResults(
+        api_key=tavily_api_key,
+        max_results=max_results,
+    )
+    results = tavily.run(query)
+    # Extract snippets from results
+    snippets = []
+    for result in results:
+        snippet = result.get("content") or result.get("snippet") or result.get("title")
+        if snippet:
+            snippets.append(snippet)
+    return snippets
@@ -0,0 +1,125 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+from langchain_ollama import OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from qdrant_client import QdrantClient
+
+
+def create_vectorstore(
+    persist_directory: str = "./qdrant_db",
+    collection_name: str = "documents",
+    embedding_model: str = "nomic-embed-text",
+) -> QdrantVectorStore:
+    """
+    Create a Qdrant vector store backed by Ollama embeddings.
+
+    Parameters
+    ----------
+    persist_directory : str
+        Path to the directory where Qdrant will store its data.
+    collection_name : str
+        Name of the collection to use.
+    embedding_model : str
+        Ollama model name for embeddings.
+
+    Returns
+    -------
+    QdrantVectorStore
+        Initialized vector store.
+    """
+    # Ensure the persistence directory exists
+    Path(persist_directory).mkdir(parents=True, exist_ok=True)
+
+    # Create a local Qdrant client that stores data on disk
+    client = QdrantClient(path=persist_directory)
+
+    # Initialize embeddings
+    embeddings = OllamaEmbeddings(model=embedding_model)
+
+    # Create or load the collection
+    vectorstore = QdrantVectorStore(
+        client=client,
+        embeddings=embeddings,
+        collection_name=collection_name,
+    )
+    return vectorstore
+
+
+def load_documents(
+    directory: str,
+    vectorstore: QdrantVectorStore,
+    chunk_size: int = 1000,
+    chunk_overlap: int = 200,
+) -> None:
+    """
+    Load all .txt and .md files from a directory, split them into chunks,
+    embed, and add to the vector store.
+
+    Parameters
+    ----------
+    directory : str
+        Path to the directory containing documents.
+    vectorstore : QdrantVectorStore
+        The vector store to populate.
+    chunk_size : int
+        Maximum size of each chunk in characters.
+    chunk_overlap : int
+        Number of overlapping characters between chunks.
+    """
+    from langchain.text_splitter import RecursiveCharacterTextSplitter
+    from langchain.docstore.document import Document
+
+    # Prepare text splitter
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+    )
+
+    # Collect all documents
+    docs = []
+    for root, _, files in os.walk(directory):
+        for file in files:
+            if file.lower().endswith((".txt", ".md")):
+                file_path = os.path.join(root, file)
+                with open(file_path, "r", encoding="utf-8") as f:
+                    content = f.read()
+                # Split content into chunks
+                chunks = splitter.split_text(content)
+                # Create Document objects with metadata
+                for i, chunk in enumerate(chunks):
+                    doc = Document(
+                        page_content=chunk,
+                        metadata={
+                            "source": file_path,
+                            "chunk_index": i,
+                        },
+                    )
+                    docs.append(doc)
+
+    if docs:
+        vectorstore.add_documents(docs)
+    else:
+        print("No documents found to load.")
+
+
+def collection_exists(vectorstore: QdrantVectorStore) -> bool:
+    """
+    Check if the collection already exists in Qdrant.
+
+    Parameters
+    ----------
+    vectorstore : QdrantVectorStore
+        The vector store to check.
+
+    Returns
+    -------
+    bool
+        True if collection exists, False otherwise.
+    """
+    try:
+        vectorstore.client.get_collection(vectorstore.collection_name)
+        return True
+    except Exception:
+        return False