feat: solution for 6a02e23da6fe2e4ac16acf65

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 <file>, /search <query>, /load <dir>, /quit")
+    print("RAG Agent CLI. Commands: /add <title> <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()