feat: solution for 'Агент с RAG-памятью'
CI / build (push) Has been cancelled

This commit is contained in:
2026-07-01 14:05:59 +03:00
parent bd49075b6e
commit 39d55136ad
14 changed files with 332 additions and 406 deletions
+31 -29
View File
@@ -1,35 +1,37 @@
const VectorStore = require('./vectorStore');
const { OpenAI } = require('openai');
const dotenv = require('dotenv');
dotenv.config();
import { OpenAI } from "@langchain/openai";
import { RetrievalQAChain } from "langchain/chains";
import { getEmbeddings } from "./embeddings.js";
import { getVectorStore } from "./vectorStore.js";
import { config } from "dotenv";
config();
const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
/**
* Creates a Retrieval QA chain using Ollama embeddings and Qdrant vector store.
* @returns {Promise<RetrievalQAChain>}
*/
export async function createAgent() {
const embeddings = getEmbeddings();
const vectorStore = await getVectorStore(embeddings);
class Agent {
constructor() {
this.vectorStore = new VectorStore();
}
const llm = new OpenAI({
temperature: 0,
modelName: "gpt-3.5-turbo",
});
async init() {
await this.vectorStore.init();
}
const chain = RetrievalQAChain.fromLLM(llm, vectorStore.asRetriever(), {
returnSourceDocuments: true,
});
async ingest(text, metadata = {}) {
await this.vectorStore.addDocument(text, metadata);
}
async ask(question) {
const results = await this.vectorStore.query(question, 3);
const context = results.documents
.map((doc, idx) => `Source ${idx + 1}:\n${doc}`)
.join('\n\n');
const prompt = `You are a helpful assistant. Use the following context to answer the question.\n\n${context}\n\nQuestion: ${question}\nAnswer:`;
const completion = await openai.chat.completions.create({
model: 'gpt-3.5-turbo',
messages: [{ role: 'user', content: prompt }],
});
return completion.choices[0].message.content.trim();
}
return chain;
}
module.exports = Agent;
/**
* Runs the agent with a given query.
* @param {string} query
* @returns {Promise<object>} The chain's output.
*/
export async function runAgent(query) {
const chain = await createAgent();
const result = await chain.call({ query });
return result;
}
+41 -13
View File
@@ -1,22 +1,50 @@
from langchain_ollama import Ollama, OllamaEmbeddings
"""
Agent implementation that performs RAG memory retrieval and response generation.
"""
from langchain import PromptTemplate, LLMChain
from langchain.chains import RetrievalQA
import config
from src.vector_store import QdrantVectorStore
from langchain.memory import ConversationBufferMemory
from langchain.llms import OpenAI
from langchain.vectorstores import Qdrant
from config import OPENAI_API_KEY, OPENAI_MODEL
from vector_store import get_vector_store
def create_agent(vector_store: QdrantVectorStore):
def build_agent() -> RetrievalQA:
"""
Create a RetrievalQA agent that uses Ollama for both embeddings and LLM.
Builds and returns a RetrievalQA chain configured with:
- OpenAI LLM for generation
- Qdrant vector store for retrieval
- ConversationBufferMemory for context
"""
# Embeddings for the vector store
embeddings = OllamaEmbeddings(model=config.OLLAMA_MODEL)
# LLM for generation
llm = OpenAI(
temperature=0,
openai_api_key=OPENAI_API_KEY,
model_name=OPENAI_MODEL
)
# LLM for generating answers
llm = Ollama(model=config.OLLAMA_MODEL)
# Vector store and retriever
vector_store: Qdrant = get_vector_store()
retriever = vector_store.as_retriever(search_kwargs={"k": 5})
# Build the RetrievalQA chain
qa_chain = RetrievalQA.from_chain_type(
# Memory to keep conversation context
memory = ConversationBufferMemory(
memory_key="chat_history",
return_messages=True
)
# RetrievalQA chain
chain = RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=vector_store.get_retriever(),
retriever=retriever,
memory=memory
)
return qa_chain
return chain
def ask_question(chain: RetrievalQA, question: str) -> str:
"""
Utility function to ask a question using the provided chain.
"""
return chain.run(question)
+17 -8
View File
@@ -1,9 +1,18 @@
import yaml
import os
# Configuration constants for the agent.
# Adjust these values if your local Ollama or Qdrant instances are running on different hosts/ports.
def load_config(path: str) -> dict:
if not os.path.exists(path):
raise FileNotFoundError(f"Config file {path} not found.")
with open(path, "r", encoding="utf-8") as f:
cfg = yaml.safe_load(f)
return cfg
# Ollama configuration
OLLAMA_HOST: str = "http://localhost"
OLLAMA_PORT: int = 11434 # Default Ollama port
# Qdrant configuration
QDRANT_HOST: str = "http://localhost"
QDRANT_PORT: int = 6333 # Default Qdrant port
QDRANT_COLLECTION_NAME: str = "rag_collection"
# OpenAI configuration (used for LLM generation)
OPENAI_API_KEY: str | None = None # Set via environment variable or .env file
OPENAI_MODEL: str = "gpt-3.5-turbo"
# Embedding model name for Ollama
OLLAMA_EMBEDDING_MODEL: str = "llama2" # Change if you use a different model
+14 -16
View File
@@ -1,21 +1,19 @@
import { OllamaEmbeddings } from 'ollama-embeddings';
import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama";
import { config } from "dotenv";
config();
/**
* Singleton instance of OllamaEmbeddings.
* The model name can be overridden via the OLLAMA_MODEL environment variable.
* Returns an instance of OllamaEmbeddings configured with the Ollama endpoint.
* @returns {OllamaEmbeddings}
*/
const modelName = process.env.OLLAMA_MODEL || 'all-minilm';
export const embeddings = new OllamaEmbeddings({
model: modelName,
// Optional: specify the Ollama host if not default
host: process.env.OLLAMA_HOST || 'http://localhost:11434'
});
export function getEmbeddings() {
const ollamaUrl = process.env.OLLAMA_URL;
if (!ollamaUrl) {
throw new Error("Environment variable OLLAMA_URL is not set.");
}
/**
* Utility to embed a single string.
* @param {string} text
* @returns {Promise<number[]>} embedding vector
*/
export async function embedText(text) {
return await embeddings.embedQuery(text);
return new OllamaEmbeddings({
model: "llama2",
baseUrl: ollamaUrl,
});
}
+9 -69
View File
@@ -1,75 +1,15 @@
"""
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.
Embeddings module using OllamaEmbeddings from langchain-community.
"""
import json
import os
from typing import List, Dict
from langchain_community.embeddings import OllamaEmbeddings
from config import OLLAMA_EMBEDDING_MODEL, OLLAMA_HOST, OLLAMA_PORT
import ollama
import numpy as np
# Cache to avoid repeated calls for the same text
_EMBED_CACHE: Dict[str, List[float]] = {}
def embed(text: str, model: str = "llama2") -> List[float]:
def get_ollama_embeddings() -> OllamaEmbeddings:
"""
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.
Returns an OllamaEmbeddings instance configured to use the local Ollama server.
"""
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)
return OllamaEmbeddings(
model=OLLAMA_EMBEDDING_MODEL,
base_url=f"{OLLAMA_HOST}:{OLLAMA_PORT}"
)
+2 -66
View File
@@ -1,67 +1,3 @@
import dotenv from 'dotenv';
import readline from 'readline';
import { search_knowledge_base } from './tools/searchKnowledgeBase.js';
import { add_to_knowledge_base } from './tools/addToKnowledgeBase.js';
import { runAgent } from "./agent.js";
dotenv.config();
/**
* Simple command-line agent that supports two commands:
* 1. /search <query> - searches the knowledge base
* 2. /add <content> - adds content to the knowledge base
* Any other input is treated as a normal message and the agent echoes it back.
*/
const rl = readline.createInterface({
input: process.stdin,
output: process.stdout,
prompt: 'Agent> '
});
console.log('Agent with RAG memory using Ollama embeddings.');
console.log('Commands:');
console.log(' /search <query> - Search knowledge base');
console.log(' /add <content> - Add content to knowledge base');
console.log(' /exit - Exit');
rl.prompt();
rl.on('line', async (line) => {
const trimmed = line.trim();
if (trimmed === '/exit') {
rl.close();
return;
}
if (trimmed.startsWith('/search ')) {
const query = trimmed.slice(8).trim();
if (!query) {
console.log('Please provide a query.');
} else {
console.log(`Searching for "${query}"...`);
const results = await search_knowledge_base(query);
if (results.length === 0) {
console.log('No relevant documents found.');
} else {
console.log('Top results:');
results.forEach((res, idx) => {
console.log(`${idx + 1}. [${res.id}] (${res.score.toFixed(4)})`);
console.log(` ${res.content}`);
});
}
}
} else if (trimmed.startsWith('/add ')) {
const content = trimmed.slice(5).trim();
if (!content) {
console.log('Please provide content to add.');
} else {
const { id } = await add_to_knowledge_base(content);
console.log(`Content added with id ${id}.`);
}
} else {
// Echo back the message (placeholder for more complex agent logic)
console.log(`You said: ${trimmed}`);
}
rl.prompt();
}).on('close', () => {
console.log('Goodbye!');
process.exit(0);
});
export { runAgent };
+24 -34
View File
@@ -1,42 +1,32 @@
"""
Entry point for the RAG agent.
"""
import os
from langchain.schema import Document
from src.vector_store import QdrantVectorStore
from src.agent import create_agent
import config
from dotenv import load_dotenv
from agent import build_agent, ask_question
from config import OPENAI_API_KEY
def main():
# Ensure Qdrant is reachable
os.environ["QDRANT_URL"] = f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}"
if config.QDRANT_API_KEY:
os.environ["QDRANT_API_KEY"] = config.QDRANT_API_KEY
def main() -> None:
# Load environment variables from .env if present
load_dotenv()
# Initialize embeddings and vector store
from langchain_ollama import OllamaEmbeddings
embeddings = OllamaEmbeddings(model=config.OLLAMA_MODEL)
vector_store = QdrantVectorStore(embeddings)
# Ensure OpenAI API key is available
if not OPENAI_API_KEY:
raise RuntimeError("OPENAI_API_KEY is not set. Please set it in environment or .env file.")
# Add sample documents (only if collection is empty)
# In a real scenario, you would load your corpus here
sample_docs = [
Document(page_content="Hello world! This is a test document.", metadata={"source": "test"}),
Document(page_content="LangChain is a powerful framework for building LLM applications.", metadata={"source": "test"}),
]
# Check if collection already has documents
try:
# Attempt to retrieve a document to see if collection is populated
vector_store.get_retriever().get_relevant_documents("test")
except Exception:
# If retrieval fails, add documents
vector_store.add_documents(sample_docs)
# Build the agent
chain = build_agent()
# Create the agent
agent = create_agent(vector_store)
# Run a sample query
query = "What is LangChain?"
print(f"Query: {query}")
result = agent.run(query)
print(f"Answer: {result}")
# Simple interactive loop
print("RAG Agent is ready. Type 'exit' to quit.")
while True:
user_input = input("\nYou: ")
if user_input.lower() in {"exit", "quit"}:
print("Goodbye!")
break
response = ask_question(chain, user_input)
print(f"Agent: {response}")
if __name__ == "__main__":
main()
+38 -47
View File
@@ -1,54 +1,45 @@
const chroma = require('./chromaClient');
const { OpenAI } = require('openai');
const dotenv = require('dotenv');
dotenv.config();
import { QdrantStore } from "@langchain/community/vectorstores/qdrant";
import { QdrantClient } from "@qdrant/js-client";
import { config } from "dotenv";
config();
const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
/**
* Initializes a Qdrant vector store with the provided embeddings instance.
* @param {OllamaEmbeddings} embeddings
* @returns {Promise<QdrantStore>}
*/
export async function getVectorStore(embeddings) {
const qdrantUrl = process.env.QDRANT_URL;
const qdrantApiKey = process.env.QDRANT_API_KEY;
class VectorStore {
constructor(collectionName = 'documents') {
this.collectionName = collectionName;
this.collection = null;
if (!qdrantUrl) {
throw new Error("Environment variable QDRANT_URL is not set.");
}
async init() {
this.collection = await chroma.getCollection({
name: this.collectionName,
metadata: { type: 'vector' },
const client = new QdrantClient({
url: qdrantUrl,
apiKey: qdrantApiKey,
});
const collectionName = "rag_collection";
// Ensure the collection exists; create if missing
const collections = await client.getCollections();
const exists = collections.collections.some(
(c) => c.name === collectionName
);
if (!exists) {
await client.createCollection({
collection_name: collectionName,
vectors_config: {
size: 768, // typical size for Llama2 embeddings
distance: "Cosine",
},
});
}
async addDocument(text, metadata = {}) {
if (!this.collection) {
await this.init();
}
const embedding = await this.getEmbedding(text);
await this.collection.add({
documents: [text],
embeddings: [embedding],
metadatas: [metadata],
});
}
async query(queryText, k = 5) {
if (!this.collection) {
await this.init();
}
const embedding = await this.getEmbedding(queryText);
const results = await this.collection.query({
queryEmbeddings: [embedding],
nResults: k,
});
return results;
}
async getEmbedding(text) {
const res = await openai.embeddings.create({
model: 'text-embedding-ada-002',
input: text,
});
return res.data[0].embedding;
}
}
module.exports = VectorStore;
return QdrantStore.fromExistingCollection(client, collectionName, {
embeddings,
});
}
+37 -24
View File
@@ -1,29 +1,42 @@
from langchain_qdrant import Qdrant
from langchain.schema import Document
import config
"""
Vector store implementation using Qdrant.
"""
class QdrantVectorStore:
from qdrant_client import QdrantClient
from qdrant_client.http import models as qdrant_models
from langchain.vectorstores import Qdrant
from config import QDRANT_HOST, QDRANT_PORT, QDRANT_COLLECTION_NAME
from embeddings import get_ollama_embeddings
def get_qdrant_client() -> QdrantClient:
"""
Wrapper around langchain_qdrant.Qdrant to provide a simple interface
for adding documents and retrieving a retriever.
Creates a Qdrant client connected to the local Qdrant instance.
"""
def __init__(self, embeddings, collection_name: str = None):
self.collection_name = collection_name or config.QDRANT_COLLECTION
self.qdrant = Qdrant(
url=f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}",
api_key=config.QDRANT_API_KEY,
collection_name=self.collection_name,
embeddings=embeddings,
return QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
def ensure_collection(client: QdrantClient, collection_name: str, vector_size: int = 768) -> None:
"""
Ensures that the specified collection exists in Qdrant.
If it does not exist, it will be created with the given vector size.
"""
if not client.has_collection(collection_name):
client.recreate_collection(
collection_name=collection_name,
vectors_config=qdrant_models.VectorParams(
size=vector_size,
distance="Cosine"
)
)
def add_documents(self, documents: list[Document]):
"""
Add a list of langchain Document objects to the Qdrant collection.
"""
self.qdrant.add_documents(documents)
def get_retriever(self):
"""
Return a retriever that can be used with LangChain chains.
"""
return self.qdrant.as_retriever()
def get_vector_store() -> Qdrant:
"""
Returns a Qdrant vector store instance ready for use with LangChain.
"""
client = get_qdrant_client()
ensure_collection(client, QDRANT_COLLECTION_NAME)
embeddings = get_ollama_embeddings()
return Qdrant(
client=client,
collection_name=QDRANT_COLLECTION_NAME,
embeddings=embeddings
)