feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'
This commit is contained in:
+32
-21
@@ -1,26 +1,37 @@
|
||||
import { Ollama } from 'langchain/llms/ollama';
|
||||
import { RetrievalQAChain } from 'langchain/chains/retrieval_qa';
|
||||
import { BaseRetriever } from 'langchain/schema';
|
||||
import { webSearch } from './webSearch.js';
|
||||
const { OpenAI } = require('openai');
|
||||
const dotenv = require('dotenv');
|
||||
dotenv.config();
|
||||
|
||||
export function createAgent(vectorStore) {
|
||||
const llm = new Ollama({
|
||||
model: process.env.OLLAMA_MODEL || 'llama3',
|
||||
baseUrl: process.env.OLLAMA_HOST || 'http://localhost:11434',
|
||||
});
|
||||
const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
|
||||
|
||||
class CombinedRetriever extends BaseRetriever {
|
||||
async getRelevantDocuments(query) {
|
||||
const chromaDocs = await vectorStore.similaritySearch(query, 3);
|
||||
const webDocs = await webSearch(query, 3);
|
||||
return [...chromaDocs, ...webDocs];
|
||||
}
|
||||
class Agent {
|
||||
constructor(vectorStore) {
|
||||
this.vectorStore = vectorStore;
|
||||
}
|
||||
|
||||
const retriever = new CombinedRetriever();
|
||||
const chain = RetrievalQAChain.fromLLMAndRetriever(llm, retriever, {
|
||||
returnSourceDocuments: true,
|
||||
});
|
||||
async ask(question) {
|
||||
const contextDocs = await this.vectorStore.query(question, 5);
|
||||
const context = contextDocs.join('\n\n');
|
||||
const prompt = `
|
||||
You are a helpful assistant. Use the following context to answer the question. If the context does not contain the answer, say you don't know.
|
||||
|
||||
return chain;
|
||||
}
|
||||
Context:
|
||||
${context}
|
||||
|
||||
Question:
|
||||
${question}
|
||||
Answer:
|
||||
`;
|
||||
|
||||
const completion = await openai.chat.completions.create({
|
||||
model: 'gpt-3.5-turbo',
|
||||
messages: [{ role: 'user', content: prompt }],
|
||||
temperature: 0.2,
|
||||
max_tokens: 300,
|
||||
});
|
||||
|
||||
return completion.choices[0].message.content.trim();
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = Agent;
|
||||
+42
-38
@@ -1,50 +1,54 @@
|
||||
import dotenv from 'dotenv';
|
||||
const readline = require('readline');
|
||||
const VectorStore = require('./vectorStore');
|
||||
const Agent = require('./agent');
|
||||
const { searchAndChunk } = require('./webSearch');
|
||||
const dotenv = require('dotenv');
|
||||
dotenv.config();
|
||||
|
||||
import { createVectorStore } from './vectorStore.js';
|
||||
import { createAgent } from './agent.js';
|
||||
import { Document } from 'langchain/document';
|
||||
|
||||
async function main() {
|
||||
const vectorStore = await createVectorStore();
|
||||
console.log('Initializing RAG agent...');
|
||||
const vectorStore = new VectorStore();
|
||||
|
||||
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.');
|
||||
}
|
||||
// Example URLs to index
|
||||
const urls = [
|
||||
'https://en.wikipedia.org/wiki/Artificial_intelligence',
|
||||
'https://en.wikipedia.org/wiki/ChromaDB',
|
||||
'https://en.wikipedia.org/wiki/OpenAI',
|
||||
];
|
||||
|
||||
const agent = createAgent(vectorStore);
|
||||
console.log('Fetching and indexing web pages...');
|
||||
const chunks = await searchAndChunk(urls);
|
||||
await vectorStore.addDocuments(chunks);
|
||||
console.log(`Indexed ${chunks.length} chunks.`);
|
||||
|
||||
const query = process.argv[2];
|
||||
if (!query) {
|
||||
console.error('Please provide a query as a command line argument.');
|
||||
process.exit(1);
|
||||
}
|
||||
const agent = new Agent(vectorStore);
|
||||
|
||||
const result = await agent.invoke({ input: query });
|
||||
console.log('Answer:', result.output);
|
||||
console.log(
|
||||
'Sources:',
|
||||
result.sourceDocuments.map((d) => d.metadata.source).join(', ')
|
||||
);
|
||||
const rl = readline.createInterface({
|
||||
input: process.stdin,
|
||||
output: process.stdout,
|
||||
});
|
||||
|
||||
const askQuestion = () => {
|
||||
rl.question('\nEnter your question (or type "exit" to quit): ', async (answer) => {
|
||||
if (answer.trim().toLowerCase() === 'exit') {
|
||||
rl.close();
|
||||
return;
|
||||
}
|
||||
console.log('\nGenerating answer...');
|
||||
try {
|
||||
const response = await agent.ask(answer);
|
||||
console.log(`\nAnswer:\n${response}`);
|
||||
} catch (err) {
|
||||
console.error('Error generating answer:', err);
|
||||
}
|
||||
askQuestion();
|
||||
});
|
||||
};
|
||||
|
||||
askQuestion();
|
||||
}
|
||||
|
||||
main().catch((err) => {
|
||||
console.error(err);
|
||||
console.error('Fatal error:', err);
|
||||
process.exit(1);
|
||||
});
|
||||
+41
-36
@@ -1,47 +1,52 @@
|
||||
import { ChromaClient } from 'chromadb';
|
||||
import { OllamaEmbeddings } from 'langchain/embeddings/ollama';
|
||||
import { Document } from 'langchain/document';
|
||||
const chromadb = require('chromadb');
|
||||
const { OpenAI } = require('openai');
|
||||
const dotenv = require('dotenv');
|
||||
dotenv.config();
|
||||
|
||||
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);
|
||||
}
|
||||
const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
|
||||
|
||||
class VectorStore {
|
||||
constructor(collection, embeddings) {
|
||||
this.collection = collection;
|
||||
this.embeddings = embeddings;
|
||||
constructor() {
|
||||
this.client = new chromadb.Client({ path: './chromadb' });
|
||||
this.collection = this.client.getCollection('rag_collection');
|
||||
}
|
||||
|
||||
async addDocuments(docs) {
|
||||
const texts = docs.map((d) => d.pageContent);
|
||||
const embeddings = await this.embeddings.embedDocuments(texts);
|
||||
await this.collection.addDocuments({
|
||||
documents: docs,
|
||||
async getEmbedding(text) {
|
||||
const response = await openai.embeddings.create({
|
||||
model: 'text-embedding-ada-002',
|
||||
input: text,
|
||||
});
|
||||
return response.data[0].embedding;
|
||||
}
|
||||
|
||||
async addDocuments(chunks) {
|
||||
const documents = [];
|
||||
const embeddings = [];
|
||||
const ids = [];
|
||||
|
||||
for (const chunk of chunks) {
|
||||
const embedding = await this.getEmbedding(chunk);
|
||||
documents.push(chunk);
|
||||
embeddings.push(embedding);
|
||||
ids.push(`${Date.now()}-${Math.random()}`);
|
||||
}
|
||||
|
||||
await this.collection.add({
|
||||
documents,
|
||||
embeddings,
|
||||
ids,
|
||||
});
|
||||
}
|
||||
|
||||
async similaritySearch(query, k = 4) {
|
||||
const embedding = await this.embeddings.embedQuery(query);
|
||||
const results = await this.collection.getNearestNeighbors({
|
||||
queryEmbeddings: [embedding],
|
||||
n: k,
|
||||
async query(queryText, k = 5) {
|
||||
const queryEmbedding = await this.getEmbedding(queryText);
|
||||
const results = await this.collection.query({
|
||||
queryEmbeddings: [queryEmbedding],
|
||||
nResults: k,
|
||||
});
|
||||
const ids = results.ids[0];
|
||||
const docs = await this.collection.getDocuments({ ids });
|
||||
return docs.map(
|
||||
(doc) =>
|
||||
new Document({
|
||||
pageContent: doc.document,
|
||||
metadata: doc.metadata,
|
||||
})
|
||||
);
|
||||
|
||||
return results.documents[0];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = VectorStore;
|
||||
+39
-21
@@ -1,24 +1,42 @@
|
||||
import fetch from 'node-fetch';
|
||||
import { Document } from 'langchain/document';
|
||||
const fetch = require('node-fetch');
|
||||
|
||||
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 },
|
||||
})
|
||||
);
|
||||
async function fetchPage(url) {
|
||||
try {
|
||||
const res = await fetch(url);
|
||||
if (!res.ok) {
|
||||
console.warn(`Failed to fetch ${url}: ${res.statusText}`);
|
||||
return '';
|
||||
}
|
||||
if (docs.length >= limit) break;
|
||||
const html = await res.text();
|
||||
// Strip HTML tags
|
||||
const text = html.replace(/<[^>]*>/g, ' ');
|
||||
// Collapse whitespace
|
||||
const cleaned = text.replace(/\s+/g, ' ').trim();
|
||||
return cleaned;
|
||||
} catch (err) {
|
||||
console.error(`Error fetching ${url}:`, err);
|
||||
return '';
|
||||
}
|
||||
return docs.slice(0, limit);
|
||||
}
|
||||
}
|
||||
|
||||
function chunkText(text, size = 500) {
|
||||
const chunks = [];
|
||||
for (let i = 0; i < text.length; i += size) {
|
||||
chunks.push(text.slice(i, i + size));
|
||||
}
|
||||
return chunks;
|
||||
}
|
||||
|
||||
async function searchAndChunk(urls) {
|
||||
const allChunks = [];
|
||||
for (const url of urls) {
|
||||
const pageText = await fetchPage(url);
|
||||
if (pageText) {
|
||||
const chunks = chunkText(pageText);
|
||||
allChunks.push(...chunks);
|
||||
}
|
||||
}
|
||||
return allChunks;
|
||||
}
|
||||
|
||||
module.exports = { searchAndChunk };
|
||||
Reference in New Issue
Block a user