370 lines
14 KiB
Python
370 lines
14 KiB
Python
"""
|
|
ResearchOwl Processor
|
|
Chunking → Quality scoring via Ollama → Embeddings → RAG synthesis
|
|
"""
|
|
import asyncio
|
|
import json
|
|
import math
|
|
import re
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
import structlog
|
|
|
|
from src.config import settings
|
|
from src.db.database import ResearchDB
|
|
|
|
logger = structlog.get_logger()
|
|
|
|
|
|
class OllamaClient:
|
|
"""Async client for Ollama API"""
|
|
|
|
def __init__(self):
|
|
self.base_url = settings.ollama_url.rstrip("/")
|
|
self.model = settings.ollama_model
|
|
|
|
async def generate(self, prompt: str, system: str = None,
|
|
timeout: int = 120, temperature: float = 0.7) -> str:
|
|
payload = {
|
|
"model": self.model,
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
"options": {
|
|
"temperature": temperature,
|
|
"num_predict": 2048,
|
|
"repeat_penalty": 1.15,
|
|
"repeat_last_n": 128,
|
|
}
|
|
}
|
|
if system:
|
|
payload["system"] = system
|
|
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
resp = await client.post(f"{self.base_url}/api/generate", json=payload)
|
|
resp.raise_for_status()
|
|
return resp.json().get("response", "").strip()
|
|
|
|
async def embed(self, text: str) -> Optional[list[float]]:
|
|
"""Get embedding vector for a text"""
|
|
payload = {"model": self.model, "prompt": text}
|
|
try:
|
|
async with httpx.AsyncClient(timeout=60) as client:
|
|
resp = await client.post(f"{self.base_url}/api/embeddings", json=payload)
|
|
resp.raise_for_status()
|
|
return resp.json().get("embedding")
|
|
except Exception as e:
|
|
logger.warning("Embedding failed", error=str(e))
|
|
return None
|
|
|
|
async def is_available(self) -> bool:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=5) as client:
|
|
resp = await client.get(f"{self.base_url}/api/tags")
|
|
return resp.status_code == 200
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def simple_chunk(text: str, chunk_size: int = 800, overlap: int = 100) -> list[str]:
|
|
"""
|
|
Split text into overlapping chunks by approximate word count.
|
|
Respects paragraph boundaries when possible.
|
|
"""
|
|
paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
|
|
chunks = []
|
|
current = []
|
|
current_words = 0
|
|
|
|
for para in paragraphs:
|
|
para_words = len(para.split())
|
|
if current_words + para_words > chunk_size and current:
|
|
chunks.append("\n\n".join(current))
|
|
# overlap: keep last paragraph
|
|
if overlap > 0 and current:
|
|
current = [current[-1]]
|
|
current_words = len(current[0].split())
|
|
else:
|
|
current = []
|
|
current_words = 0
|
|
current.append(para)
|
|
current_words += para_words
|
|
|
|
if current:
|
|
chunks.append("\n\n".join(current))
|
|
|
|
return chunks
|
|
|
|
|
|
def cosine_similarity(a: list[float], b: list[float]) -> float:
|
|
"""Simple cosine similarity"""
|
|
if not a or not b or len(a) != len(b):
|
|
return 0.0
|
|
dot = sum(x * y for x, y in zip(a, b))
|
|
norm_a = math.sqrt(sum(x * x for x in a))
|
|
norm_b = math.sqrt(sum(x * x for x in b))
|
|
if norm_a == 0 or norm_b == 0:
|
|
return 0.0
|
|
return dot / (norm_a * norm_b)
|
|
|
|
|
|
class ContentProcessor:
|
|
"""
|
|
Processes scraped sources:
|
|
1. Chunks content
|
|
2. Scores quality with Ollama
|
|
3. Generates embeddings
|
|
4. Stores high-quality chunks
|
|
"""
|
|
|
|
def __init__(self, db: ResearchDB, ollama: OllamaClient):
|
|
self.db = db
|
|
self.ollama = ollama
|
|
|
|
async def process_session(self, session_id: int, topic: str,
|
|
progress_callback=None) -> dict:
|
|
"""Process all scraped sources for a session"""
|
|
from src.db.database import ResearchDB
|
|
sources = await self.db.get_all_sources(session_id)
|
|
scraped = [s for s in sources if s["status"] == "scraped"]
|
|
|
|
logger.info("Processing sources", total=len(scraped))
|
|
scraped = await self._dedup_sources(session_id, scraped)
|
|
logger.info("After dedup", unique=len(scraped))
|
|
total_chunks = 0
|
|
total_words = 0
|
|
|
|
semaphore = asyncio.Semaphore(3) # process 3 sources at once
|
|
|
|
async def process_one(source):
|
|
async with semaphore:
|
|
n = await self._process_source(session_id, topic, source)
|
|
return n
|
|
|
|
results = await asyncio.gather(*[process_one(s) for s in scraped],
|
|
return_exceptions=True)
|
|
|
|
for i, r in enumerate(results):
|
|
if isinstance(r, Exception):
|
|
logger.error("Source processing raised exception",
|
|
source_id=scraped[i]["id"], error=str(r), exc_info=r)
|
|
elif isinstance(r, int):
|
|
total_chunks += r
|
|
|
|
total_words = sum(s.get("word_count", 0) for s in scraped)
|
|
await self.db.update_session(
|
|
session_id,
|
|
total_chunks=total_chunks,
|
|
total_words=total_words
|
|
)
|
|
|
|
if progress_callback:
|
|
await progress_callback(total_chunks=total_chunks, total_words=total_words)
|
|
|
|
return {"total_chunks": total_chunks, "total_words": total_words}
|
|
|
|
async def _dedup_sources(self, session_id: int,
|
|
scraped: list[dict]) -> list[dict]:
|
|
try:
|
|
import hashlib
|
|
seen_hashes: set = set()
|
|
seen_prefixes: list = []
|
|
unique: list = []
|
|
duplicates = 0
|
|
|
|
for source in scraped:
|
|
content = await self.db.get_source_content(source["id"])
|
|
if not content:
|
|
unique.append(source)
|
|
continue
|
|
|
|
content_hash = hashlib.md5(content[:2000].encode()).hexdigest()
|
|
if content_hash in seen_hashes:
|
|
duplicates += 1
|
|
await self.db.update_source(source["id"], status="skipped")
|
|
continue
|
|
seen_hashes.add(content_hash)
|
|
|
|
prefix = content[:300].strip().lower()
|
|
prefix_words = set(prefix.split())
|
|
is_dup = False
|
|
if len(prefix_words) >= 10:
|
|
for seen_prefix_words in seen_prefixes:
|
|
intersection = len(prefix_words & seen_prefix_words)
|
|
union = len(prefix_words | seen_prefix_words)
|
|
if intersection / max(union, 1) > 0.85:
|
|
is_dup = True
|
|
break
|
|
|
|
if is_dup:
|
|
duplicates += 1
|
|
await self.db.update_source(source["id"], status="skipped")
|
|
continue
|
|
|
|
seen_prefixes.append(prefix_words)
|
|
unique.append(source)
|
|
|
|
if duplicates > 0:
|
|
logger.info("Dedup complete", session_id=session_id,
|
|
original=len(scraped), duplicates=duplicates,
|
|
unique=len(unique))
|
|
return unique
|
|
except Exception as e:
|
|
logger.warning("Dedup failed, processing all sources", error=str(e))
|
|
return scraped
|
|
|
|
async def _process_source(self, session_id: int, topic: str, source: dict) -> int:
|
|
"""Chunk, score, embed and store a single source. Returns chunk count."""
|
|
source_id = source["id"]
|
|
|
|
content = await self.db.get_source_content(source_id)
|
|
if not content:
|
|
logger.warning("No content in source_contents", source_id=source_id)
|
|
return 0
|
|
|
|
chunks = simple_chunk(content, settings.chunk_size, settings.chunk_overlap)
|
|
logger.info("Processing source", source_id=source_id,
|
|
content_len=len(content), num_chunks=len(chunks),
|
|
quality_threshold=settings.quality_threshold)
|
|
stored = 0
|
|
filtered_quality = 0
|
|
|
|
for i, chunk in enumerate(chunks):
|
|
words = len(chunk.split())
|
|
if words < 30:
|
|
continue
|
|
|
|
quality = await self._score_quality(chunk, topic, session_id)
|
|
if quality < settings.quality_threshold:
|
|
filtered_quality += 1
|
|
logger.debug("Chunk filtered by quality", source_id=source_id,
|
|
chunk_index=i, quality=round(quality, 2),
|
|
threshold=settings.quality_threshold, words=words)
|
|
continue
|
|
|
|
embedding = await self.ollama.embed(chunk[:1000])
|
|
|
|
await self.db.add_chunk(
|
|
session_id=session_id,
|
|
source_id=source_id,
|
|
content=chunk,
|
|
chunk_index=i,
|
|
token_count=words,
|
|
quality_score=quality,
|
|
embedding=embedding
|
|
)
|
|
stored += 1
|
|
|
|
if filtered_quality > 0 and stored == 0:
|
|
logger.warning(
|
|
"All chunks filtered by quality — consider lowering QUALITY_THRESHOLD "
|
|
"(currently %.1f) or set QUALITY_THRESHOLD=0 to disable",
|
|
settings.quality_threshold,
|
|
source_id=source_id, chunks_total=len(chunks),
|
|
chunks_filtered=filtered_quality
|
|
)
|
|
|
|
logger.info("Source processed", source_id=source_id, stored=stored)
|
|
return stored
|
|
|
|
async def _score_quality(self, chunk: str, topic: str,
|
|
session_id: int | None = None) -> float:
|
|
"""Score 0-1 relevance to topic. Uses Claude Haiku if API key set, else Ollama."""
|
|
if settings.anthropic_api_key:
|
|
return await self._score_with_claude(chunk, topic, session_id)
|
|
return await self._score_with_ollama(chunk, topic)
|
|
|
|
async def _score_with_claude(self, chunk: str, topic: str,
|
|
session_id: int | None = None) -> float:
|
|
import anthropic
|
|
prompt = (
|
|
f'Rate 0-10 how relevant this text is to the topic "{topic}". '
|
|
f'Be generous — if the text is tangentially related, score 4+. '
|
|
f'Only score below 3 if completely unrelated. '
|
|
f'Reply with only a number.\n\nText:\n{chunk[:500]}'
|
|
)
|
|
try:
|
|
client = anthropic.AsyncAnthropic(api_key=settings.anthropic_api_key)
|
|
msg = await client.messages.create(
|
|
model=settings.claude_model,
|
|
max_tokens=10,
|
|
messages=[{"role": "user", "content": prompt}]
|
|
)
|
|
if session_id is not None:
|
|
try:
|
|
await self.db.log_api_call(
|
|
session_id, "scoring", settings.claude_model,
|
|
msg.usage.input_tokens, msg.usage.output_tokens
|
|
)
|
|
except Exception as log_err:
|
|
logger.warning("Failed to log API usage", error=str(log_err))
|
|
response = msg.content[0].text.strip()
|
|
numbers = re.findall(r'\b(\d+(?:\.\d+)?)\b', response)
|
|
if numbers:
|
|
score = float(numbers[0])
|
|
normalized = min(1.0, score / 10.0)
|
|
logger.debug("Claude relevance score", raw=score, normalized=round(normalized, 2))
|
|
return normalized
|
|
return 0.5
|
|
except Exception as e:
|
|
logger.warning("Claude scoring failed, falling back to Ollama", error=str(e))
|
|
return await self._score_with_ollama(chunk, topic)
|
|
|
|
async def _score_with_ollama(self, chunk: str, topic: str) -> float:
|
|
prompt = (
|
|
f'Score 0-10: how relevant is this text to the topic "{topic}"?\n'
|
|
f"0 = completely unrelated, 10 = directly and specifically about this topic.\n\n"
|
|
f"Text:\n{chunk[:500]}\n\n"
|
|
f"Reply with ONLY a single integer 0-10. No explanation."
|
|
)
|
|
try:
|
|
response = await self.ollama.generate(prompt, temperature=0.1)
|
|
numbers = re.findall(r'\b(\d+(?:\.\d+)?)\b', response)
|
|
if numbers:
|
|
score = float(numbers[0])
|
|
normalized = min(1.0, score / 10.0)
|
|
logger.debug("Ollama relevance score", raw=score, normalized=round(normalized, 2))
|
|
return normalized
|
|
logger.debug("No number in Ollama relevance response", response=response[:80])
|
|
return 0.6
|
|
except Exception as e:
|
|
logger.warning("Ollama relevance scoring failed", error=str(e))
|
|
return 0.6
|
|
|
|
async def rag_query(self, session_id: int, query: str, top_k: int = 20) -> str:
|
|
"""
|
|
Retrieve most relevant chunks for a query using embeddings + keyword fallback
|
|
"""
|
|
# Get query embedding
|
|
query_embedding = await self.ollama.embed(query)
|
|
|
|
# Get top quality chunks
|
|
chunks = await self.db.get_top_chunks(session_id, limit=300)
|
|
|
|
if query_embedding and chunks:
|
|
# Rank by embedding similarity
|
|
scored = []
|
|
for chunk in chunks:
|
|
emb = chunk.get("embedding")
|
|
if emb and isinstance(emb, str):
|
|
try:
|
|
emb = json.loads(emb)
|
|
except Exception:
|
|
emb = None
|
|
sim = cosine_similarity(query_embedding, emb) if emb else 0.5
|
|
scored.append((sim * 0.7 + chunk["quality_score"] * 0.3, chunk))
|
|
|
|
scored.sort(key=lambda x: x[0], reverse=True)
|
|
top_chunks = [c for _, c in scored[:top_k]]
|
|
else:
|
|
# Fallback: just use quality score
|
|
top_chunks = chunks[:top_k]
|
|
|
|
# Build context
|
|
context_parts = []
|
|
for chunk in top_chunks:
|
|
source_label = f"[{chunk.get('source_type', 'web').upper()}] {chunk.get('title', 'Unknown')}"
|
|
context_parts.append(f"{source_label}:\n{chunk['content']}")
|
|
|
|
return "\n\n---\n\n".join(context_parts)
|