modified: duck_duck_go_hermes_unturf.py

duck_duck_go_hermes_unturf.py

@@ -7,10 +7,7 @@ from urllib.parse import quote_plus, urlparse, parse_qs, unquote
 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,
@@ -23,6 +20,9 @@ from sqlalchemy import (
 from sqlalchemy.orm import declarative_base, sessionmaker
 from openai import OpenAI
 from bs4 import BeautifulSoup
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
 
 # Configure logging
 logging.basicConfig(
@@ -34,20 +34,25 @@ logger = logging.getLogger(__name__)
 
 # Default Hermes endpoints
 DEFAULT_HERMES_ENDPOINTS = [
-    "https://hermes.ai.unturf.com/v1",  # 80k context window
+    "https://hermes.ai.unturf.com/v1",
-    "https://hermes2.ai.unturf.com/v1",  # Smaller context window
+    "https://hermes2.ai.unturf.com/v1",
 ]
-# Max characters to send to extraction to avoid context overflow
+# Max characters to send to extraction
 MAX_HTML_INPUT_CHARS = 50000
-# Default crawl delay if robots.txt is missing or doesn't specify
+# Default crawl delay
-DEFAULT_CRAWL_DELAY = 2.0  # seconds
+DEFAULT_CRAWL_DELAY = 2.0
-# Token limits for Hermes endpoints
+# Token limits for Hermes
 HERMES_TOKEN_LIMIT = 80000
-HERMES2_TOKEN_LIMIT = 48000
 HERMES_INPUT_TOKENS = 72000
 HERMES_COMPLETION_TOKENS = 8000
+HERMES2_TOKEN_LIMIT = 48000
 HERMES2_INPUT_TOKENS = 40000
 HERMES2_COMPLETION_TOKENS = 8000
+# Max words for chunking (when needed)
+MAX_WORDS_LONG = 10000
+MAX_WORDS_SHORT = 5000
+# Context window limit (words, ~4 chars per token)
+CONTEXT_WINDOW_WORDS = 20000  # ~80K tokens
 
 Base = declarative_base()
 
@@ -59,7 +64,7 @@ class Article(Base):
     title = Column(String, nullable=False)
     raw_html = Column(Text, nullable=False)
     extracted_content = Column(Text, nullable=False)
-    summary = Column(Text, nullable=False)  # Stores extracted_content unless summarized
+    summary = Column(Text, nullable=False)
     fetched_at = Column(
         DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
     )
@@ -75,7 +80,6 @@ class SQLAlchemyDuckDuckGoCrawler:
         hermes_endpoints=None,
         user_agent_append="",
     ):
-        # Setup Hermes clients
         self.hermes_endpoints = hermes_endpoints or DEFAULT_HERMES_ENDPOINTS
         self.clients = [
             OpenAI(base_url=ep, api_key=api_key) for ep in self.hermes_endpoints
@@ -83,16 +87,12 @@ class SQLAlchemyDuckDuckGoCrawler:
         self.api_key = api_key
         self.model = model
         self.db_path = db_path
-
-        # Configure user agent
         default_requests_ua = f"python-requests/{requests.__version__}"
         self.user_agent = (
             f"{default_requests_ua} unturf-deep-research {user_agent_append}".strip()
         )
         self.session = requests.Session()
         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}",
@@ -104,29 +104,19 @@ class SQLAlchemyDuckDuckGoCrawler:
         self.SessionLocal = sessionmaker(
             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 = {}
         self.domain_crawl_delays = {}
         self.robots_txt_content = {}
 
     def _get_domain(self, url):
-        """Extract the domain from a URL."""
         parsed = urlparse(url)
         return parsed.netloc
 
     def _fetch_robots_txt(self, domain):
-        """Fetch and parse robots.txt for a domain."""
         if domain in self.robot_parsers:
             logger.info(f"Using cached robots.txt for {domain}")
             return self.robot_parsers[domain]
-
         logger.info(f"Fetching robots.txt for {domain}")
         robots_url = f"https://{domain}/robots.txt"
         parser = urllib.robotparser.RobotFileParser()
@@ -149,12 +139,10 @@ class SQLAlchemyDuckDuckGoCrawler:
         return self.robot_parsers[domain]
 
     def _can_fetch(self, url):
-        """Check if crawling the URL is allowed per robots.txt."""
         domain = self._get_domain(url)
         parser = self._fetch_robots_txt(domain)
         if parser is None:
             return True
-
         can_fetch = parser.can_fetch(self.user_agent, url)
         if not can_fetch:
             logger.warning(f"Blocked by robots.txt: {url}")
@@ -192,7 +180,6 @@ class SQLAlchemyDuckDuckGoCrawler:
         return can_fetch
 
     def _enforce_crawl_delay(self, domain):
-        """Enforce crawl delay for the domain."""
         delay = self.domain_crawl_delays.get(domain, DEFAULT_CRAWL_DELAY)
         last_fetched = self.domain_last_fetched.get(domain, 0)
         elapsed = time.time() - last_fetched
@@ -205,7 +192,6 @@ class SQLAlchemyDuckDuckGoCrawler:
         self.domain_last_fetched[domain] = time.time()
 
     def search_duckduckgo(self, query, max_results=10):
-        """Search DuckDuckGo and extract target URLs from redirect links."""
         encoded = quote_plus(query)
         url = f"https://html.duckduckgo.com/html/?q={encoded}"
         resp = self.session.get(url, timeout=10)
@@ -242,7 +228,6 @@ class SQLAlchemyDuckDuckGoCrawler:
         return results
 
     def fetch_webpage(self, url):
-        """Fetch a webpage, validating the URL first."""
         if not url.startswith(("http://", "https://")):
             logger.error(f"Invalid URL scheme: {url}")
             return None
@@ -259,8 +244,6 @@ class SQLAlchemyDuckDuckGoCrawler:
             return None
 
     def _fanout_call(self, messages, max_tokens, prefer_hermes=False):
-        """Call Hermes endpoints, preferring hermes.ai.unturf.com for aggregate_and_answer."""
-
         def call_client(client, endpoint):
             try:
                 adjusted_max_tokens = max_tokens
@@ -300,7 +283,9 @@ class SQLAlchemyDuckDuckGoCrawler:
                 if result is not None:
                     logger.info(f"Success using {used_endpoint}")
                     return result
-                logger.warning(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)
@@ -333,7 +318,6 @@ class SQLAlchemyDuckDuckGoCrawler:
         raise RuntimeError("All Hermes endpoints failed.")
 
     def extract_with_hermes(self, html):
-        """Extract main article content using Hermes."""
         text = BeautifulSoup(html, "html.parser").get_text(separator="\n")
         if len(text) > MAX_HTML_INPUT_CHARS:
             text = text[:MAX_HTML_INPUT_CHARS]
@@ -341,8 +325,9 @@ class SQLAlchemyDuckDuckGoCrawler:
             {
                 "role": "system",
                 "content": (
-                    "Extract the main article content, preserving formatting and structure, "
+                    "Extract the main article content, preserving all formatting, structure, and relevant details, "
-                    "excluding ads and navigation. Return only the full text."
+                    "including headers, paragraphs, lists, and key text. Exclude only ads, navigation, and unrelated boilerplate. "
+                    "Maximize content retention to capture comprehensive information, ensuring no relevant text is omitted."
                 ),
             },
             {"role": "user", "content": text},
@@ -355,7 +340,6 @@ class SQLAlchemyDuckDuckGoCrawler:
             return self._fanout_call(messages, max_tokens=HERMES_COMPLETION_TOKENS // 2)
 
     def summarize_with_hermes(self, content):
-        """Summarize content as bullet points, used only when token limit is exceeded."""
         messages = [
             {
                 "role": "system",
@@ -367,13 +351,13 @@ class SQLAlchemyDuckDuckGoCrawler:
         ]
         return self._fanout_call(messages, max_tokens=HERMES_COMPLETION_TOKENS)
 
-    def _split_into_chunks(self, text, max_words=500):
+    def _split_into_chunks(self, text, content_length):
-        """Split text into chunks of approximately max_words."""
+        max_words = MAX_WORDS_LONG if content_length > 5000 else MAX_WORDS_SHORT
+        logger.info(f"Using max_words={max_words} for content length={content_length}")
         words = text.split()
         chunks = []
         current_chunk = []
         current_word_count = 0
-
         for word in words:
             current_chunk.append(word)
             current_word_count += 1
@@ -383,42 +367,28 @@ class SQLAlchemyDuckDuckGoCrawler:
                 current_word_count = 0
         if current_chunk:
             chunks.append(" ".join(current_chunk))
-        return chunks
+        return chunks if len(chunks) > 1 else [text]
 
     def cache_article(self, url, title, html, extracted, summary):
-        """Cache article in database and store chunks in ChromaDB."""
         with self.SessionLocal() as db:
             art = Article(
                 url=url,
                 title=title,
                 raw_html=html,
                 extracted_content=extracted,
-                summary=summary,  # Stores extracted_content unless summarized
+                summary=summary,
                 fetched_at=datetime.now(timezone.utc),
             )
             db.add(art)
             try:
                 db.commit()
+                logger.info(f"Cached article: {url}")
             except Exception as e:
                 db.rollback()
                 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():
@@ -429,61 +399,191 @@ class SQLAlchemyDuckDuckGoCrawler:
             if not html:
                 return
             extracted = self.extract_with_hermes(html)
-            # Skip summarization; use extracted content as summary
             self.cache_article(url, title, html, extracted, extracted)
 
-    def aggregate_and_answer(self, query):
+    def _extract_keywords_with_hermes(self, query):
-        """Aggregate relevant article chunks from ChromaDB and generate an answer."""
+        messages = [
-        # Embed the query
+            {
-        query_embedding = self.embedder.encode([query], show_progress_bar=False)[0]
+                "role": "system",
-
+                "content": (
-        # Query ChromaDB for relevant chunks
+                    "Analyze the following query and extract a list of up to 5 relevant keywords or phrases "
-        results = self.collection.query(
+                    "that capture the main topics or entities. Focus on nouns, proper nouns, and key concepts. "
-            query_embeddings=[query_embedding.tolist()],
+                    "Avoid generic terms like 'what', 'is', or redundant variations. "
-            n_results=50,  # Retrieve more to filter by token count
+                    "Return the keywords as a comma-separated string."
+                ),
+            },
+            {"role": "user", "content": query},
+        ]
+        try:
+            result = self._fanout_call(messages, max_tokens=100)
+            keywords = [k.strip() for k in result.split(",") if k.strip()]
+            if len(keywords) > 5:
+                keywords = keywords[:5]
+            logger.info(f"Extracted keywords for query '{query}': {keywords}")
+            return keywords
+        except RuntimeError:
+            logger.warning(
+                "Failed to extract keywords with Hermes, falling back to query split"
             )
+            return query.lower().split()
+
+    def _search_sqlite(self, query, keywords, max_results=10):
+        with self.SessionLocal() as db:
+            articles = db.query(Article).all()
+            if not articles:
+                logger.warning("No articles found in SQLite cache.")
+                return []
+
+            # Keyword-based filtering
+            core_keywords = keywords + query.lower().split()
+            relevant_articles = []
+            for article in articles:
+                content_lower = article.extracted_content.lower()
+                keyword_score = sum(
+                    1 for kw in core_keywords if kw.lower() in content_lower
+                )
+                if keyword_score > 0:
+                    relevant_articles.append((article, keyword_score))
+
+            # Sort by keyword score
+            relevant_articles.sort(key=lambda x: x[1], reverse=True)
+            top_articles = relevant_articles[:max_results]
+
+            # TF-IDF similarity for refined ranking
+            if top_articles:
+                documents = [article.extracted_content for article, _ in top_articles]
+                vectorizer = TfidfVectorizer(stop_words="english")
+                try:
+                    tfidf_matrix = vectorizer.fit_transform(documents + [query])
+                    similarities = cosine_similarity(
+                        tfidf_matrix[-1], tfidf_matrix[:-1]
+                    )[0]
+                    scored_articles = [
+                        (article, score, keyword_score)
+                        for (article, keyword_score), score in zip(
+                            top_articles, similarities
+                        )
+                    ]
+                    scored_articles.sort(
+                        key=lambda x: 0.5 * x[1]
+                        + 0.5 * (x[2] / max(1, max(s[2] for s in scored_articles))),
+                        reverse=True,
+                    )
+                    return [
+                        {
+                            "url": article.url,
+                            "title": article.title,
+                            "content": article.extracted_content,
+                        }
+                        for article, _, _ in scored_articles
+                    ]
+                except ValueError as e:
+                    logger.warning(
+                        f"TF-IDF failed: {e}, falling back to keyword ranking"
+                    )
+                    return [
+                        {
+                            "url": article.url,
+                            "title": article.title,
+                            "content": article.extracted_content,
+                        }
+                        for article, _ in top_articles
+                    ]
+            return []
+
+    def aggregate_and_answer(self, query, search_results=None, query_keywords=None):
+        # Search SQLite cache
+        cache_results = self._search_sqlite(query, query_keywords)
+        logger.info(f"Retrieved {len(cache_results)} articles from SQLite cache")
 
-        # Collect chunks up to token limit
         def estimate_tokens(text):
             return len(text) // 4 + 1
 
         combined_content = []
-        total_tokens = 0
+        total_words = 0
-        prompt_template = (
+        prompt_template = f"Based on these article excerpts:\n{{}}\n\nProvide a comprehensive answer to: {query}"
-            f"Based on these article excerpts:\n{{}}\n\nProvide a comprehensive answer to: {query}"
-        )
         template_tokens = estimate_tokens(prompt_template.format(""))
 
-        for doc, metadata in zip(results["documents"][0], results["metadatas"][0]):
+        # Add cached content
-            chunk = doc
+        for result in cache_results:
-            chunk_tokens = estimate_tokens(chunk)
+            content = result["content"]
-            if total_tokens + chunk_tokens + template_tokens <= HERMES_INPUT_TOKENS:
+            word_count = len(content.split())
+            content_tokens = estimate_tokens(content)
+            if total_words + word_count + template_tokens <= CONTEXT_WINDOW_WORDS:
                 combined_content.append(
-                    f"From {metadata['url']} (Title: {metadata['title']}):\n{chunk}"
+                    f"From {result['url']} (Title: {result['title']}):\n{content}"
+                )
+                total_words += word_count
+                logger.info(
+                    f"Included article {result['url']} (words: {word_count}, tokens: {content_tokens})"
                 )
-                total_tokens += chunk_tokens
             else:
+                # Chunk if content exceeds context window
+                chunks = self._split_into_chunks(content, len(content))
+                for chunk in chunks:
+                    chunk_words = len(chunk.split())
+                    chunk_tokens = estimate_tokens(chunk)
+                    if (
+                        total_words + chunk_words + template_tokens
+                        <= CONTEXT_WINDOW_WORDS
+                    ):
+                        combined_content.append(
+                            f"From {result['url']} (Title: {result['title']}):\n{chunk}"
+                        )
+                        total_words += chunk_words
+                        logger.info(
+                            f"Included chunk from {result['url']} (words: {chunk_words}, tokens: {chunk_tokens})"
+                        )
+                    else:
+                        logger.info(
+                            f"Skipping chunk from {result['url']} (exceeds context window: {total_words + chunk_words})"
+                        )
                         break
 
+        if not combined_content and search_results:
+            logger.warning(
+                "No relevant content found in cache, using search results as fallback"
+            )
+            fallback_content = "\n".join(
+                f"- {title}: {url}" for title, url in search_results
+            )
+            messages = [
+                {
+                    "role": "user",
+                    "content": f"Based on these search results:\n{fallback_content}\n\nAnswer: {query}",
+                }
+            ]
+            try:
+                return self._fanout_call(
+                    messages, max_tokens=HERMES_COMPLETION_TOKENS, prefer_hermes=True
+                )
+            except RuntimeError:
+                return "No relevant information found, and fallback answer generation failed."
+
         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")
+        total_tokens = estimate_tokens(combined) + template_tokens
+        logger.info(
+            f"Using {total_tokens} tokens from {len(combined_content)} articles/chunks"
+        )
         messages = [{"role": "user", "content": prompt_template.format(combined)}]
         return self._fanout_call(
             messages, max_tokens=HERMES_COMPLETION_TOKENS, prefer_hermes=True
         )
 
     def run(self, query, max_results=10):
+        query_keywords = self._extract_keywords_with_hermes(query)
         hits = self.search_duckduckgo(query, max_results)
         with ThreadPoolExecutor(max_workers=4) as executor:
             futures = [executor.submit(self.process_url, hit) for hit in hits]
             for future in as_completed(futures):
                 future.result()
         logger.info("Generating comprehensive answer...")
-        result = self.aggregate_and_answer(query)
+        result = self.aggregate_and_answer(
+            query, search_results=hits, query_keywords=query_keywords
+        )
         print(result)
 
 
@@ -505,7 +605,6 @@ if __name__ == "__main__":
         "--user-agent-append", default="", help="String to append to default user agent"
     )
     args = parser.parse_args()
-
     crawler = SQLAlchemyDuckDuckGoCrawler(
         api_key=args.api_key, model=args.model, user_agent_append=args.user_agent_append
     )