diff --git a/Dockerfile b/Dockerfile index 1399ef7..9e674a7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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"] \ No newline at end of file +# Start Ollama server in background and run FastAPI +CMD ["sh", "-c", "ollama serve & uvicorn src.main:app --host 0.0.0.0 --port 8000"] \ No newline at end of file diff --git a/README.md b/README.md index eacd3b6..c72d24b 100644 --- a/README.md +++ b/README.md @@ -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 \ No newline at end of file +MIT License +--- +This implementation follows the assignment constraints: **only Ollama** is used for embeddings and LLM, no OpenAI services are involved. \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 214ecb0..08e42e3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ -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 +ollama==0.1.0 +fastapi==0.110.0 +uvicorn==0.29.0 +numpy==1.26.4 +pydantic==2.7.1 \ No newline at end of file diff --git a/src/agent.py b/src/agent.py index 09d7bf7..941efed 100644 --- a/src/agent.py +++ b/src/agent.py @@ -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() \ No newline at end of file + 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() \ No newline at end of file diff --git a/src/embeddings.py b/src/embeddings.py index 33bd57c..c210d47 100644 --- a/src/embeddings.py +++ b/src/embeddings.py @@ -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) \ No newline at end of file + 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) \ No newline at end of file diff --git a/src/llm.py b/src/llm.py new file mode 100644 index 0000000..5256468 --- /dev/null +++ b/src/llm.py @@ -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 \ No newline at end of file diff --git a/src/main.py b/src/main.py index 9090959..e5facd4 100644 --- a/src/main.py +++ b/src/main.py @@ -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() \ No newline at end of file +@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} \ No newline at end of file