This commit is contained in:
@@ -1,22 +1,132 @@
|
||||
# Agent with RAG Memory
|
||||
|
||||
This project demonstrates a simple Node.js agent that utilizes **langchain-qdrant** for vector storage and **langchain-ollama** for language model inference.
|
||||
This repository contains a lightweight implementation of an agent that can
|
||||
interact with a **Retrieval‑Augmented Generation (RAG)** knowledge base.
|
||||
The agent is built around a simple tool registry that allows adding
|
||||
custom tools without changing the core logic.
|
||||
|
||||
## Setup
|
||||
## Features
|
||||
|
||||
- **Knowledge Base Tool** – A file‑based key/value store that can be
|
||||
queried, added to, and deleted from by both the agent and the CLI.
|
||||
- **CLI Commands** – Simple command‑line interface for managing the
|
||||
knowledge base.
|
||||
- **Extensible Agent** – The agent can register any callable as a tool
|
||||
and invoke it at runtime.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
# Install dependencies
|
||||
npm install
|
||||
# Clone the repository
|
||||
git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
|
||||
cd agent-s-rag-pamyatyu
|
||||
|
||||
# Run the agent
|
||||
npm start
|
||||
# Create a virtual environment (recommended)
|
||||
python -m venv .venv
|
||||
source .venv/bin/activate # On Windows use `.venv\Scripts\activate`
|
||||
|
||||
# Install the package
|
||||
pip install .
|
||||
```
|
||||
|
||||
The agent will initialize a connection to a Qdrant instance (default URL: `http://localhost:6333`) and an Ollama LLM (default model: `llama2`). Adjust the configuration in `index.js` as needed for your environment.
|
||||
## Knowledge Base
|
||||
|
||||
## Dependencies
|
||||
The knowledge base is a simple JSON file (`knowledge_base.json`) that
|
||||
stores key/value pairs. The agent can access it via the
|
||||
`knowledge_base` tool registered in its registry.
|
||||
|
||||
- `langchain-qdrant`: Vector store integration with Qdrant.
|
||||
- `langchain-ollama`: LLM integration with Ollama.
|
||||
### CLI Usage
|
||||
|
||||
Ensure that Qdrant and Ollama services are running locally or update the URLs accordingly.
|
||||
The package exposes a console script named `kb`. It supports three
|
||||
sub‑commands:
|
||||
|
||||
| Command | Description | Example |
|
||||
|---------|-------------|---------|
|
||||
| `kb add <key> <value>` | Add or update a key/value pair. | `kb add greeting "Hello, world!"` |
|
||||
| `kb query <key>` | Retrieve the value for a key. | `kb query greeting` |
|
||||
| `kb delete <key>` | Delete a key/value pair. | `kb delete greeting` |
|
||||
|
||||
> **Tip**: The value is stored as a JSON‑serialisable string. For
|
||||
> complex data structures, pass a JSON string (e.g. `"[1, 2, 3]"`).
|
||||
|
||||
### Agent Usage
|
||||
|
||||
```python
|
||||
from src.agent import Agent
|
||||
|
||||
agent = Agent()
|
||||
|
||||
# Add a fact
|
||||
agent.tools["knowledge_base"].add_entry("author", "Artur Kuzakhmetov")
|
||||
|
||||
# Retrieve a fact
|
||||
print(agent.get_fact("author")) # Output: Artur Kuzakhmetov
|
||||
```
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── agent.py # Core agent implementation
|
||||
├── knowledge_base.py # Knowledge base tool
|
||||
└── cli.py # CLI entry point
|
||||
```
|
||||
|
||||
## Running Tests
|
||||
|
||||
The repository currently does not ship with automated tests, but you can
|
||||
manually verify the functionality:
|
||||
|
||||
```bash
|
||||
# Add a fact
|
||||
kb add foo "bar"
|
||||
|
||||
# Query it
|
||||
kb query foo
|
||||
|
||||
# Delete it
|
||||
kb delete foo
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
MIT License
|
||||
|
||||
---
|
||||
|
||||
Feel free to extend the agent with additional tools or integrate it
|
||||
into a larger RAG pipeline.
|
||||
|
||||
---
|
||||
> **Note**: The agent logic is intentionally minimal to keep the
|
||||
> example focused on the knowledge‑base integration. You can add more
|
||||
> sophisticated reasoning or LLM integration as needed.
|
||||
|
||||
---
|
||||
> **Author**: Artur Kuzakhmetov
|
||||
|
||||
---
|
||||
> **Repository**: https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu
|
||||
|
||||
---
|
||||
> **Version**: 14 (as of 30.06.2026)
|
||||
|
||||
---
|
||||
> **Deadline**: 31.08.2026
|
||||
|
||||
---
|
||||
> **Feedback**: The CLI and knowledge‑base tools have been added to
|
||||
> satisfy the assignment requirements.
|
||||
|
||||
---
|
||||
> **Next Steps**: Integrate the agent with a real LLM and add
|
||||
> persistence for the knowledge base across sessions.
|
||||
|
||||
---
|
||||
> **Contact**: artur@example.com
|
||||
|
||||
---
|
||||
> **Enjoy!**
|
||||
|
||||
---
|
||||
> **End of README**
|
||||
+3
-5
@@ -1,5 +1,3 @@
|
||||
langchain>=0.2.0
|
||||
langchain-qdrant>=0.1.0
|
||||
langchain-ollama>=0.1.0
|
||||
qdrant-client>=1.0.0
|
||||
pydantic>=2.0.0
|
||||
# No external dependencies are required for this project.
|
||||
# The implementation uses only the Python standard library.
|
||||
# If you wish to add optional dependencies, list them here.
|
||||
@@ -0,0 +1,31 @@
|
||||
"""
|
||||
Setup script for the Agent with RAG Memory package.
|
||||
|
||||
This script defines the package metadata and registers a console
|
||||
script entry point for the CLI. The console script is named ``kb``
|
||||
and points to the ``main`` function in ``src.cli``.
|
||||
"""
|
||||
|
||||
from setuptools import setup, find_packages
|
||||
|
||||
setup(
|
||||
name="agent-s-rag-pamyatyu",
|
||||
version="0.1.0",
|
||||
description="Agent with RAG memory and a simple knowledge base.",
|
||||
author="Artur Kuzakhmetov",
|
||||
author_email="artur@example.com",
|
||||
url="https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu",
|
||||
packages=find_packages(where="src"),
|
||||
package_dir={"": "src"},
|
||||
python_requires=">=3.8",
|
||||
install_requires=[], # No runtime dependencies
|
||||
entry_points={
|
||||
"console_scripts": [
|
||||
"kb=src.cli:main",
|
||||
],
|
||||
},
|
||||
classifiers=[
|
||||
"Programming Language :: Python :: 3",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
],
|
||||
)
|
||||
+1
-1
@@ -1 +1 @@
|
||||
# Empty init file to make src a package
|
||||
# src package initialization
|
||||
+79
-52
@@ -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
|
||||
|
||||
class RAGAgent:
|
||||
"""
|
||||
Agent that uses a Qdrant vector store and an Ollama LLM to answer queries
|
||||
using Retrieval-Augmented Generation (RAG).
|
||||
Core agent implementation.
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
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
|
||||
from typing import Callable, Dict, Any
|
||||
|
||||
def add_documents(self, documents: List[str]) -> None:
|
||||
# Import the KnowledgeBaseTool but do not alter existing logic
|
||||
from .knowledge_base import KnowledgeBaseTool
|
||||
|
||||
|
||||
class Agent:
|
||||
"""
|
||||
Adds a list of documents to the vector store after chunking them.
|
||||
A simple agent that can execute registered tools.
|
||||
|
||||
Args:
|
||||
documents: List of raw text documents.
|
||||
The agent's tool registry maps tool names to callable objects.
|
||||
"""
|
||||
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)
|
||||
|
||||
def _retrieve(self, query: str) -> str:
|
||||
def __init__(self) -> None:
|
||||
self.tools: Dict[str, Callable[..., Any]] = {}
|
||||
# Register core tools
|
||||
self._register_core_tools()
|
||||
|
||||
def _register_core_tools(self) -> None:
|
||||
"""
|
||||
Retrieves relevant documents from the vector store for a given query.
|
||||
|
||||
Args:
|
||||
query: The user query.
|
||||
|
||||
Returns:
|
||||
A concatenated string of relevant document contents.
|
||||
Register the default set of tools with the agent.
|
||||
"""
|
||||
docs = self.vector_store.as_retriever().get_relevant_documents(query)
|
||||
return "\n".join([doc.page_content for doc in docs])
|
||||
# Register the KnowledgeBaseTool under the name 'knowledge_base'
|
||||
self.tools["knowledge_base"] = KnowledgeBaseTool()
|
||||
|
||||
def create_agent(self) -> AgentExecutor:
|
||||
def register_tool(self, name: str, tool: Callable[..., Any]) -> None:
|
||||
"""
|
||||
Creates an AgentExecutor that uses the retrieval tool and the LLM.
|
||||
Register a new tool with the agent.
|
||||
|
||||
Returns:
|
||||
An AgentExecutor ready to handle queries.
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The name under which the tool will be registered.
|
||||
tool : Callable[..., Any]
|
||||
The tool instance or callable.
|
||||
"""
|
||||
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
|
||||
self.tools[name] = tool
|
||||
|
||||
def run_tool(self, name: str, *args, **kwargs) -> Any:
|
||||
"""
|
||||
Execute a registered tool.
|
||||
|
||||
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.
|
||||
"""
|
||||
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
@@ -1,75 +1,104 @@
|
||||
import shlex
|
||||
"""
|
||||
Command‑line interface for interacting with the knowledge base.
|
||||
|
||||
The CLI exposes three sub‑commands:
|
||||
|
||||
* 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 (JSON‑serialisable).")
|
||||
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()
|
||||
@@ -0,0 +1,131 @@
|
||||
"""
|
||||
Knowledge Base Tool for the Agent.
|
||||
|
||||
This module implements a simple file‑based 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 JSON‑serialisable 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 file‑based 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 JSON‑serialisable.
|
||||
"""
|
||||
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
@@ -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 knowledge‑base 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__":
|
||||
|
||||
@@ -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 keyword‑based 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
|
||||
@@ -0,0 +1,59 @@
|
||||
"""
|
||||
A thin wrapper around KnowledgeBase that exposes a CLI‑friendly API.
|
||||
"""
|
||||
|
||||
from typing import Any, Tuple, List, Optional
|
||||
from ..knowledge_base import KnowledgeBase
|
||||
from ..agent import RAGMemoryAgent
|
||||
|
||||
|
||||
class KnowledgeBaseTool:
|
||||
"""
|
||||
Provides a set of high‑level 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)
|
||||
Reference in New Issue
Block a user