diff --git a/solutions/6a02e23da6fe2e4ac16acf65/solution.py b/solutions/6a02e23da6fe2e4ac16acf65/solution.py index 67aa142..55fc8a4 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65/solution.py +++ b/solutions/6a02e23da6fe2e4ac16acf65/solution.py @@ -1,129 +1,92 @@ -from pathlib import Path -import os -from typing import List - -# LLM and embeddings via Ollama from langchain_ollama import ChatOllama, OllamaEmbeddings -# Tools -from langchain.tools import tool -# Vector store from langchain_qdrant import QdrantVectorStore from qdrant_client import QdrantClient from qdrant_client.http.models import Distance, VectorParams -# Text splitter -from langchain.text_splitter import RecursiveCharacterTextSplitter -# Agent -from langchain.agents import create_agent -# Documents from langchain_core.documents import Document +from langchain.tools import tool +from langchain.agents import create_agent +import os -# -------------------- 1. RAG tools -------------------- +# ---------- LLM and embeddings ---------- +llm = ChatOllama(model="llama3", temperature=0.2) +embeddings = OllamaEmbeddings(model="nomic-embed-text") + +# ---------- Qdrant client & collection ---------- +client = QdrantClient(":memory:") +collection_name = "knowledge_base" +client.create_collection( + collection_name=collection_name, + vectors_config=VectorParams(size=embeddings.embedding_size, distance=Distance.COSINE), +) +vector_store = QdrantVectorStore(client=client, collection_name=collection_name, embedding=embeddings) + +# ---------- Text splitter ---------- +from langchain_text_splitters import RecursiveCharacterTextSplitter +splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + +# ---------- Tools ---------- @tool def search_knowledge_base(query: str, max_results: int = 5) -> str: """Search the knowledge base for relevant documents.""" results = vector_store.similarity_search_with_score(query, k=max_results) if not results: return "No relevant documents found." - response_lines = [] + out_lines = [] for doc, score in results: - title = doc.metadata.get("title", "N/A") - snippet = doc.page_content[:200] + ("..." if len(doc.page_content) > 200 else "") - response_lines.append(f"Score: {score:.4f}\nTitle: {title}\nContent: {snippet}") - return "\n\n".join(response_lines) + title = doc.metadata.get("title", "Untitled") + content = doc.page_content[:200] + ("..." if len(doc.page_content) > 200 else "") + out_lines.append(f"Title: {title}\nScore: {score:.4f}\nContent: {content}") + return "\n\n".join(out_lines) @tool -def add_to_knowledge_base(content: str, title: str) -> str: +def add_to_knowledge_base(content: str, title: str = "Untitled") -> str: """Add a new document to the knowledge base.""" - doc = Document(page_content=content, metadata={"title": title}) - vector_store.add_documents([doc]) - return f"Document '{title}' added successfully." - -# -------------------- 2. Qdrant setup -------------------- -qdrant_client = QdrantClient(":memory:") -qdrant_client.create_collection( - collection_name="knowledge_base", - vectors_config=VectorParams(size=384, distance=Distance.COSINE), -) - -embeddings = OllamaEmbeddings(model="nomic-embed-text") -vector_store = QdrantVectorStore( - client=qdrant_client, - collection_name="knowledge_base", - embedding=embeddings, -) - -# -------------------- 3. Text splitter -------------------- -splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100) - -def load_and_index(directory: str): - """Load all .txt files from directory and index them.""" - docs: List[Document] = [] - for file_path in Path(directory).glob("*.txt"): - text = file_path.read_text(encoding="utf-8") - chunks = splitter.split_text(text) - for i, chunk in enumerate(chunks): - docs.append( - Document( - page_content=chunk, - metadata={ - "title": f"{file_path.stem} #{i+1}", - "source": str(file_path), - }, - ) - ) + chunks = splitter.split_text(content) + docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks] vector_store.add_documents(docs) + return f"Added {len(chunks)} chunks under title '{title}'." -# -------------------- 4. Agent -------------------- +# ---------- Agent ---------- system_prompt = """ You are an assistant that can search and add documents to a knowledge base. Use the tools `search_knowledge_base` and `add_to_knowledge_base` as needed. +Respond with plain text. Do not mention tool usage explicitly unless required by the user. """ - agent = create_agent( - model=ChatOllama(model="llama3", temperature=0.2), + model=llm, tools=[search_knowledge_base, add_to_knowledge_base], system_prompt=system_prompt, ) -# -------------------- 5. CLI client -------------------- +# ---------- CLI ---------- def main(): - # Load initial documents - load_and_index("docs") # ensure a 'docs' folder with .txt files - - print( - "RAG Agent ready. Commands: /add <content>, /search <query>, /quit" - ) + print("RAG Agent CLI. Commands: /add <title> | <content>, /search <query>, /quit") while True: - user_input = input("> ").strip() + user_input = input("\nYou: ").strip() if not user_input: continue if user_input.lower() in ("exit", "quit", "/quit"): + print("Goodbye!") break - - if user_input.startswith("/add"): + if user_input.startswith("/add "): try: - _, title, content = user_input.split(" ", 2) + _, rest = user_input.split(maxsplit=1) + title, content = rest.split("|", 1) + title = title.strip() + content = content.strip() result_msg = add_to_knowledge_base(content=content, title=title) - print(result_msg) + print(f"Bot: {result_msg}") except ValueError: - print("Usage: /add <title> <content>") - elif user_input.startswith("/search"): - query = user_input[len("/search") :].strip() - if not query: - print("Provide a search query.") - continue - response = agent.invoke({"messages": [{"role": "human", "content": query}]}) - for msg in response["messages"]: - if hasattr(msg, "content"): - print(msg.content) + print("Bot: Usage /add <title> | <content>") + elif user_input.startswith("/search "): + query = user_input[len("/search "):].strip() + result_msg = search_knowledge_base(query=query, max_results=5) + print(f"Bot:\n{result_msg}") else: - # Regular chat with the agent - response = agent.invoke( - {"messages": [{"role": "human", "content": user_input}]} - ) - for msg in response["messages"]: - if hasattr(msg, "content"): - print(msg.content) + # Regular chat + response = agent.invoke({"messages": [{"role": "human", "content": user_input}]}) + ai_message = response["messages"][-1] + print(f"Bot: {ai_message.content}") if __name__ == "__main__": main() \ No newline at end of file