diff --git a/main.py b/main.py index 2dc9a82..7fc1645 100644 --- a/main.py +++ b/main.py @@ -1,8 +1,4 @@ -"""Main module implementing a RAG-enabled agent with Qdrant and Ollama. - -The code follows the assignment specification and uses only the required libraries. -""" - +# Main script implementing Qdrant-based RAG agent import os import sys import json @@ -12,162 +8,120 @@ from typing import List, Dict, Any from langchain_ollama import ChatOllama, OllamaEmbeddings from langchain_qdrant import QdrantVectorStore from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain.tools import tool -from langchain.agents import create_agent, AgentExecutor -from langchain.schema import AgentAction, AgentFinish -from langchain_core.prompts import ChatPromptTemplate +from langchain.tools import tool, BaseTool +from langchain.agents import create_agent, AgentExecutor, AgentType -# --------------------------------------------------------------------------- # Configuration -# --------------------------------------------------------------------------- -QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333") -QDRANT_COLLECTION = "knowledge" +QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost") +QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333")) +COLLECTION_NAME = "knowledge" EMBEDDING_MODEL = "nomic-embed-text" LLM_MODEL = "llama3" -CHUNK_SIZE = 1000 # characters -CHUNK_OVERLAP = 200 -# --------------------------------------------------------------------------- -# Vector store helper -# --------------------------------------------------------------------------- -class KnowledgeBase: - """Wrapper around QdrantVectorStore providing add/search helpers.""" - - def __init__(self, url: str = QDRANT_URL, collection: str = QDRANT_COLLECTION): - self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) - self.store = QdrantVectorStore.from_existing_collection( +# Vector store wrapper +class QdrantStore: + def __init__(self, host: str, port: int, collection: str): + self.store = QdrantVectorStore( + url=f"http://{host}:{port}", collection_name=collection, - url=url, - embedding=self.embeddings, + embedding=OllamaEmbeddings(model=EMBEDDING_MODEL), ) - def add_documents(self, documents: List[str], titles: List[str]): - """Adds a list of documents with corresponding titles to the store. + def add_documents(self, documents: List[str], metadatas: List[Dict[str, Any]]): + self.store.add_texts(documents, metadatas=metadatas) - Each document is split into chunks before being stored. - """ - splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP) - for doc, title in zip(documents, titles): - chunks = splitter.split_text(doc) - metadatas = [{"title": title, "chunk_index": i} for i in range(len(chunks))] - self.store.add_texts(chunks, metadatas=metadatas) - - def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]: - """Semantic search in the knowledge base. - - Returns a list of dicts with keys: text, title, distance. - """ - results = self.store.similarity_search_with_score(query, k=max_results) + def similarity_search(self, query: str, k: int = 5) -> List[Dict[str, Any]]: + results = self.store.similarity_search(query, k=k) return [ { - "text": text, - "title": meta.get("title", "unknown"), - "distance": score, + "content": doc.page_content, + "metadata": doc.metadata, + "score": doc.metadata.get("score", 0), } - for text, score, meta in results + for doc in results ] -# --------------------------------------------------------------------------- -# Global knowledge base instance -# --------------------------------------------------------------------------- -kb = KnowledgeBase() +# Text splitter +text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + +# Store instance +store = QdrantStore(QDRANT_HOST, QDRANT_PORT, COLLECTION_NAME) -# --------------------------------------------------------------------------- # Tools -# --------------------------------------------------------------------------- -@tool("search_knowledge_base") +@tool("search_knowledge_base", "Semantic search in the knowledge base") def search_knowledge_base(query: str, max_results: int = 5) -> str: - """Search the knowledge base for a query and return formatted results.""" - results = kb.search(query, max_results) - if not results: - return "No relevant documents found." - formatted = [f"Title: {r['title']}\nSnippet: {r['text'][:200]}...\nDistance: {r['distance']:.4f}" for r in results] - return "\n\n".join(formatted) + results = store.similarity_search(query, k=max_results) + return json.dumps(results, ensure_ascii=False) -@tool("add_to_knowledge_base") -def add_to_knowledge_base(content: str, title: str) -> str: - """Add a new document to the knowledge base.""" - kb.add_documents([content], [title]) - return f"Document '{title}' added successfully." - -# --------------------------------------------------------------------------- -# Agent definition -# --------------------------------------------------------------------------- -SYSTEM_PROMPT = ( - "You are an assistant that can search and add documents to a local knowledge base. " - "Use the provided tools to perform semantic search and store new information. " - "When answering user queries, first determine if the user needs a search or an addition. " - "If no relevant information is found, suggest adding new content." -) - -prompt = ChatPromptTemplate.from_messages([ - ("system", SYSTEM_PROMPT), - ("user", "{input}"), -]) +@tool("add_to_knowledge_base", "Add a document to the knowledge base") +def add_to_knowledge_base(content: str, title: str = "Untitled") -> str: + chunks = text_splitter.split_text(content) + metadatas = [{"title": title, "chunk_idx": i} for i in range(len(chunks))] + store.add_documents(chunks, metadatas) + return f"Added {len(chunks)} chunks titled '{title}'." +# Agent +llm = ChatOllama(model=LLM_MODEL) agent = create_agent( - llm=ChatOllama(model=LLM_MODEL), + llm=llm, tools=[search_knowledge_base, add_to_knowledge_base], - prompt=prompt, - verbose=True, + agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, + system_message="You are an assistant that can search and add information to a local knowledge base. Use the tools when appropriate.", ) +executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True) -agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base]) +# CLI helpers -# --------------------------------------------------------------------------- -# CLI client -# --------------------------------------------------------------------------- -def load_documents_from_dir(directory: str): - """Load all .txt files from a directory and add them to the knowledge base.""" - docs = [] - titles = [] - for path in Path(directory).glob("**/*.txt"): - text = path.read_text(encoding="utf-8") - docs.append(text) - titles.append(path.stem) - if docs: - kb.add_documents(docs, titles) - print(f"Loaded {len(docs)} documents from {directory}.") - else: - print("No .txt files found.") +def load_documents_from_dir(dir_path: str): + for path in Path(dir_path).glob("**/*"): + if path.suffix.lower() in {".txt", ".md"}: + content = path.read_text(encoding="utf-8") + title = path.stem + add_to_knowledge_base(content, title) + print("Loading complete.") - -def interactive_loop(): - print("RAG Agent CLI. Commands: /add <content>, /search <query>, /load <dir>, /quit") +def main(): + if len(sys.argv) > 1 and sys.argv[1] == "load": + if len(sys.argv) < 3: + print("Usage: python main.py load <directory>") + sys.exit(1) + load_documents_from_dir(sys.argv[2]) + sys.exit(0) + print("Interactive mode. Commands: /add <title> <file>, /search <query>, /quit") while True: try: - user_input = input("> ") - except (KeyboardInterrupt, EOFError): - print("\nExiting.") + user_input = input("\n> ") + except (EOFError, KeyboardInterrupt): break if not user_input: continue if user_input.startswith("/quit"): - print("Goodbye!") break - elif user_input.startswith("/add "): + if user_input.startswith("/add"): + parts = user_input.split(maxsplit=2) + if len(parts) != 3: + print("Usage: /add <title> <file_path>") + continue + title, file_path = parts[1], parts[2] try: - _, rest = user_input.split("/add ", 1) - title, content = rest.split(" ", 1) - response = add_to_knowledge_base(content, title) - print(response) - except ValueError: - print("Usage: /add <title> <content>") - elif user_input.startswith("/search "): - query = user_input.split("/search ", 1)[1] + content = Path(file_path).read_text(encoding="utf-8") + except Exception as e: + print(f"Error reading file: {e}") + continue + print(add_to_knowledge_base(content, title)) + continue + if user_input.startswith("/search"): + query = user_input[len("/search"):].strip() + if not query: + print("Provide a query.") + continue results = search_knowledge_base(query) - print(results) - elif user_input.startswith("/load "): - dir_path = user_input.split("/load ", 1)[1] - load_documents_from_dir(dir_path) - else: - # Treat as normal user query - result = agent_executor.invoke({"input": user_input}) - print(result.get("output", "")) + print("Search results:") + for r in json.loads(results): + print(f"- {r['metadata'].get('title', 'Untitled')} (chunk {r['metadata'].get('chunk_idx')})\n {r['content'][:200]}...") + continue + response = executor.invoke({"input": user_input}) + print(response["output"]) if __name__ == "__main__": - if len(sys.argv) > 1 and sys.argv[1] == "cli": - interactive_loop() - else: - print("Usage: python main.py cli") - print("Run the interactive CLI with: python main.py cli") + main()