From 9001ea3fa8cdc6e21db7b712e5e510fe2e07cd32 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:27:58 +0000 Subject: [PATCH] feat: solution for 6a02e23da6fe2e4ac16acf65 --- .../6a02e23da6fe2e4ac16acf65/solution.py | 87 ++++++++++--------- 1 file changed, 45 insertions(+), 42 deletions(-) diff --git a/solutions/6a02e23da6fe2e4ac16acf65/solution.py b/solutions/6a02e23da6fe2e4ac16acf65/solution.py index 377d54a..b12048d 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65/solution.py +++ b/solutions/6a02e23da6fe2e4ac16acf65/solution.py @@ -8,7 +8,6 @@ from langchain_ollama import OllamaEmbeddings from langchain_core.documents import Document from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain.agents import create_agent -import os # ---------- LLM ---------- llm = ChatOpenAI( @@ -18,53 +17,49 @@ llm = ChatOpenAI( temperature=0.7, ) -# ---------- Vector Store ---------- +# ---------- Embeddings ---------- +embeddings = OllamaEmbeddings(model="nomic-embed-text") + +# ---------- Qdrant client ---------- client = QdrantClient(":memory:") client.create_collection( collection_name="knowledge_base", - vectors_config=VectorParams(size=1024, distance=Distance.COSINE), + vectors_config=VectorParams(size=embeddings.embedding_size, distance=Distance.COSINE), ) -embeddings = OllamaEmbeddings(model="nomic-embed-text") vector_store = QdrantVectorStore( client=client, collection_name="knowledge_base", embedding=embeddings, ) -# ---------- Text Splitter ---------- +# ---------- Text splitter ---------- splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100) # ---------- Tools ---------- @tool -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}'." +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." + result_lines = [ + f"{i+1}. {doc.page_content[:200]}..." for i, (doc, _) in enumerate(docs_with_score) + ] + return "\n".join(result_lines) @tool -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) +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": title}) for chunk in chunks] + vector_store.add_documents(documents) + return f"Added {len(documents)} chunks from '{title}'." # ---------- Agent ---------- -system_prompt = """ -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=[add_to_knowledge_base, search_knowledge_base], - system_prompt=system_prompt, + tools=[search_knowledge_base, add_to_knowledge_base], + system_prompt="You are an assistant that can search and add documents to the knowledge base.", ) # ---------- CLI ---------- @@ -77,29 +72,37 @@ def main(): break if not inp: continue - if inp.lower() in ("exit", "quit", "/quit"): + if inp.lower() in ("quit", "exit"): print("Bye!") break + + # Add document if inp.startswith("/add "): parts = inp[5:].split(None, 1) if len(parts) != 2: print("Usage: /add <content>") continue title, content = parts - res = add_to_knowledge_base(content=content, title=title) - print(res) - elif inp.startswith("/search "): + res = agent.invoke({"messages": [{"role": "human", "content": f"Add document '{title}'"}]}) + for msg in res["messages"]: + if hasattr(msg, "tool_calls"): + print(msg.content) + continue + + # Search query + if inp.startswith("/search "): 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 chat with agent - result = agent.invoke({"messages": [{"role": "human", "content": inp}]}) - ai_msg = result["messages"][-1] - print(ai_msg.content) + res = agent.invoke({"messages": [{"role": "human", "content": f"Search: {query}"}]}) + for msg in res["messages"]: + if hasattr(msg, "tool_calls"): + print(msg.content) + continue + + # General chat + res = agent.invoke({"messages": [{"role": "human", "content": inp}]}) + for msg in res["messages"]: + if hasattr(msg, "content"): + print(msg.content) if __name__ == "__main__": main() \ No newline at end of file