feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'
This commit is contained in:
+69
-32
@@ -1,42 +1,79 @@
|
||||
import argparse
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
from src.vectorstore import create_vectorstore, load_documents
|
||||
from src.agent import create_agent
|
||||
import sys
|
||||
from typing import List
|
||||
|
||||
import openai
|
||||
|
||||
from vector_store import ingest_documents, get_relevant_chunks
|
||||
from web_search import search_web
|
||||
|
||||
# Load OpenAI API key
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
if not OPENAI_API_KEY:
|
||||
print("Error: OPENAI_API_KEY environment variable not set.")
|
||||
sys.exit(1)
|
||||
openai.api_key = OPENAI_API_KEY
|
||||
|
||||
def generate_answer(context: str, question: str) -> str:
|
||||
"""
|
||||
Generate an answer using OpenAI ChatCompletion with the provided context.
|
||||
"""
|
||||
system_prompt = "You are a helpful assistant. Use the provided context to answer the question."
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
|
||||
]
|
||||
try:
|
||||
response = openai.ChatCompletion.create(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
max_tokens=512
|
||||
)
|
||||
return response["choices"][0]["message"]["content"].strip()
|
||||
except Exception as e:
|
||||
print(f"OpenAI request failed: {e}")
|
||||
return ""
|
||||
|
||||
def ingest_mode(file_paths: List[str]) -> None:
|
||||
ingest_documents(file_paths)
|
||||
|
||||
def query_mode(question: str) -> None:
|
||||
# Retrieve relevant chunks from local vector store
|
||||
local_chunks = get_relevant_chunks(question, k=5)
|
||||
local_context = "\n\n".join([chunk for _, chunk in local_chunks])
|
||||
|
||||
# Perform web search for up-to-date info
|
||||
web_snippets = search_web(question, num_results=3)
|
||||
web_context = "\n\n".join(web_snippets)
|
||||
|
||||
# Combine contexts
|
||||
combined_context = f"Local documents:\n{local_context}\n\nWeb results:\n{web_context}"
|
||||
|
||||
# Generate answer
|
||||
answer = generate_answer(combined_context, question)
|
||||
print("\nAnswer:\n")
|
||||
print(answer)
|
||||
|
||||
def main():
|
||||
load_dotenv()
|
||||
# Initialize or load the vector store
|
||||
vectorstore = create_vectorstore(persist_directory="./chroma_db")
|
||||
parser = argparse.ArgumentParser(description="RAG Agent with ChromaDB and Web Search")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
# Load documents into the vector store if not already loaded
|
||||
# (Chroma will load existing data automatically)
|
||||
load_documents("./documents", vectorstore)
|
||||
ingest_parser = subparsers.add_parser("ingest", help="Ingest documents into the vector store")
|
||||
ingest_parser.add_argument("files", nargs="+", help="Paths to text files to ingest")
|
||||
|
||||
# Create the agent
|
||||
agent = create_agent(vectorstore)
|
||||
query_parser = subparsers.add_parser("query", help="Ask a question to the RAG agent")
|
||||
query_parser.add_argument("question", help="The question to ask")
|
||||
|
||||
print("\n=== RAG Agent with ChromaDB and Tavily ===")
|
||||
print("Type your question (or 'exit' to quit):")
|
||||
args = parser.parse_args()
|
||||
|
||||
while True:
|
||||
try:
|
||||
user_input = input("\nYou: ").strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\nGoodbye!")
|
||||
break
|
||||
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
if user_input.lower() in {"exit", "quit"}:
|
||||
print("Goodbye!")
|
||||
break
|
||||
|
||||
try:
|
||||
response = agent.run(user_input)
|
||||
print(f"\nAgent: {response}")
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
if args.command == "ingest":
|
||||
ingest_mode(args.files)
|
||||
elif args.command == "query":
|
||||
query_mode(args.question)
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,85 @@
|
||||
import os
|
||||
from typing import List, Tuple
|
||||
|
||||
import chromadb
|
||||
from chromadb import PersistentClient
|
||||
from chromadb.config import Settings
|
||||
import openai
|
||||
|
||||
# Load OpenAI API key from environment
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
if not OPENAI_API_KEY:
|
||||
raise RuntimeError("OPENAI_API_KEY environment variable not set.")
|
||||
openai.api_key = OPENAI_API_KEY
|
||||
|
||||
# ChromaDB persistent client settings
|
||||
CHROMA_DB_PATH = os.getenv("CHROMA_DB_PATH", "./chromadb")
|
||||
CHROMA_COLLECTION_NAME = os.getenv("CHROMA_COLLECTION_NAME", "rag_collection")
|
||||
|
||||
# Initialize Chroma client
|
||||
client = PersistentClient(path=CHROMA_DB_PATH, settings=Settings(chroma_api_impl="chromadb.api.fastapi.FastAPI"))
|
||||
collection = client.get_or_create_collection(name=CHROMA_COLLECTION_NAME)
|
||||
|
||||
|
||||
def _split_text(text: str, chunk_size: int = 500, overlap: int = 50) -> List[str]:
|
||||
"""
|
||||
Split text into chunks of approximately chunk_size characters with overlap.
|
||||
"""
|
||||
chunks = []
|
||||
start = 0
|
||||
text_length = len(text)
|
||||
while start < text_length:
|
||||
end = min(start + chunk_size, text_length)
|
||||
chunk = text[start:end]
|
||||
chunks.append(chunk)
|
||||
start += chunk_size - overlap
|
||||
return chunks
|
||||
|
||||
|
||||
def _embed_text(text: str) -> List[float]:
|
||||
"""
|
||||
Generate embedding for a single text string using OpenAI embeddings.
|
||||
"""
|
||||
response = openai.Embedding.create(
|
||||
model="text-embedding-ada-002",
|
||||
input=text
|
||||
)
|
||||
return response["data"][0]["embedding"]
|
||||
|
||||
|
||||
def ingest_documents(file_paths: List[str]) -> None:
|
||||
"""
|
||||
Ingest a list of file paths into the Chroma collection.
|
||||
Each file is read, split into chunks, embedded, and stored.
|
||||
"""
|
||||
for file_path in file_paths:
|
||||
if not os.path.isfile(file_path):
|
||||
print(f"Skipping non-existent file: {file_path}")
|
||||
continue
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
chunks = _split_text(content)
|
||||
embeddings = [_embed_text(chunk) for chunk in chunks]
|
||||
ids = [f"{os.path.basename(file_path)}_{i}" for i in range(len(chunks))]
|
||||
collection.add(
|
||||
ids=ids,
|
||||
documents=chunks,
|
||||
embeddings=embeddings
|
||||
)
|
||||
print(f"Ingested {len(chunks)} chunks from {file_path}.")
|
||||
|
||||
|
||||
def get_relevant_chunks(query: str, k: int = 5) -> List[Tuple[str, str]]:
|
||||
"""
|
||||
Retrieve top-k relevant chunks for a query.
|
||||
Returns a list of tuples (chunk_id, chunk_text).
|
||||
"""
|
||||
query_embedding = _embed_text(query)
|
||||
results = collection.query(
|
||||
query_embeddings=[query_embedding],
|
||||
n_results=k,
|
||||
include=["documents", "ids"]
|
||||
)
|
||||
ids = results["ids"][0]
|
||||
docs = results["documents"][0]
|
||||
return list(zip(ids, docs))
|
||||
@@ -0,0 +1,65 @@
|
||||
import os
|
||||
import re
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
from typing import List
|
||||
|
||||
# DuckDuckGo search URL
|
||||
DDG_SEARCH_URL = "https://duckduckgo.com/html/"
|
||||
|
||||
def _extract_text_from_html(html: str) -> str:
|
||||
"""
|
||||
Extract visible text from HTML, removing scripts and styles.
|
||||
"""
|
||||
soup = BeautifulSoup(html, "html.parser")
|
||||
for script in soup(["script", "style"]):
|
||||
script.decompose()
|
||||
text = soup.get_text(separator="\n")
|
||||
lines = (line.strip() for line in text.splitlines())
|
||||
chunks = [phrase.strip() for phrase in lines if phrase.strip()]
|
||||
return "\n".join(chunks)
|
||||
|
||||
def search_web(query: str, num_results: int = 3) -> List[str]:
|
||||
"""
|
||||
Perform a web search using DuckDuckGo and return the top num_results snippets.
|
||||
"""
|
||||
params = {
|
||||
"q": query,
|
||||
"s": "0",
|
||||
"dc": "0",
|
||||
"kl": "us-en",
|
||||
"kp": "-2",
|
||||
"kp": "-2",
|
||||
"kp": "-2",
|
||||
"kp": "-2",
|
||||
}
|
||||
headers = {
|
||||
"User-Agent": "Mozilla/5.0 (compatible; RAG-Agent/1.0; +https://example.com/bot)"
|
||||
}
|
||||
try:
|
||||
response = requests.get(DDG_SEARCH_URL, params=params, headers=headers, timeout=10)
|
||||
response.raise_for_status()
|
||||
except requests.RequestException as e:
|
||||
print(f"Web search request failed: {e}")
|
||||
return []
|
||||
|
||||
soup = BeautifulSoup(response.text, "html.parser")
|
||||
results = []
|
||||
for a in soup.select("a.result__a"):
|
||||
href = a.get("href")
|
||||
if href:
|
||||
results.append(href)
|
||||
if len(results) >= num_results:
|
||||
break
|
||||
|
||||
snippets = []
|
||||
for url in results:
|
||||
try:
|
||||
page_resp = requests.get(url, headers=headers, timeout=10)
|
||||
page_resp.raise_for_status()
|
||||
snippet = _extract_text_from_html(page_resp.text)[:500] # limit snippet size
|
||||
snippets.append(snippet)
|
||||
except requests.RequestException:
|
||||
continue
|
||||
|
||||
return snippets
|
||||
Reference in New Issue
Block a user