diff --git a/README.md b/README.md index bdae846..4c8e3ad 100644 --- a/README.md +++ b/README.md @@ -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 и веб-п \ No newline at end of file +Практическое задание: RAG-аге \ No newline at end of file diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..fadceaf --- /dev/null +++ b/agent.py @@ -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 \ No newline at end of file diff --git a/main.py b/main.py new file mode 100644 index 0000000..1f3e161 --- /dev/null +++ b/main.py @@ -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() \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index b0b31f9..722e91d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,8 @@ langchain -langchain-chroma -langchain-tavily langchain-ollama +langchain-qdrant +langchain-tavily tavily-python chromadb -python-dotenv \ No newline at end of file +python-dotenv +qdrant-client \ No newline at end of file diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..cb02c6f --- /dev/null +++ b/tools.py @@ -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 \ No newline at end of file diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..c5353c2 --- /dev/null +++ b/vectorstore.py @@ -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 \ No newline at end of file