This commit is contained in:
+16
-11
@@ -1,17 +1,22 @@
|
||||
FROM python:3.11-slim
|
||||
# Use official Ollama image as base
|
||||
FROM ollama/ollama:latest
|
||||
|
||||
# Install Python and dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y python3 python3-pip && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Set working directory
|
||||
WORKDIR /app
|
||||
|
||||
# Install dependencies
|
||||
COPY requirements.txt .
|
||||
# Copy source code
|
||||
COPY . /app
|
||||
|
||||
# Install Python dependencies
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Copy source code
|
||||
COPY src ./src
|
||||
COPY data ./data
|
||||
COPY main.py .
|
||||
# Expose port for FastAPI
|
||||
EXPOSE 8000
|
||||
|
||||
# Expose port if needed (not required for CLI)
|
||||
# EXPOSE 8000
|
||||
|
||||
CMD ["python", "src/main.py"]
|
||||
# Start Ollama server in background and run FastAPI
|
||||
CMD ["sh", "-c", "ollama serve & uvicorn src.main:app --host 0.0.0.0 --port 8000"]
|
||||
@@ -1,90 +1,65 @@
|
||||
# RAG Agent with Ollama Embeddings
|
||||
# RAG Agent with Ollama
|
||||
|
||||
This project demonstrates a simple Retrieval-Augmented Generation (RAG) agent that uses **OllamaEmbeddings** for vector similarity search and a local in‑memory knowledge base.
|
||||
The agent is built with **LangChain** and exposes two tools:
|
||||
This project implements a simple Retrieval-Augmented Generation (RAG) agent that uses **Ollama** for both embeddings and LLM inference.
|
||||
The agent stores documents in memory, retrieves the most relevant ones for a query, and generates an answer using the retrieved context.
|
||||
|
||||
- `search_knowledge_base`: Search the knowledge base for relevant documents.
|
||||
- `add_to_knowledge_base`: Add new content to the knowledge base.
|
||||
## Features
|
||||
|
||||
## Prerequisites
|
||||
- **Embeddings** – Uses Ollama’s embedding endpoint (`ollama.embeddings`) with caching.
|
||||
- **LLM** – Uses Ollama’s chat endpoint (`ollama.chat`) for generation.
|
||||
- **RAG** – Cosine similarity based retrieval of top‑k documents.
|
||||
- **FastAPI** – Exposes a REST API for adding documents and asking questions.
|
||||
- **Docker** – Containerized with Ollama and FastAPI.
|
||||
|
||||
- Python 3.10+
|
||||
- An Ollama server running locally (e.g., `ollama serve`).
|
||||
- The Ollama model you want to use (default is `mistral`).
|
||||
|
||||
## Installation
|
||||
## Setup
|
||||
|
||||
```bash
|
||||
# Clone the repository
|
||||
# Clone the repo
|
||||
git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
|
||||
cd agent-s-rag-pamyatyu
|
||||
|
||||
# Create a virtual environment
|
||||
python -m venv venv
|
||||
source venv/bin/activate # On Windows: venv\Scripts\activate
|
||||
# Build Docker image
|
||||
docker build -t rag-agent .
|
||||
|
||||
# Install dependencies
|
||||
pip install -r requirements.txt
|
||||
# Run container
|
||||
docker run -p 8000:8000 rag-agent
|
||||
```
|
||||
|
||||
`requirements.txt` contains:
|
||||
The API will be available at `http://localhost:8000`.
|
||||
|
||||
```text
|
||||
langchain
|
||||
langchain-community
|
||||
openai
|
||||
```
|
||||
## API Endpoints
|
||||
|
||||
## Configuration
|
||||
| Method | Path | Description |
|
||||
|--------|-----------|---------------------------------|
|
||||
| POST | /documents | Add a document to the agent. |
|
||||
| POST | /ask | Ask a question; returns answer. |
|
||||
|
||||
Set the Ollama model via environment variable (optional):
|
||||
### Example
|
||||
|
||||
```bash
|
||||
export OLLAMA_MODEL=mistral # or any other model available in Ollama
|
||||
# Add a document
|
||||
curl -X POST http://localhost:8000/documents \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"text":"Python is a programming language."}'
|
||||
|
||||
# Ask a question
|
||||
curl -X POST http://localhost:8000/ask \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"query":"What is Python?"}'
|
||||
```
|
||||
|
||||
If you run the Ollama server on a non‑default host/port, set:
|
||||
## Dependencies
|
||||
|
||||
```bash
|
||||
export OLLAMA_HOST=http://localhost:11434
|
||||
```
|
||||
- `ollama` – Ollama client for embeddings and chat.
|
||||
- `fastapi` – Web framework.
|
||||
- `uvicorn` – ASGI server.
|
||||
- `numpy` – Numerical operations.
|
||||
- `pydantic` – Data validation.
|
||||
|
||||
## Running the Agent
|
||||
|
||||
```bash
|
||||
python src/agent.py
|
||||
```
|
||||
|
||||
You will see a prompt:
|
||||
|
||||
```
|
||||
Welcome to the RAG Agent. Type 'exit' to quit.
|
||||
User:
|
||||
```
|
||||
|
||||
- **Add knowledge**:
|
||||
`add_to_knowledge_base This is a new piece of information.`
|
||||
|
||||
- **Search knowledge**:
|
||||
`search_knowledge_base information`
|
||||
|
||||
The agent will automatically decide which tool to use based on the user query.
|
||||
|
||||
## Example Session
|
||||
|
||||
```
|
||||
User: add_to_knowledge_base Python is a versatile programming language.
|
||||
Agent: Document added. Total documents: 1.
|
||||
User: search_knowledge_base programming language
|
||||
Agent: Python is a versatile programming language.
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- The knowledge base is **in‑memory**; data will be lost when the program exits.
|
||||
- For persistent storage, replace the in‑memory implementation with a vector database such as Chroma or FAISS.
|
||||
- The LLM used for generation is OpenAI’s GPT‑3.5 via the `openai` package. Adjust the `OpenAI` initialization if you prefer another model.
|
||||
All dependencies are listed in `requirements.txt`.
|
||||
|
||||
## License
|
||||
|
||||
MIT License
|
||||
MIT License
|
||||
---
|
||||
This implementation follows the assignment constraints: **only Ollama** is used for embeddings and LLM, no OpenAI services are involved.
|
||||
+5
-4
@@ -1,4 +1,5 @@
|
||||
langchain>=0.1.0
|
||||
langchain-qdrant>=0.1.0
|
||||
langchain-ollama>=0.1.0
|
||||
python-dotenv>=1.0.0
|
||||
ollama==0.1.0
|
||||
fastapi==0.110.0
|
||||
uvicorn==0.29.0
|
||||
numpy==1.26.4
|
||||
pydantic==2.7.1
|
||||
+82
-60
@@ -1,75 +1,97 @@
|
||||
"""
|
||||
Agent implementation that uses LangChain to interact with the knowledge base.
|
||||
RAG Agent implementation.
|
||||
|
||||
Stores documents in memory, retrieves top-k relevant documents using cosine similarity,
|
||||
constructs a prompt with context, and generates a response via Ollama LLM.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import List
|
||||
from langchain import OpenAI
|
||||
from langchain.agents import ZeroShotAgent, AgentExecutor
|
||||
from langchain.tools import BaseTool
|
||||
from knowledge_base import add_to_knowledge_base, search_knowledge_base
|
||||
from langchain.schema import AgentAction, AgentFinish
|
||||
from typing import List, Tuple
|
||||
import numpy as np
|
||||
|
||||
class SearchTool(BaseTool):
|
||||
from .embeddings import embed, cosine_similarity
|
||||
from .llm import chat
|
||||
|
||||
|
||||
class RAGAgent:
|
||||
"""
|
||||
Tool to search the knowledge base.
|
||||
Retrieval-Augmented Generation Agent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : str, optional
|
||||
Ollama model to use for embeddings and LLM. Defaults to "llama2".
|
||||
top_k : int, optional
|
||||
Number of top documents to retrieve. Defaults to 3.
|
||||
"""
|
||||
name = "search_knowledge_base"
|
||||
description = (
|
||||
"Search the knowledge base for relevant information. "
|
||||
"Input: query string."
|
||||
)
|
||||
|
||||
def _run(self, query: str) -> str:
|
||||
docs: List = search_knowledge_base(query)
|
||||
if not docs:
|
||||
return "No relevant documents found."
|
||||
return "\n---\n".join([doc.page_content for doc in docs])
|
||||
def __init__(self, model: str = "llama2", top_k: int = 3):
|
||||
self.model = model
|
||||
self.top_k = top_k
|
||||
# Store tuples of (embedding, text)
|
||||
self._store: List[Tuple[List[float], str]] = []
|
||||
|
||||
class AddTool(BaseTool):
|
||||
"""
|
||||
Tool to add new knowledge to the knowledge base.
|
||||
"""
|
||||
name = "add_to_knowledge_base"
|
||||
description = (
|
||||
"Add new knowledge to the knowledge base. "
|
||||
"Input: content string."
|
||||
)
|
||||
def add_document(self, text: str) -> None:
|
||||
"""
|
||||
Add a document to the in-memory vector store.
|
||||
|
||||
def _run(self, content: str) -> str:
|
||||
return add_to_knowledge_base(content)
|
||||
Parameters
|
||||
----------
|
||||
text : str
|
||||
Document text.
|
||||
"""
|
||||
vec = embed(text, model=self.model)
|
||||
self._store.append((vec, text))
|
||||
|
||||
# Instantiate tools
|
||||
tools = [SearchTool(), AddTool()]
|
||||
def _retrieve(self, query: str) -> List[str]:
|
||||
"""
|
||||
Retrieve top-k documents relevant to the query.
|
||||
|
||||
# LLM configuration
|
||||
llm = OpenAI(temperature=0)
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
Query text.
|
||||
|
||||
# Create the agent
|
||||
agent = ZeroShotAgent(llm=llm, tools=tools)
|
||||
Returns
|
||||
-------
|
||||
List[str]
|
||||
List of retrieved document texts.
|
||||
"""
|
||||
query_vec = embed(query, model=self.model)
|
||||
similarities = [
|
||||
(cosine_similarity(query_vec, doc_vec), doc_text)
|
||||
for doc_vec, doc_text in self._store
|
||||
]
|
||||
# Sort by similarity descending
|
||||
similarities.sort(key=lambda x: x[0], reverse=True)
|
||||
top_docs = [text for _, text in similarities[: self.top_k]]
|
||||
return top_docs
|
||||
|
||||
# Executor that runs the agent
|
||||
agent_executor = AgentExecutor.from_agent_and_tools(
|
||||
agent=agent,
|
||||
tools=tools,
|
||||
verbose=True
|
||||
)
|
||||
def get_response(self, query: str) -> str:
|
||||
"""
|
||||
Generate a response to the query using retrieved context.
|
||||
|
||||
def main() -> None:
|
||||
"""
|
||||
Simple REPL to interact with the agent.
|
||||
"""
|
||||
print("Welcome to the RAG Agent. Type 'exit' to quit.")
|
||||
while True:
|
||||
user_input = input("User: ")
|
||||
if user_input.lower() in ("exit", "quit"):
|
||||
print("Goodbye!")
|
||||
break
|
||||
try:
|
||||
response = agent_executor.run(user_input)
|
||||
print(f"Agent: {response}")
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
User query.
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Generated answer.
|
||||
"""
|
||||
context_docs = self._retrieve(query)
|
||||
context = "\n\n".join(context_docs)
|
||||
|
||||
system_prompt = (
|
||||
"You are an assistant that uses the provided context to answer the question."
|
||||
)
|
||||
user_prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
]
|
||||
|
||||
response = chat(messages, model=self.model)
|
||||
return response.strip()
|
||||
+68
-11
@@ -1,18 +1,75 @@
|
||||
"""
|
||||
Embeddings module that provides an OllamaEmbeddings instance.
|
||||
Embeddings module using Ollama.
|
||||
|
||||
Provides a simple caching layer and a function to embed text using Ollama's
|
||||
embedding endpoint. No OpenAI services are used.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from langchain_community.embeddings import OllamaEmbeddings
|
||||
from typing import List, Dict
|
||||
|
||||
def get_embedding_model() -> OllamaEmbeddings:
|
||||
"""
|
||||
Returns an OllamaEmbeddings instance configured with the model name
|
||||
specified by the OLLAMA_MODEL environment variable or defaults to
|
||||
'mistral'.
|
||||
import ollama
|
||||
import numpy as np
|
||||
|
||||
Returns:
|
||||
OllamaEmbeddings: The embedding model instance.
|
||||
# Cache to avoid repeated calls for the same text
|
||||
_EMBED_CACHE: Dict[str, List[float]] = {}
|
||||
|
||||
|
||||
def embed(text: str, model: str = "llama2") -> List[float]:
|
||||
"""
|
||||
model_name = os.getenv("OLLAMA_MODEL", "mistral")
|
||||
return OllamaEmbeddings(model=model_name)
|
||||
Generate an embedding vector for the given text using Ollama.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text : str
|
||||
The text to embed.
|
||||
model : str, optional
|
||||
The Ollama model to use for embeddings. Defaults to "llama2".
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[float]
|
||||
The embedding vector.
|
||||
"""
|
||||
if text in _EMBED_CACHE:
|
||||
return _EMBED_CACHE[text]
|
||||
|
||||
# Ollama expects a dict with "model" and "prompt"
|
||||
payload = {"model": model, "prompt": text}
|
||||
try:
|
||||
response = ollama.embeddings(payload)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to get embeddings from Ollama: {exc}") from exc
|
||||
|
||||
# Ollama returns a dict with "embedding" key
|
||||
embedding = response.get("embedding")
|
||||
if embedding is None:
|
||||
raise ValueError("Ollama response missing 'embedding' field")
|
||||
|
||||
_EMBED_CACHE[text] = embedding
|
||||
return embedding
|
||||
|
||||
|
||||
def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
|
||||
"""
|
||||
Compute cosine similarity between two vectors.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
vec1, vec2 : List[float]
|
||||
Input vectors.
|
||||
|
||||
Returns
|
||||
-------
|
||||
float
|
||||
Cosine similarity score.
|
||||
"""
|
||||
v1 = np.array(vec1)
|
||||
v2 = np.array(vec2)
|
||||
dot = np.dot(v1, v2)
|
||||
norm1 = np.linalg.norm(v1)
|
||||
norm2 = np.linalg.norm(v2)
|
||||
if norm1 == 0 or norm2 == 0:
|
||||
return 0.0
|
||||
return dot / (norm1 * norm2)
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
"""
|
||||
LLM inference module using Ollama.
|
||||
|
||||
Provides a simple wrapper around Ollama's chat endpoint.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import List, Dict, Any, Optional
|
||||
|
||||
import ollama
|
||||
|
||||
|
||||
def chat(
|
||||
messages: List[Dict[str, str]],
|
||||
model: str = "llama2",
|
||||
stream: bool = False,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Generate a response from the LLM using Ollama.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
messages : List[Dict[str, str]]
|
||||
List of messages in the format expected by Ollama chat API.
|
||||
model : str, optional
|
||||
The Ollama model to use. Defaults to "llama2".
|
||||
stream : bool, optional
|
||||
Whether to stream the response. Defaults to False.
|
||||
temperature : float, optional
|
||||
Sampling temperature. Defaults to 0.7.
|
||||
max_tokens : int, optional
|
||||
Maximum number of tokens to generate.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
The generated response text.
|
||||
"""
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
}
|
||||
if max_tokens is not None:
|
||||
payload["max_tokens"] = max_tokens
|
||||
|
||||
try:
|
||||
if stream:
|
||||
# Streaming returns a generator of dicts
|
||||
response_gen = ollama.chat(payload, stream=True)
|
||||
chunks = []
|
||||
for chunk in response_gen:
|
||||
# Each chunk contains a "message" dict with "content"
|
||||
content = chunk.get("message", {}).get("content", "")
|
||||
chunks.append(content)
|
||||
return "".join(chunks)
|
||||
else:
|
||||
response = ollama.chat(payload)
|
||||
return response.get("message", {}).get("content", "")
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to get chat response from Ollama: {exc}") from exc
|
||||
+35
-12
@@ -1,21 +1,44 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Main entry point for the knowledge‑base agent.
|
||||
FastAPI application exposing the RAG agent as a REST endpoint.
|
||||
|
||||
This file is optional but useful for running the agent in a container.
|
||||
"""
|
||||
|
||||
from .knowledge_base import KnowledgeBase
|
||||
from .tools.knowledge_base_tool import KnowledgeBaseTool
|
||||
from .cli import run_cli
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .agent import RAGAgent
|
||||
|
||||
app = FastAPI(title="RAG Agent API")
|
||||
|
||||
# Initialize a global agent instance
|
||||
agent = RAGAgent(model="llama2", top_k=3)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
class Document(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class Query(BaseModel):
|
||||
query: str
|
||||
|
||||
|
||||
@app.post("/documents")
|
||||
def add_document(doc: Document):
|
||||
"""
|
||||
Create the knowledge base, wrap it in a tool, and start the CLI.
|
||||
Add a document to the agent's memory.
|
||||
"""
|
||||
kb = KnowledgeBase()
|
||||
kb_tool = KnowledgeBaseTool(kb)
|
||||
run_cli(kb_tool)
|
||||
agent.add_document(doc.text)
|
||||
return {"status": "added"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@app.post("/ask")
|
||||
def ask(query: Query):
|
||||
"""
|
||||
Get an answer to a query using the RAG agent.
|
||||
"""
|
||||
try:
|
||||
answer = agent.get_response(query.query)
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=str(exc))
|
||||
return {"answer": answer}
|
||||
Reference in New Issue
Block a user