feat: solution for 'Практическое задание: Агент с RAG-памятью'
This commit is contained in:
+6
-12
@@ -1,17 +1,11 @@
|
||||
[project]
|
||||
name = "rag-agent"
|
||||
version = "0.1.0"
|
||||
description = "AI agent with local RAG memory using Qdrant and Ollama"
|
||||
authors = [{name = "Your Name", email = "you@example.com"}]
|
||||
description = "RAG agent with Qdrant and Ollama"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = [
|
||||
"langchain>=0.1.0",
|
||||
"langchain-qdrant>=0.1.0",
|
||||
"langchain-ollama>=0.1.0",
|
||||
"qdrant-client>=1.0.0",
|
||||
"python-dotenv>=1.0.0"
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=61.0"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
"langchain>=0.2.0",
|
||||
"langchain-qdrant",
|
||||
"langchain-ollama",
|
||||
"qdrant-client"
|
||||
]
|
||||
+4
-5
@@ -1,5 +1,4 @@
|
||||
langchain>=0.1.0
|
||||
langchain-qdrant>=0.1.0
|
||||
langchain-ollama>=0.1.0
|
||||
qdrant-client>=1.0.0
|
||||
python-dotenv>=1.0.0
|
||||
langchain>=0.2.0
|
||||
langchain-qdrant
|
||||
langchain-ollama
|
||||
qdrant-client
|
||||
+28
-52
@@ -1,59 +1,35 @@
|
||||
from typing import List
|
||||
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 import LLMChain
|
||||
from langchain.chat_models import ChatOllama
|
||||
from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
|
||||
from langchain.agents import AgentExecutor, Tool
|
||||
|
||||
from .tools import search_knowledge_base, add_to_knowledge_base
|
||||
|
||||
|
||||
def create_agent(
|
||||
llm_model: str = "llama3",
|
||||
tools: List[Tool] = None,
|
||||
verbose: bool = True,
|
||||
) -> AgentExecutor:
|
||||
def create_agent():
|
||||
"""
|
||||
Create an AgentExecutor that uses the provided tools and a system prompt
|
||||
instructing the agent to use the knowledge base.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
llm_model : str
|
||||
The Ollama model to use.
|
||||
tools : List[Tool]
|
||||
List of LangChain tools to expose to the agent.
|
||||
verbose : bool
|
||||
Whether to enable verbose output.
|
||||
|
||||
Returns
|
||||
-------
|
||||
AgentExecutor
|
||||
Configured agent executor.
|
||||
Create a LangChain agent configured to use the knowledge base tools.
|
||||
"""
|
||||
if tools is None:
|
||||
tools = [search_knowledge_base, add_to_knowledge_base]
|
||||
|
||||
# System prompt instructing the agent to use the knowledge base
|
||||
system_prompt = SystemMessagePromptTemplate.from_template(
|
||||
"""
|
||||
You are an AI assistant that has access to a knowledge base. Use the provided tools to search the knowledge base or add new documents. When answering user queries, first decide if you need to search the knowledge base. If so, use the `search_knowledge_base` tool. If you need to add new information, use the `add_to_knowledge_base` tool. Always provide a concise answer after retrieving relevant information.
|
||||
"""
|
||||
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."
|
||||
)
|
||||
]
|
||||
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."
|
||||
)
|
||||
|
||||
human_prompt = HumanMessagePromptTemplate.from_template("{input}")
|
||||
|
||||
chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])
|
||||
|
||||
llm = ChatOllama(model=llm_model)
|
||||
|
||||
llm_chain = LLMChain(llm=llm, prompt=chat_prompt)
|
||||
|
||||
agent = AgentExecutor.from_llm_and_tools(
|
||||
llm=llm_chain,
|
||||
agent = initialize_agent(
|
||||
tools=tools,
|
||||
verbose=verbose,
|
||||
agent="zero-shot-react-description",
|
||||
llm=llm,
|
||||
agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
||||
verbose=True,
|
||||
agent_kwargs={"system_message": system_prompt}
|
||||
)
|
||||
|
||||
return agent
|
||||
+14
-4
@@ -1,6 +1,16 @@
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain.schema import Document
|
||||
from typing import List
|
||||
|
||||
|
||||
def chunk_text(text: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> list[str]:
|
||||
def chunk_document(content: str, title: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> List[Document]:
|
||||
"""
|
||||
Split a document into chunks suitable for vector storage.
|
||||
Each chunk is stored as a Document with metadata.
|
||||
"""
|
||||
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||||
return splitter.split_text(text)
|
||||
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
|
||||
+26
-77
@@ -1,50 +1,26 @@
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from src.loader import load_documents_from_directory
|
||||
from src.tools import add_to_knowledge_base, search_knowledge_base
|
||||
from src.agent import create_agent
|
||||
|
||||
from .vector_store import QdrantVectorStore
|
||||
from .agent import create_agent
|
||||
from .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()
|
||||
|
||||
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.")
|
||||
|
||||
def load_documents_from_directory(
|
||||
directory: Path, vector_store: QdrantVectorStore
|
||||
) -> None:
|
||||
"""
|
||||
Load all text files from the specified directory into the vector store.
|
||||
"""
|
||||
if not directory.is_dir():
|
||||
print(f"Directory {directory} does not exist.")
|
||||
return
|
||||
|
||||
documents = []
|
||||
for file_path in directory.rglob("*"):
|
||||
if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md"}:
|
||||
content = file_path.read_text(encoding="utf-8")
|
||||
title = file_path.stem
|
||||
documents.append({"content": content, "title": title})
|
||||
if documents:
|
||||
vector_store.add_documents(documents)
|
||||
print(f"Loaded {len(documents)} documents into the knowledge base.")
|
||||
else:
|
||||
print("No documents found in the directory.")
|
||||
|
||||
|
||||
def interactive_loop(agent, max_results: int = 5):
|
||||
"""
|
||||
Simple interactive CLI loop for the agent.
|
||||
Commands:
|
||||
/add <title> <content> - Add a new document
|
||||
/search <query> - Search the knowledge base
|
||||
/quit - Exit
|
||||
"""
|
||||
print("Welcome to the RAG Agent CLI. Type /quit to exit.")
|
||||
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("\nExiting.")
|
||||
print("\nGoodbye!")
|
||||
break
|
||||
|
||||
if not user_input:
|
||||
@@ -54,56 +30,29 @@ def interactive_loop(agent, max_results: int = 5):
|
||||
print("Goodbye!")
|
||||
break
|
||||
|
||||
if user_input.startswith("/add"):
|
||||
parts = user_input.split(maxsplit=2)
|
||||
if user_input.lower().startswith("/add"):
|
||||
parts = user_input.split(" ", 2)
|
||||
if len(parts) < 3:
|
||||
print("Usage: /add <title> <content>")
|
||||
continue
|
||||
title, content = parts[1], parts[2]
|
||||
response = agent.run({"input": f"/add {title} {content}"})
|
||||
_, title, content = parts
|
||||
response = add_to_knowledge_base(content, title)
|
||||
print(response)
|
||||
continue
|
||||
|
||||
if user_input.startswith("/search"):
|
||||
query = user_input[len("/search"):].strip()
|
||||
if not query:
|
||||
if user_input.lower().startswith("/search"):
|
||||
parts = user_input.split(" ", 1)
|
||||
if len(parts) < 2:
|
||||
print("Usage: /search <query>")
|
||||
continue
|
||||
response = agent.run({"input": f"/search {query}"})
|
||||
_, query = parts
|
||||
response = search_knowledge_base(query)
|
||||
print(response)
|
||||
continue
|
||||
|
||||
# Default: treat as normal query
|
||||
response = agent.run({"input": user_input})
|
||||
# Treat as normal query for the agent
|
||||
response = agent.run(user_input)
|
||||
print(response)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="RAG Agent CLI")
|
||||
parser.add_argument(
|
||||
"--load-dir",
|
||||
type=str,
|
||||
help="Directory containing documents to load into the knowledge base",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-results",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Maximum number of search results to return",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
vector_store = QdrantVectorStore()
|
||||
# Override global vector store used by tools
|
||||
global_vector_store.__dict__.update(vector_store.__dict__)
|
||||
|
||||
if args.load_dir:
|
||||
load_documents_from_directory(Path(args.load_dir), vector_store)
|
||||
|
||||
agent = create_agent(verbose=True)
|
||||
|
||||
interactive_loop(agent, max_results=args.max_results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+14
-27
@@ -1,35 +1,22 @@
|
||||
import os
|
||||
import uuid
|
||||
from typing import List, Dict
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
from langchain.schema import Document
|
||||
|
||||
from .vector_store import ChromaVectorStore
|
||||
from .chunking import chunk_text
|
||||
from .config import CHROMA_DB_PATH
|
||||
|
||||
store = ChromaVectorStore(CHROMA_DB_PATH)
|
||||
|
||||
|
||||
def load_documents_from_directory(directory_path: str):
|
||||
def load_documents_from_directory(directory: str) -> List[Document]:
|
||||
"""
|
||||
Load all .txt files from a directory into the vector store.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
directory_path : str
|
||||
Path to the directory containing text files.
|
||||
Recursively load all text-based files from a directory into Documents.
|
||||
Supported extensions: .txt, .md, .py, .json, .csv
|
||||
"""
|
||||
for root, _, files in os.walk(directory_path):
|
||||
docs = []
|
||||
for root, _, files in os.walk(directory):
|
||||
for file in files:
|
||||
if file.lower().endswith(".txt"):
|
||||
file_path = os.path.join(root, file)
|
||||
if file.lower().endswith(('.txt', '.md', '.py', '.json', '.csv')):
|
||||
path = Path(root) / file
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
title = file
|
||||
chunks = chunk_text(content)
|
||||
ids = [str(uuid.uuid4()) for _ in chunks]
|
||||
metadatas = [{"title": title} for _ in chunks]
|
||||
store.add_documents(chunks, metadatas, ids)
|
||||
print(f"Loaded {len(chunks)} chunks from {file_path}")
|
||||
docs.append(Document(page_content=content, metadata={"source": str(path)}))
|
||||
except Exception as e:
|
||||
print(f"Failed to load {file_path}: {e}")
|
||||
print(f"Failed to read {path}: {e}")
|
||||
return docs
|
||||
+4
-3
@@ -1,3 +1,4 @@
|
||||
# This file is intentionally left empty.
|
||||
# The application entry point is defined in src/cli.py.
|
||||
# Importing src.cli.main will start the CLI when executed as a script.
|
||||
from src.cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+20
-39
@@ -1,49 +1,30 @@
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from langchain.tools import tool
|
||||
|
||||
from .vector_store import QdrantVectorStore
|
||||
|
||||
# Instantiate a global vector store (will be overridden in main)
|
||||
vector_store = QdrantVectorStore()
|
||||
|
||||
from src.vector_store import vector_store
|
||||
from src.chunking import chunk_document
|
||||
from typing import List
|
||||
|
||||
@tool
|
||||
def search_knowledge_base(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
|
||||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||||
"""
|
||||
Search the knowledge base for the given query.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
The search query.
|
||||
max_results : int, optional
|
||||
Maximum number of results to return.
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[Dict[str, Any]]
|
||||
List of search results with title, content, and score.
|
||||
Search the knowledge base for relevant information.
|
||||
Returns a formatted string of results.
|
||||
"""
|
||||
return vector_store.semantic_search(query, max_results)
|
||||
|
||||
results = vector_store.semantic_search(query, k=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}...")
|
||||
return "\n\n".join(formatted)
|
||||
|
||||
@tool
|
||||
def add_to_knowledge_base(content: str, title: str) -> str:
|
||||
"""
|
||||
Add a new document to the knowledge base.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
content : str
|
||||
The full text content of the document.
|
||||
title : str
|
||||
The title of the document.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Confirmation message.
|
||||
Add new content to the knowledge base.
|
||||
The content is split into chunks before being stored.
|
||||
"""
|
||||
vector_store.add_documents([{"content": content, "title": title}])
|
||||
return f"Document '{title}' added to the knowledge base."
|
||||
docs = chunk_document(content, title)
|
||||
vector_store.add_documents(docs)
|
||||
return f"Added {len(docs)} chunks from '{title}' to the knowledge base."
|
||||
+30
-83
@@ -1,104 +1,51 @@
|
||||
import os
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from typing import List
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
from langchain_qdrant import Qdrant
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from langchain_qdrant import QdrantVectorStore
|
||||
from langchain.schema import Document
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http import models as qdrant_models
|
||||
|
||||
|
||||
class QdrantVectorStore:
|
||||
class QdrantStore:
|
||||
"""
|
||||
Wrapper around Qdrant vector store with Ollama embeddings.
|
||||
Wrapper around QdrantVectorStore that handles collection creation,
|
||||
embedding generation via Ollama, and semantic search.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 6333,
|
||||
collection_name: str = "rag_collection",
|
||||
chunk_size: int = 1000,
|
||||
chunk_overlap: int = 200,
|
||||
):
|
||||
def __init__(self, collection_name: str = "knowledge_base", host: str = "localhost", port: int = 6333):
|
||||
self.client = QdrantClient(host=host, port=port)
|
||||
self.collection_name = collection_name
|
||||
self.chunk_size = chunk_size
|
||||
self.chunk_overlap = chunk_overlap
|
||||
self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
|
||||
# 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.create_collection(
|
||||
self.client.recreate_collection(
|
||||
collection_name=collection_name,
|
||||
vectors_config=qdrant_models.VectorParams(
|
||||
size=384, # size of nomic-embed-text embeddings
|
||||
distance=qdrant_models.Distance.COSINE,
|
||||
),
|
||||
size=vector_size,
|
||||
distance="Cosine"
|
||||
)
|
||||
)
|
||||
|
||||
self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
self.text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=self.chunk_size,
|
||||
chunk_overlap=self.chunk_overlap,
|
||||
self.store = QdrantVectorStore(
|
||||
client=self.client,
|
||||
collection_name=collection_name,
|
||||
embeddings=self.embeddings
|
||||
)
|
||||
|
||||
def add_documents(self, documents: List[Dict[str, str]]) -> None:
|
||||
def add_documents(self, documents: List[Document]):
|
||||
"""
|
||||
Add documents to the vector store.
|
||||
Add a list of Documents to the vector store.
|
||||
"""
|
||||
self.store.add_documents(documents)
|
||||
|
||||
Each document dict must contain 'content' and 'title'.
|
||||
def semantic_search(self, query: str, k: int = 5):
|
||||
"""
|
||||
for doc in documents:
|
||||
content = doc.get("content", "")
|
||||
title = doc.get("title", "Untitled")
|
||||
# Split into chunks
|
||||
chunks = self.text_splitter.split_text(content)
|
||||
# Embed each chunk
|
||||
embeddings = self.embeddings.embed_documents(chunks)
|
||||
# Prepare payloads
|
||||
payloads = [
|
||||
{
|
||||
"title": title,
|
||||
"chunk_index": idx,
|
||||
"content": chunk,
|
||||
}
|
||||
for idx, chunk in enumerate(chunks)
|
||||
]
|
||||
# Upsert into Qdrant
|
||||
self.client.upsert(
|
||||
collection_name=self.collection_name,
|
||||
points=[
|
||||
qdrant_models.PointStruct(
|
||||
id=None,
|
||||
vector=emb,
|
||||
payload=payload,
|
||||
)
|
||||
for emb, payload in zip(embeddings, payloads)
|
||||
],
|
||||
)
|
||||
Perform a semantic similarity search and return top-k Documents.
|
||||
"""
|
||||
return self.store.similarity_search(query, k=k)
|
||||
|
||||
def semantic_search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Perform semantic search in the vector store.
|
||||
|
||||
Returns a list of dicts with 'title', 'content', and 'score'.
|
||||
"""
|
||||
query_embedding = self.embeddings.embed_query(query)
|
||||
search_result = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query_embedding,
|
||||
limit=max_results,
|
||||
with_payload=True,
|
||||
score_threshold=0.0,
|
||||
)
|
||||
results = []
|
||||
for point in search_result:
|
||||
payload = point.payload
|
||||
results.append(
|
||||
{
|
||||
"title": payload.get("title", "Untitled"),
|
||||
"content": payload.get("content", ""),
|
||||
"score": point.score,
|
||||
}
|
||||
)
|
||||
return results
|
||||
# Global instance used by tools and agent
|
||||
vector_store = QdrantStore()
|
||||
Reference in New Issue
Block a user