diff --git a/solution.py b/solution.py index 4b764ca..94b5003 100644 --- a/solution.py +++ b/solution.py @@ -1,8 +1,145 @@ +import os import sys -from main import run_cli +from pathlib import Path +from typing import List + +import requests +from bs4 import BeautifulSoup +from langchain.embeddings.openai import OpenAIEmbeddings +from langchain.vectorstores import Chroma +from langchain.llms import OpenAI +from langchain.chains.question_answering import load_qa_chain +from rich.console import Console +from rich.prompt import Prompt + +# --------------------------------------------------------------------------- # +# Конфигурация и глобальные объекты +# --------------------------------------------------------------------------- # + +console = Console() + +CHROMA_DIR = Path("./chroma_db") +EMBEDDINGS_MODEL = "text-embedding-ada-002" +LLM_MODEL = "gpt-3.5-turbo" + +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +if not OPENAI_API_KEY: + console.print("[red]Ошибка: переменная окружения OPENAI_API_KEY не задана[/red]") + sys.exit(1) + +embeddings = OpenAIEmbeddings(model=EMBEDDINGS_MODEL, openai_api_key=OPENAI_API_KEY) +llm = OpenAI(model_name=LLM_MODEL, temperature=0.2, openai_api_key=OPENAI_API_KEY) + + +def init_vector_store() -> Chroma: + """ + Инициализирует или загружает существующую базу данных Chroma. + """ + return Chroma( + persist_directory=str(CHROMA_DIR), + embedding_function=embeddings, + ) + + +# --------------------------------------------------------------------------- # +# Веб‑парсинг и загрузка контента +# --------------------------------------------------------------------------- # + +def fetch_text_from_url(url: str) -> str: + """ + Загружает страницу по URL, удаляет скрипты/стили и возвращает чистый текст. + """ + try: + resp = requests.get(url, timeout=10) + resp.raise_for_status() + except Exception as e: + console.print(f"[yellow]Не удалось загрузить {url}: {e}[/yellow]") + return "" + + soup = BeautifulSoup(resp.text, "html.parser") + + # Удаляем скрипты и стили + for tag in soup(["script", "style"]): + tag.decompose() + + text = soup.get_text(separator="\n") + lines = [line.strip() for line in text.splitlines()] + cleaned = "\n".join([l for l in lines if l]) + return cleaned + + +def ingest_urls(urls: List[str], vector_store: Chroma) -> None: + """ + Загружает тексты по списку URL и сохраняет их в векторный хранилище. + """ + docs = [] + metadatas = [] + + for url in urls: + console.print(f"[cyan]Обрабатываем {url}[/cyan]") + content = fetch_text_from_url(url) + if not content: + continue + # Разбиваем на части по 1000 символов для удобства векторизации + chunk_size = 1000 + for i in range(0, len(content), chunk_size): + chunk = content[i : i + chunk_size] + docs.append(chunk) + metadatas.append({"source": url}) + + if docs: + vector_store.add_texts(docs, metadatas=metadatas) + console.print(f"[green]Добавлено {len(docs)} фрагментов в базу[/green]") + else: + console.print("[red]Нет данных для добавления[/red]") + + +# --------------------------------------------------------------------------- # +# RAG‑обработчик запросов +# --------------------------------------------------------------------------- # + +def answer_query(query: str, vector_store: Chroma) -> str: + """ + Получает ответ на вопрос, используя векторное хранение и LLM. + """ + # Поиск похожих документов + docs = vector_store.similarity_search(query, k=5) + if not docs: + return "Извините, я не нашёл подходящей информации." + + chain = load_qa_chain(llm, chain_type="stuff") + result = chain.run(input_documents=docs, question=query) + return result + + +# --------------------------------------------------------------------------- # +# Основная логика программы +# --------------------------------------------------------------------------- # + +def main() -> None: + vector_store = init_vector_store() + + # Если база пуста – запрашиваем URL для загрузки + if not list(vector_store.get_all_ids()): + console.print("[yellow]База данных пустая. Пожалуйста, укажите URL-адреса для индексации[/yellow]") + urls_input = Prompt.ask("Введите URL через запятую") + urls = [u.strip() for u in urls_input.split(",") if u.strip()] + ingest_urls(urls, vector_store) + vector_store.persist() + + console.print("[bold green]RAG‑агент готов к работе![/bold green]") + while True: + try: + query = Prompt.ask("\nВведите ваш вопрос (или 'выход' для завершения)") + if query.lower() in ("выход", "exit", "quit"): + console.print("[blue]До свидания![/blue]") + break + answer = answer_query(query, vector_store) + console.print(f"\n[bold]Ответ:[/bold]\n{answer}\n") + except KeyboardInterrupt: + console.print("\n[blue]Прервано пользователем. Завершение работы.[/blue]") + break + if __name__ == "__main__": - try: - run_cli() - except KeyboardInterrupt: - sys.exit(0) \ No newline at end of file + main() \ No newline at end of file