Update app.py

app.py

@@ -4,6 +4,7 @@ import requests
 import trafilatura
 from ddgs import DDGS
 import uvicorn
+import re
 import numpy as np
 
 from sentence_transformers import SentenceTransformer
@@ -60,6 +61,19 @@ def clean_text(text: str):
 
     return text.strip()
 
+def split_sentences(text):
+
+    sentences = re.split(
+        r'(?<=[.!?])\s+',
+        text
+    )
+
+    return [
+        s.strip()
+        for s in sentences
+        if len(s.strip()) > 40
+    ]
+
 # --------------------------------------------------
 # CHUNKING
 # --------------------------------------------------
@@ -167,34 +181,67 @@ def rank_chunks(query, chunks):
 # COMPRESS TEXT
 # --------------------------------------------------
 
-def compress_chunks(chunks):
+def compress_chunks(query, chunks):
+
+    try:
 
-    compressed = []
+        all_sentences = []
 
-    total_chars = 0
+        for chunk in chunks:
 
-    for chunk in chunks:
+            sentences = split_sentences(chunk)
 
-        chunk = chunk.strip()
+            all_sentences.extend(sentences)
 
-        if not chunk:
-            continue
+        if not all_sentences:
+            return ""
 
-        remaining = (
-            MAX_FINAL_CONTEXT_CHARS
-            - total_chars
+        query_embedding = embedding_model.encode(
+            [query]
         )
 
-        if remaining <= 0:
-            break
+        sentence_embeddings = embedding_model.encode(
+            all_sentences
+        )
 
-        chunk = chunk[:remaining]
+        similarities = cosine_similarity(
+            query_embedding,
+            sentence_embeddings
+        )[0]
 
-        compressed.append(chunk)
+        ranked_sentences = sorted(
+            zip(all_sentences, similarities),
+            key=lambda x: x[1],
+            reverse=True
+        )
 
-        total_chars += len(chunk)
+        final_sentences = []
 
-    return "\n\n".join(compressed)
+        total_chars = 0
+
+        used = set()
+
+        for sentence, score in ranked_sentences:
+
+            normalized = sentence.lower().strip()
+
+            if normalized in used:
+                continue
+
+            used.add(normalized)
+
+            if total_chars + len(sentence) > MAX_FINAL_CONTEXT_CHARS:
+                break
+
+            final_sentences.append(sentence)
+
+            total_chars += len(sentence)
+
+        return "\n\n".join(final_sentences)
+
+    except Exception:
+
+        return "\n\n".join(chunks[:2])
 
 # --------------------------------------------------
 # RESEARCH PIPELINE
@@ -245,6 +292,7 @@ def perform_research(query):
                 )
 
                 compressed_content = compress_chunks(
+                    query,
                     ranked_chunks
                 )