This commit is contained in:
@@ -1,95 +1,98 @@
|
||||
# Agent with RAG Memory (ChromaDB)
|
||||
# Agent with RAG Memory
|
||||
|
||||
This project implements a Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** as its sole vector store. The agent can ingest documents, store their embeddings, retrieve relevant passages, and generate answers using OpenAI’s GPT models.
|
||||
This project implements a retrieval‑augmented generation (RAG) agent that uses **Qdrant** as the vector store and **Ollama** for local LLM inference.
|
||||
The agent follows the latest LangChain API and is fully configurable via environment variables.
|
||||
|
||||
## Features
|
||||
|
||||
- **Vector Store** – Uses ChromaDB for storing and querying embeddings.
|
||||
- **Embeddings** – Generated with OpenAI’s `text-embedding-ada-002`.
|
||||
- **Chat** – Generates responses with OpenAI’s `gpt-3.5-turbo`.
|
||||
- **Public API** – The `Agent` class exposes `init`, `ingest`, and `ask` methods, keeping the original interface unchanged.
|
||||
- **Qdrant** vector store (via `langchain-qdrant`)
|
||||
- **Ollama** local LLM integration (via `langchain-ollama`)
|
||||
- Retrieval‑augmented generation with conversation memory
|
||||
- Simple CLI interface for quick testing
|
||||
|
||||
## Setup
|
||||
## Installation
|
||||
|
||||
1. **Clone the repository**
|
||||
```bash
|
||||
# Clone the repository
|
||||
git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
|
||||
cd agent-s-rag-pamyatyu
|
||||
|
||||
```bash
|
||||
git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
|
||||
cd agent-s-rag-pamyatyu
|
||||
```
|
||||
# Create a virtual environment (optional but recommended)
|
||||
python -m venv .venv
|
||||
source .venv/bin/activate # On Windows: .venv\\Scripts\\activate
|
||||
|
||||
2. **Install dependencies**
|
||||
# Install dependencies
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
```bash
|
||||
npm install
|
||||
```
|
||||
## Configuration
|
||||
|
||||
3. **Configure environment variables**
|
||||
Create a `.env` file in the project root (or set environment variables directly):
|
||||
|
||||
Create a `.env` file in the project root (or export the variables in your shell):
|
||||
```dotenv
|
||||
# Qdrant
|
||||
QDRANT_HOST=localhost
|
||||
QDRANT_PORT=6333
|
||||
QDRANT_API_KEY= # leave empty if no key
|
||||
|
||||
```dotenv
|
||||
# ChromaDB
|
||||
CHROMA_URL=localhost
|
||||
CHROMA_PORT=8000
|
||||
# Ollama
|
||||
OLLAMA_HOST=localhost
|
||||
OLLAMA_PORT=11434
|
||||
OLLAMA_MODEL=llama3
|
||||
|
||||
# OpenAI
|
||||
OPENAI_API_KEY=YOUR_OPENAI_API_KEY
|
||||
```
|
||||
# Optional: collection name
|
||||
QDRANT_COLLECTION=documents
|
||||
```
|
||||
|
||||
- `CHROMA_URL` and `CHROMA_PORT` point to your ChromaDB instance.
|
||||
- `OPENAI_API_KEY` is required for embeddings and chat completions.
|
||||
|
||||
4. **Run ChromaDB**
|
||||
|
||||
Ensure a ChromaDB server is running on the specified host/port. You can start a local instance with Docker:
|
||||
|
||||
```bash
|
||||
docker run -d -p 8000:8000 chromadb/chroma
|
||||
```
|
||||
> **Note**: The Qdrant instance must be running and accessible at the specified host/port.
|
||||
> The Ollama server must be running locally and expose the chosen model.
|
||||
|
||||
## Usage
|
||||
|
||||
```js
|
||||
const { Agent } = require('./src');
|
||||
### Adding Documents
|
||||
|
||||
(async () => {
|
||||
const agent = new Agent();
|
||||
await agent.init();
|
||||
```python
|
||||
from src.agent import Agent
|
||||
from langchain_core.documents import Document
|
||||
|
||||
// Ingest documents
|
||||
await agent.ingest('The quick brown fox jumps over the lazy dog.', { source: 'example.txt' });
|
||||
agent = Agent()
|
||||
|
||||
// Ask a question
|
||||
const answer = await agent.ask('What did the fox do?');
|
||||
console.log(answer);
|
||||
})();
|
||||
docs = [
|
||||
Document(page_content="Python is a programming language.", metadata={"source": "python.txt"}),
|
||||
Document(page_content="LangChain is a framework for LLM applications.", metadata={"source": "langchain.txt"}),
|
||||
]
|
||||
|
||||
agent.add_documents(docs)
|
||||
```
|
||||
|
||||
## API
|
||||
|
||||
| Method | Description |
|
||||
|--------|-------------|
|
||||
| `init()` | Initializes the vector store (creates collection if needed). |
|
||||
| `ingest(text, metadata)` | Adds a document to the vector store. |
|
||||
| `ask(question)` | Retrieves relevant passages and generates an answer. |
|
||||
|
||||
## Testing
|
||||
|
||||
If you have a test suite, run:
|
||||
### Querying the Agent
|
||||
|
||||
```bash
|
||||
npm test
|
||||
python -m src.agent "What is LangChain?"
|
||||
```
|
||||
|
||||
All tests should pass after the ChromaDB integration.
|
||||
or from Python:
|
||||
|
||||
## Notes
|
||||
```python
|
||||
response = agent.run("What is LangChain?")
|
||||
print(response["answer"])
|
||||
```
|
||||
|
||||
- The agent’s public API remains unchanged; only the underlying vector store implementation has been swapped to ChromaDB.
|
||||
- No new external services are introduced beyond ChromaDB and the existing OpenAI usage.
|
||||
- Ensure that the ChromaDB server is reachable; otherwise, the agent will throw connection errors.
|
||||
The response will include the answer and the source documents used.
|
||||
|
||||
---
|
||||
## Project Structure
|
||||
|
||||
Happy coding!
|
||||
```
|
||||
agent-s-rag-pamyatyu/
|
||||
├── src/
|
||||
│ ├── agent.py
|
||||
│ ├── config.py
|
||||
│ └── vector_store.py
|
||||
├── requirements.txt
|
||||
├── pyproject.toml
|
||||
└── README.md
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
MIT License
|
||||
+7
-11
@@ -1,19 +1,15 @@
|
||||
[project]
|
||||
name = "agent-s-rag-pamyatyu"
|
||||
version = "0.1.0"
|
||||
description = "A simple RAG agent with interactive CLI using LangChain tools."
|
||||
authors = [
|
||||
{ name = "Artur Kuzakhmetov", email = "artur@example.com" }
|
||||
]
|
||||
description = "Agent with RAG memory using Qdrant and Ollama"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.8"
|
||||
dependencies = [
|
||||
"langchain>=0.0.0",
|
||||
"python-dotenv>=0.21.0"
|
||||
]
|
||||
requires-python = ">=3.9"
|
||||
|
||||
[project.scripts]
|
||||
agent-s-rag = "src.main:main"
|
||||
[project.dependencies]
|
||||
langchain = ">=0.1.0"
|
||||
langchain-qdrant = ">=0.1.0"
|
||||
langchain-ollama = ">=0.1.0"
|
||||
python-dotenv = ">=1.0.0"
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=42", "wheel"]
|
||||
|
||||
+4
-3
@@ -1,3 +1,4 @@
|
||||
# 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.
|
||||
langchain>=0.1.0
|
||||
langchain-qdrant>=0.1.0
|
||||
langchain-ollama>=0.1.0
|
||||
python-dotenv>=1.0.0
|
||||
+73
-72
@@ -1,93 +1,94 @@
|
||||
"""
|
||||
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.
|
||||
Agent implementation using the latest LangChain API.
|
||||
The agent uses a RetrievalQA chain backed by the Qdrant vector store
|
||||
and the Ollama LLM for local inference.
|
||||
"""
|
||||
|
||||
from typing import Callable, Dict, Any
|
||||
from __future__ import annotations
|
||||
|
||||
# Import the KnowledgeBaseTool but do not alter existing logic
|
||||
from .knowledge_base import KnowledgeBaseTool
|
||||
from typing import Any
|
||||
|
||||
from langchain_ollama import Ollama
|
||||
from langchain.chains import RetrievalQA
|
||||
from langchain.memory import ConversationBufferMemory
|
||||
from langchain.prompts import PromptTemplate
|
||||
|
||||
from .config import OLLAMA_HOST, OLLAMA_PORT, OLLAMA_MODEL
|
||||
from .vector_store import QdrantVectorStore
|
||||
from langchain_core.documents import Document
|
||||
|
||||
|
||||
class Agent:
|
||||
"""
|
||||
A simple agent that can execute registered tools.
|
||||
|
||||
The agent's tool registry maps tool names to callable objects.
|
||||
A simple retrieval-based agent that answers user queries
|
||||
using documents stored in Qdrant and an Ollama LLM.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.tools: Dict[str, Callable[..., Any]] = {}
|
||||
# Register core tools
|
||||
self._register_core_tools()
|
||||
def __init__(self, collection_name: str = "documents"):
|
||||
# Initialize LLM
|
||||
self.llm = Ollama(
|
||||
model=OLLAMA_MODEL,
|
||||
base_url=f"http://{OLLAMA_HOST}:{OLLAMA_PORT}",
|
||||
)
|
||||
|
||||
def _register_core_tools(self) -> None:
|
||||
# Initialize vector store
|
||||
self.vector_store = QdrantVectorStore(collection_name=collection_name)
|
||||
|
||||
# Memory to keep conversation context
|
||||
self.memory = ConversationBufferMemory(memory_key="chat_history")
|
||||
|
||||
# Prompt template
|
||||
self.prompt = PromptTemplate(
|
||||
input_variables=["chat_history", "question"],
|
||||
template=(
|
||||
"You are a helpful assistant. Use the following context to answer the question.\n"
|
||||
"Context:\n{chat_history}\n"
|
||||
"Question: {question}\n"
|
||||
"Answer:"
|
||||
),
|
||||
)
|
||||
|
||||
# RetrievalQA chain
|
||||
self.chain = RetrievalQA.from_chain_type(
|
||||
llm=self.llm,
|
||||
chain_type="stuff",
|
||||
retriever=self.vector_store.as_retriever(),
|
||||
memory=self.memory,
|
||||
return_source_documents=True,
|
||||
chain_type_kwargs={"prompt": self.prompt},
|
||||
)
|
||||
|
||||
def add_documents(self, documents: list[Document]) -> None:
|
||||
"""
|
||||
Register the default set of tools with the agent.
|
||||
Add documents to the underlying vector store.
|
||||
"""
|
||||
# Register the KnowledgeBaseTool under the name 'knowledge_base'
|
||||
self.tools["knowledge_base"] = KnowledgeBaseTool()
|
||||
self.vector_store.add_documents(documents)
|
||||
|
||||
def register_tool(self, name: str, tool: Callable[..., Any]) -> None:
|
||||
def run(self, question: str) -> Any:
|
||||
"""
|
||||
Register a new tool with the agent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The name under which the tool will be registered.
|
||||
tool : Callable[..., Any]
|
||||
The tool instance or callable.
|
||||
Run the agent on a user question.
|
||||
Returns the LLM's answer and the source documents.
|
||||
"""
|
||||
self.tools[name] = tool
|
||||
result = self.chain({"question": question})
|
||||
return result
|
||||
|
||||
def run_tool(self, name: str, *args, **kwargs) -> Any:
|
||||
"""
|
||||
Execute a registered tool.
|
||||
def __call__(self, question: str) -> Any:
|
||||
return self.run(question)
|
||||
|
||||
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.
|
||||
if __name__ == "__main__":
|
||||
import json
|
||||
import sys
|
||||
|
||||
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)
|
||||
# Simple CLI usage
|
||||
agent = Agent()
|
||||
if len(sys.argv) > 1:
|
||||
query = " ".join(sys.argv[1:])
|
||||
else:
|
||||
query = input("Enter your question: ")
|
||||
|
||||
# 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)
|
||||
# ...
|
||||
response = agent.run(query)
|
||||
print("\nAnswer:\n", response["answer"])
|
||||
print("\nSources:")
|
||||
for doc in response["source_documents"]:
|
||||
print(f"- {doc.metadata.get('source', 'unknown')}")
|
||||
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
Configuration helper for the agent project.
|
||||
Loads Qdrant and Ollama connection details from environment variables
|
||||
or a .env file. Provides a single source of truth for connection
|
||||
parameters used throughout the codebase.
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables from .env if present
|
||||
load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env")
|
||||
|
||||
# Qdrant configuration
|
||||
QDRANT_HOST: str = os.getenv("QDRANT_HOST", "localhost")
|
||||
QDRANT_PORT: int = int(os.getenv("QDRANT_PORT", "6333"))
|
||||
QDRANT_API_KEY: str | None = os.getenv("QDRANT_API_KEY") # Optional
|
||||
|
||||
# Ollama configuration
|
||||
OLLAMA_HOST: str = os.getenv("OLLAMA_HOST", "localhost")
|
||||
OLLAMA_PORT: int = int(os.getenv("OLLAMA_PORT", "11434"))
|
||||
OLLAMA_MODEL: str = os.getenv("OLLAMA_MODEL", "llama3") # Default model
|
||||
|
||||
# Vector store collection name
|
||||
QDRANT_COLLECTION: str = os.getenv("QDRANT_COLLECTION", "documents")
|
||||
@@ -0,0 +1,97 @@
|
||||
"""
|
||||
Vector store implementation using Qdrant via langchain-qdrant.
|
||||
Provides a simple interface for adding documents and performing
|
||||
similarity search. Embeddings are generated using OpenAIEmbeddings
|
||||
by default, but can be overridden by passing a custom embedding
|
||||
function.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Iterable, List, Optional
|
||||
|
||||
from langchain.embeddings import OpenAIEmbeddings
|
||||
from langchain_qdrant import Qdrant
|
||||
from langchain.vectorstores import VectorStore
|
||||
from langchain_core.documents import Document
|
||||
|
||||
from .config import (
|
||||
QDRANT_HOST,
|
||||
QDRANT_PORT,
|
||||
QDRANT_API_KEY,
|
||||
QDRANT_COLLECTION,
|
||||
)
|
||||
|
||||
|
||||
class QdrantVectorStore(VectorStore):
|
||||
"""
|
||||
A wrapper around langchain_qdrant.Qdrant that implements the
|
||||
VectorStore interface expected by LangChain chains.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embeddings: Optional[OpenAIEmbeddings] = None,
|
||||
collection_name: str = QDRANT_COLLECTION,
|
||||
):
|
||||
self.embeddings = embeddings or OpenAIEmbeddings()
|
||||
self.collection_name = collection_name
|
||||
|
||||
# Initialize Qdrant client
|
||||
self.client = Qdrant(
|
||||
host=QDRANT_HOST,
|
||||
port=QDRANT_PORT,
|
||||
api_key=QDRANT_API_KEY,
|
||||
collection_name=self.collection_name,
|
||||
)
|
||||
|
||||
def add_documents(self, documents: Iterable[Document]) -> None:
|
||||
"""
|
||||
Add a collection of documents to the Qdrant store.
|
||||
"""
|
||||
texts = [doc.page_content for doc in documents]
|
||||
metadatas = [doc.metadata for doc in documents]
|
||||
ids = [doc.id for doc in documents if doc.id is not None]
|
||||
|
||||
# Embed the documents
|
||||
embeddings = self.embeddings.embed_documents(texts)
|
||||
|
||||
# Upsert into Qdrant
|
||||
self.client.upsert(
|
||||
embeddings=embeddings,
|
||||
documents=texts,
|
||||
metadatas=metadatas,
|
||||
ids=ids,
|
||||
)
|
||||
|
||||
def similarity_search(
|
||||
self,
|
||||
query: str,
|
||||
k: int = 5,
|
||||
filter: Optional[dict] = None,
|
||||
) -> List[Document]:
|
||||
"""
|
||||
Perform a similarity search against the Qdrant store.
|
||||
"""
|
||||
query_embedding = self.embeddings.embed_query(query)
|
||||
results = self.client.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=k,
|
||||
filter=filter,
|
||||
)
|
||||
# Convert results to Document objects
|
||||
return [
|
||||
Document(
|
||||
page_content=result["payload"]["text"],
|
||||
metadata=result["payload"],
|
||||
id=result["id"],
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
|
||||
# The following methods are required by the VectorStore interface
|
||||
def embed_query(self, query: str) -> List[float]:
|
||||
return self.embeddings.embed_query(query)
|
||||
|
||||
def embed_documents(self, documents: List[str]) -> List[List[float]]:
|
||||
return self.embeddings.embed_documents(documents)
|
||||
Reference in New Issue
Block a user