feat: solution for 'Агент с RAG-памятью'

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 source code
-COPY requirements.txt .
+COPY . /app
+
+# Install Python dependencies
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Copy source code
+# Expose port for FastAPI
-COPY src ./src
+EXPOSE 8000
-COPY data ./data
-COPY main.py .
 
-# Expose port if needed (not required for CLI)
+# Start Ollama server in background and run FastAPI
-# EXPOSE 8000
+CMD ["sh", "-c", "ollama serve & uvicorn src.main:app --host 0.0.0.0 --port 8000"]
-
-CMD ["python", "src/main.py"]

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.  
+This project implements a simple Retrieval-Augmented Generation (RAG) agent that uses **Ollama** for both embeddings and LLM inference.  
-The agent is built with **LangChain** and exposes two tools:
+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.
+## Features
-- `add_to_knowledge_base`: Add new content to the knowledge base.
 
-## 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+
+## Setup
-- An Ollama server running locally (e.g., `ollama serve`).
-- The Ollama model you want to use (default is `mistral`).
-
-## Installation
 
 ```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
+# Build Docker image
-python -m venv venv
+docker build -t rag-agent .
-source venv/bin/activate  # On Windows: venv\Scripts\activate
 
-# Install dependencies
+# Run container
-pip install -r requirements.txt
+docker run -p 8000:8000 rag-agent
 ```
 
-`requirements.txt` contains:
+The API will be available at `http://localhost:8000`.
 
-```text
+## API Endpoints
-langchain
-langchain-community
-openai
-```
 
-## 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
+- `ollama` – Ollama client for embeddings and chat.
-export OLLAMA_HOST=http://localhost:11434
+- `fastapi` – Web framework.
-```
+- `uvicorn` – ASGI server.
+- `numpy` – Numerical operations.
+- `pydantic` – Data validation.
 
-## Running the Agent
+All dependencies are listed in `requirements.txt`.
-
-```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.
 
 ## License
 
 MIT License
+--- 
+This implementation follows the assignment constraints: **only Ollama** is used for embeddings and LLM, no OpenAI services are involved.

requirements.txt

@@ -1,4 +1,5 @@
-langchain>=0.1.0
+ollama==0.1.0
-langchain-qdrant>=0.1.0
+fastapi==0.110.0
-langchain-ollama>=0.1.0
+uvicorn==0.29.0
-python-dotenv>=1.0.0
+numpy==1.26.4
+pydantic==2.7.1

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, Tuple
-from typing import List
+import numpy as np
-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
 
-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 = (
+    def __init__(self, model: str = "llama2", top_k: int = 3):
-        "Search the knowledge base for relevant information. "
+        self.model = model
-        "Input: query string."
+        self.top_k = top_k
+        # Store tuples of (embedding, text)
+        self._store: List[Tuple[List[float], str]] = []
+
+    def add_document(self, text: str) -> None:
+        """
+        Add a document to the in-memory vector store.
+
+        Parameters
+        ----------
+        text : str
+            Document text.
+        """
+        vec = embed(text, model=self.model)
+        self._store.append((vec, text))
+
+    def _retrieve(self, query: str) -> List[str]:
+        """
+        Retrieve top-k documents relevant to the query.
+
+        Parameters
+        ----------
+        query : str
+            Query text.
+
+        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
+
+    def get_response(self, query: str) -> str:
+        """
+        Generate a response to the query using retrieved context.
+
+        Parameters
+        ----------
+        query : str
+            User query.
+
+        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:"
 
-    def _run(self, query: str) -> str:
+        messages = [
-        docs: List = search_knowledge_base(query)
+            {"role": "system", "content": system_prompt},
-        if not docs:
+            {"role": "user", "content": user_prompt},
-            return "No relevant documents found."
+        ]
-        return "\n---\n".join([doc.page_content for doc in docs])
 
-class AddTool(BaseTool):
+        response = chat(messages, model=self.model)
-    """
+        return response.strip()
-    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 _run(self, content: str) -> str:
-        return add_to_knowledge_base(content)
-
-# Instantiate tools
-tools = [SearchTool(), AddTool()]
-
-# LLM configuration
-llm = OpenAI(temperature=0)
-
-# Create the agent
-agent = ZeroShotAgent(llm=llm, tools=tools)
-
-# Executor that runs the agent
-agent_executor = AgentExecutor.from_agent_and_tools(
-    agent=agent,
-    tools=tools,
-    verbose=True
-)
-
-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}")
-
-if __name__ == "__main__":
-    main()

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:
+import ollama
-    """
+import numpy as np
-    Returns an OllamaEmbeddings instance configured with the model name
-    specified by the OLLAMA_MODEL environment variable or defaults to
-    'mistral'.
 
-    Returns:
+# Cache to avoid repeated calls for the same text
-        OllamaEmbeddings: The embedding model instance.
+_EMBED_CACHE: Dict[str, List[float]] = {}
+
+
+def embed(text: str, model: str = "llama2") -> List[float]:
     """
-    model_name = os.getenv("OLLAMA_MODEL", "mistral")
+    Generate an embedding vector for the given text using Ollama.
-    return OllamaEmbeddings(model=model_name)
+
+    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)

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

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 fastapi import FastAPI, HTTPException
-from .tools.knowledge_base_tool import KnowledgeBaseTool
+from pydantic import BaseModel
-from .cli import run_cli
+
+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()
+    agent.add_document(doc.text)
-    kb_tool = KnowledgeBaseTool(kb)
+    return {"status": "added"}
-    run_cli(kb_tool)
 
 
-if __name__ == "__main__":
+@app.post("/ask")
-    main()
+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}