feat: solution for 6a02e23da6fe2e4ac16acf65

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 <title> <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()