From 160dcf1a17e734d576e82ffbb3ddfb117464943c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D0=B5=D0=BC=20=D0=92=D0=BB=D0=B0=D0=B4?= =?UTF-8?q?=D0=B8=D0=BC=D0=B8=D1=80=D0=BE=D0=B2=D0=B8=D1=87=20=D0=91=D0=B0?= =?UTF-8?q?=D0=B1=D0=B0=D0=B9=D0=BA=D0=B8=D0=BD?= Date: Thu, 28 May 2026 16:05:04 +0000 Subject: [PATCH] feat: solution for 6a02e23da6fe2e4ac16acf65 --- .../6a02e23da6fe2e4ac16acf65/solution.py | 115 ++++++++---------- 1 file changed, 49 insertions(+), 66 deletions(-) diff --git a/solutions/6a02e23da6fe2e4ac16acf65/solution.py b/solutions/6a02e23da6fe2e4ac16acf65/solution.py index 988b5ec..377d54a 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65/solution.py +++ b/solutions/6a02e23da6fe2e4ac16acf65/solution.py @@ -18,75 +18,58 @@ llm = ChatOpenAI( temperature=0.7, ) -# ---------- Embeddings ---------- -embeddings = OllamaEmbeddings(model="nomic-embed-text") - -# ---------- Qdrant ---------- +# ---------- Vector Store ---------- client = QdrantClient(":memory:") -collection_name = "knowledge_base" - -try: - client.get_collection(collection_name) -except Exception: - # Use a typical embedding size for nomic-embed-text (768) - client.create_collection( - collection_name=collection_name, - vectors_config=VectorParams(size=768, distance=Distance.COSINE), - ) - +client.create_collection( + collection_name="knowledge_base", + vectors_config=VectorParams(size=1024, distance=Distance.COSINE), +) +embeddings = OllamaEmbeddings(model="nomic-embed-text") vector_store = QdrantVectorStore( client=client, - collection_name=collection_name, + collection_name="knowledge_base", embedding=embeddings, ) -# ---------- Text splitter ---------- -splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) +# ---------- Text Splitter ---------- +splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100) # ---------- Tools ---------- @tool -def search_knowledge_base(query: str, max_results: int = 5) -> str: - """Search the knowledge base for relevant documents.""" - docs_with_score = vector_store.similarity_search_with_score(query, k=max_results) - if not docs_with_score: - return "No results found." - return "\n".join( - f"{i+1}. {doc.page_content[:200]}..." - for i, (doc, _) in enumerate(docs_with_score) - ) +def add_to_knowledge_base(content: str, title: str) -> str: + """Add a document to the knowledge base.""" + docs = splitter.split_text(content) + documents = [Document(page_content=c, metadata={"title": title}) for c in docs] + vector_store.add_documents(documents) + return f"Added {len(docs)} chunks under title '{title}'." @tool -def add_to_knowledge_base(content: str, title: str = "") -> str: - """Add a new document to the knowledge base.""" - chunks = splitter.split_text(content) - docs = [Document(page_content=c, metadata={"title": title}) for c in chunks] - vector_store.add_documents(docs) - return f"Added {len(chunks)} chunks under title '{title}'." +def search_knowledge_base(query: str, max_results: int = 5) -> str: + """Search the knowledge base for relevant information.""" + results = vector_store.similarity_search_with_score(query, k=max_results) + if not results: + return "No relevant documents found." + out_lines = [] + for doc, score in results: + title = doc.metadata.get("title", "Untitled") + out_lines.append(f"[{score:.2f}] {title}: {doc.page_content[:200]}...") + return "\n".join(out_lines) # ---------- Agent ---------- system_prompt = """ -You are an assistant that can search and add information to a knowledge base. -Use the tools `search_knowledge_base` and `add_to_knowledge_base` as needed. +You are an assistant with access to a knowledge base. +Use the tools `add_to_knowledge_base` and `search_knowledge_base` as needed. +Respond concisely. If you need more info, ask the user. """ - agent = create_agent( model=llm, - tools=[search_knowledge_base, add_to_knowledge_base], + tools=[add_to_knowledge_base, search_knowledge_base], system_prompt=system_prompt, ) # ---------- CLI ---------- -def load_directory(path: str): - """Load all text files from a directory into the knowledge base.""" - for root, _, files in os.walk(path): - for file in files: - if file.lower().endswith(".txt"): - with open(os.path.join(root, file), encoding="utf-8") as f: - content = f.read() - add_to_knowledge_base(content=content, title=file) - def main(): - print("Welcome to the RAG agent. Commands: /add , /search , /load , /quit") + print("RAG Agent CLI. Commands: /add <content>, /search <query>, /quit") while True: try: inp = input("> ").strip() @@ -94,29 +77,29 @@ def main(): break if not inp: continue - if inp.lower() in ("/quit", "exit"): - print("Goodbye!") + if inp.lower() in ("exit", "quit", "/quit"): + print("Bye!") break if inp.startswith("/add "): - _, file_path = inp.split(maxsplit=1) - try: - with open(file_path, encoding="utf-8") as f: - content = f.read() - print(add_to_knowledge_base(content=content, title=os.path.basename(file_path))) - except Exception as e: - print(f"Error adding file: {e}") + parts = inp[5:].split(None, 1) + if len(parts) != 2: + print("Usage: /add <title> <content>") + continue + title, content = parts + res = add_to_knowledge_base(content=content, title=title) + print(res) elif inp.startswith("/search "): - _, query = inp.split(maxsplit=1) - print(search_knowledge_base(query=query)) - elif inp.startswith("/load "): - _, dir_path = inp.split(maxsplit=1) - load_directory(dir_path) - print(f"Loaded documents from {dir_path}") + query = inp[8:].strip() + if not query: + print("Usage: /search <query>") + continue + res = search_knowledge_base(query=query, max_results=5) + print(res) else: - # Regular conversation - response = agent.invoke({"messages": [{"role": "human", "content": inp}]}) - msg = response["messages"][-1] - print(msg.content) + # Regular chat with agent + result = agent.invoke({"messages": [{"role": "human", "content": inp}]}) + ai_msg = result["messages"][-1] + print(ai_msg.content) if __name__ == "__main__": main() \ No newline at end of file