add agent.py

agent.py

@@ -0,0 +1,140 @@
+"""
+Agent and vector store setup for the RAG task.
+
+This module defines:
+* `QdrantVectorStore` wrapper that uses Ollama embeddings.
+* Two tools – ``search_knowledge_base`` and ``add_to_knowledge_base``.
+* A helper to create the agent via :func:`langchain.agents.create_agent`.
+"""
+
+import os
+from pathlib import Path
+from typing import List, Dict
+
+from langchain_ollama import ChatOllama, OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from langchain.tools import tool
+from langchain.agents import create_agent
+from langchain_core.messages import HumanMessage
+
+# ---------------------------------------------------------------------------
+# Vector store configuration
+# ---------------------------------------------------------------------------
+
+QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
+QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
+COLLECTION_NAME = "knowledge"
+
+# Initialize embeddings and vector store.  The client is created lazily on first use.
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+vector_store: QdrantVectorStore | None = None
+
+
+def get_vector_store() -> QdrantVectorStore:
+    """Return a singleton Qdrant vector store instance.
+
+    The collection is created automatically if it does not exist.
+    """
+    global vector_store
+    if vector_store is None:
+        from qdrant_client import QdrantClient
+        client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
+        vector_store = QdrantVectorStore(
+            client=client,
+            collection_name=COLLECTION_NAME,
+            embedding=embeddings,
+        )
+    return vector_store
+
+# ---------------------------------------------------------------------------
+# Tools
+# ---------------------------------------------------------------------------
+
+@tool("search_knowledge_base")
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
+    """Semantic search in the knowledge base.
+
+    Parameters
+    ----------
+    query: str
+        Search query.
+    max_results: int, optional
+        Number of top results to return. Defaults to 5.
+
+    Returns
+    -------
+    str
+        A numbered list of passages or a message if nothing was found.
+    """
+    store = get_vector_store()
+    docs = store.similarity_search(query, k=max_results)
+    if not docs:
+        return "No relevant documents found."
+    return "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs))
+
+@tool("add_to_knowledge_base")
+def add_to_knowledge_base(content: str, title: str = "document") -> str:
+    """Add a document to the knowledge base.
+
+    Parameters
+    ----------
+    content: str
+        Text of the document.
+    title: str, optional
+        Title used as metadata. Defaults to ``"document"``.
+
+    Returns
+    -------
+    str
+        Confirmation message.
+    """
+    store = get_vector_store()
+    from langchain_core.documents import Document
+    doc = Document(page_content=content, metadata={"title": title})
+    store.add_documents([doc])
+    return f"Added '{title}' to knowledge base."
+
+# ---------------------------------------------------------------------------
+# Agent creation helper
+# ---------------------------------------------------------------------------
+
+llm = ChatOllama(model="llama3", temperature=0.0)
+
+SYSTEM_PROMPT = (
+    "You are an assistant that can search and add information to a local knowledge base.\n"
+    "Use the tools `search_knowledge_base` and `add_to_knowledge_base`.\n"
+    "When searching, return the most relevant passages. When adding, confirm success."
+)
+
+
+def create_rag_agent():
+    """Return a LangChain agent configured with the RAG tools."""
+    agent = create_agent(
+        llm=llm,
+        tools=[search_knowledge_base, add_to_knowledge_base],
+        system_prompt=SYSTEM_PROMPT,
+    )
+    return agent
+
+# ---------------------------------------------------------------------------
+# Example usage (can be imported by main.py)
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    # Simple demo: add a short doc and search it.
+    agent = create_rag_agent()
+    print("Adding sample document...")
+    res = agent.invoke(
+        {"messages": [HumanMessage(content="Add to knowledge base: content='Python is great' title='Python intro'")]},
+        {"configurable": {"thread_id": "demo-1"}},
+    )
+    print(res["messages"][-1].content)
+
+    print("Searching for Python...")
+    res = agent.invoke(
+        {"messages": [HumanMessage(content="Search for Python')"],},
+        {"configurable": {"thread_id": "demo-2"}},
+    )
+    print(res["messages"][-1].content)
+
+# End of agent.py