Updated main.py with Qdrant-based implementation
This commit is contained in:
@@ -1,8 +1,4 @@
|
||||
"""Main module implementing a RAG-enabled agent with Qdrant and Ollama.
|
||||
|
||||
The code follows the assignment specification and uses only the required libraries.
|
||||
"""
|
||||
|
||||
# Main script implementing Qdrant-based RAG agent
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
@@ -12,162 +8,120 @@ 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
|
||||
from langchain.agents import create_agent, AgentExecutor
|
||||
from langchain.schema import AgentAction, AgentFinish
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain.tools import tool, BaseTool
|
||||
from langchain.agents import create_agent, AgentExecutor, AgentType
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
|
||||
QDRANT_COLLECTION = "knowledge"
|
||||
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"
|
||||
CHUNK_SIZE = 1000 # characters
|
||||
CHUNK_OVERLAP = 200
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Vector store helper
|
||||
# ---------------------------------------------------------------------------
|
||||
class KnowledgeBase:
|
||||
"""Wrapper around QdrantVectorStore providing add/search helpers."""
|
||||
|
||||
def __init__(self, url: str = QDRANT_URL, collection: str = QDRANT_COLLECTION):
|
||||
self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||||
self.store = QdrantVectorStore.from_existing_collection(
|
||||
# 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,
|
||||
url=url,
|
||||
embedding=self.embeddings,
|
||||
embedding=OllamaEmbeddings(model=EMBEDDING_MODEL),
|
||||
)
|
||||
|
||||
def add_documents(self, documents: List[str], titles: List[str]):
|
||||
"""Adds a list of documents with corresponding titles to the store.
|
||||
def add_documents(self, documents: List[str], metadatas: List[Dict[str, Any]]):
|
||||
self.store.add_texts(documents, metadatas=metadatas)
|
||||
|
||||
Each document is split into chunks before being stored.
|
||||
"""
|
||||
splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
|
||||
for doc, title in zip(documents, titles):
|
||||
chunks = splitter.split_text(doc)
|
||||
metadatas = [{"title": title, "chunk_index": i} for i in range(len(chunks))]
|
||||
self.store.add_texts(chunks, metadatas=metadatas)
|
||||
|
||||
def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
|
||||
"""Semantic search in the knowledge base.
|
||||
|
||||
Returns a list of dicts with keys: text, title, distance.
|
||||
"""
|
||||
results = self.store.similarity_search_with_score(query, k=max_results)
|
||||
def similarity_search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
|
||||
results = self.store.similarity_search(query, k=k)
|
||||
return [
|
||||
{
|
||||
"text": text,
|
||||
"title": meta.get("title", "unknown"),
|
||||
"distance": score,
|
||||
"content": doc.page_content,
|
||||
"metadata": doc.metadata,
|
||||
"score": doc.metadata.get("score", 0),
|
||||
}
|
||||
for text, score, meta in results
|
||||
for doc in results
|
||||
]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Global knowledge base instance
|
||||
# ---------------------------------------------------------------------------
|
||||
kb = KnowledgeBase()
|
||||
# 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")
|
||||
@tool("search_knowledge_base", "Semantic search in the knowledge base")
|
||||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||||
"""Search the knowledge base for a query and return formatted results."""
|
||||
results = kb.search(query, max_results)
|
||||
if not results:
|
||||
return "No relevant documents found."
|
||||
formatted = [f"Title: {r['title']}\nSnippet: {r['text'][:200]}...\nDistance: {r['distance']:.4f}" for r in results]
|
||||
return "\n\n".join(formatted)
|
||||
results = store.similarity_search(query, k=max_results)
|
||||
return json.dumps(results, ensure_ascii=False)
|
||||
|
||||
@tool("add_to_knowledge_base")
|
||||
def add_to_knowledge_base(content: str, title: str) -> str:
|
||||
"""Add a new document to the knowledge base."""
|
||||
kb.add_documents([content], [title])
|
||||
return f"Document '{title}' added successfully."
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Agent definition
|
||||
# ---------------------------------------------------------------------------
|
||||
SYSTEM_PROMPT = (
|
||||
"You are an assistant that can search and add documents to a local knowledge base. "
|
||||
"Use the provided tools to perform semantic search and store new information. "
|
||||
"When answering user queries, first determine if the user needs a search or an addition. "
|
||||
"If no relevant information is found, suggest adding new content."
|
||||
)
|
||||
|
||||
prompt = ChatPromptTemplate.from_messages([
|
||||
("system", SYSTEM_PROMPT),
|
||||
("user", "{input}"),
|
||||
])
|
||||
@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=ChatOllama(model=LLM_MODEL),
|
||||
llm=llm,
|
||||
tools=[search_knowledge_base, add_to_knowledge_base],
|
||||
prompt=prompt,
|
||||
verbose=True,
|
||||
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)
|
||||
|
||||
agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base])
|
||||
# CLI helpers
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI client
|
||||
# ---------------------------------------------------------------------------
|
||||
def load_documents_from_dir(directory: str):
|
||||
"""Load all .txt files from a directory and add them to the knowledge base."""
|
||||
docs = []
|
||||
titles = []
|
||||
for path in Path(directory).glob("**/*.txt"):
|
||||
text = path.read_text(encoding="utf-8")
|
||||
docs.append(text)
|
||||
titles.append(path.stem)
|
||||
if docs:
|
||||
kb.add_documents(docs, titles)
|
||||
print(f"Loaded {len(docs)} documents from {directory}.")
|
||||
else:
|
||||
print("No .txt files found.")
|
||||
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 interactive_loop():
|
||||
print("RAG Agent CLI. Commands: /add <title> <content>, /search <query>, /load <dir>, /quit")
|
||||
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("> ")
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
print("\nExiting.")
|
||||
user_input = input("\n> ")
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
break
|
||||
if not user_input:
|
||||
continue
|
||||
if user_input.startswith("/quit"):
|
||||
print("Goodbye!")
|
||||
break
|
||||
elif user_input.startswith("/add "):
|
||||
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:
|
||||
_, rest = user_input.split("/add ", 1)
|
||||
title, content = rest.split(" ", 1)
|
||||
response = add_to_knowledge_base(content, title)
|
||||
print(response)
|
||||
except ValueError:
|
||||
print("Usage: /add <title> <content>")
|
||||
elif user_input.startswith("/search "):
|
||||
query = user_input.split("/search ", 1)[1]
|
||||
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(results)
|
||||
elif user_input.startswith("/load "):
|
||||
dir_path = user_input.split("/load ", 1)[1]
|
||||
load_documents_from_dir(dir_path)
|
||||
else:
|
||||
# Treat as normal user query
|
||||
result = agent_executor.invoke({"input": user_input})
|
||||
print(result.get("output", ""))
|
||||
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__":
|
||||
if len(sys.argv) > 1 and sys.argv[1] == "cli":
|
||||
interactive_loop()
|
||||
else:
|
||||
print("Usage: python main.py cli")
|
||||
print("Run the interactive CLI with: python main.py cli")
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user