feat: solution for 'Практическое задание: Агент с RAG-памятью'
This commit is contained in:
+7
-4
@@ -1,4 +1,7 @@
|
||||
langchain>=0.2.0
|
||||
langchain-qdrant
|
||||
langchain-ollama
|
||||
qdrant-client
|
||||
langchain==0.2.0
|
||||
langchain-ollama==0.2.0
|
||||
langchain-qdrant==0.2.0
|
||||
qdrant-client==1.7.0
|
||||
python-dotenv==1.0.0
|
||||
tqdm==4.66.1
|
||||
pydantic==2.7.0
|
||||
+19
-23
@@ -1,35 +1,31 @@
|
||||
from langchain_ollama import ChatOllama
|
||||
from langchain.agents import initialize_agent, AgentType
|
||||
from langchain.tools import Tool
|
||||
from src.tools import search_knowledge_base, add_to_knowledge_base
|
||||
from langchain_ollama import Ollama
|
||||
from langchain.agents import initialize_agent, AgentExecutor, AgentType
|
||||
from langchain.tools import BaseTool
|
||||
from typing import List
|
||||
|
||||
def create_agent():
|
||||
def create_agent(tools: List[BaseTool]) -> AgentExecutor:
|
||||
"""
|
||||
Create a LangChain agent configured to use the knowledge base tools.
|
||||
Create a LangChain agent that can use the provided tools.
|
||||
|
||||
Parameters:
|
||||
tools (List[BaseTool]): List of tools for the agent.
|
||||
|
||||
Returns:
|
||||
AgentExecutor: Configured agent.
|
||||
"""
|
||||
llm = ChatOllama(model="llama3")
|
||||
tools = [
|
||||
Tool(
|
||||
name="search_knowledge_base",
|
||||
func=search_knowledge_base,
|
||||
description="Search the knowledge base for relevant information."
|
||||
),
|
||||
Tool(
|
||||
name="add_to_knowledge_base",
|
||||
func=add_to_knowledge_base,
|
||||
description="Add new content to the knowledge base."
|
||||
)
|
||||
]
|
||||
llm = Ollama(model="llama3")
|
||||
system_prompt = (
|
||||
"You are an AI assistant that helps users by searching and adding information to a knowledge base. "
|
||||
"Use the provided tools to answer queries. If you need to add new information, call add_to_knowledge_base. "
|
||||
"If you need to retrieve information, call search_knowledge_base. Provide concise answers."
|
||||
"You are an AI assistant with access to a knowledge base. "
|
||||
"Use the following tools to answer user queries:\n"
|
||||
"- search_knowledge_base: Search the knowledge base.\n"
|
||||
"- add_to_knowledge_base: Add a new document to the knowledge base.\n"
|
||||
"When you need to use a tool, call it with the appropriate arguments."
|
||||
)
|
||||
agent = initialize_agent(
|
||||
tools=tools,
|
||||
llm=llm,
|
||||
agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
||||
verbose=True,
|
||||
agent_kwargs={"system_message": system_prompt}
|
||||
agent_kwargs={"system_message": system_prompt},
|
||||
)
|
||||
return agent
|
||||
+17
-13
@@ -1,16 +1,20 @@
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain.schema import Document
|
||||
from typing import List
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
|
||||
def chunk_document(content: str, title: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> List[Document]:
|
||||
def chunk_document(text: str, chunk_size: int = 1000, chunk_overlap: int = 200):
|
||||
"""
|
||||
Split a document into chunks suitable for vector storage.
|
||||
Each chunk is stored as a Document with metadata.
|
||||
Split a document into chunks using RecursiveCharacterTextSplitter.
|
||||
|
||||
Args:
|
||||
text: The full text of the document.
|
||||
chunk_size: Maximum size of each chunk in characters.
|
||||
chunk_overlap: Number of characters to overlap between chunks.
|
||||
|
||||
Returns:
|
||||
List of chunk strings.
|
||||
"""
|
||||
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||||
texts = splitter.split_text(content)
|
||||
documents = []
|
||||
for i, text in enumerate(texts):
|
||||
metadata = {"source": title, "chunk_id": i}
|
||||
documents.append(Document(page_content=text, metadata=metadata))
|
||||
return documents
|
||||
splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
separators=["\n\n", "\n", " ", ""],
|
||||
)
|
||||
return splitter.split_text(text)
|
||||
+74
-39
@@ -1,58 +1,93 @@
|
||||
import argparse
|
||||
from src.loader import load_documents_from_directory
|
||||
from src.tools import add_to_knowledge_base, search_knowledge_base
|
||||
import sys
|
||||
from typing import List
|
||||
from src.agent import create_agent
|
||||
from src.tools import search_knowledge_base, add_to_knowledge_base
|
||||
from src.loader import load_documents_from_directory
|
||||
from src.vector_store import QdrantVectorStore
|
||||
from src.tools import vector_store as global_vector_store
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="RAG Agent CLI")
|
||||
parser.add_argument("--init-dir", type=str, help="Directory to load documents from")
|
||||
args = parser.parse_args()
|
||||
# Initialize vector store
|
||||
vector_store = QdrantVectorStore()
|
||||
# Override global reference for tools
|
||||
global global_vector_store
|
||||
global_vector_store = vector_store
|
||||
|
||||
if args.init_dir:
|
||||
docs = load_documents_from_directory(args.init_dir)
|
||||
for doc in docs:
|
||||
add_to_knowledge_base(doc.page_content, doc.metadata.get("source", "unknown"))
|
||||
print(f"Loaded {len(docs)} documents into the knowledge base.")
|
||||
# Prepare tools
|
||||
tools = [search_knowledge_base, add_to_knowledge_base]
|
||||
|
||||
# Create agent
|
||||
agent = create_agent(tools)
|
||||
|
||||
print("Welcome to the RAG Agent CLI.")
|
||||
print("Commands:")
|
||||
print(" /load <directory> - Load all .txt and .md files from directory into the knowledge base.")
|
||||
print(" /add <title> - Add a new document. You will be prompted for content.")
|
||||
print(" /search <query> - Search the knowledge base.")
|
||||
print(" /quit - Exit the program.")
|
||||
print("")
|
||||
|
||||
agent = create_agent()
|
||||
print("Agent ready. Type your queries. Use /add <title> <content>, /search <query>, or /quit to exit.")
|
||||
while True:
|
||||
try:
|
||||
user_input = input(">> ").strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\nGoodbye!")
|
||||
print("\nExiting.")
|
||||
break
|
||||
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
if user_input.lower() == "/quit":
|
||||
print("Goodbye!")
|
||||
break
|
||||
if user_input.startswith("/"):
|
||||
parts = user_input.split(maxsplit=1)
|
||||
command = parts[0][1:].lower()
|
||||
arg = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
if user_input.lower().startswith("/add"):
|
||||
parts = user_input.split(" ", 2)
|
||||
if len(parts) < 3:
|
||||
print("Usage: /add <title> <content>")
|
||||
continue
|
||||
_, title, content = parts
|
||||
response = add_to_knowledge_base(content, title)
|
||||
if command == "quit":
|
||||
print("Goodbye!")
|
||||
break
|
||||
|
||||
elif command == "load":
|
||||
if not arg:
|
||||
print("Please provide a directory path.")
|
||||
continue
|
||||
docs = load_documents_from_directory(arg)
|
||||
if not docs:
|
||||
print("No documents found.")
|
||||
continue
|
||||
vector_store.add_documents(docs)
|
||||
print(f"Loaded {len(docs)} documents into the knowledge base.")
|
||||
|
||||
elif command == "add":
|
||||
title = arg.strip()
|
||||
if not title:
|
||||
print("Please provide a title for the document.")
|
||||
continue
|
||||
print("Enter document content. Finish with a single line containing only 'END'.")
|
||||
lines = []
|
||||
while True:
|
||||
line = input()
|
||||
if line.strip() == "END":
|
||||
break
|
||||
lines.append(line)
|
||||
content = "\n".join(lines)
|
||||
result = add_to_knowledge_base(content=content, title=title)
|
||||
print(result)
|
||||
|
||||
elif command == "search":
|
||||
query = arg.strip()
|
||||
if not query:
|
||||
print("Please provide a search query.")
|
||||
continue
|
||||
result = search_knowledge_base(query=query, max_results=5)
|
||||
print(result)
|
||||
|
||||
else:
|
||||
print(f"Unknown command: {command}")
|
||||
|
||||
else:
|
||||
# Treat as a natural language query for the agent
|
||||
response = agent.run(user_input)
|
||||
print(response)
|
||||
continue
|
||||
|
||||
if user_input.lower().startswith("/search"):
|
||||
parts = user_input.split(" ", 1)
|
||||
if len(parts) < 2:
|
||||
print("Usage: /search <query>")
|
||||
continue
|
||||
_, query = parts
|
||||
response = search_knowledge_base(query)
|
||||
print(response)
|
||||
continue
|
||||
|
||||
# Treat as normal query for the agent
|
||||
response = agent.run(user_input)
|
||||
print(response)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+14
-15
@@ -1,22 +1,21 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
from langchain.schema import Document
|
||||
from typing import List, Dict
|
||||
import re
|
||||
|
||||
def load_documents_from_directory(directory: str) -> List[Document]:
|
||||
def load_documents_from_directory(directory: str) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Recursively load all text-based files from a directory into Documents.
|
||||
Supported extensions: .txt, .md, .py, .json, .csv
|
||||
Load all .txt and .md files from the given directory.
|
||||
|
||||
Returns a list of dicts with keys: title, content.
|
||||
"""
|
||||
docs = []
|
||||
for root, _, files in os.walk(directory):
|
||||
for file in files:
|
||||
if file.lower().endswith(('.txt', '.md', '.py', '.json', '.csv')):
|
||||
path = Path(root) / file
|
||||
try:
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
docs.append(Document(page_content=content, metadata={"source": str(path)}))
|
||||
except Exception as e:
|
||||
print(f"Failed to read {path}: {e}")
|
||||
for file_path in Path(directory).rglob("*"):
|
||||
if file_path.suffix.lower() in {".txt", ".md"}:
|
||||
try:
|
||||
content = file_path.read_text(encoding="utf-8")
|
||||
title = file_path.stem
|
||||
docs.append({"title": title, "content": content})
|
||||
except Exception as e:
|
||||
print(f"Failed to read {file_path}: {e}")
|
||||
return docs
|
||||
+4
-4
@@ -1,4 +1,4 @@
|
||||
from src.cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
# This file is intentionally left blank.
|
||||
# The CLI functionality is provided in src/cli.py.
|
||||
# To run the application, execute:
|
||||
# python -m src.cli
|
||||
+39
-17
@@ -1,30 +1,52 @@
|
||||
from langchain.tools import tool
|
||||
from src.vector_store import vector_store
|
||||
from src.chunking import chunk_document
|
||||
from typing import List
|
||||
from typing import List, Dict, Any
|
||||
from src.vector_store import QdrantVectorStore
|
||||
|
||||
@tool
|
||||
# Instantiate a global vector store (will be overridden in main)
|
||||
vector_store: QdrantVectorStore = None
|
||||
|
||||
@tool("search_knowledge_base")
|
||||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||||
"""
|
||||
Search the knowledge base for relevant information.
|
||||
Returns a formatted string of results.
|
||||
Search the knowledge base for the given query and return top results.
|
||||
|
||||
Parameters:
|
||||
query (str): The search query.
|
||||
max_results (int): Number of results to return.
|
||||
|
||||
Returns:
|
||||
str: Formatted search results.
|
||||
"""
|
||||
results = vector_store.semantic_search(query, k=max_results)
|
||||
if vector_store is None:
|
||||
raise ValueError("Vector store not initialized.")
|
||||
results = vector_store.search(query, max_results)
|
||||
if not results:
|
||||
return "No relevant documents found."
|
||||
formatted = []
|
||||
for i, doc in enumerate(results, 1):
|
||||
source = doc.metadata.get("source", "unknown")
|
||||
snippet = doc.page_content[:500].replace("\n", " ")
|
||||
formatted.append(f"{i}. Source: {source}\n{snippet}...")
|
||||
for idx, res in enumerate(results, 1):
|
||||
formatted.append(
|
||||
f"{idx}. Title: {res['title']}\n Score: {res['score']:.4f}\n Snippet: {res['content'][:200]}..."
|
||||
)
|
||||
return "\n\n".join(formatted)
|
||||
|
||||
@tool
|
||||
@tool("add_to_knowledge_base")
|
||||
def add_to_knowledge_base(content: str, title: str) -> str:
|
||||
"""
|
||||
Add new content to the knowledge base.
|
||||
The content is split into chunks before being stored.
|
||||
Add a new document to the knowledge base.
|
||||
|
||||
Parameters:
|
||||
content (str): Full text of the document.
|
||||
title (str): Title of the document.
|
||||
|
||||
Returns:
|
||||
str: Confirmation message.
|
||||
"""
|
||||
docs = chunk_document(content, title)
|
||||
vector_store.add_documents(docs)
|
||||
return f"Added {len(docs)} chunks from '{title}' to the knowledge base."
|
||||
if vector_store is None:
|
||||
raise ValueError("Vector store not initialized.")
|
||||
document = {
|
||||
"content": content,
|
||||
"title": title,
|
||||
"metadata": {},
|
||||
}
|
||||
vector_store.add_documents([document])
|
||||
return f"Document '{title}' added to the knowledge base."
|
||||
+79
-36
@@ -1,51 +1,94 @@
|
||||
from typing import List
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
from langchain_qdrant import QdrantVectorStore
|
||||
from langchain.schema import Document
|
||||
import os
|
||||
from typing import List, Dict, Any
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http import models as qdrant_models
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
from langchain_qdrant import Qdrant
|
||||
from src.chunking import chunk_document
|
||||
|
||||
class QdrantStore:
|
||||
"""
|
||||
Wrapper around QdrantVectorStore that handles collection creation,
|
||||
embedding generation via Ollama, and semantic search.
|
||||
"""
|
||||
def __init__(self, collection_name: str = "knowledge_base", host: str = "localhost", port: int = 6333):
|
||||
self.client = QdrantClient(host=host, port=port)
|
||||
class QdrantVectorStore:
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: str = "knowledge_base",
|
||||
host: str = "localhost",
|
||||
port: int = 6333,
|
||||
embedding_model: str = "nomic-embed-text",
|
||||
):
|
||||
self.collection_name = collection_name
|
||||
self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
self.client = QdrantClient(host=host, port=port)
|
||||
self.embedding_model = embedding_model
|
||||
self.embeddings = OllamaEmbeddings(model=self.embedding_model)
|
||||
|
||||
# Determine vector size from a sample embedding
|
||||
sample_vector = self.embeddings.embed_query("")
|
||||
vector_size = len(sample_vector)
|
||||
|
||||
# Create collection if it does not exist
|
||||
if not self.client.has_collection(collection_name):
|
||||
self.client.recreate_collection(
|
||||
collection_name=collection_name,
|
||||
# Ensure collection exists
|
||||
if not self.client.collection_exists(self.collection_name):
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=qdrant_models.VectorParams(
|
||||
size=vector_size,
|
||||
distance="Cosine"
|
||||
)
|
||||
size=self.embeddings.embedding_size,
|
||||
distance=qdrant_models.Distance.COSINE,
|
||||
),
|
||||
)
|
||||
|
||||
self.store = QdrantVectorStore(
|
||||
self.qdrant = Qdrant(
|
||||
client=self.client,
|
||||
collection_name=collection_name,
|
||||
embeddings=self.embeddings
|
||||
collection_name=self.collection_name,
|
||||
embeddings=self.embeddings,
|
||||
)
|
||||
|
||||
def add_documents(self, documents: List[Document]):
|
||||
def add_documents(self, documents: List[Dict[str, Any]]):
|
||||
"""
|
||||
Add a list of Documents to the vector store.
|
||||
Add documents to the vector store. Each document dict should contain:
|
||||
- content: full text
|
||||
- title: document title
|
||||
- metadata: optional dict
|
||||
"""
|
||||
self.store.add_documents(documents)
|
||||
points = []
|
||||
for doc in documents:
|
||||
content = doc.get("content", "")
|
||||
title = doc.get("title", "")
|
||||
metadata = doc.get("metadata", {})
|
||||
chunks = chunk_document(content)
|
||||
for idx, chunk in enumerate(chunks):
|
||||
point = qdrant_models.PointStruct(
|
||||
id=None,
|
||||
vector=self.embeddings.embed_query(chunk),
|
||||
payload={
|
||||
"title": title,
|
||||
"chunk_index": idx,
|
||||
"content": chunk,
|
||||
**metadata,
|
||||
},
|
||||
)
|
||||
points.append(point)
|
||||
|
||||
def semantic_search(self, query: str, k: int = 5):
|
||||
"""
|
||||
Perform a semantic similarity search and return top-k Documents.
|
||||
"""
|
||||
return self.store.similarity_search(query, k=k)
|
||||
if points:
|
||||
self.client.upsert(
|
||||
collection_name=self.collection_name,
|
||||
points=points,
|
||||
)
|
||||
|
||||
# Global instance used by tools and agent
|
||||
vector_store = QdrantStore()
|
||||
def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Perform a semantic search in the vector store.
|
||||
|
||||
Returns a list of dicts with keys: title, content, score.
|
||||
"""
|
||||
query_vector = self.embeddings.embed_query(query)
|
||||
search_result = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query_vector,
|
||||
limit=max_results,
|
||||
with_payload=True,
|
||||
with_vector=False,
|
||||
)
|
||||
results = []
|
||||
for hit in search_result:
|
||||
payload = hit.payload
|
||||
results.append(
|
||||
{
|
||||
"title": payload.get("title", ""),
|
||||
"content": payload.get("content", ""),
|
||||
"score": hit.score,
|
||||
}
|
||||
)
|
||||
return results
|
||||
Reference in New Issue
Block a user