128 lines
4.6 KiB
Python
128 lines
4.6 KiB
Python
# Main script implementing Qdrant-based RAG agent
|
|
import os
|
|
import sys
|
|
import json
|
|
from pathlib import Path
|
|
from typing import List, Dict, Any
|
|
|
|
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
|
from langchain_qdrant import QdrantVectorStore
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain.tools import tool, BaseTool
|
|
from langchain.agents import create_agent, AgentExecutor, AgentType
|
|
|
|
# Configuration
|
|
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
|
|
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
|
|
COLLECTION_NAME = "knowledge"
|
|
EMBEDDING_MODEL = "nomic-embed-text"
|
|
LLM_MODEL = "llama3"
|
|
|
|
# Vector store wrapper
|
|
class QdrantStore:
|
|
def __init__(self, host: str, port: int, collection: str):
|
|
self.store = QdrantVectorStore(
|
|
url=f"http://{host}:{port}",
|
|
collection_name=collection,
|
|
embedding=OllamaEmbeddings(model=EMBEDDING_MODEL),
|
|
)
|
|
|
|
def add_documents(self, documents: List[str], metadatas: List[Dict[str, Any]]):
|
|
self.store.add_texts(documents, metadatas=metadatas)
|
|
|
|
def similarity_search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
|
|
results = self.store.similarity_search(query, k=k)
|
|
return [
|
|
{
|
|
"content": doc.page_content,
|
|
"metadata": doc.metadata,
|
|
"score": doc.metadata.get("score", 0),
|
|
}
|
|
for doc in results
|
|
]
|
|
|
|
# Text splitter
|
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
|
|
|
# Store instance
|
|
store = QdrantStore(QDRANT_HOST, QDRANT_PORT, COLLECTION_NAME)
|
|
|
|
# Tools
|
|
@tool("search_knowledge_base", "Semantic search in the knowledge base")
|
|
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
|
results = store.similarity_search(query, k=max_results)
|
|
return json.dumps(results, ensure_ascii=False)
|
|
|
|
@tool("add_to_knowledge_base", "Add a document to the knowledge base")
|
|
def add_to_knowledge_base(content: str, title: str = "Untitled") -> str:
|
|
chunks = text_splitter.split_text(content)
|
|
metadatas = [{"title": title, "chunk_idx": i} for i in range(len(chunks))]
|
|
store.add_documents(chunks, metadatas)
|
|
return f"Added {len(chunks)} chunks titled '{title}'."
|
|
|
|
# Agent
|
|
llm = ChatOllama(model=LLM_MODEL)
|
|
agent = create_agent(
|
|
llm=llm,
|
|
tools=[search_knowledge_base, add_to_knowledge_base],
|
|
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
|
system_message="You are an assistant that can search and add information to a local knowledge base. Use the tools when appropriate.",
|
|
)
|
|
executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)
|
|
|
|
# CLI helpers
|
|
|
|
def load_documents_from_dir(dir_path: str):
|
|
for path in Path(dir_path).glob("**/*"):
|
|
if path.suffix.lower() in {".txt", ".md"}:
|
|
content = path.read_text(encoding="utf-8")
|
|
title = path.stem
|
|
add_to_knowledge_base(content, title)
|
|
print("Loading complete.")
|
|
|
|
def main():
|
|
if len(sys.argv) > 1 and sys.argv[1] == "load":
|
|
if len(sys.argv) < 3:
|
|
print("Usage: python main.py load <directory>")
|
|
sys.exit(1)
|
|
load_documents_from_dir(sys.argv[2])
|
|
sys.exit(0)
|
|
print("Interactive mode. Commands: /add <title> <file>, /search <query>, /quit")
|
|
while True:
|
|
try:
|
|
user_input = input("\n> ")
|
|
except (EOFError, KeyboardInterrupt):
|
|
break
|
|
if not user_input:
|
|
continue
|
|
if user_input.startswith("/quit"):
|
|
break
|
|
if user_input.startswith("/add"):
|
|
parts = user_input.split(maxsplit=2)
|
|
if len(parts) != 3:
|
|
print("Usage: /add <title> <file_path>")
|
|
continue
|
|
title, file_path = parts[1], parts[2]
|
|
try:
|
|
content = Path(file_path).read_text(encoding="utf-8")
|
|
except Exception as e:
|
|
print(f"Error reading file: {e}")
|
|
continue
|
|
print(add_to_knowledge_base(content, title))
|
|
continue
|
|
if user_input.startswith("/search"):
|
|
query = user_input[len("/search"):].strip()
|
|
if not query:
|
|
print("Provide a query.")
|
|
continue
|
|
results = search_knowledge_base(query)
|
|
print("Search results:")
|
|
for r in json.loads(results):
|
|
print(f"- {r['metadata'].get('title', 'Untitled')} (chunk {r['metadata'].get('chunk_idx')})\n {r['content'][:200]}...")
|
|
continue
|
|
response = executor.invoke({"input": user_input})
|
|
print(response["output"])
|
|
|
|
if __name__ == "__main__":
|
|
main()
|