diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..a38d25c --- /dev/null +++ b/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