Правки solution.py: Создать модуль vectorstore.py, tools.py, agent.py, init_know
This commit is contained in:
+5
-142
@@ -1,145 +1,8 @@
|
|||||||
import os
|
|
||||||
import sys
|
import sys
|
||||||
from pathlib import Path
|
from main import run_cli
|
||||||
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__":
|
if __name__ == "__main__":
|
||||||
main()
|
try:
|
||||||
|
run_cli()
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
sys.exit(0)
|
||||||
Reference in New Issue
Block a user