engineering/spider.unturf.com
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add chromadb so we can do RAG

Browse source 
modified:   .gitignore
	modified:   duck_duck_go_hermes_unturf.py
This commit is contained in:
Russell Ballestrini 2025-04-26 00:22:43 -04:00
parent dff8089f7c
commit 95c94db6f3
2 changed files with 111 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

2
.gitignore vendored
View file

@ -1 +1,3 @@
data/
chroma_db/
crawler.log

174
duck_duck_go_hermes_unturf.py
View file

@ -8,7 +8,9 @@ from datetime import datetime, timezone
from concurrent.futures import ThreadPoolExecutor, as_completed
import urllib.robotparser
import threading
import logging
import chromadb
from sentence_transformers import SentenceTransformer
from sqlalchemy import (
create_engine,
Column,
@ -22,6 +24,14 @@ from sqlalchemy.orm import declarative_base, sessionmaker
from openai import OpenAI
from bs4 import BeautifulSoup
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
handlers=[logging.StreamHandler(), logging.FileHandler("crawler.log")],
)
logger = logging.getLogger(__name__)
# Default Hermes endpoints
DEFAULT_HERMES_ENDPOINTS = [
"https://hermes.ai.unturf.com/v1", # 80k context window
@ -32,8 +42,8 @@ MAX_HTML_INPUT_CHARS = 50000
# Default crawl delay if robots.txt is missing or doesn't specify
DEFAULT_CRAWL_DELAY = 2.0 # seconds
# Token limits for Hermes endpoints
HERMES_TOKEN_LIMIT = 80000 # hermes.ai.unturf.com
HERMES2_TOKEN_LIMIT = 48000 # hermes2.ai.unturf.com
HERMES_TOKEN_LIMIT = 80000
HERMES2_TOKEN_LIMIT = 48000
HERMES_INPUT_TOKENS = 72000
HERMES_COMPLETION_TOKENS = 8000
HERMES2_INPUT_TOKENS = 40000
@ -83,6 +93,7 @@ class SQLAlchemyDuckDuckGoCrawler:
self.session.headers.update({"User-Agent": self.user_agent})
# Initialize DB
os.makedirs(os.path.dirname(db_path), exist_ok=True)
self.engine = create_engine(
f"sqlite:///{self.db_path}",
echo=False,
@ -94,6 +105,11 @@ class SQLAlchemyDuckDuckGoCrawler:
bind=self.engine, autoflush=False, autocommit=False
)
# Initialize ChromaDB
self.chroma_client = chromadb.PersistentClient(path="./chroma_db")
self.collection = self.chroma_client.get_or_create_collection(name="articles")
self.embedder = SentenceTransformer("all-MiniLM-L6-v2")
# Robots.txt parser and domain tracking
self.robot_parsers = {}
self.domain_last_fetched = {}
@ -108,10 +124,10 @@ class SQLAlchemyDuckDuckGoCrawler:
def _fetch_robots_txt(self, domain):
"""Fetch and parse robots.txt for a domain."""
if domain in self.robot_parsers:
print(f"Using cached robots.txt for {domain}")
logger.info(f"Using cached robots.txt for {domain}")
return self.robot_parsers[domain]
print(f"Fetching robots.txt for {domain}")
logger.info(f"Fetching robots.txt for {domain}")
robots_url = f"https://{domain}/robots.txt"
parser = urllib.robotparser.RobotFileParser()
parser.set_url(robots_url)
@ -126,7 +142,7 @@ class SQLAlchemyDuckDuckGoCrawler:
delay if delay is not None else DEFAULT_CRAWL_DELAY
)
except Exception as e:
print(f"Could not fetch robots.txt for {domain}: {e}")
logger.error(f"Could not fetch robots.txt for {domain}: {e}")
self.robot_parsers[domain] = None
self.robots_txt_content[domain] = None
self.domain_crawl_delays[domain] = DEFAULT_CRAWL_DELAY
@ -141,10 +157,10 @@ class SQLAlchemyDuckDuckGoCrawler:
can_fetch = parser.can_fetch(self.user_agent, url)
if not can_fetch:
print(f"Blocked by robots.txt: {url}")
logger.warning(f"Blocked by robots.txt: {url}")
robots_content = self.robots_txt_content.get(domain, "")
if robots_content:
print(f"Relevant robots.txt rules for {domain}:")
logger.info(f"Relevant robots.txt rules for {domain}:")
current_user_agent = None
relevant_rules = []
for line in robots_content.splitlines():
@ -166,13 +182,13 @@ class SQLAlchemyDuckDuckGoCrawler:
f" User-agent: {current_user_agent}\n Disallow: {rule}"
)
if relevant_rules:
print("\n".join(relevant_rules))
logger.info("\n".join(relevant_rules))
else:
print(
logger.info(
" No specific Disallow rules found; may be blocked by a broad rule."
)
else:
print(" No robots.txt content available to display rules.")
logger.info(" No robots.txt content available to display rules.")
return can_fetch
def _enforce_crawl_delay(self, domain):
@ -182,7 +198,7 @@ class SQLAlchemyDuckDuckGoCrawler:
elapsed = time.time() - last_fetched
if elapsed < delay:
sleep_time = delay - elapsed
print(
logger.info(
f"Delaying crawl for {domain} by {sleep_time:.2f} seconds due to crawl delay"
)
time.sleep(sleep_time)
@ -213,7 +229,7 @@ class SQLAlchemyDuckDuckGoCrawler:
target_url = "https://" + target_url
results.append((title, target_url))
else:
print(f"Skipping invalid redirect URL: {href}")
logger.warning(f"Skipping invalid redirect URL: {href}")
else:
if href and not href.startswith(("javascript:", "#")):
if href.startswith("//"):
@ -222,13 +238,13 @@ class SQLAlchemyDuckDuckGoCrawler:
href = "https://" + href
results.append((title, href))
else:
print(f"Skipping invalid URL: {href}")
logger.warning(f"Skipping invalid URL: {href}")
return results
def fetch_webpage(self, url):
"""Fetch a webpage, validating the URL first."""
if not url.startswith(("http://", "https://")):
print(f"Invalid URL scheme: {url}")
logger.error(f"Invalid URL scheme: {url}")
return None
try:
if not self._can_fetch(url):
@ -239,7 +255,7 @@ class SQLAlchemyDuckDuckGoCrawler:
r.raise_for_status()
return r.text
except requests.exceptions.RequestException as e:
print(f"Error fetching {url}: {e}")
logger.error(f"Error fetching {url}: {e}")
return None
def _fanout_call(self, messages, max_tokens, prefer_hermes=False):
@ -254,7 +270,7 @@ class SQLAlchemyDuckDuckGoCrawler:
adjusted_max_tokens = max(
1000, HERMES2_TOKEN_LIMIT - input_tokens
)
print(
logger.info(
f"Adjusted max_tokens to {adjusted_max_tokens} for {endpoint} due to context limit"
)
response = client.chat.completions.create(
@ -265,7 +281,7 @@ class SQLAlchemyDuckDuckGoCrawler:
)
return response.choices[0].message.content, endpoint
except Exception as e:
print(f"Error calling endpoint {endpoint}: {e}")
logger.error(f"Error calling endpoint {endpoint}: {e}")
return None, endpoint
if prefer_hermes:
@ -279,12 +295,12 @@ class SQLAlchemyDuckDuckGoCrawler:
None,
)
if hermes_client:
print(f"Attempting {hermes_endpoint} for final answer")
logger.info(f"Attempting {hermes_endpoint} for final answer")
result, used_endpoint = call_client(hermes_client, hermes_endpoint)
if result is not None:
print(f"Success using {used_endpoint}")
logger.info(f"Success using {used_endpoint}")
return result
print(f"Fallback: {hermes_endpoint} failed, trying other endpoints")
logger.warning(f"Fallback: {hermes_endpoint} failed, trying other endpoints")
other_clients = [
(c, e)
for c, e in zip(self.clients, self.hermes_endpoints)
@ -297,10 +313,10 @@ class SQLAlchemyDuckDuckGoCrawler:
for future in as_completed(futures):
result, used_endpoint = future.result()
if result is not None:
print(f"Success using {used_endpoint}")
logger.info(f"Success using {used_endpoint}")
return result
else:
print(
logger.warning(
f"Warning: {hermes_endpoint} not in endpoints, using parallel execution"
)
else:
@ -312,7 +328,7 @@ class SQLAlchemyDuckDuckGoCrawler:
for future in as_completed(futures):
result, used_endpoint = future.result()
if result is not None:
print(f"Success using {used_endpoint}")
logger.info(f"Success using {used_endpoint}")
return result
raise RuntimeError("All Hermes endpoints failed.")
@ -351,8 +367,26 @@ class SQLAlchemyDuckDuckGoCrawler:
]
return self._fanout_call(messages, max_tokens=HERMES_COMPLETION_TOKENS)
def _split_into_chunks(self, text, max_words=500):
"""Split text into chunks of approximately max_words."""
words = text.split()
chunks = []
current_chunk = []
current_word_count = 0
for word in words:
current_chunk.append(word)
current_word_count += 1
if current_word_count >= max_words:
chunks.append(" ".join(current_chunk))
current_chunk = []
current_word_count = 0
if current_chunk:
chunks.append(" ".join(current_chunk))
return chunks
def cache_article(self, url, title, html, extracted, summary):
"""Cache article in database, using extracted_content as summary."""
"""Cache article in database and store chunks in ChromaDB."""
with self.SessionLocal() as db:
art = Article(
url=url,
@ -367,17 +401,30 @@ class SQLAlchemyDuckDuckGoCrawler:
db.commit()
except Exception as e:
db.rollback()
print(f"Error caching article {url}: {e}")
logger.error(f"Error caching article {url}: {e}")
raise
# Store chunks in ChromaDB
chunks = self._split_into_chunks(extracted)
embeddings = self.embedder.encode(chunks, show_progress_bar=False)
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
chunk_id = f"{url}_chunk_{i}"
self.collection.upsert(
ids=[chunk_id],
embeddings=[embedding.tolist()],
documents=[chunk],
metadatas=[{"url": url, "title": title, "chunk_index": i}],
)
logger.info(f"Stored {len(chunks)} chunks for {url} in ChromaDB")
def process_url(self, title_url):
"""Process a URL, extracting content but not summarizing."""
title, url = title_url
with self.SessionLocal() as db:
if db.query(Article).filter_by(url=url).first():
print(f"Already cached: {url}")
logger.info(f"Already cached: {url}")
return
print(f"Fetching: {url}")
logger.info(f"Fetching: {url}")
html = self.fetch_webpage(url)
if not html:
return
@ -386,52 +433,44 @@ class SQLAlchemyDuckDuckGoCrawler:
self.cache_article(url, title, html, extracted, extracted)
def aggregate_and_answer(self, query):
"""Aggregate article content and generate an answer, summarizing if needed."""
with self.SessionLocal() as db:
articles = db.query(Article).all()
content_list = [
(art.summary, art.extracted_content) for art in articles
] # summary is extracted_content
"""Aggregate relevant article chunks from ChromaDB and generate an answer."""
# Embed the query
query_embedding = self.embedder.encode([query], show_progress_bar=False)[0]
# Query ChromaDB for relevant chunks
results = self.collection.query(
query_embeddings=[query_embedding.tolist()],
n_results=50, # Retrieve more to filter by token count
)
# Collect chunks up to token limit
def estimate_tokens(text):
return len(text) // 4 + 1
# Try with full extracted content
combined = "\n\n".join(content for content, _ in content_list)
prompt_template = f"Based on these article contents:\n{{}}\n\nProvide a comprehensive answer to: {query}"
combined_content = []
total_tokens = 0
prompt_template = (
f"Based on these article excerpts:\n{{}}\n\nProvide a comprehensive answer to: {query}"
)
template_tokens = estimate_tokens(prompt_template.format(""))
combined_tokens = estimate_tokens(combined)
if combined_tokens + template_tokens <= HERMES_INPUT_TOKENS:
print(f"Using full content ({combined_tokens} tokens)")
messages = [{"role": "user", "content": prompt_template.format(combined)}]
return self._fanout_call(
messages, max_tokens=HERMES_COMPLETION_TOKENS, prefer_hermes=True
for doc, metadata in zip(results["documents"][0], results["metadatas"][0]):
chunk = doc
chunk_tokens = estimate_tokens(chunk)
if total_tokens + chunk_tokens + template_tokens <= HERMES_INPUT_TOKENS:
combined_content.append(
f"From {metadata['url']} (Title: {metadata['title']}):\n{chunk}"
)
total_tokens += chunk_tokens
else:
break
# Summarize if token limit exceeded
print(
f"Warning: Input exceeds token limit ({combined_tokens + template_tokens} tokens). Summarizing content."
)
summaries = []
for _, extracted_content in content_list:
summary = self.summarize_with_hermes(extracted_content)
summaries.append(summary)
combined = "\n\n".join(summaries)
combined_tokens = estimate_tokens(combined)
if combined_tokens + template_tokens > HERMES_INPUT_TOKENS:
print(
f"Warning: Summarized input still exceeds token limit ({combined_tokens + template_tokens} tokens). Truncating."
)
max_chars = (HERMES_INPUT_TOKENS - template_tokens) * 4
combined = combined[:max_chars]
last_newline = combined.rfind("\n\n", 0, max_chars)
if last_newline != -1:
combined = combined[:last_newline]
combined_tokens = estimate_tokens(combined)
print(f"Truncated to {combined_tokens} tokens.")
if not combined_content:
logger.warning("No relevant chunks found within token limit.")
return "No relevant information found to answer the query."
combined = "\n\n".join(combined_content)
logger.info(f"Using {total_tokens} tokens from {len(combined_content)} chunks")
messages = [{"role": "user", "content": prompt_template.format(combined)}]
return self._fanout_call(
messages, max_tokens=HERMES_COMPLETION_TOKENS, prefer_hermes=True
@ -443,8 +482,9 @@ class SQLAlchemyDuckDuckGoCrawler:
futures = [executor.submit(self.process_url, hit) for hit in hits]
for future in as_completed(futures):
future.result()
print("Generating comprehensive answer...")
print(self.aggregate_and_answer(query))
logger.info("Generating comprehensive answer...")
result = self.aggregate_and_answer(query)
print(result)
if __name__ == "__main__":