RAG на своих документах

import chromadb
from chromadb.utils import embedding_functions
import requests
import os

class LocalRAG:
    """RAG-система на локальных моделях через Ollama."""
    
    def __init__(self, collection_name: str = "my_docs"):
        self.client = chromadb.Client()
        
        # Используем Ollama для эмбеддингов
        self.ollama_ef = embedding_functions.OllamaEmbeddingFunction(
            url="http://localhost:11434/api/embeddings",
            model_name="nomic-embed-text"
        )
        
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=self.ollama_ef
        )
    
    def add_document(self, text: str, metadata: dict = None, doc_id: str = None):
        """Добавить документ в базу."""
        self.collection.add(
            documents=[text],
            metadatas=[metadata or {}],
            ids=[doc_id or str(hash(text))]
        )
        print(f"  ✓ Добавлен документ ({len(text)} символов)")
    
    def add_directory(self, path: str):
        """Добавить все .txt и .md файлы из директории."""
        for filename in os.listdir(path):
            if filename.endswith((".txt", ".md", ".py", ".json")):
                filepath = os.path.join(path, filename)
                with open(filepath, "r") as f:
                    content = f.read()
                self.add_document(
                    content,
                    metadata={"source": filename},
                    doc_id=filename
                )
    
    def query(self, question: str, n_results: int = 3) -> str:
        """Задать вопрос по документам."""
        
        # 1. Ищем похожие куски
        results = self.collection.query(
            query_texts=[question],
            n_results=n_results
        )
        
        if not results["documents"][0]:
            return "Нет релевантных документов в базе"
        
        # 2. Собираем контекст
        context = "\n\n".join([
            f"[{meta.get('source', 'unknown')}]: {doc}"
            for doc, meta in zip(results["documents"][0], results["metadatas"][0])
        ])
        
        # 3. Отправляем модели с контекстом
        response = requests.post("http://localhost:11434/api/chat", json={
            "model": "qwen3.5:4b",
            "messages": [
                {"role": "system", "content": (
                    "Ты — Q&A ассистент. Отвечай на вопросы ТОЛЬКО на основе "
                    "предоставленных документов. Если в документах нет ответа — "
                    "так и скажи. Не придумывай факты."
                )},
                {"role": "user", "content": (
                    f"Контекст из документов:\n{context}\n\n"
                    f"Вопрос: {question}"
                )}
            ],
            "stream": False
        })
        
        answer = response.json()["message"]["content"]
        
        # 4. Показываем источники
        sources = [meta.get('source', 'unknown') 
                   for meta in results["metadatas"][0]]
        
        return f"{answer}\n\n📚 Источники: {', '.join(sources)}"


# === Пример использования ===
if __name__ == "__main__":
    rag = LocalRAG("handbook")
    
    # Индексируем документы
    print("Индексация документов...")
    rag.add_document(
        "Ollama — программа для запуска LLM локально. "
        "Устанавливается через brew install ollama.",
        {"source": "getting-started.md", "topic": "installation"}
    )
    rag.add_document(
        "Для RAG используйте Open WebUI или GPT4All. "
        "Эмбеддинги: nomic-embed-text.",
        {"source": "rag.md", "topic": "rag"}
    )
    
    # Задаём вопрос
    print("\nВопрос: Как установить Ollama?")
    print(rag.query("Как установить Ollama?"))