diff --git a/requirements.txt b/requirements.txt index a69b6f2..d1b51f1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,7 @@ -langchain>=0.2.0 -langchain-qdrant -langchain-ollama -qdrant-client \ No newline at end of file +langchain==0.2.0 +langchain-ollama==0.2.0 +langchain-qdrant==0.2.0 +qdrant-client==1.7.0 +python-dotenv==1.0.0 +tqdm==4.66.1 +pydantic==2.7.0 \ No newline at end of file diff --git a/src/agent.py b/src/agent.py index d3f927b..1f248aa 100644 --- a/src/agent.py +++ b/src/agent.py @@ -1,35 +1,31 @@ -from langchain_ollama import ChatOllama -from langchain.agents import initialize_agent, AgentType -from langchain.tools import Tool -from src.tools import search_knowledge_base, add_to_knowledge_base +from langchain_ollama import Ollama +from langchain.agents import initialize_agent, AgentExecutor, AgentType +from langchain.tools import BaseTool +from typing import List -def create_agent(): +def create_agent(tools: List[BaseTool]) -> AgentExecutor: """ - Create a LangChain agent configured to use the knowledge base tools. + Create a LangChain agent that can use the provided tools. + + Parameters: + tools (List[BaseTool]): List of tools for the agent. + + Returns: + AgentExecutor: Configured agent. """ - llm = ChatOllama(model="llama3") - tools = [ - Tool( - name="search_knowledge_base", - func=search_knowledge_base, - description="Search the knowledge base for relevant information." - ), - Tool( - name="add_to_knowledge_base", - func=add_to_knowledge_base, - description="Add new content to the knowledge base." - ) - ] + llm = Ollama(model="llama3") system_prompt = ( - "You are an AI assistant that helps users by searching and adding information to a knowledge base. " - "Use the provided tools to answer queries. If you need to add new information, call add_to_knowledge_base. " - "If you need to retrieve information, call search_knowledge_base. Provide concise answers." + "You are an AI assistant with access to a knowledge base. " + "Use the following tools to answer user queries:\n" + "- search_knowledge_base: Search the knowledge base.\n" + "- add_to_knowledge_base: Add a new document to the knowledge base.\n" + "When you need to use a tool, call it with the appropriate arguments." ) agent = initialize_agent( tools=tools, llm=llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, - agent_kwargs={"system_message": system_prompt} + agent_kwargs={"system_message": system_prompt}, ) return agent \ No newline at end of file diff --git a/src/chunking.py b/src/chunking.py index c0c0b09..3a9cd36 100644 --- a/src/chunking.py +++ b/src/chunking.py @@ -1,16 +1,20 @@ -from langchain.text_splitter import RecursiveCharacterTextSplitter -from langchain.schema import Document -from typing import List +from langchain_text_splitters import RecursiveCharacterTextSplitter -def chunk_document(content: str, title: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> List[Document]: +def chunk_document(text: str, chunk_size: int = 1000, chunk_overlap: int = 200): """ - Split a document into chunks suitable for vector storage. - Each chunk is stored as a Document with metadata. + Split a document into chunks using RecursiveCharacterTextSplitter. + + Args: + text: The full text of the document. + chunk_size: Maximum size of each chunk in characters. + chunk_overlap: Number of characters to overlap between chunks. + + Returns: + List of chunk strings. """ - splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) - texts = splitter.split_text(content) - documents = [] - for i, text in enumerate(texts): - metadata = {"source": title, "chunk_id": i} - documents.append(Document(page_content=text, metadata=metadata)) - return documents \ No newline at end of file + splitter = RecursiveCharacterTextSplitter( + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + separators=["\n\n", "\n", " ", ""], + ) + return splitter.split_text(text) \ No newline at end of file diff --git a/src/cli.py b/src/cli.py index 46666f7..4282515 100644 --- a/src/cli.py +++ b/src/cli.py @@ -1,58 +1,93 @@ -import argparse -from src.loader import load_documents_from_directory -from src.tools import add_to_knowledge_base, search_knowledge_base +import sys +from typing import List from src.agent import create_agent +from src.tools import search_knowledge_base, add_to_knowledge_base +from src.loader import load_documents_from_directory +from src.vector_store import QdrantVectorStore +from src.tools import vector_store as global_vector_store def main(): - parser = argparse.ArgumentParser(description="RAG Agent CLI") - parser.add_argument("--init-dir", type=str, help="Directory to load documents from") - args = parser.parse_args() + # Initialize vector store + vector_store = QdrantVectorStore() + # Override global reference for tools + global global_vector_store + global_vector_store = vector_store - if args.init_dir: - docs = load_documents_from_directory(args.init_dir) - for doc in docs: - add_to_knowledge_base(doc.page_content, doc.metadata.get("source", "unknown")) - print(f"Loaded {len(docs)} documents into the knowledge base.") + # Prepare tools + tools = [search_knowledge_base, add_to_knowledge_base] + + # Create agent + agent = create_agent(tools) + + print("Welcome to the RAG Agent CLI.") + print("Commands:") + print(" /load - Load all .txt and .md files from directory into the knowledge base.") + print(" /add - Add a new document. You will be prompted for content.") + print(" /search <query> - Search the knowledge base.") + print(" /quit - Exit the program.") + print("") - agent = create_agent() - print("Agent ready. Type your queries. Use /add <title> <content>, /search <query>, or /quit to exit.") while True: try: user_input = input(">> ").strip() except (EOFError, KeyboardInterrupt): - print("\nGoodbye!") + print("\nExiting.") break if not user_input: continue - if user_input.lower() == "/quit": - print("Goodbye!") - break + if user_input.startswith("/"): + parts = user_input.split(maxsplit=1) + command = parts[0][1:].lower() + arg = parts[1] if len(parts) > 1 else "" - if user_input.lower().startswith("/add"): - parts = user_input.split(" ", 2) - if len(parts) < 3: - print("Usage: /add <title> <content>") - continue - _, title, content = parts - response = add_to_knowledge_base(content, title) + if command == "quit": + print("Goodbye!") + break + + elif command == "load": + if not arg: + print("Please provide a directory path.") + continue + docs = load_documents_from_directory(arg) + if not docs: + print("No documents found.") + continue + vector_store.add_documents(docs) + print(f"Loaded {len(docs)} documents into the knowledge base.") + + elif command == "add": + title = arg.strip() + if not title: + print("Please provide a title for the document.") + continue + print("Enter document content. Finish with a single line containing only 'END'.") + lines = [] + while True: + line = input() + if line.strip() == "END": + break + lines.append(line) + content = "\n".join(lines) + result = add_to_knowledge_base(content=content, title=title) + print(result) + + elif command == "search": + query = arg.strip() + if not query: + print("Please provide a search query.") + continue + result = search_knowledge_base(query=query, max_results=5) + print(result) + + else: + print(f"Unknown command: {command}") + + else: + # Treat as a natural language query for the agent + response = agent.run(user_input) print(response) - continue - - if user_input.lower().startswith("/search"): - parts = user_input.split(" ", 1) - if len(parts) < 2: - print("Usage: /search <query>") - continue - _, query = parts - response = search_knowledge_base(query) - print(response) - continue - - # Treat as normal query for the agent - response = agent.run(user_input) - print(response) if __name__ == "__main__": main() \ No newline at end of file diff --git a/src/loader.py b/src/loader.py index f90479a..9bf322e 100644 --- a/src/loader.py +++ b/src/loader.py @@ -1,22 +1,21 @@ import os from pathlib import Path -from typing import List -from langchain.schema import Document +from typing import List, Dict +import re -def load_documents_from_directory(directory: str) -> List[Document]: +def load_documents_from_directory(directory: str) -> List[Dict[str, str]]: """ - Recursively load all text-based files from a directory into Documents. - Supported extensions: .txt, .md, .py, .json, .csv + Load all .txt and .md files from the given directory. + + Returns a list of dicts with keys: title, content. """ docs = [] - for root, _, files in os.walk(directory): - for file in files: - if file.lower().endswith(('.txt', '.md', '.py', '.json', '.csv')): - path = Path(root) / file - try: - with open(path, 'r', encoding='utf-8') as f: - content = f.read() - docs.append(Document(page_content=content, metadata={"source": str(path)})) - except Exception as e: - print(f"Failed to read {path}: {e}") + for file_path in Path(directory).rglob("*"): + if file_path.suffix.lower() in {".txt", ".md"}: + try: + content = file_path.read_text(encoding="utf-8") + title = file_path.stem + docs.append({"title": title, "content": content}) + except Exception as e: + print(f"Failed to read {file_path}: {e}") return docs \ No newline at end of file diff --git a/src/main.py b/src/main.py index 1c5ce43..07a7909 100644 --- a/src/main.py +++ b/src/main.py @@ -1,4 +1,4 @@ -from src.cli import main - -if __name__ == "__main__": - main() \ No newline at end of file +# This file is intentionally left blank. +# The CLI functionality is provided in src/cli.py. +# To run the application, execute: +# python -m src.cli \ No newline at end of file diff --git a/src/tools.py b/src/tools.py index dbe4066..e7ade8e 100644 --- a/src/tools.py +++ b/src/tools.py @@ -1,30 +1,52 @@ from langchain.tools import tool -from src.vector_store import vector_store -from src.chunking import chunk_document -from typing import List +from typing import List, Dict, Any +from src.vector_store import QdrantVectorStore -@tool +# Instantiate a global vector store (will be overridden in main) +vector_store: QdrantVectorStore = None + +@tool("search_knowledge_base") def search_knowledge_base(query: str, max_results: int = 5) -> str: """ - Search the knowledge base for relevant information. - Returns a formatted string of results. + Search the knowledge base for the given query and return top results. + + Parameters: + query (str): The search query. + max_results (int): Number of results to return. + + Returns: + str: Formatted search results. """ - results = vector_store.semantic_search(query, k=max_results) + if vector_store is None: + raise ValueError("Vector store not initialized.") + results = vector_store.search(query, max_results) if not results: return "No relevant documents found." formatted = [] - for i, doc in enumerate(results, 1): - source = doc.metadata.get("source", "unknown") - snippet = doc.page_content[:500].replace("\n", " ") - formatted.append(f"{i}. Source: {source}\n{snippet}...") + for idx, res in enumerate(results, 1): + formatted.append( + f"{idx}. Title: {res['title']}\n Score: {res['score']:.4f}\n Snippet: {res['content'][:200]}..." + ) return "\n\n".join(formatted) -@tool +@tool("add_to_knowledge_base") def add_to_knowledge_base(content: str, title: str) -> str: """ - Add new content to the knowledge base. - The content is split into chunks before being stored. + Add a new document to the knowledge base. + + Parameters: + content (str): Full text of the document. + title (str): Title of the document. + + Returns: + str: Confirmation message. """ - docs = chunk_document(content, title) - vector_store.add_documents(docs) - return f"Added {len(docs)} chunks from '{title}' to the knowledge base." \ No newline at end of file + if vector_store is None: + raise ValueError("Vector store not initialized.") + document = { + "content": content, + "title": title, + "metadata": {}, + } + vector_store.add_documents([document]) + return f"Document '{title}' added to the knowledge base." \ No newline at end of file diff --git a/src/vector_store.py b/src/vector_store.py index 7237ad3..080bee9 100644 --- a/src/vector_store.py +++ b/src/vector_store.py @@ -1,51 +1,94 @@ -from typing import List -from langchain_ollama import OllamaEmbeddings -from langchain_qdrant import QdrantVectorStore -from langchain.schema import Document +import os +from typing import List, Dict, Any from qdrant_client import QdrantClient from qdrant_client.http import models as qdrant_models +from langchain_ollama import OllamaEmbeddings +from langchain_qdrant import Qdrant +from src.chunking import chunk_document -class QdrantStore: - """ - Wrapper around QdrantVectorStore that handles collection creation, - embedding generation via Ollama, and semantic search. - """ - def __init__(self, collection_name: str = "knowledge_base", host: str = "localhost", port: int = 6333): - self.client = QdrantClient(host=host, port=port) +class QdrantVectorStore: + def __init__( + self, + collection_name: str = "knowledge_base", + host: str = "localhost", + port: int = 6333, + embedding_model: str = "nomic-embed-text", + ): self.collection_name = collection_name - self.embeddings = OllamaEmbeddings(model="nomic-embed-text") + self.client = QdrantClient(host=host, port=port) + self.embedding_model = embedding_model + self.embeddings = OllamaEmbeddings(model=self.embedding_model) - # Determine vector size from a sample embedding - sample_vector = self.embeddings.embed_query("") - vector_size = len(sample_vector) - - # Create collection if it does not exist - if not self.client.has_collection(collection_name): - self.client.recreate_collection( - collection_name=collection_name, + # Ensure collection exists + if not self.client.collection_exists(self.collection_name): + self.client.create_collection( + collection_name=self.collection_name, vectors_config=qdrant_models.VectorParams( - size=vector_size, - distance="Cosine" - ) + size=self.embeddings.embedding_size, + distance=qdrant_models.Distance.COSINE, + ), ) - self.store = QdrantVectorStore( + self.qdrant = Qdrant( client=self.client, - collection_name=collection_name, - embeddings=self.embeddings + collection_name=self.collection_name, + embeddings=self.embeddings, ) - def add_documents(self, documents: List[Document]): + def add_documents(self, documents: List[Dict[str, Any]]): """ - Add a list of Documents to the vector store. + Add documents to the vector store. Each document dict should contain: + - content: full text + - title: document title + - metadata: optional dict """ - self.store.add_documents(documents) + points = [] + for doc in documents: + content = doc.get("content", "") + title = doc.get("title", "") + metadata = doc.get("metadata", {}) + chunks = chunk_document(content) + for idx, chunk in enumerate(chunks): + point = qdrant_models.PointStruct( + id=None, + vector=self.embeddings.embed_query(chunk), + payload={ + "title": title, + "chunk_index": idx, + "content": chunk, + **metadata, + }, + ) + points.append(point) - def semantic_search(self, query: str, k: int = 5): - """ - Perform a semantic similarity search and return top-k Documents. - """ - return self.store.similarity_search(query, k=k) + if points: + self.client.upsert( + collection_name=self.collection_name, + points=points, + ) -# Global instance used by tools and agent -vector_store = QdrantStore() \ No newline at end of file + def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]: + """ + Perform a semantic search in the vector store. + + Returns a list of dicts with keys: title, content, score. + """ + query_vector = self.embeddings.embed_query(query) + search_result = self.client.search( + collection_name=self.collection_name, + query_vector=query_vector, + limit=max_results, + with_payload=True, + with_vector=False, + ) + results = [] + for hit in search_result: + payload = hit.payload + results.append( + { + "title": payload.get("title", ""), + "content": payload.get("content", ""), + "score": hit.score, + } + ) + return results \ No newline at end of file