diff --git a/README.md b/README.md index faff64b..6879689 100644 --- a/README.md +++ b/README.md @@ -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! \ No newline at end of file +``` +agent-s-rag-pamyatyu/ +├── src/ +│ ├── agent.py +│ ├── config.py +│ └── vector_store.py +├── requirements.txt +├── pyproject.toml +└── README.md +``` + +## License + +MIT License \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 5c98e77..c48c7da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/requirements.txt b/requirements.txt index ac3059e..214ecb0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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. \ No newline at end of file +langchain>=0.1.0 +langchain-qdrant>=0.1.0 +langchain-ollama>=0.1.0 +python-dotenv>=1.0.0 \ No newline at end of file diff --git a/src/agent.py b/src/agent.py index ac9ef31..3068779 100644 --- a/src/agent.py +++ b/src/agent.py @@ -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) - # ... \ No newline at end of file + response = agent.run(query) + print("\nAnswer:\n", response["answer"]) + print("\nSources:") + for doc in response["source_documents"]: + print(f"- {doc.metadata.get('source', 'unknown')}") \ No newline at end of file diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000..3bd6292 --- /dev/null +++ b/src/config.py @@ -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") \ No newline at end of file diff --git a/src/vector_store.py b/src/vector_store.py new file mode 100644 index 0000000..6e35a38 --- /dev/null +++ b/src/vector_store.py @@ -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) \ No newline at end of file