feat: solution for 'Практическое задание: Агент с RAG-памятью'
This commit is contained in:
@@ -0,0 +1,17 @@
|
||||
[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"}]
|
||||
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"
|
||||
+5
-4
@@ -1,4 +1,5 @@
|
||||
langchain
|
||||
langchain-ollama
|
||||
langchain-text-splitters
|
||||
chromadb
|
||||
langchain>=0.1.0
|
||||
langchain-qdrant>=0.1.0
|
||||
langchain-ollama>=0.1.0
|
||||
qdrant-client>=1.0.0
|
||||
python-dotenv>=1.0.0
|
||||
+43
-29
@@ -1,45 +1,59 @@
|
||||
from langchain_ollama import Ollama
|
||||
from langchain.agents import initialize_agent, AgentType
|
||||
from langchain.tools import Tool
|
||||
from typing import List
|
||||
|
||||
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
|
||||
from .config import LLM_MODEL
|
||||
|
||||
|
||||
def create_agent():
|
||||
def create_agent(
|
||||
llm_model: str = "llama3",
|
||||
tools: List[Tool] = None,
|
||||
verbose: bool = True,
|
||||
) -> AgentExecutor:
|
||||
"""
|
||||
Create an RAG-enabled agent that can search and add to a knowledge base.
|
||||
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
|
||||
The configured agent.
|
||||
Configured agent executor.
|
||||
"""
|
||||
llm = Ollama(model=LLM_MODEL)
|
||||
if tools is None:
|
||||
tools = [search_knowledge_base, add_to_knowledge_base]
|
||||
|
||||
tools = [
|
||||
Tool(
|
||||
name="search_knowledge_base",
|
||||
func=search_knowledge_base,
|
||||
description="Search the knowledge base for relevant documents."
|
||||
),
|
||||
Tool(
|
||||
name="add_to_knowledge_base",
|
||||
func=add_to_knowledge_base,
|
||||
description="Add a new document to the knowledge base."
|
||||
),
|
||||
]
|
||||
|
||||
system_prompt = (
|
||||
"You are an AI assistant that can search and add information to a knowledge base. "
|
||||
"Use the provided tools to answer user queries."
|
||||
# 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.
|
||||
"""
|
||||
)
|
||||
|
||||
agent = initialize_agent(
|
||||
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,
|
||||
tools=tools,
|
||||
llm=llm,
|
||||
agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
|
||||
verbose=True,
|
||||
system_message=system_prompt,
|
||||
verbose=verbose,
|
||||
agent="zero-shot-react-description",
|
||||
)
|
||||
|
||||
return agent
|
||||
+86
-44
@@ -1,66 +1,108 @@
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from .vector_store import QdrantVectorStore
|
||||
from .agent import create_agent
|
||||
from .tools import add_to_knowledge_base, search_knowledge_base
|
||||
from .tools import vector_store as global_vector_store
|
||||
|
||||
|
||||
def main():
|
||||
agent = create_agent()
|
||||
print("Welcome to the RAG Agent CLI.")
|
||||
print("Commands:")
|
||||
print(" /add <file_path> - Add a text file to the knowledge base.")
|
||||
print(" /search <query> - Search the knowledge base.")
|
||||
print(" /quit - Exit the program.")
|
||||
print("Any other input will be sent to the agent for general processing.\n")
|
||||
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.")
|
||||
while True:
|
||||
try:
|
||||
user_input = input(">> ").strip()
|
||||
except EOFError:
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\nExiting.")
|
||||
break
|
||||
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
if user_input.startswith("/add"):
|
||||
parts = user_input.split(maxsplit=1)
|
||||
if len(parts) < 2:
|
||||
print("Usage: /add <file_path>")
|
||||
continue
|
||||
file_path = parts[1]
|
||||
path_obj = Path(file_path)
|
||||
if not path_obj.is_file():
|
||||
print(f"File not found: {file_path}")
|
||||
continue
|
||||
try:
|
||||
content = path_obj.read_text(encoding="utf-8")
|
||||
title = path_obj.name
|
||||
result = add_to_knowledge_base(content, title)
|
||||
print(result)
|
||||
except Exception as e:
|
||||
print(f"Error reading file: {e}")
|
||||
|
||||
elif user_input.startswith("/search"):
|
||||
parts = user_input.split(maxsplit=1)
|
||||
if len(parts) < 2:
|
||||
print("Usage: /search <query>")
|
||||
continue
|
||||
query = parts[1]
|
||||
result = search_knowledge_base(query, max_results=5)
|
||||
print(result)
|
||||
|
||||
elif user_input.startswith("/quit"):
|
||||
if user_input.lower() == "/quit":
|
||||
print("Goodbye!")
|
||||
break
|
||||
|
||||
else:
|
||||
# General query to the agent
|
||||
try:
|
||||
response = agent.run(user_input)
|
||||
if user_input.startswith("/add"):
|
||||
parts = user_input.split(maxsplit=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}"})
|
||||
print(response)
|
||||
except Exception as e:
|
||||
print(f"Agent error: {e}")
|
||||
continue
|
||||
|
||||
if user_input.startswith("/search"):
|
||||
query = user_input[len("/search"):].strip()
|
||||
if not query:
|
||||
print("Usage: /search <query>")
|
||||
continue
|
||||
response = agent.run({"input": f"/search {query}"})
|
||||
print(response)
|
||||
continue
|
||||
|
||||
# Default: treat as normal query
|
||||
response = agent.run({"input": 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__":
|
||||
|
||||
+3
-4
@@ -1,4 +1,3 @@
|
||||
from .cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
# 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.
|
||||
+16
-32
@@ -1,47 +1,34 @@
|
||||
import uuid
|
||||
from typing import List, Dict
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from langchain.tools import tool
|
||||
from .vector_store import ChromaVectorStore
|
||||
from .chunking import chunk_text
|
||||
from .config import CHROMA_DB_PATH
|
||||
|
||||
# Initialize a single vector store instance
|
||||
store = ChromaVectorStore(CHROMA_DB_PATH)
|
||||
from .vector_store import QdrantVectorStore
|
||||
|
||||
# Instantiate a global vector store (will be overridden in main)
|
||||
vector_store = QdrantVectorStore()
|
||||
|
||||
|
||||
@tool("search_knowledge_base")
|
||||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||||
@tool
|
||||
def search_knowledge_base(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search the knowledge base for the most relevant documents.
|
||||
Search the knowledge base for the given query.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
The search query.
|
||||
max_results : int, optional
|
||||
Number of top results to return (default is 5).
|
||||
Maximum number of results to return.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Formatted search results.
|
||||
List[Dict[str, Any]]
|
||||
List of search results with title, content, and score.
|
||||
"""
|
||||
results = store.search(query, limit=max_results)
|
||||
docs = results["documents"][0]
|
||||
distances = results["distances"][0]
|
||||
metadatas = results["metadatas"][0]
|
||||
output = []
|
||||
for doc, dist, meta in zip(docs, distances, metadatas):
|
||||
output.append(
|
||||
f"Title: {meta.get('title', 'N/A')}\n"
|
||||
f"Distance: {dist:.4f}\n"
|
||||
f"Content: {doc}\n"
|
||||
)
|
||||
return "\n".join(output)
|
||||
return vector_store.semantic_search(query, max_results)
|
||||
|
||||
|
||||
@tool("add_to_knowledge_base")
|
||||
@tool
|
||||
def add_to_knowledge_base(content: str, title: str) -> str:
|
||||
"""
|
||||
Add a new document to the knowledge base.
|
||||
@@ -51,15 +38,12 @@ def add_to_knowledge_base(content: str, title: str) -> str:
|
||||
content : str
|
||||
The full text content of the document.
|
||||
title : str
|
||||
A title or identifier for the document.
|
||||
The title of the document.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Confirmation message.
|
||||
"""
|
||||
chunks = chunk_text(content)
|
||||
ids = [str(uuid.uuid4()) for _ in chunks]
|
||||
metadatas = [{"title": title} for _ in chunks]
|
||||
store.add_documents(chunks, metadatas, ids)
|
||||
return f"Added {len(chunks)} chunks to the knowledge base."
|
||||
vector_store.add_documents([{"content": content, "title": title}])
|
||||
return f"Document '{title}' added to the knowledge base."
|
||||
+94
-24
@@ -1,34 +1,104 @@
|
||||
import uuid
|
||||
from typing import List, Dict
|
||||
import os
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from chromadb import Client
|
||||
from chromadb.config import Settings
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
|
||||
from .config import CHROMA_DB_PATH, EMBEDDING_MODEL
|
||||
from langchain_qdrant import Qdrant
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http import models as qdrant_models
|
||||
|
||||
|
||||
class ChromaVectorStore:
|
||||
def __init__(self, db_path: str = CHROMA_DB_PATH):
|
||||
self.client = Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory=db_path))
|
||||
self.collection_name = "knowledge_base"
|
||||
self.collection = self.client.get_or_create_collection(name=self.collection_name)
|
||||
self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||||
class QdrantVectorStore:
|
||||
"""
|
||||
Wrapper around Qdrant vector store with Ollama embeddings.
|
||||
"""
|
||||
|
||||
def add_documents(self, documents: List[str], metadatas: List[Dict], ids: List[str]):
|
||||
embeddings = self.embeddings.embed_documents(documents)
|
||||
self.collection.add(
|
||||
documents=documents,
|
||||
embeddings=embeddings,
|
||||
metadatas=metadatas,
|
||||
ids=ids
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 6333,
|
||||
collection_name: str = "rag_collection",
|
||||
chunk_size: int = 1000,
|
||||
chunk_overlap: int = 200,
|
||||
):
|
||||
self.client = QdrantClient(host=host, port=port)
|
||||
self.collection_name = collection_name
|
||||
self.chunk_size = chunk_size
|
||||
self.chunk_overlap = chunk_overlap
|
||||
|
||||
# Create collection if it does not exist
|
||||
if not self.client.has_collection(collection_name):
|
||||
self.client.create_collection(
|
||||
collection_name=collection_name,
|
||||
vectors_config=qdrant_models.VectorParams(
|
||||
size=384, # size of nomic-embed-text embeddings
|
||||
distance=qdrant_models.Distance.COSINE,
|
||||
),
|
||||
)
|
||||
|
||||
def search(self, query: str, limit: int = 5):
|
||||
self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
self.text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=self.chunk_size,
|
||||
chunk_overlap=self.chunk_overlap,
|
||||
)
|
||||
|
||||
def add_documents(self, documents: List[Dict[str, str]]) -> None:
|
||||
"""
|
||||
Add documents to the vector store.
|
||||
|
||||
Each document dict must contain 'content' and 'title'.
|
||||
"""
|
||||
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)
|
||||
],
|
||||
)
|
||||
|
||||
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)
|
||||
results = self.collection.query(
|
||||
query_embeddings=[query_embedding],
|
||||
n_results=limit,
|
||||
include=["documents", "distances", "metadatas"]
|
||||
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
|
||||
Reference in New Issue
Block a user