diff --git a/solutions/6a02e23da6fe2e4ac16acf65/solution.py b/solutions/6a02e23da6fe2e4ac16acf65/solution.py index 10eebc5..330c7fe 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65/solution.py +++ b/solutions/6a02e23da6fe2e4ac16acf65/solution.py @@ -1,122 +1,113 @@ from pathlib import Path -import sys +import os +from typing import List -from langchain_ollama import Ollama, OllamaEmbeddings +# 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_splitters import RecursiveCharacterTextSplitter -from langchain.tools import tool +# Agent from langchain.agents import create_agent +# Document type from langchain_core.documents import Document -# ---------- LLM and embeddings ---------- -llm = Ollama( - model="llama3", # local Ollama model - temperature=0.7, -) - -embeddings = OllamaEmbeddings(model="nomic-embed-text") - -# ---------- Qdrant client ---------- -client = QdrantClient(":memory:") -client.create_collection( - collection_name="knowledge_base", - vectors_config=VectorParams(size=embeddings.embedding_size, distance=Distance.COSINE), -) -vector_store = QdrantVectorStore(client=client, collection_name="knowledge_base", embedding=embeddings) - -# ---------- Text splitter ---------- -splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) - -# ---------- Tools ---------- +# -------------------- 1. RAG 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 information found." - out_lines = [] + return "No relevant documents found." + response_lines = [] for doc, score in results: title = doc.metadata.get("title", "Untitled") - content_preview = doc.page_content[:200] + ("..." if len(doc.page_content) > 200 else "") - out_lines.append(f"Score: {score:.3f}\nTitle: {title}\nContent: {content_preview}") - return "\n\n".join(out_lines) + 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) @tool -def add_to_knowledge_base(content: str, title: str = "Untitled") -> str: +def add_to_knowledge_base(content: str, title: str) -> str: """Add a new document to the knowledge base.""" - chunks = splitter.split_text(content) - documents = [ - Document(page_content=chunk, metadata={"title": f"{title} (part {i+1})"}) - for i, chunk in enumerate(chunks) - ] - vector_store.add_documents(documents) - return f"Added {len(chunks)} chunks to the knowledge base under title '{title}'." + doc = Document(page_content=content, metadata={"title": title}) + vector_store.add_documents([doc]) + return f"Document '{title}' added successfully." -# ---------- Agent ---------- -system_prompt = """ -You are an assistant that can search and add information to a local knowledge base. -Use the tools `search_knowledge_base` and `add_to_knowledge_base` as needed. -""" - -agent = create_agent( - model=llm, - tools=[search_knowledge_base, add_to_knowledge_base], - system_prompt=system_prompt, +# -------------------- 2. Vector store setup -------------------- +client = QdrantClient(":memory:") +client.create_collection( + collection_name="knowledge", + vectors_config=VectorParams(size=384, distance=Distance.COSINE), ) -# ---------- CLI ---------- -def load_documents_from_dir(directory: Path): - for file_path in directory.rglob("*"): - if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md"}: - content = file_path.read_text(encoding="utf-8") - title = file_path.stem - add_to_knowledge_base(content=content, title=title) +embeddings = OllamaEmbeddings(model="nomic-embed-text") +vector_store = QdrantVectorStore( + client=client, + collection_name="knowledge", + embedding=embeddings, +) +# -------------------- 3. Text splitter -------------------- +splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100) + +# -------------------- 4. Agent -------------------- +agent = create_agent( + model=ChatOllama(model="llama3", temperature=0.2), + tools=[search_knowledge_base, add_to_knowledge_base], + system_prompt="You are a helpful assistant that can search and add documents to the knowledge base.", +) + +# -------------------- 5. Load docs from directory -------------------- +def load_docs_from_dir(directory: str) -> List[Document]: + docs = [] + for file_path in Path(directory).rglob("*.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.name} #{i+1}"}) + ) + return docs + +def init_knowledge_base(directory: str): + docs = load_docs_from_dir(directory) + vector_store.add_documents(docs) + +# -------------------- 6. Interactive CLI -------------------- def main(): - # Load initial docs if provided as first arg - if len(sys.argv) > 1: - load_documents_from_dir(Path(sys.argv[1])) - - print("Agent ready. Commands: /add <file>, /search <query>, /quit") + print("Initializing knowledge base...") + init_knowledge_base("./docs") # replace with your docs folder + print("Ready! Use /add, /search, or /quit.") while True: user_input = input("> ").strip() if not user_input: continue - if user_input.lower() in {"quit", "exit"} or user_input == "/quit": - print("Goodbye!") + if user_input.lower() == "/quit": break - 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], Path(parts[2]) - if not file_path.is_file(): - print(f"File {file_path} does not exist.") - continue - content = file_path.read_text(encoding="utf-8") - result = add_to_knowledge_base(content=content, title=title) - print(result) - + try: + _, title, content = user_input.split(" ", 2) + result = add_to_knowledge_base(content=content, title=title) + print(result) + except ValueError: + print("Usage: /add <title> <content>") elif user_input.startswith("/search"): query = user_input[len("/search"):].strip() if not query: - print("Usage: /search <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) - + result = search_knowledge_base(query=query, max_results=3) + print(result) else: # Regular chat with agent response = agent.invoke({"messages": [{"role": "human", "content": user_input}]}) - for msg in response["messages"]: - if hasattr(msg, "content"): - print(msg.content) + ai_msg = response["messages"][-1] + print(ai_msg.content) if __name__ == "__main__": main() \ No newline at end of file