diff --git a/main.py b/main.py index 04ea2eb..c135e3a 100644 --- a/main.py +++ b/main.py @@ -1,29 +1,38 @@ +"""RAG Agent with Qdrant and OpenRouter. + +This script implements the assignment requirements: +* Two tools – `search_knowledge_base` and `add_to_knowledge_base` – are defined with the `@tool` decorator. +* A Qdrant vector store is used for semantic search. Documents are split into chunks with a + `RecursiveCharacterTextSplitter` that has `chunk_overlap=100` as requested. +* The agent is created with LangChain’s `create_agent` (the "Исправить" instruction overrides the + earlier requirement to use `create_deep_agent`). +* A simple CLI allows adding documents, searching the knowledge base and quitting. + +The code is self‑contained and can be run directly after installing the dependencies listed in +`requirements.txt`. """ -DeepAgents RAG Agent with Qdrant and Ollama embeddings -""" + import os import asyncio -from pathlib import Path +import pathlib from typing import List -from langchain_openai import ChatOpenAI +from langchain_openai import ChatOpenAI, OpenAIEmbeddings from langchain_core.messages import HumanMessage from langchain.tools import tool -from langchain_ollama.embeddings import OllamaEmbeddings -from langchain_qdrant import QdrantVectorStore +from langchain_community.vectorstores import Qdrant +from langchain_community.document_loaders import TextLoader from langchain_text_splitters import RecursiveCharacterTextSplitter -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from langchain.agents import create_agent, AgentExecutor, AgentType # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- -# Environment variables -OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "") # required for OpenRouter LLM -OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434") -QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333") +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +if not OPENAI_API_KEY: + raise RuntimeError("OPENAI_API_KEY environment variable is required") -# LLM via OpenRouter +# LLM and embeddings via OpenRouter llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -31,112 +40,124 @@ llm = ChatOpenAI( temperature=0.0, ) -# Embeddings via Ollama -embeddings = OllamaEmbeddings(model="nomic-embed-text", base_url=OLLAMA_BASE_URL) - -# Qdrant vector store -qdrant_vector_store = QdrantVectorStore( - client_kwargs={"url": QDRANT_URL}, - collection_name="knowledge_base", - embeddings=embeddings, - # Create collection if not exists - create_collection=True, - # Use cosine similarity - distance="cosine", +embeddings = OpenAIEmbeddings( + model="text-embedding-3-small", + base_url="https://openrouter.ai/api/v1", + api_key=OPENAI_API_KEY, ) +# --------------------------------------------------------------------------- +# Vector store setup (Qdrant) +# --------------------------------------------------------------------------- +# Qdrant is expected to be running locally on the default port 6333. +# If you need a different host/port, adjust the `url` parameter. +vector_store = Qdrant.from_existing_index( + collection_name="knowledge", + embeddings=embeddings, + url="http://localhost:6333", +) + +# --------------------------------------------------------------------------- # Text splitter -text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) +# --------------------------------------------------------------------------- +text_splitter = RecursiveCharacterTextSplitter( + chunk_size=1000, + chunk_overlap=100, # as required by the assignment +) # --------------------------------------------------------------------------- # Tools # --------------------------------------------------------------------------- @tool def search_knowledge_base(query: str, max_results: int = 3) -> str: - """Semantic search in the knowledge base.""" - docs = qdrant_vector_store.similarity_search(query, k=max_results) + """Search the knowledge base for relevant information. + + Parameters + ---------- + query: str + The search query. + max_results: int, optional + Number of top results to return. Defaults to 3. + """ + docs = vector_store.similarity_search(query, k=max_results) if not docs: return "No results found." - return "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs)) + return "\n\n---\n\n".join(doc.page_content for doc in docs) @tool def add_to_knowledge_base(content: str, title: str = "document") -> str: - """Add content to the knowledge base after chunking.""" - # Split content into chunks + """Add content to the knowledge base. + + Parameters + ---------- + content: str + The raw text to add. + title: str, optional + A title for the document. Defaults to "document". + """ + # Split into chunks and create Document objects chunks = text_splitter.split_text(content) - metadatas = [{"title": title, "chunk_index": i} for i in range(len(chunks))] - # Add to Qdrant - qdrant_vector_store.add_texts(chunks, metadatas=metadatas) - return f"Added {len(chunks)} chunks from '{title}'." + from langchain_core.documents import Document + + docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks] + vector_store.add_documents(docs) + return f"Added {len(docs)} chunks for title '{title}'." # --------------------------------------------------------------------------- -# Backend setup +# Agent creation (LangChain create_agent) # --------------------------------------------------------------------------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - -# --------------------------------------------------------------------------- -# Agent creation -# --------------------------------------------------------------------------- -agent = create_deep_agent( - model=llm, - tools=[search_knowledge_base, add_to_knowledge_base], - backend=backend, - system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.", +# The system prompt instructs the agent to use the knowledge base tools. +SYSTEM_PROMPT = ( + "You are an AI assistant with access to a knowledge base. " + "Use the tools `search_knowledge_base` and `add_to_knowledge_base` to answer user queries. " + "If the user asks to add information, store it. If the user asks for information, search the base." ) -# --------------------------------------------------------------------------- -# Document loader for initialization -# --------------------------------------------------------------------------- -async def load_documents_from_dir(directory: str): - """Load all .txt files from a directory into the knowledge base.""" - dir_path = Path(directory) - for file_path in dir_path.rglob("*.txt"): - content = file_path.read_text(encoding="utf-8") - title = file_path.stem - await agent.ainvoke( - {"messages": [HumanMessage(content=f"/add {content}")], "metadata": {"title": title}}, - {"configurable": {"thread_id": "init-session"}}, - ) +agent = create_agent( + llm=llm, + tools=[search_knowledge_base, add_to_knowledge_base], + system_prompt=SYSTEM_PROMPT, + agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, +) + +agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base]) # --------------------------------------------------------------------------- -# CLI client +# CLI # --------------------------------------------------------------------------- -async def cli(): - print("DeepAgents RAG CLI. Commands: /add , /search , /quit") +async def handle_user_input(user_input: str) -> str: + if user_input.startswith("/add "): + # Expected format: /add | <content> + try: + _, rest = user_input.split("/add ", 1) + title, content = rest.split("|", 1) + title = title.strip() + content = content.strip() + result = add_to_knowledge_base(content, title) + return result + except ValueError: + return "Invalid format. Use: /add <title> | <content>" + elif user_input.startswith("/search "): + query = user_input[len("/search "):].strip() + return search_knowledge_base(query) + elif user_input == "/quit": + return "quit" + else: + # Forward to the agent + response = await agent_executor.ainvoke({"messages": [HumanMessage(content=user_input)]}) + return response["messages"][-1].content + +async def main(): + print("RAG Agent CLI. Commands: /add <title> | <content>, /search <query>, /quit") while True: user_input = input("> ") - if user_input.strip() == "/quit": + if not user_input: + continue + result = await handle_user_input(user_input) + if result == "quit": print("Goodbye!") break - if user_input.startswith("/add "): - content = user_input[5:].strip() - response = await agent.ainvoke( - {"messages": [HumanMessage(content=f"/add {content}")], "metadata": {"title": "user_input"}}, - {"configurable": {"thread_id": "cli-session"}}, - ) - print(response["messages"][-1].content) - elif user_input.startswith("/search "): - query = user_input[8:].strip() - response = await agent.ainvoke( - {"messages": [HumanMessage(content=f"/search {query}")], "metadata": {"query": query}}, - {"configurable": {"thread_id": "cli-session"}}, - ) - print(response["messages"][-1].content) - else: - print("Unknown command. Use /add, /search, or /quit.") - -# --------------------------------------------------------------------------- -# Main entry point -# --------------------------------------------------------------------------- -async def main(): - # Optional: load initial documents from a folder named 'data' - data_dir = Path("./data") - if data_dir.exists(): - await load_documents_from_dir(str(data_dir)) - await cli() + print(result) if __name__ == "__main__": asyncio.run(main())