514 lines
21 KiB
Python
514 lines
21 KiB
Python
"""
|
|
title: Obsidian RAG Pipeline
|
|
author: Daniel Henry
|
|
version: 0.17
|
|
description: Updated for llama-swap with llama.cpp (OpenAI-compatible API)
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import time
|
|
import math
|
|
import urllib.parse
|
|
from typing import AsyncGenerator
|
|
|
|
import aiohttp
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class Pipe:
|
|
|
|
class Valves(BaseModel):
|
|
# Endpoints
|
|
llamacpp_url: str = Field(default="http://ollama.internal.henryhosted.com:9292")
|
|
qdrant_url: str = Field(default="http://app-01.internal.henryhosted.com:6333")
|
|
|
|
# Qdrant
|
|
collection_name: str = Field(default="obsidian_vault")
|
|
retrieve_count: int = Field(
|
|
default=50, description="Candidates to fetch from Qdrant"
|
|
)
|
|
qdrant_score_threshold: float = Field(
|
|
default=0.3, description="Minimum similarity score"
|
|
)
|
|
|
|
# Reranker
|
|
rerank_enabled: bool = Field(
|
|
default=True, description="Set to False to skip reranking"
|
|
)
|
|
rerank_logit: bool = Field(
|
|
default=False, description="Enable if reranker outputs logits"
|
|
)
|
|
rerank_debug: bool = Field(
|
|
default=False, description="Output all rerank values into think"
|
|
)
|
|
rerank_model: str = Field(
|
|
default="bge-reranker-v2-m3-q8_0",
|
|
description="Reranker model name",
|
|
)
|
|
rerank_timeout: float = Field(default=60.0)
|
|
min_rerank_score: float = Field(
|
|
default=0.01, description="Minimum rerank score to keep"
|
|
)
|
|
final_top_k: int = Field(
|
|
default=10, description="Chunks to keep after reranking"
|
|
)
|
|
|
|
# LLM
|
|
embedding_model: str = Field(
|
|
default="nomic-embed-text-v1.5.f16",
|
|
description="Embedding model name",
|
|
)
|
|
llm_model: str = Field(
|
|
default="qwen2.5-3b-instruct-q4_k_m",
|
|
description="LLM model name",
|
|
)
|
|
llm_max_tokens: int = Field(
|
|
default=2048, description="Max tokens for LLM response"
|
|
)
|
|
llm_timeout: float = Field(default=300.0)
|
|
query_rewrite_model: str = Field(
|
|
default="",
|
|
description="Model for query rewriting. Leave empty to use llm_model.",
|
|
)
|
|
|
|
# Obsidian
|
|
vault_name: str = Field(
|
|
default="Main", description="For generating obsidian:// links"
|
|
)
|
|
|
|
# Display
|
|
show_thinking: bool = Field(default=True)
|
|
show_sources: bool = Field(default=True)
|
|
show_stats: bool = Field(default=True)
|
|
token_warning_threshold: int = Field(
|
|
default=6000, description="Warn if context exceeds this"
|
|
)
|
|
|
|
def __init__(self):
|
|
self.valves = self.Valves()
|
|
|
|
def _estimate_tokens(self, text: str) -> int:
|
|
"""Rough token estimate: ~4 chars per token for English text."""
|
|
return len(text) // 4
|
|
|
|
async def pipe(self, body: dict) -> AsyncGenerator[str, None]:
|
|
messages = body.get("messages", [])
|
|
if not messages:
|
|
yield "No messages provided."
|
|
return
|
|
|
|
query = messages[-1].get("content", "").strip()
|
|
if not query:
|
|
yield "Empty query."
|
|
return
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
async for chunk in self._execute(session, query, messages):
|
|
yield chunk
|
|
|
|
async def _execute(
|
|
self,
|
|
session: aiohttp.ClientSession,
|
|
query: str,
|
|
messages: list[dict],
|
|
) -> AsyncGenerator[str, None]:
|
|
|
|
think = self.valves.show_thinking
|
|
total_prompt_tokens = 0
|
|
|
|
# Start thinking block
|
|
if think:
|
|
yield "<think>\n"
|
|
yield f"**LLM Model:** {self.valves.llm_model}\n"
|
|
yield f"**Query:** {query}\n\n"
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Step 1: Rewrite query with conversation context
|
|
# ─────────────────────────────────────────────
|
|
if think:
|
|
yield "**Step 1: Query Rewriting**\n"
|
|
|
|
t0 = time.time()
|
|
rewrite_model = self.valves.query_rewrite_model or self.valves.llm_model
|
|
current_question = messages[-1].get("content", "")
|
|
|
|
# Build conversation context for rewriting (only if there's prior conversation)
|
|
conversation_for_rewrite = []
|
|
for m in messages[:-1]:
|
|
role = m.get("role", "")
|
|
content = m.get("content", "")
|
|
if role == "user":
|
|
conversation_for_rewrite.append(f"User: {content}")
|
|
elif role == "assistant":
|
|
truncated = content[:500] + "..." if len(content) > 500 else content
|
|
conversation_for_rewrite.append(f"Assistant: {truncated}")
|
|
|
|
if conversation_for_rewrite:
|
|
rewrite_prompt = f"""Do not interpret or answer the question. Simply add enough context from the conversation so the question makes sense on its own.
|
|
|
|
Conversation:
|
|
{chr(10).join(conversation_for_rewrite)}
|
|
|
|
Latest question: {current_question}
|
|
|
|
Rewrite the question to be standalone (respond with ONLY the rewritten question, nothing else):"""
|
|
|
|
try:
|
|
async with session.post(
|
|
f"{self.valves.llamacpp_url}/v1/chat/completions",
|
|
json={
|
|
"model": rewrite_model,
|
|
"messages": [{"role": "user", "content": rewrite_prompt}],
|
|
"stream": False,
|
|
"max_tokens": 256,
|
|
},
|
|
timeout=aiohttp.ClientTimeout(total=30),
|
|
) as resp:
|
|
if resp.status == 200:
|
|
data = await resp.json()
|
|
rewritten = (
|
|
data.get("choices", [{}])[0]
|
|
.get("message", {})
|
|
.get("content", "")
|
|
.strip()
|
|
)
|
|
if rewritten and len(rewritten) < 1000:
|
|
search_query = rewritten
|
|
else:
|
|
search_query = current_question
|
|
else:
|
|
search_query = current_question
|
|
except Exception as e:
|
|
if think:
|
|
yield f" ⚠ Rewrite failed: {e}, using original query\n"
|
|
search_query = current_question
|
|
else:
|
|
search_query = current_question
|
|
|
|
if think:
|
|
yield f" Model: {rewrite_model}\n"
|
|
yield f" Original: {current_question}\n"
|
|
yield f" Search query: {search_query}\n"
|
|
yield f" ✓ Done ({time.time() - t0:.2f}s)\n\n"
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Step 2: Embed
|
|
# ─────────────────────────────────────────────
|
|
if think:
|
|
yield "**Step 2: Embedding**\n"
|
|
t0 = time.time()
|
|
|
|
try:
|
|
async with session.post(
|
|
f"{self.valves.llamacpp_url}/v1/embeddings",
|
|
json={
|
|
"model": self.valves.embedding_model,
|
|
"input": search_query,
|
|
},
|
|
timeout=aiohttp.ClientTimeout(total=30),
|
|
) as resp:
|
|
if resp.status != 200:
|
|
error_text = await resp.text()
|
|
if think:
|
|
yield f" ✗ HTTP {resp.status}: {error_text}\n</think>\n\n"
|
|
yield f"Embedding failed: HTTP {resp.status}"
|
|
return
|
|
data = await resp.json()
|
|
embedding = data.get("data", [{}])[0].get("embedding")
|
|
if not embedding:
|
|
if think:
|
|
yield " ✗ No embedding in response\n</think>\n\n"
|
|
yield "Embedding failed: No embedding returned"
|
|
return
|
|
except Exception as e:
|
|
if think:
|
|
yield f" ✗ {e}\n</think>\n\n"
|
|
yield f"Embedding failed: {e}"
|
|
return
|
|
|
|
if think:
|
|
yield f" ✓ Done ({time.time() - t0:.2f}s)\n\n"
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Step 3: Search Qdrant
|
|
# ─────────────────────────────────────────────
|
|
if think:
|
|
yield "**Step 3: Qdrant Search**\n"
|
|
t0 = time.time()
|
|
|
|
try:
|
|
async with session.post(
|
|
f"{self.valves.qdrant_url}/collections/{self.valves.collection_name}/points/search",
|
|
json={
|
|
"vector": embedding,
|
|
"limit": self.valves.retrieve_count,
|
|
"with_payload": True,
|
|
"score_threshold": self.valves.qdrant_score_threshold,
|
|
},
|
|
timeout=aiohttp.ClientTimeout(total=15),
|
|
) as resp:
|
|
if resp.status != 200:
|
|
if think:
|
|
yield f" ✗ HTTP {resp.status}\n</think>\n\n"
|
|
yield f"Qdrant search failed: HTTP {resp.status}"
|
|
return
|
|
qdrant_results = (await resp.json()).get("result", [])
|
|
except Exception as e:
|
|
if think:
|
|
yield f" ✗ {e}\n</think>\n\n"
|
|
yield f"Qdrant search failed: {e}"
|
|
return
|
|
|
|
if think:
|
|
yield f" ✓ Found {len(qdrant_results)} chunks ({time.time() - t0:.2f}s)\n"
|
|
|
|
if not qdrant_results:
|
|
if think:
|
|
yield " ✗ No results\n</think>\n\n"
|
|
yield "No relevant notes found for this query."
|
|
return
|
|
|
|
if think:
|
|
yield " Top 5:\n"
|
|
for i, r in enumerate(qdrant_results[:5]):
|
|
name = r.get("payload", {}).get("fileName", "?")
|
|
score = r.get("score", 0)
|
|
yield f" {i+1}. [{score:.4f}] {name}\n"
|
|
yield "\n"
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Step 4: Rerank (optional)
|
|
# ─────────────────────────────────────────────
|
|
if self.valves.rerank_enabled:
|
|
if think:
|
|
yield "**Step 4: Reranking**\n"
|
|
yield f"**Rerank Model:** {self.valves.rerank_model}\n"
|
|
t0 = time.time()
|
|
|
|
docs_for_rerank = [
|
|
r.get("payload", {}).get("content", "") for r in qdrant_results
|
|
]
|
|
|
|
try:
|
|
async with session.post(
|
|
f"{self.valves.llamacpp_url}/v1/rerank",
|
|
json={
|
|
"model": self.valves.rerank_model,
|
|
"query": search_query,
|
|
"documents": docs_for_rerank,
|
|
},
|
|
timeout=aiohttp.ClientTimeout(total=self.valves.rerank_timeout),
|
|
) as resp:
|
|
if resp.status != 200:
|
|
error_text = await resp.text()
|
|
if think:
|
|
yield f" ⚠ Reranker failed: HTTP {resp.status} - {error_text}, using Qdrant order\n\n"
|
|
chunks = qdrant_results[: self.valves.final_top_k]
|
|
else:
|
|
rerank_data = await resp.json()
|
|
rerank_results = rerank_data.get("results", [])
|
|
|
|
scored = []
|
|
for item in rerank_results:
|
|
idx = item["index"]
|
|
score = item["relevance_score"]
|
|
if self.valves.rerank_logit:
|
|
score = 1 / (1 + math.exp(-item["relevance_score"]))
|
|
if think and self.valves.rerank_debug:
|
|
yield f" • Debug: Doc {idx} score: {score}\n"
|
|
if score >= self.valves.min_rerank_score:
|
|
chunk = qdrant_results[idx].copy()
|
|
chunk["rerank_score"] = score
|
|
scored.append(chunk)
|
|
|
|
scored.sort(key=lambda x: x["rerank_score"], reverse=True)
|
|
chunks = scored[: self.valves.final_top_k]
|
|
|
|
if think:
|
|
yield f" ✓ Kept {len(chunks)} chunks ({time.time() - t0:.2f}s)\n"
|
|
if chunks:
|
|
yield " Top 5 after rerank:\n"
|
|
for i, c in enumerate(chunks[:5]):
|
|
name = c.get("payload", {}).get("fileName", "?")
|
|
score = c.get("rerank_score", 0)
|
|
yield f" {i+1}. [{score:.4f}] {name}\n"
|
|
yield "\n"
|
|
|
|
except Exception as e:
|
|
if think:
|
|
yield f" ⚠ Reranker error: {e}, using Qdrant order\n\n"
|
|
chunks = qdrant_results[: self.valves.final_top_k]
|
|
else:
|
|
if think:
|
|
yield "**Step 4: Reranking** (disabled)\n\n"
|
|
chunks = qdrant_results[: self.valves.final_top_k]
|
|
|
|
if not chunks:
|
|
if think:
|
|
yield " ✗ No chunks after filtering\n</think>\n\n"
|
|
yield "No relevant notes passed the relevance threshold."
|
|
return
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Step 5: Build context
|
|
# ─────────────────────────────────────────────
|
|
if think:
|
|
yield "**Step 5: Build Context**\n"
|
|
|
|
context_parts = []
|
|
for i, chunk in enumerate(chunks, 1):
|
|
payload = chunk.get("payload", {})
|
|
file_name = payload.get("fileName", "Unknown")
|
|
content = payload.get("content", "").strip()
|
|
source = payload.get("source", "")
|
|
|
|
# CHANGE: Explicit bracketed ID format
|
|
part = f"[{i}] File: {file_name}\n"
|
|
if source:
|
|
part += f"Source: {source}\n"
|
|
part += f"\n{content}"
|
|
context_parts.append(part)
|
|
|
|
context = "\n\n---\n\n".join(context_parts)
|
|
context_tokens = self._estimate_tokens(context)
|
|
|
|
if think:
|
|
yield f" ✓ {len(chunks)} chunks, ~{context_tokens:,} tokens\n"
|
|
if context_tokens > self.valves.token_warning_threshold:
|
|
yield f" ⚠ Warning: large context may affect quality\n"
|
|
yield "\n"
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Step 6: Build prompt and call LLM
|
|
# ─────────────────────────────────────────────
|
|
if think:
|
|
yield "**Step 6: Generate Response**\n"
|
|
yield "</think>\n\n"
|
|
|
|
system_prompt = f"""You are a helpful assistant. Use the provided notes to answer the user's question.
|
|
|
|
RULES:
|
|
1. Use the <notes> as your source of truth.
|
|
2. Cite facts using the bracketed ID number [1].
|
|
3. SYNTHESIS: You are encouraged to draw connections between different notes to form a complete answer.
|
|
4. INFERENCE: If the answer is not explicitly written but can be logically inferred from the notes, you may answer, but please use phrases like "The notes imply..." or "Based on [1], it suggests..."
|
|
5. If the answer is completely absent, say "I couldn't find that in your notes."
|
|
|
|
<notes>
|
|
{context}
|
|
</notes>"""
|
|
|
|
# Build conversation, stripping previous sources from assistant messages
|
|
conversation = []
|
|
for m in messages:
|
|
role = m.get("role")
|
|
if role not in ("user", "assistant"):
|
|
continue
|
|
msg = {"role": role, "content": m.get("content", "")}
|
|
if role == "assistant" and "**Sources:**" in msg["content"]:
|
|
msg["content"] = msg["content"].split("**Sources:**")[0].strip()
|
|
conversation.append(msg)
|
|
|
|
llm_messages = [{"role": "system", "content": system_prompt}] + conversation
|
|
|
|
# Estimate prompt tokens
|
|
prompt_text = system_prompt + "".join(m["content"] for m in conversation)
|
|
total_prompt_tokens = self._estimate_tokens(prompt_text)
|
|
|
|
llm_payload = {
|
|
"model": self.valves.llm_model,
|
|
"messages": llm_messages,
|
|
"stream": True,
|
|
"max_tokens": self.valves.llm_max_tokens,
|
|
}
|
|
|
|
completion_tokens = 0
|
|
completion_text = ""
|
|
|
|
try:
|
|
async with session.post(
|
|
f"{self.valves.llamacpp_url}/v1/chat/completions",
|
|
json=llm_payload,
|
|
timeout=aiohttp.ClientTimeout(total=self.valves.llm_timeout),
|
|
) as resp:
|
|
if resp.status != 200:
|
|
error_text = await resp.text()
|
|
yield f"LLM error: HTTP {resp.status} - {error_text}"
|
|
return
|
|
|
|
async for line in resp.content:
|
|
if not line:
|
|
continue
|
|
|
|
line_str = line.decode("utf-8").strip()
|
|
if not line_str or line_str.startswith(":"):
|
|
continue
|
|
|
|
if line_str.startswith("data: "):
|
|
line_str = line_str[6:]
|
|
|
|
if line_str == "[DONE]":
|
|
break
|
|
|
|
try:
|
|
data = json.loads(line_str)
|
|
delta = data.get("choices", [{}])[0].get("delta", {})
|
|
if content := delta.get("content"):
|
|
yield content
|
|
completion_text += content
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
except asyncio.TimeoutError:
|
|
yield "\n\n⚠️ LLM timed out"
|
|
return
|
|
except Exception as e:
|
|
yield f"\n\nLLM error: {e}"
|
|
return
|
|
|
|
# Estimate completion tokens
|
|
completion_tokens = self._estimate_tokens(completion_text)
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Sources
|
|
# ─────────────────────────────────────────────
|
|
if self.valves.show_sources:
|
|
# We now track 'indices' list along with the count
|
|
source_counts: dict[str, dict] = {}
|
|
|
|
# 'chunks' is still available from Step 4/Step 3
|
|
for i, chunk in enumerate(chunks, 1):
|
|
payload = chunk.get("payload", {})
|
|
path = payload.get("filePath", "")
|
|
name = payload.get("fileName", "Unknown")
|
|
|
|
if path in source_counts:
|
|
source_counts[path]["count"] += 1
|
|
source_counts[path]["indices"].append(i)
|
|
else:
|
|
source_counts[path] = {
|
|
"name": name,
|
|
"path": path,
|
|
"count": 1,
|
|
"indices": [i],
|
|
}
|
|
|
|
yield "\n\n---\n**Sources:**\n"
|
|
for src in source_counts.values():
|
|
vault = urllib.parse.quote(self.valves.vault_name)
|
|
path = urllib.parse.quote(src["path"])
|
|
uri = f"obsidian://open?vault={vault}&file={path}"
|
|
|
|
# Format indices like: [1, 2, 5]
|
|
indices_str = ", ".join(map(str, src["indices"]))
|
|
|
|
yield f"- [{src['name']}]({uri}) (Chunks: {indices_str})\n"
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Stats
|
|
# ─────────────────────────────────────────────
|
|
if self.valves.show_stats:
|
|
yield f"\n*~{total_prompt_tokens:,} in / ~{completion_tokens:,} out (estimated)*"
|
|
|