108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
"""
|
|
Agent with RAG memory using Qdrant and Ollama.
|
|
|
|
This is a minimal example that demonstrates:
|
|
1. Connecting to a local Qdrant instance.
|
|
2. Creating two tools: `search_knowledge_base` and `add_to_knowledge_base`.
|
|
3. Using LangChain's RecursiveCharacterTextSplitter to chunk documents.
|
|
4. Building an agent with the tools via `create_agent`.
|
|
5. A simple REPL that accepts `/add`, `/search` and `/quit` commands.
|
|
|
|
To run:
|
|
pip install -r requirements.txt
|
|
python main.py
|
|
|
|
Make sure a Qdrant instance is running locally (default port 6333) and Ollama is available at http://localhost:11434.
|
|
"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import List, Dict
|
|
|
|
from langchain_community.document_loaders import DirectoryLoader
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain.embeddings.ollama import OllamaEmbeddings
|
|
from langchain.vectorstores.qdrant import QdrantVectorStore
|
|
from langchain.agents import create_agent, AgentExecutor, Tool
|
|
from langchain.schema import Document
|
|
|
|
# Configuration
|
|
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
|
|
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama2")
|
|
COLLECTION_NAME = "rag_collection"
|
|
|
|
# Initialize embeddings and vector store
|
|
embeddings = OllamaEmbeddings(model=OLLAMA_MODEL)
|
|
vector_store = QdrantVectorStore(
|
|
url=QDRANT_URL,
|
|
collection_name=COLLECTION_NAME,
|
|
embedding_function=embeddings,
|
|
)
|
|
|
|
# Tool: search knowledge base
|
|
|
|
def search_knowledge_base(query: str) -> str:
|
|
"""Return the top 3 relevant snippets for a query."""
|
|
docs = vector_store.similarity_search_with_score(query, k=3)
|
|
if not docs:
|
|
return "No relevant documents found."
|
|
results = [f"{idx+1}. {doc[0].page_content[:200]}… (score: {doc[1]:.4f})" for idx, doc in enumerate(docs)]
|
|
return "\n".join(results)
|
|
|
|
# Tool: add to knowledge base
|
|
|
|
def add_to_knowledge_base(file_path: str) -> str:
|
|
"""Load a text file, chunk it and add to Qdrant."""
|
|
loader = DirectoryLoader(Path(file_path).parent.as_posix(), glob=Path(file_path).name)
|
|
docs = loader.load()
|
|
if not docs:
|
|
return f"No documents found in {file_path}."
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
chunks: List[Document] = []
|
|
for doc in docs:
|
|
chunks.extend(splitter.split_documents([doc]))
|
|
vector_store.add_documents(chunks)
|
|
return f"Added {len(chunks)} chunks from {file_path} to the knowledge base."
|
|
|
|
# Define tools list
|
|
TOOLS: List[Tool] = [
|
|
Tool(
|
|
name="search_knowledge_base",
|
|
func=search_knowledge_base,
|
|
description="Search the local Qdrant knowledge base for relevant information.",
|
|
),
|
|
Tool(
|
|
name="add_to_knowledge_base",
|
|
func=add_to_knowledge_base,
|
|
description="Add a text file to the knowledge base. Provide full path.",
|
|
),
|
|
]
|
|
|
|
# Build agent
|
|
agent = create_agent(TOOLS, llm=embeddings) # embeddings can act as LLM via Ollama
|
|
executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)
|
|
|
|
# Simple REPL
|
|
if __name__ == "__main__":
|
|
print("Welcome to the RAG agent. Commands: /add <file>, /search <query>, /quit")
|
|
while True:
|
|
try:
|
|
inp = input(">>> ")
|
|
except EOFError:
|
|
break
|
|
if not inp:
|
|
continue
|
|
if inp.startswith("/quit"):
|
|
print("Goodbye!")
|
|
break
|
|
elif inp.startswith("/add "):
|
|
path = inp.split(maxsplit=1)[1]
|
|
print(add_to_knowledge_base(path))
|
|
elif inp.startswith("/search "):
|
|
query = inp.split(maxsplit=1)[1]
|
|
print(search_knowledge_base(query))
|
|
else:
|
|
# Treat as normal agent prompt
|
|
result = executor.invoke({"input": inp})
|
|
print(result.get("output", ""))
|