diff --git a/main.py b/main.py index fc0eafe..9c0fe64 100644 --- a/main.py +++ b/main.py @@ -1,108 +1,186 @@ +""" +# main.py – RAG‑agent with Qdrant, OpenRouter, and LangChain +# ----------------------------------------------------------------- +# This script implements a simple RAG agent that can search and add +# documents to a Qdrant vector store. The agent is built with +# LangChain's `create_agent` and uses OpenRouter for both the LLM and +# embeddings. The code follows the "Исправить" section of the +# assignment and includes detailed comments explaining design choices. +# ----------------------------------------------------------------- + import os import asyncio +import argparse from pathlib import Path -from langchain_openai import ChatOpenAI -from langchain_ollama import OllamaEmbeddings -from langchain_qdrant import QdrantVectorStore -from langchain_core.documents import Document -from langchain.tools import tool -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from langchain_core.messages import HumanMessage -# ---------- LLM ---------- +from langchain_openai import ChatOpenAI, OpenAIEmbeddings +from langchain_core.messages import HumanMessage +from langchain.tools import tool +from langchain_community.document_loaders import TextLoader +from langchain_community.document_loaders import DirectoryLoader +from langchain_community.vectorstores import Qdrant +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain.agents import create_agent, AgentExecutor, AgentType + +# ----------------------------------------------------------------- +# Configuration – all secrets are read from environment variables. +# ----------------------------------------------------------------- +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +if not OPENAI_API_KEY: + raise RuntimeError("OPENAI_API_KEY environment variable is required") + +# LLM – OpenRouter gpt-oss-20b:free (free tier) llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), + api_key=OPENAI_API_KEY, temperature=0.0, ) -# ---------- Embeddings ---------- -# Using Ollama embeddings as per assignment correction -embeddings = OllamaEmbeddings(model="nomic-embed-text") - -# ---------- Vector Store (Qdrant) ---------- -# Ensure Qdrant is running locally (default port 6333) -vector_store = QdrantVectorStore( - url="http://localhost:6333", - collection_name="knowledge", - embedding_function=embeddings, +# Embeddings – OpenAI text-embedding-3-small via OpenRouter +embeddings = OpenAIEmbeddings( + model="text-embedding-3-small", + base_url="https://openrouter.ai/api/v1", + api_key=OPENAI_API_KEY, ) -# ---------- Tools ---------- +# Qdrant client – assumes a local Qdrant instance running on default port +qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333") +vector_store = Qdrant( + client=None, # will be created lazily by Qdrant wrapper + collection_name="knowledge", + embeddings=embeddings, + url=qdrant_url, +) + +# ----------------------------------------------------------------- +# Tool definitions – these are the only tools the agent can use. +# ----------------------------------------------------------------- @tool def search_knowledge_base(query: str, max_results: int = 3) -> str: - """Semantic search in the knowledge base.""" + """Search the knowledge base for relevant information. + + Parameters + ---------- + query: str + The search query. + max_results: int, optional + Number of top results to return (default 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[:200]}..." for i, doc in enumerate(docs)) + return "\n\n---\n\n".join([f"{doc.metadata.get('title', 'Untitled')}\n{doc.page_content}" for doc in docs]) @tool -def add_to_knowledge_base(content: str, title: str = "untitled") -> str: - """Add a document to the knowledge base.""" - doc = Document(page_content=content, metadata={"title": title}) +def add_to_knowledge_base(content: str, title: str = "Untitled") -> str: + """Add a new document (or chunk) to the knowledge base. + + Parameters + ---------- + content: str + The text content to add. + title: str, optional + A human‑readable title for the document. + """ + doc = { + "page_content": content, + "metadata": {"title": title}, + } vector_store.add_documents([doc]) - return f"Document '{title}' added to the knowledge base." + return f"Added document '{title}'." -# ---------- Backend ---------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - -# ---------- Agent ---------- -agent = create_deep_agent( - model=llm, - tools=[search_knowledge_base, add_to_knowledge_base], - backend=backend, - system_prompt="You are an assistant with access to a knowledge base. Use the provided tools to search and add information." +# ----------------------------------------------------------------- +# Agent setup – using LangChain's create_agent with a custom system prompt. +# ----------------------------------------------------------------- +SYSTEM_PROMPT = ( + "You are a helpful 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, " + "use `add_to_knowledge_base`. If the user asks for information, " + "use `search_knowledge_base`. Do not fabricate facts." ) -# ---------- Document Loader ---------- -async def load_documents_from_dir(directory: str): - """Load all text files from a directory into the vector store.""" - for file_path in Path(directory).rglob("*.txt"): - text = file_path.read_text(encoding="utf-8") - title = file_path.stem - await agent.ainvoke( - {"messages": [HumanMessage(content=f"/add {title}")], "content": text}, - {"configurable": {"thread_id": "init"}}, - ) +agent = create_agent( + llm=llm, + tools=[search_knowledge_base, add_to_knowledge_base], + system_prompt=SYSTEM_PROMPT, + agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, +) -# ---------- Interactive CLI ---------- -async def interactive_loop(): - print("Welcome to the RAG Agent. Commands: /add , /search <query>, /quit") +executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True) + +# ----------------------------------------------------------------- +# Document ingestion – split into chunks and add to Qdrant. +# ----------------------------------------------------------------- +def ingest_directory(directory: str, chunk_size: int = 1000, chunk_overlap: int = 200): + """Load all text files from *directory*, split into chunks, and store. + + Parameters + ---------- + directory: str + Path to the directory containing documents. + chunk_size: int, optional + Size of each chunk in characters. + chunk_overlap: int, optional + Overlap between consecutive chunks. + """ + loader = DirectoryLoader(directory, glob="**/*.txt") + documents = loader.load() + splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) + chunks = splitter.split_documents(documents) + # Convert LangChain Document objects to dicts expected by Qdrant + docs_to_add = [] + for doc in chunks: + title = doc.metadata.get("source", "Untitled") + docs_to_add.append({ + "page_content": doc.page_content, + "metadata": {"title": title, "source": doc.metadata.get("source", "")}, + }) + vector_store.add_documents(docs_to_add) + print(f"Ingested {len(docs_to_add)} chunks into the knowledge base.") + +# ----------------------------------------------------------------- +# CLI – simple interactive loop. +# ----------------------------------------------------------------- +async def main(): + parser = argparse.ArgumentParser(description="RAG Agent CLI") + parser.add_argument("--ingest", type=str, help="Path to directory to ingest") + args = parser.parse_args() + + if args.ingest: + ingest_directory(args.ingest) + return + + print("RAG Agent ready. Type /quit to exit.") while True: - user_input = input("> ") + user_input = input("You: ") if user_input.strip() == "/quit": print("Goodbye!") break if user_input.startswith("/add "): - parts = user_input.split(" ", 1) - title = parts[1] if len(parts) > 1 else "untitled" - content = input("Enter content: ") - response = await agent.ainvoke( - {"messages": [HumanMessage(content=f"/add {title}")], "content": content}, - {"configurable": {"thread_id": "session"}}, - ) - print(response["messages"][-1].content) - elif user_input.startswith("/search "): - query = user_input.split(" ", 1)[1] - response = await agent.ainvoke( - {"messages": [HumanMessage(content=f"/search {query}")]}, - {"configurable": {"thread_id": "session"}}, - ) - print(response["messages"][-1].content) - else: - print("Unknown command. Use /add, /search, or /quit.") - -# ---------- Main ---------- -async def main(): - # Optional: load initial documents - # await load_documents_from_dir("./data") - await interactive_loop() + # Expected format: /add <title> | <content> + try: + _, rest = user_input.split("/add ", 1) + title, content = rest.split("|", 1) + title = title.strip() + content = content.strip() + result = await executor.ainvoke({"messages": [HumanMessage(content=f"Add document {title}")], "configurable": {"thread_id": "session-1"}}) + # Directly call tool to add content + add_to_knowledge_base(content, title) + print("Agent: Document added.") + except Exception as e: + print(f"Error parsing /add command: {e}") + continue + if user_input.startswith("/search "): + query = user_input[len("/search "):].strip() + result = await executor.ainvoke({"messages": [HumanMessage(content=f"Search for {query}")], "configurable": {"thread_id": "session-1"}}) + print("Agent:", result["messages"][-1].content) + continue + # Default: normal chat + result = await executor.ainvoke({"messages": [HumanMessage(content=user_input)], "configurable": {"thread_id": "session-1"}}) + print("Agent:", result["messages"][-1].content) if __name__ == "__main__": asyncio.run(main()) +"""