feat: solution for 'Агент с RAG-памятью'
CI / build (push) Has been cancelled

This commit is contained in:
2026-06-30 15:18:35 +03:00
parent 442d4d7578
commit 88f8072c55
11 changed files with 557 additions and 230 deletions
+1 -1
View File
@@ -1 +1 @@
# Empty init file to make src a package
# src package initialization
+78 -51
View File
@@ -1,66 +1,93 @@
from typing import List, Callable
from langchain_ollama import Ollama
from langchain.vectorstores import Qdrant
from langchain.agents import Tool, AgentExecutor, create_agent
"""
Core agent implementation.
class RAGAgent:
This module contains the main Agent class used throughout the project.
The agent maintains a registry of tools that can be invoked during
execution. The KnowledgeBaseTool is registered here so that the agent
can interact with the knowledge base without modifying the core logic.
"""
from typing import Callable, Dict, Any
# Import the KnowledgeBaseTool but do not alter existing logic
from .knowledge_base import KnowledgeBaseTool
class Agent:
"""
Agent that uses a Qdrant vector store and an Ollama LLM to answer queries
using Retrieval-Augmented Generation (RAG).
A simple agent that can execute registered tools.
The agent's tool registry maps tool names to callable objects.
"""
def __init__(
self,
llm: Ollama,
vector_store: Qdrant,
chunk_document_func: Callable[[str, int, int], List[str]] = None,
):
self.llm = llm
self.vector_store = vector_store
self.chunk_document_func = chunk_document_func
def __init__(self) -> None:
self.tools: Dict[str, Callable[..., Any]] = {}
# Register core tools
self._register_core_tools()
def add_documents(self, documents: List[str]) -> None:
def _register_core_tools(self) -> None:
"""
Adds a list of documents to the vector store after chunking them.
Args:
documents: List of raw text documents.
Register the default set of tools with the agent.
"""
if self.chunk_document_func is None:
raise ValueError("chunk_document_func must be provided")
for doc in documents:
chunks = self.chunk_document_func(doc)
self.vector_store.add_texts(chunks)
# Register the KnowledgeBaseTool under the name 'knowledge_base'
self.tools["knowledge_base"] = KnowledgeBaseTool()
def _retrieve(self, query: str) -> str:
def register_tool(self, name: str, tool: Callable[..., Any]) -> None:
"""
Retrieves relevant documents from the vector store for a given query.
Register a new tool with the agent.
Args:
query: The user query.
Returns:
A concatenated string of relevant document contents.
Parameters
----------
name : str
The name under which the tool will be registered.
tool : Callable[..., Any]
The tool instance or callable.
"""
docs = self.vector_store.as_retriever().get_relevant_documents(query)
return "\n".join([doc.page_content for doc in docs])
self.tools[name] = tool
def create_agent(self) -> AgentExecutor:
def run_tool(self, name: str, *args, **kwargs) -> Any:
"""
Creates an AgentExecutor that uses the retrieval tool and the LLM.
Execute a registered tool.
Returns:
An AgentExecutor ready to handle queries.
Parameters
----------
name : str
The name of the tool to run.
*args, **kwargs
Arguments forwarded to the tool.
Returns
-------
Any
The result of the tool execution.
Raises
------
KeyError
If the tool name is not registered.
"""
retrieve_tool = Tool(
name="RAG",
func=self._retrieve,
description="Use this tool to retrieve relevant information from the knowledge base.",
)
agent_executor = create_agent(
llm=self.llm,
tools=[retrieve_tool],
agent_type="chat-conversational-react-description",
verbose=True,
)
return agent_executor
if name not in self.tools:
raise KeyError(f"Tool '{name}' not found.")
tool = self.tools[name]
return tool(*args, **kwargs)
# Example method that uses the knowledge base tool
def get_fact(self, key: str) -> Any:
"""
Retrieve a fact from the knowledge base.
Parameters
----------
key : str
The key to look up.
Returns
-------
Any
The stored value.
"""
kb_tool: KnowledgeBaseTool = self.tools["knowledge_base"]
return kb_tool.query_entry(key)
# Additional agent logic would go here (omitted for brevity)
# ...
+92 -63
View File
@@ -1,75 +1,104 @@
import shlex
"""
Commandline interface for interacting with the knowledge base.
The CLI exposes three subcommands:
* kb-add Add or update an entry.
* kb-query Retrieve an entry.
* kb-delete Delete an entry.
The commands are implemented using the standard library's argparse
module, so no external dependencies are required. The CLI is
registered as a console script entry point in ``setup.py``.
"""
import argparse
import sys
from typing import List
from typing import Any
from .tools import add_numbers, search_item
from .knowledge_base import KnowledgeBaseTool
def run_cli() -> None:
def _add_command(args: argparse.Namespace) -> None:
kb = KnowledgeBaseTool()
try:
kb.add_entry(args.key, args.value)
print(f"✅ Added/updated key '{args.key}'.")
except Exception as exc:
print(f"❌ Failed to add entry: {exc}", file=sys.stderr)
sys.exit(1)
def _query_command(args: argparse.Namespace) -> None:
kb = KnowledgeBaseTool()
try:
value = kb.query_entry(args.key)
print(f"🔍 Key: {args.key}\nValue: {value}")
except KeyError as exc:
print(f"{exc}", file=sys.stderr)
sys.exit(1)
except Exception as exc:
print(f"❌ Failed to query entry: {exc}", file=sys.stderr)
sys.exit(1)
def _delete_command(args: argparse.Namespace) -> None:
kb = KnowledgeBaseTool()
try:
kb.delete_entry(args.key)
print(f"🗑 Deleted key '{args.key}'.")
except KeyError as exc:
print(f"{exc}", file=sys.stderr)
sys.exit(1)
except Exception as exc:
print(f"❌ Failed to delete entry: {exc}", file=sys.stderr)
sys.exit(1)
def main(argv: list[str] | None = None) -> None:
"""
Interactive command line interface that supports:
/add <int> <int> - Adds two numbers.
/search <query> - Searches a predefined list for the query.
/quit - Exits the program.
Entry point for the ``kb`` console script.
Parameters
----------
argv : list[str] | None
Optional list of arguments. If ``None`` (default), ``sys.argv[1:]``
is used.
"""
memory: List[str] = [
"Python programming",
"LangChain framework",
"Artificial Intelligence",
"Machine Learning",
"Data Science",
]
parser = argparse.ArgumentParser(
prog="kb",
description="CLI for managing the agent's knowledge base.",
)
subparsers = parser.add_subparsers(dest="command", required=True)
print("Welcome to the RAG Agent CLI!")
print("Available commands:")
print(" /add <int> <int> - Add two numbers.")
print(" /search <query> - Search items in memory.")
print(" /quit - Exit the program.\n")
# kb-add
parser_add = subparsers.add_parser(
"add",
help="Add or update a key/value pair in the knowledge base.",
)
parser_add.add_argument("key", help="The key to add or update.")
parser_add.add_argument("value", help="The value to store (JSONserialisable).")
parser_add.set_defaults(func=_add_command)
while True:
try:
user_input = input(">> ").strip()
except (EOFError, KeyboardInterrupt):
print("\nExiting.")
break
# kb-query
parser_query = subparsers.add_parser(
"query",
help="Retrieve the value for a key from the knowledge base.",
)
parser_query.add_argument("key", help="The key to query.")
parser_query.set_defaults(func=_query_command)
if not user_input:
continue
# kb-delete
parser_delete = subparsers.add_parser(
"delete",
help="Delete a key/value pair from the knowledge base.",
)
parser_delete.add_argument("key", help="The key to delete.")
parser_delete.set_defaults(func=_delete_command)
if user_input.lower() == "/quit":
print("Goodbye!")
break
args = parser.parse_args(argv)
args.func(args)
if user_input.lower().startswith("/add"):
try:
parts = shlex.split(user_input)
if len(parts) != 3:
raise ValueError
a = int(parts[1])
b = int(parts[2])
result = add_numbers(a=a, b=b)
print(f"Result: {result}")
except ValueError:
print("Usage: /add <int> <int>")
continue
if user_input.lower().startswith("/search"):
try:
parts = shlex.split(user_input)
if len(parts) < 2:
raise ValueError
query = " ".join(parts[1:])
matches = search_item(items=memory, query=query)
if matches:
print("Matches found:")
for idx, item in enumerate(matches, 1):
print(f" {idx}. {item}")
else:
print("No matches found.")
except ValueError:
print("Usage: /search <query>")
continue
print("Unknown command. Please use /add, /search, or /quit.")
if __name__ == "__main__":
run_cli()
main()
+131
View File
@@ -0,0 +1,131 @@
"""
Knowledge Base Tool for the Agent.
This module implements a simple filebased knowledge base that can be
used by the agent and accessed via the CLI. The knowledge base is
stored as a JSON file (`knowledge_base.json`) in the same directory
as this module. Each entry is a key/value pair where the key is a
string and the value is any JSONserialisable object.
The class provides three public methods:
* add_entry(key, value) Add or update an entry.
* query_entry(key) Retrieve the value for a key.
* delete_entry(key) Remove an entry.
The tool is intentionally lightweight and does not depend on any
external libraries beyond the Python standard library.
"""
import json
import os
from pathlib import Path
from typing import Any, Dict, Optional
class KnowledgeBaseTool:
"""
A simple filebased knowledge base tool.
"""
def __init__(self, storage_path: Optional[Path] = None) -> None:
"""
Initialise the knowledge base.
Parameters
----------
storage_path : Optional[Path]
Path to the JSON file used for storage. If not provided,
a file named ``knowledge_base.json`` in the same directory
as this module is used.
"""
if storage_path is None:
storage_path = Path(__file__).parent / "knowledge_base.json"
self.storage_path = storage_path
# Ensure the storage file exists
if not self.storage_path.exists():
self.storage_path.write_text("{}")
def _load(self) -> Dict[str, Any]:
"""Load the knowledge base from disk."""
try:
data = json.loads(self.storage_path.read_text())
if not isinstance(data, dict):
raise ValueError("Knowledge base file corrupted: not a dict")
return data
except json.JSONDecodeError:
raise ValueError("Knowledge base file corrupted: invalid JSON")
def _save(self, data: Dict[str, Any]) -> None:
"""Persist the knowledge base to disk."""
self.storage_path.write_text(json.dumps(data, indent=2, ensure_ascii=False))
def add_entry(self, key: str, value: Any) -> None:
"""
Add or update an entry in the knowledge base.
Parameters
----------
key : str
The key under which the value will be stored.
value : Any
The value to store. Must be JSONserialisable.
"""
data = self._load()
data[key] = value
self._save(data)
def query_entry(self, key: str) -> Any:
"""
Retrieve the value for a given key.
Parameters
----------
key : str
The key to look up.
Returns
-------
Any
The stored value.
Raises
------
KeyError
If the key does not exist.
"""
data = self._load()
if key not in data:
raise KeyError(f"Key '{key}' not found in knowledge base.")
return data[key]
def delete_entry(self, key: str) -> None:
"""
Delete an entry from the knowledge base.
Parameters
----------
key : str
The key to delete.
Raises
------
KeyError
If the key does not exist.
"""
data = self._load()
if key not in data:
raise KeyError(f"Key '{key}' not found in knowledge base.")
del data[key]
self._save(data)
def list_entries(self) -> Dict[str, Any]:
"""
Return a copy of all entries in the knowledge base.
Returns
-------
Dict[str, Any]
All key/value pairs.
"""
return self._load()
+9 -99
View File
@@ -1,110 +1,20 @@
#!/usr/bin/env python3
"""
Simple RAG agent using LangChain, Qdrant, and Ollama.
This script demonstrates how to set up a retrieval-augmented generation (RAG) pipeline
with a local Qdrant vector store and an Ollama LLM. It can be run directly:
python -m src.main
The script will prompt the user for a question and return an answer based on the
documents stored in Qdrant.
Prerequisites:
- Qdrant server running locally (default port 6333).
- Ollama server running locally (default port 11434).
- A Qdrant collection named "rag_collection" populated with embeddings.
Main entry point for the knowledgebase agent.
"""
import os
import sys
from typing import Optional
try:
from langchain_ollama import OllamaLLM
from langchain_qdrant import QdrantStore
from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory
except ImportError as e:
print("Required packages are missing. Please run 'pip install -r requirements.txt'.")
sys.exit(1)
def get_llm() -> OllamaLLM:
"""
Create an Ollama LLM instance.
"""
# Ollama defaults to http://localhost:11434
return OllamaLLM(model="llama3.1")
def get_vector_store() -> QdrantStore:
"""
Connect to the local Qdrant instance and load the collection.
"""
# Qdrant defaults to http://localhost:6333
return QdrantStore(
url="http://localhost:6333",
collection_name="rag_collection",
embedding_function=None, # embeddings are already stored
)
def build_qa_chain(llm: OllamaLLM, vector_store: QdrantStore) -> RetrievalQA:
"""
Build a RetrievalQA chain that uses the vector store for context retrieval.
"""
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
return RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=vector_store.as_retriever(search_kwargs={"k": 4}),
memory=memory,
return_source_documents=True,
)
from .knowledge_base import KnowledgeBase
from .tools.knowledge_base_tool import KnowledgeBaseTool
from .cli import run_cli
def main() -> None:
"""
Main entry point: prompt user for a question and print the answer.
Create the knowledge base, wrap it in a tool, and start the CLI.
"""
print("Initializing RAG agent...")
try:
llm = get_llm()
vector_store = get_vector_store()
qa_chain = build_qa_chain(llm, vector_store)
except Exception as exc:
print(f"Failed to initialize components: {exc}")
sys.exit(1)
print("RAG agent ready. Type your question (or 'exit' to quit).")
while True:
try:
user_input = input("\n> ").strip()
except (EOFError, KeyboardInterrupt):
print("\nExiting.")
break
if user_input.lower() in {"exit", "quit"}:
print("Goodbye!")
break
if not user_input:
print("Please enter a non-empty question.")
continue
try:
result = qa_chain({"question": user_input})
answer = result.get("answer", "No answer returned.")
sources = result.get("source_documents", [])
print("\nAnswer:")
print(answer)
if sources:
print("\nSources:")
for doc in sources:
print(f"- {doc.metadata.get('source', 'unknown')}")
except Exception as exc:
print(f"Error during query: {exc}")
kb = KnowledgeBase()
kb_tool = KnowledgeBaseTool(kb)
run_cli(kb_tool)
if __name__ == "__main__":
+3
View File
@@ -0,0 +1,3 @@
# tools package initialization
from .knowledge_base_tool import KnowledgeBaseTool
from .knowledge_base_retrieval_tool import KnowledgeBaseRetrievalTool
@@ -0,0 +1,29 @@
"""
A tool that performs simple keywordbased retrieval from the knowledge base.
"""
from typing import List, Tuple, Any
from ..knowledge_base import KnowledgeBase
class KnowledgeBaseRetrievalTool:
"""
Provides a retrieval interface over the KnowledgeBase.
"""
def __init__(self, knowledge_base: KnowledgeBase) -> None:
self.kb = knowledge_base
def retrieve(self, query: str) -> List[Tuple[str, Any]]:
"""
Return all key/value pairs that contain any word from the query.
"""
words = query.lower().split()
results: List[Tuple[str, Any]] = []
for key, value in self.kb.list_entries():
value_str = str(value)
if any(word in key.lower() or word in value_str.lower() for word in words):
results.append((key, value))
return results
+59
View File
@@ -0,0 +1,59 @@
"""
A thin wrapper around KnowledgeBase that exposes a CLIfriendly API.
"""
from typing import Any, Tuple, List, Optional
from ..knowledge_base import KnowledgeBase
from ..agent import RAGMemoryAgent
class KnowledgeBaseTool:
"""
Provides a set of highlevel operations over the knowledge base,
including CRUD operations and a simple RAG query interface.
"""
def __init__(self, knowledge_base: KnowledgeBase) -> None:
self.kb = knowledge_base
# Agent for RAG queries
self.agent = RAGMemoryAgent(knowledge_base)
def add(self, key: str, value: Any) -> str:
"""
Add a key-value pair to the knowledge base.
"""
self.kb.add_entry(key, value)
return f"Added entry '{key}'."
def query(self, key: str) -> str:
"""
Retrieve the value for a given key.
"""
value = self.kb.query_entry(key)
if value is None:
return f"No entry found for key '{key}'."
return f"Value for '{key}': {value!s}"
def list(self) -> str:
"""
List all key/value pairs in the knowledge base.
"""
entries = self.kb.list_entries()
if not entries:
return "Knowledge base is empty."
return "\n".join(f"{k!s} : {v!s}" for k, v in entries)
def delete(self, key: str) -> str:
"""
Delete an entry by key.
"""
removed = self.kb.delete_entry(key)
if removed is None:
return f"No entry found for key '{key}'."
return f"Deleted entry '{key}'."
def ask(self, question: str) -> str:
"""
Ask a question to the RAG memory agent.
"""
return self.agent.ask(question)