feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'
This commit is contained in:
+21
-23
@@ -1,28 +1,26 @@
|
||||
import { VectorStore } from "./vectorStore.js";
|
||||
import { webSearch } from "./webSearch.js";
|
||||
import { Ollama } from 'langchain/llms/ollama';
|
||||
import { RetrievalQAChain } from 'langchain/chains/retrieval_qa';
|
||||
import { BaseRetriever } from 'langchain/schema';
|
||||
import { webSearch } from './webSearch.js';
|
||||
|
||||
/**
|
||||
* A simple RAG agent that retrieves relevant documents from the vector store
|
||||
* and optionally performs a web search if no relevant documents are found.
|
||||
*/
|
||||
export class Agent {
|
||||
constructor(vectorStore) {
|
||||
this.vectorStore = vectorStore;
|
||||
}
|
||||
export function createAgent(vectorStore) {
|
||||
const llm = new Ollama({
|
||||
model: process.env.OLLAMA_MODEL || 'llama3',
|
||||
baseUrl: process.env.OLLAMA_HOST || 'http://localhost:11434',
|
||||
});
|
||||
|
||||
/**
|
||||
* Processes a user query and returns the best answer.
|
||||
* @param {string} query
|
||||
* @returns {Promise<string>}
|
||||
*/
|
||||
async answer(query) {
|
||||
const results = await this.vectorStore.query(query, 3);
|
||||
if (results.length > 0 && results[0].score < 0.5) {
|
||||
// Return the most relevant document text
|
||||
return results[0].text;
|
||||
class CombinedRetriever extends BaseRetriever {
|
||||
async getRelevantDocuments(query) {
|
||||
const chromaDocs = await vectorStore.similaritySearch(query, 3);
|
||||
const webDocs = await webSearch(query, 3);
|
||||
return [...chromaDocs, ...webDocs];
|
||||
}
|
||||
// Fallback to web search
|
||||
const html = await webSearch(query);
|
||||
return `No relevant local documents found. Here is the raw web search result:\n${html}`;
|
||||
}
|
||||
|
||||
const retriever = new CombinedRetriever();
|
||||
const chain = RetrievalQAChain.fromLLMAndRetriever(llm, retriever, {
|
||||
returnSourceDocuments: true,
|
||||
});
|
||||
|
||||
return chain;
|
||||
}
|
||||
+38
-17
@@ -1,26 +1,47 @@
|
||||
import { VectorStore } from "./vectorStore.js";
|
||||
import { Agent } from "./agent.js";
|
||||
import dotenv from 'dotenv';
|
||||
dotenv.config();
|
||||
|
||||
import { createVectorStore } from './vectorStore.js';
|
||||
import { createAgent } from './agent.js';
|
||||
import { Document } from 'langchain/document';
|
||||
|
||||
async function main() {
|
||||
const store = new VectorStore();
|
||||
await store.init();
|
||||
const vectorStore = await createVectorStore();
|
||||
|
||||
// Sample documents to index
|
||||
const docs = [
|
||||
{ id: "1", text: "ChromaDB is a fast, lightweight vector database." },
|
||||
{ id: "2", text: "It supports in-memory and persistent storage." },
|
||||
{ id: "3", text: "You can use it with various embedding models." },
|
||||
];
|
||||
if (process.env.ADD_SAMPLE_DOCS === 'true') {
|
||||
const sampleDocs = [
|
||||
new Document({
|
||||
pageContent:
|
||||
'LangChain is a framework for building applications powered by language models.',
|
||||
metadata: { source: 'LangChain Docs' },
|
||||
}),
|
||||
new Document({
|
||||
pageContent: 'ChromaDB is a vector database for storing embeddings.',
|
||||
metadata: { source: 'ChromaDB Docs' },
|
||||
}),
|
||||
new Document({
|
||||
pageContent: 'Ollama is a lightweight LLM server that can run locally.',
|
||||
metadata: { source: 'Ollama Docs' },
|
||||
}),
|
||||
];
|
||||
await vectorStore.addDocuments(sampleDocs);
|
||||
console.log('Sample documents added to ChromaDB.');
|
||||
}
|
||||
|
||||
await store.addDocuments(docs);
|
||||
const agent = createAgent(vectorStore);
|
||||
|
||||
const agent = new Agent(store);
|
||||
const query = process.argv[2];
|
||||
if (!query) {
|
||||
console.error('Please provide a query as a command line argument.');
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
const query = process.argv[2] || "What is ChromaDB?";
|
||||
console.log(`Query: ${query}`);
|
||||
const answer = await agent.answer(query);
|
||||
console.log("\nAnswer:");
|
||||
console.log(answer);
|
||||
const result = await agent.invoke({ input: query });
|
||||
console.log('Answer:', result.output);
|
||||
console.log(
|
||||
'Sources:',
|
||||
result.sourceDocuments.map((d) => d.metadata.source).join(', ')
|
||||
);
|
||||
}
|
||||
|
||||
main().catch((err) => {
|
||||
|
||||
+34
-56
@@ -1,69 +1,47 @@
|
||||
import { Client } from "@chromadb/chromadb";
|
||||
import { ChromaClient } from 'chromadb';
|
||||
import { OllamaEmbeddings } from 'langchain/embeddings/ollama';
|
||||
import { Document } from 'langchain/document';
|
||||
|
||||
/**
|
||||
* Simple embedding function that converts text into a fixed-length numeric vector.
|
||||
* This is a placeholder and should be replaced with a real embedding model for production use.
|
||||
*/
|
||||
function embed(text) {
|
||||
const vector = Array.from(text)
|
||||
.map((c) => c.charCodeAt(0))
|
||||
.slice(0, 10);
|
||||
while (vector.length < 10) {
|
||||
vector.push(0);
|
||||
}
|
||||
return vector;
|
||||
export async function createVectorStore() {
|
||||
const chroma = new ChromaClient({ path: process.env.CHROMA_URL });
|
||||
const collection = await chroma.getOrCreateCollection({
|
||||
name: process.env.CHROMA_COLLECTION,
|
||||
});
|
||||
const embeddings = new OllamaEmbeddings({
|
||||
model: process.env.OLLAMA_EMBEDDING_MODEL || 'nomic-embed-text',
|
||||
});
|
||||
return new VectorStore(collection, embeddings);
|
||||
}
|
||||
|
||||
export class VectorStore {
|
||||
constructor() {
|
||||
this.client = new Client();
|
||||
this.collection = null;
|
||||
class VectorStore {
|
||||
constructor(collection, embeddings) {
|
||||
this.collection = collection;
|
||||
this.embeddings = embeddings;
|
||||
}
|
||||
|
||||
async init() {
|
||||
this.collection = await this.client.getOrCreateCollection({
|
||||
name: "rag_collection",
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds an array of documents to the collection.
|
||||
* @param {Array<{id: string, text: string}>} docs
|
||||
*/
|
||||
async addDocuments(docs) {
|
||||
if (!this.collection) {
|
||||
throw new Error("VectorStore not initialized. Call init() first.");
|
||||
}
|
||||
const ids = docs.map((d) => d.id);
|
||||
const embeddings = docs.map((d) => embed(d.text));
|
||||
const documents = docs.map((d) => d.text);
|
||||
await this.collection.add({
|
||||
ids,
|
||||
const texts = docs.map((d) => d.pageContent);
|
||||
const embeddings = await this.embeddings.embedDocuments(texts);
|
||||
await this.collection.addDocuments({
|
||||
documents: docs,
|
||||
embeddings,
|
||||
documents,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Queries the collection for the most relevant documents.
|
||||
* @param {string} queryText
|
||||
* @param {number} nResults
|
||||
* @returns {Promise<Array<{id: string, text: string, score: number}>>}
|
||||
*/
|
||||
async query(queryText, nResults = 3) {
|
||||
if (!this.collection) {
|
||||
throw new Error("VectorStore not initialized. Call init() first.");
|
||||
}
|
||||
const queryEmbedding = embed(queryText);
|
||||
const results = await this.collection.query({
|
||||
queryEmbeddings: [queryEmbedding],
|
||||
nResults,
|
||||
async similaritySearch(query, k = 4) {
|
||||
const embedding = await this.embeddings.embedQuery(query);
|
||||
const results = await this.collection.getNearestNeighbors({
|
||||
queryEmbeddings: [embedding],
|
||||
n: k,
|
||||
});
|
||||
// results is an array of objects with ids, documents, and scores
|
||||
return results[0].ids.map((id, idx) => ({
|
||||
id,
|
||||
text: results[0].documents[idx],
|
||||
score: results[0].distances[idx],
|
||||
}));
|
||||
const ids = results.ids[0];
|
||||
const docs = await this.collection.getDocuments({ ids });
|
||||
return docs.map(
|
||||
(doc) =>
|
||||
new Document({
|
||||
pageContent: doc.document,
|
||||
metadata: doc.metadata,
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
+21
-14
@@ -1,17 +1,24 @@
|
||||
import fetch from "node-fetch";
|
||||
import fetch from 'node-fetch';
|
||||
import { Document } from 'langchain/document';
|
||||
|
||||
/**
|
||||
* Performs a simple web search using DuckDuckGo's HTML interface.
|
||||
* This is a lightweight example and does not use an official API.
|
||||
* @param {string} query
|
||||
* @returns {Promise<string>} The raw HTML of the search results page.
|
||||
*/
|
||||
export async function webSearch(query) {
|
||||
const url = `https://duckduckgo.com/html/?q=${encodeURIComponent(query)}`;
|
||||
const response = await fetch(url);
|
||||
if (!response.ok) {
|
||||
throw new Error(`Web search failed with status ${response.status}`);
|
||||
export async function webSearch(query, limit = 3) {
|
||||
const url = `https://api.duckduckgo.com/?q=${encodeURIComponent(
|
||||
query
|
||||
)}&format=json&pretty=1`;
|
||||
const res = await fetch(url);
|
||||
const data = await res.json();
|
||||
const topics = data.RelatedTopics || [];
|
||||
const docs = [];
|
||||
for (const topic of topics) {
|
||||
if (topic.Text) {
|
||||
docs.push(
|
||||
new Document({
|
||||
pageContent: topic.Text,
|
||||
metadata: { source: 'DuckDuckGo', url: topic.FirstURL },
|
||||
})
|
||||
);
|
||||
}
|
||||
if (docs.length >= limit) break;
|
||||
}
|
||||
const html = await response.text();
|
||||
return html;
|
||||
return docs.slice(0, limit);
|
||||
}
|
||||
Reference in New Issue
Block a user