feat: solution for 6a02e23da6fe2e4ac16acf65

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 <title> <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 "):
             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>")
+                print("Bot: 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)
+            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()