admin/OpenJarvis
1
0
Fork 0
mirror of https://github.com/open-jarvis/OpenJarvis.git synced 2026-07-30 02:42:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Deep Research Granola integration + citation fixes (#366)

Browse Source
This commit is contained in:
Robby Manihani
2026-05-20 15:56:42 -07:00
committed by GitHub
parent 75cb70526e
commit 363bdbfa41
7 changed files with 779 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/openjarvis/agents/research_loop.py
+182 -23
View File
@@ -17,10 +17,11 @@ from __future__ import annotations
import json
import logging
import re
import sys
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional
from typing import Any, Callable, Dict, List, Optional, Tuple
from openjarvis.connectors.hybrid_search import HybridSearch, SearchHit
from openjarvis.core.types import Message, Role, ToolCall
@@ -102,7 +103,15 @@ SEARCH_TOOL_SPEC: Dict[str, Any] = {
},
"sources": {
"type": "array",
"description": "Restrict to these connectors (e.g. ['gmail']).",
"description": (
"Restrict the search to one or more connectors. Use this "
"whenever the user names a data source (e.g. \"in my "
"Granola notes\" → ['granola']; \"check Slack and Gmail\" "
"→ ['slack', 'gmail']). Valid IDs include: gmail, slack, "
"granola, notion, obsidian, gcalendar, gdrive, gmail_imap, "
"outlook, imessage, whatsapp, apple_notes, apple_contacts, "
"gcontacts, google_tasks, github_notifications."
),
"items": {"type": "string"},
},
"limit": {
@@ -117,7 +126,12 @@ SEARCH_TOOL_SPEC: Dict[str, Any] = {
}
SYSTEM_PROMPT = """You are a research assistant with access to the user's personal knowledge corpus (their email, notes, calendar). You answer questions by calling two tools:
SYSTEM_PROMPT = """You are a research assistant with access to the user's personal knowledge corpus.
The user's corpus contains data from these sources only:
{available_sources}
You answer questions by calling two tools:
search(query, person=None, time_range=None, sources=None, limit=20)
clarify(question)
@@ -126,10 +140,12 @@ Strategy:
1. If the user names a person, ALWAYS pass `person=` rather than relying on lexical match. Hybrid search will fuzzy-match name or address fragments.
2. When the user mentions ANY time window — "this past week", "recently", "last month", "past few days", "yesterday" — you MUST translate it to a `time_range` parameter. Today is {today}.
3. The `time_range` argument is a JSON object: `{{"start": "<ISO 8601>", "end": "<ISO 8601>"}}`. Either bound may be omitted, but pass at least one whenever the user gave you a temporal cue.
4. If the first structured search returns nothing useful, broaden with a semantic query and drop filters one at a time.
5. You have a clarify tool. Only use it AFTER at least one search attempt. Use it when: you found multiple ambiguous matches (e.g. 3 different people named John), search returned zero results and the query might need reframing, or the scope is too broad to synthesize meaningfully. Never use clarify before searching — always try first.
6. After receiving a clarify response, use the information to construct a precise search with the correct person, time_range, and query parameters. Never send an empty query or a query with no parameters — extract every concrete signal from the user's reply (names, dates, topics) and put it on the call.
7. Tool calls — search AND clarify — share a budget of 5 total. Spend wisely.
4. When the user names a specific data source — "my Granola notes", "in Slack", "from my email" — you MUST pass `sources=[...]` with the matching connector ID. Only use IDs that appear in the connected-sources list above; do NOT invent or assume sources that are not connected. Common synonyms: "meeting notes"/"meetings"/"transcripts" → granola; "email"/"inbox" → gmail; "DMs"/"channels" → slack. Without this filter the search returns mail/messages ABOUT a tool instead of records FROM that tool.
4a. Never apologize about sources that aren't in the connected-sources list — if the user asks about "Notion" but Notion isn't connected, just say "Notion isn't connected, but here's what I found in {available_sources}" and answer from what is available.
5. If the first structured search returns nothing useful, broaden with a semantic query and drop filters one at a time.
6. You have a clarify tool. Only use it AFTER at least one search attempt. Use it when: you found multiple ambiguous matches (e.g. 3 different people named John), search returned zero results and the query might need reframing, or the scope is too broad to synthesize meaningfully. Never use clarify before searching — always try first.
7. After receiving a clarify response, use the information to construct a precise search with the correct person, time_range, sources, and query parameters. Never send an empty query or a query with no parameters — extract every concrete signal from the user's reply (names, dates, topics, sources) and put it on the call.
8. Tool calls — search AND clarify — share a budget of 5 total. Spend wisely.
Synthesis rules:
- Cite sources as individual numbers in square brackets. Always separate — write [4] [7] [20], never [4, 7, 20]. Never format citations as markdown links. Just the number in brackets: [1]. The `ref` field on each hit is the citation number.
@@ -161,21 +177,24 @@ def shape_results_for_model(
detailed_top: int = 5,
thread_ctx_per_hit: int = 3,
total_cap: int = 20,
ref_offset: int = 0,
) -> Dict[str, Any]:
"""Compact a hit list into a JSON payload the planner can chew through.
The first ``detailed_top`` rows keep their content snippet and trimmed
thread context; the remainder are summarised to title + sender + date so
the planner still sees the breadth of what's available without blowing the
context window. Each hit gets a 1-indexed numeric ``ref`` so the synthesis
can cite it as ``[N]``.
context window. Each hit gets a numeric ``ref`` (1-indexed, plus
``ref_offset``) so the synthesis can cite it as ``[N]``. The offset lets
multi-search runs hand the planner globally unique refs across calls so
a later renumbering pass can dedupe by first appearance.
"""
out_hits: List[Dict[str, Any]] = []
visible = hits[:total_cap]
for i, h in enumerate(visible):
sender = h.participants[0] if h.participants else ""
base = {
"ref": i + 1,
"ref": i + 1 + ref_offset,
"title": h.title,
"sender": sender,
"timestamp": h.timestamp,
@@ -223,11 +242,15 @@ def _bare_doc_id(source: str, document_id: str) -> str:
def _hit_url(source: str, document_id: str) -> str:
"""Build a clickable URL for a hit, when we know how to link it.
"""Reconstruct a clickable URL from a hit's ``doc_id`` alone.
Gmail and Slack are the linkable sources today; anything else falls
back to an empty string and the client renders a non-clickable
citation chip.
Used as a *fallback* when the connector didn't persist a URL on the
chunk (``SearchHit.url`` is empty). Reconstruction only works for
sources whose doc_id encodes everything the permalink needs — Gmail
and Slack today. Sources whose doc_id is just an opaque ID (e.g.
Granola, where the web URL uses a different UUID than the API note_id)
must populate ``Document.url`` at ingest time; we can't make a working
link from the doc_id alone.
Gmail ids land here in two flavors:
@@ -275,10 +298,76 @@ def _hit_url(source: str, document_id: str) -> str:
return ""
_CITE_RE = re.compile(r"\[(\d+)\]")
def renumber_citations(
text: str,
ref_to_source: Dict[int, Dict[str, Any]],
) -> Tuple[str, List[Dict[str, Any]]]:
"""Renumber ``[N]`` citations in ``text`` by first-appearance order.
The planner sees globally-offset refs across multiple search calls
(search 1 returns 1..20, search 2 returns 21..40, …). When the
synthesis arrives, the first ref the model actually cited becomes
``[1]``, the second unique one becomes ``[2]``, and so on. Repeats
map to the same new ref. Refs the synthesis never cites are dropped
from the returned ``sources`` list — only the ones the user can
actually click on get carried through.
Parameters
----------
text:
Synthesis text containing inline ``[N]`` references.
ref_to_source:
Mapping from the original (offset) ref to the source dict that
``build_sources_for_client`` produced for that hit.
Returns
-------
(new_text, ordered_sources)
``new_text`` has every cited ``[N]`` rewritten to its new
sequence number. ``ordered_sources`` is the deduped list of
source dicts in the order they appear in the synthesis, each
with its ``ref`` field set to the new sequence number.
"""
old_to_new: Dict[int, int] = {}
ordered: List[Dict[str, Any]] = []
for m in _CITE_RE.finditer(text):
try:
old = int(m.group(1))
except ValueError:
continue
if old in old_to_new:
continue
src = ref_to_source.get(old)
if src is None:
# Synthesis cited a ref that doesn't exist in the corpus —
# leave the literal text alone, drop the source entry.
continue
new_ref = len(ordered) + 1
old_to_new[old] = new_ref
renumbered_src = dict(src)
renumbered_src["ref"] = new_ref
ordered.append(renumbered_src)
def _replace(match: "re.Match[str]") -> str:
try:
old = int(match.group(1))
except ValueError:
return match.group(0)
new = old_to_new.get(old)
return f"[{new}]" if new is not None else match.group(0)
new_text = _CITE_RE.sub(_replace, text)
return new_text, ordered
def build_sources_for_client(
hits: List[SearchHit],
*,
total_cap: int = 20,
ref_offset: int = 0,
) -> List[Dict[str, Any]]:
"""Produce the citation-friendly sources list streamed to the frontend.
@@ -291,14 +380,21 @@ def build_sources_for_client(
out: List[Dict[str, Any]] = []
for i, h in enumerate(hits[:total_cap]):
sender = h.participants[0] if h.participants else ""
# Prefer the URL the connector stored at ingest time (Granola's
# ``web_url``, Notion's page URL, etc.) — it's the only reliable
# link for sources whose web URL doesn't derive from the doc_id.
# Fall back to the doc_id-based reconstruction for sources where
# that still works (Slack, Gmail).
url = h.url or _hit_url(h.source, h.document_id)
out.append(
{
"ref": i + 1,
"ref": i + 1 + ref_offset,
"title": h.title,
"sender": sender,
"date": _hit_date(h.timestamp),
"source": h.source,
"source_id": _bare_doc_id(h.source, h.document_id),
"url": _hit_url(h.source, h.document_id),
"url": url,
}
)
return out
@@ -389,6 +485,7 @@ class ResearchAgent:
num_ctx: int = 16384,
clarify_handler: Optional[Callable[[str], str]] = None,
on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
available_sources: Optional[List[str]] = None,
) -> None:
self._engine = engine
self._search = search
@@ -399,6 +496,10 @@ class ResearchAgent:
self._num_ctx = int(num_ctx)
self._clarify_handler = clarify_handler or _default_clarify_handler
self._on_event = on_event
# Explicit list wins; otherwise we'll discover sources from the
# KnowledgeStore on each run() call so the prompt stays accurate
# even as the user connects new connectors mid-session.
self._available_sources_override = available_sources
def _emit(self, event: Dict[str, Any]) -> None:
"""Fire ``self._on_event`` if set; swallow callback errors."""
@@ -485,17 +586,56 @@ class ResearchAgent:
# Loop
# ------------------------------------------------------------------
def _resolve_available_sources(self) -> List[str]:
"""Return the source IDs the user actually has data for.
Override > live query of the KnowledgeStore. Failure to read the
store (e.g. no _store attribute on the search backend) returns
``[]`` so the prompt still formats — better empty than crashing.
"""
if self._available_sources_override is not None:
return list(self._available_sources_override)
store = getattr(self._search, "_store", None)
if store is None:
return []
try:
return list(store.distinct_sources())
except Exception as exc: # noqa: BLE001
logger.debug("distinct_sources() failed: %s", exc)
return []
def run(self, query: str) -> ResearchResult:
"""Run the loop end-to-end and return the synthesis plus a trace."""
sources_list = self._resolve_available_sources()
if sources_list:
sources_blurb = ", ".join(sources_list)
else:
sources_blurb = (
"(no connected sources — tell the user to connect a "
"connector before searching)"
)
sys_msg = Message(
role=Role.SYSTEM,
content=SYSTEM_PROMPT.format(today=datetime.now().isoformat(timespec="minutes")),
content=SYSTEM_PROMPT.format(
today=datetime.now().isoformat(timespec="minutes"),
available_sources=sources_blurb,
),
)
messages: List[Message] = [sys_msg, Message(role=Role.USER, content=query)]
invocations: List[ToolInvocation] = []
total_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
# Global ref counter: each search increments by the number of hits
# it returned so the planner sees unique refs across calls. The
# accumulator lets us renumber whatever the synthesis cites at the
# end into a single deduped client-facing sources list.
next_ref: int = 1
ref_to_source: Dict[int, Dict[str, Any]] = {}
def _finalize(text: str) -> Tuple[str, List[Dict[str, Any]]]:
return renumber_citations(text, ref_to_source)
iterations = 0
for _ in range(self._max_iterations + 1):
iterations += 1
@@ -520,8 +660,14 @@ class ResearchAgent:
if not tool_calls_raw:
if content.strip():
answer = content.strip()
self._emit({"type": "final_answer", "text": answer})
answer, final_sources = _finalize(content.strip())
self._emit(
{
"type": "final_answer",
"text": answer,
"sources": final_sources,
}
)
return ResearchResult(
answer=answer,
iterations=iterations,
@@ -542,7 +688,9 @@ class ResearchAgent:
)
continue
fallback = "(model returned no content and no tool calls)"
self._emit({"type": "final_answer", "text": fallback})
self._emit(
{"type": "final_answer", "text": fallback, "sources": []}
)
return ResearchResult(
answer=fallback,
iterations=iterations,
@@ -579,16 +727,24 @@ class ResearchAgent:
self._emit({"type": "search_call", "arguments": args})
inv = self._execute_search(args)
invocations.append(inv)
offset = next_ref - 1
sources_for_search = build_sources_for_client(
inv.raw_hits, ref_offset=offset
)
self._emit(
{
"type": "search_result",
"num_hits": inv.num_results,
"top_titles": inv.top_titles,
"sources": build_sources_for_client(inv.raw_hits),
"sources": sources_for_search,
}
)
for src in sources_for_search:
ref_to_source[int(src["ref"])] = src
next_ref += len(sources_for_search)
tool_output = json.dumps(
shape_results_for_model(inv.raw_hits), ensure_ascii=False
shape_results_for_model(inv.raw_hits, ref_offset=offset),
ensure_ascii=False,
)
elif name == "clarify":
# Enforce the "search first" rule at runtime so we don't
@@ -686,7 +842,10 @@ class ResearchAgent:
"(no synthesis available — the search budget was exhausted "
"and the model returned no text response)"
)
self._emit({"type": "final_answer", "text": answer})
answer, final_sources = _finalize(answer)
self._emit(
{"type": "final_answer", "text": answer, "sources": final_sources}
)
return ResearchResult(
answer=answer,
iterations=iterations,
src/openjarvis/connectors/granola.py
+26 -8
View File
@@ -10,6 +10,7 @@ Settings → API (requires Business or Enterprise plan).
from __future__ import annotations
import logging
from datetime import datetime, timezone
from typing import Any, Dict, Iterator, List, Optional
@@ -21,6 +22,8 @@ from openjarvis.core.config import DEFAULT_CONFIG_DIR
from openjarvis.core.registry import ConnectorRegistry
from openjarvis.tools._stubs import ToolSpec
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
@@ -286,6 +289,7 @@ class GranolaConnector(BaseConnector):
created_after=created_after,
)
notes: List[Dict[str, Any]] = list_resp.get("notes", [])
logger.info("Granola: Found %d notes on this page", len(notes))
for note_summary in notes:
note_id: str = note_summary.get("id", "")
@@ -297,11 +301,18 @@ class GranolaConnector(BaseConnector):
title: str = note.get("title", "")
owner: Dict[str, Any] = note.get("owner") or {}
author: str = owner.get("email", "")
author: str = (owner.get("email") or "").lower()
attendees: List[Dict[str, Any]] = note.get("attendees") or []
participants: List[str] = [
a.get("email", "") for a in attendees if a.get("email")
(a.get("email") or "").lower()
for a in attendees
if a.get("email")
]
participants_raw: List[str] = [
a.get("name") or a.get("email") or ""
for a in attendees
if a.get("name") or a.get("email")
]
created_at_str: str = note.get("created_at", "")
@@ -309,11 +320,14 @@ class GranolaConnector(BaseConnector):
content = _format_note_content(note)
# Build URL from calendar event if available, else None
cal_event: Optional[Dict[str, Any]] = note.get("calendar_event")
url: Optional[str] = None
if cal_event:
url = note.get("url")
cal_event: Dict[str, Any] = note.get("calendar_event") or {}
channel: str = cal_event.get("event_title") or "meeting"
# ``web_url`` (e.g. https://notes.granola.ai/d/{uuid}) is the
# only reliable way to deep-link to a Granola note — the API
# ``note_id`` and the web UUID are different, so we must store
# what the API gives us here.
web_url: Optional[str] = note.get("web_url") or None
doc = Document(
doc_id=f"granola:{note_id}",
@@ -323,8 +337,11 @@ class GranolaConnector(BaseConnector):
title=title,
author=author,
participants=participants,
participants_raw=participants_raw,
channel=channel,
thread_id=note_id,
timestamp=timestamp,
url=url,
url=web_url,
metadata={
"note_id": note_id,
"owner_name": owner.get("name", ""),
@@ -343,6 +360,7 @@ class GranolaConnector(BaseConnector):
self._items_synced = synced
self._last_sync = datetime.now(tz=timezone.utc)
logger.info("Granola: Sync complete, %d notes total", synced)
def sync_status(self) -> SyncStatus:
"""Return sync progress from the most recent :meth:`sync` call."""
src/openjarvis/connectors/hybrid_search.py
+7 -1
View File
@@ -54,6 +54,11 @@ class SearchHit:
vector_score: float
thread_id: str = ""
thread_context: List[Dict[str, Any]] = field(default_factory=list)
# ``url`` is the connector-provided deep-link, persisted on
# ``knowledge_chunks.url``. Empty when the source didn't supply one — in
# that case callers may fall back to a doc_id-based reconstruction (Slack,
# Gmail), or render the citation as non-clickable.
url: str = ""
def to_dict(self) -> Dict[str, Any]:
return {
@@ -435,7 +440,7 @@ class HybridSearch:
meta_rows = self._store._conn.execute(
f"""
SELECT id, doc_id, content, source, title, author, participants,
timestamp, thread_id, chunk_index
timestamp, thread_id, chunk_index, url
FROM knowledge_chunks
WHERE id IN ({placeholders})
""",
@@ -463,6 +468,7 @@ class HybridSearch:
vector_score=vec_score,
thread_id=r["thread_id"] or "",
thread_context=self._thread_context(r["thread_id"] or "", chunk_id),
url=r["url"] or "",
)
)
return hits
src/openjarvis/connectors/store.py
+15
View File
@@ -470,6 +470,21 @@ class KnowledgeStore(MemoryBackend):
row = self._conn.execute("SELECT COUNT(*) FROM knowledge_chunks").fetchone()
return row[0] if row else 0
def distinct_sources(self) -> List[str]:
"""Return the sorted list of distinct ``source`` values currently indexed.
Used by the research agent to populate the system prompt with the
sources the user actually has connected — so the model doesn't
mention "Notion" or "Apple Notes" when nothing from those sources
is in the corpus.
"""
rows = self._conn.execute(
"SELECT DISTINCT source FROM knowledge_chunks "
"WHERE source IS NOT NULL AND source != '' "
"ORDER BY source"
).fetchall()
return [r[0] for r in rows]
def close(self) -> None:
"""Close the underlying SQLite connection."""
try:
src/openjarvis/server/research_router.py
+18 -4
View File
@@ -25,7 +25,7 @@ import logging
import re
import threading
import time
from typing import Any, AsyncGenerator, Callable, Dict, Optional
from typing import Any, AsyncGenerator, Callable, Dict, List, Optional
from fastapi import APIRouter
from fastapi.responses import StreamingResponse
@@ -393,6 +393,7 @@ async def _stream_research(query: str, model: str) -> AsyncGenerator[str, None]:
final_answer: Optional[str] = None
final_usage: Dict[str, int] = {}
final_sources: List[Dict[str, Any]] = []
try:
while True:
event = await queue.get()
@@ -408,18 +409,29 @@ async def _stream_research(query: str, model: str) -> AsyncGenerator[str, None]:
continue
# We translate the agent's `final_answer` event into a stream of
# `synthesis` chunks so the client sees the answer materialize
# incrementally rather than as a single blob.
# incrementally rather than as a single blob. The accompanying
# ``sources`` array is the renumbered, deduped citation list the
# frontend should render under the final answer.
if etype == "final_answer":
final_answer = event.get("text", "")
final_sources = list(event.get("sources") or [])
for piece in _chunk_synthesis(final_answer or ""):
yield _sse({"type": "synthesis", "text": piece})
if final_sources:
yield _sse(
{"type": "final_sources", "sources": final_sources}
)
continue
yield _sse(event)
# If the agent thread crashed before producing a final answer, the
# client still gets the error frame (emitted above) followed by done.
yield _sse({"type": "done", "usage": final_usage})
# The done frame also carries the deduped sources so a client that
# only listens for ``done`` still gets the canonical citation list.
yield _sse(
{"type": "done", "usage": final_usage, "sources": final_sources}
)
except Exception as exc: # noqa: BLE001
# Consumer loop crashed unexpectedly (e.g. JSON serialization fault,
# logic bug). Surface a clean error frame rather than letting the
@@ -431,7 +443,9 @@ async def _stream_research(query: str, model: str) -> AsyncGenerator[str, None]:
"message": f"Research failed: {type(exc).__name__}: {exc}",
}
)
yield _sse({"type": "done", "usage": final_usage})
yield _sse(
{"type": "done", "usage": final_usage, "sources": final_sources}
)
finally:
# The worker may still be cleaning up (rarely) — make sure we don't
# leak a dangling task. Swallow any straggler exception so a worker
tests/agents/test_research_loop.py
+397 -1
View File
@@ -12,7 +12,16 @@ from unittest.mock import MagicMock
import pytest
from openjarvis.agents.research_loop import ResearchAgent, _hit_url
from openjarvis.agents.research_loop import (
SEARCH_TOOL_SPEC,
SYSTEM_PROMPT,
ResearchAgent,
_hit_url,
build_sources_for_client,
renumber_citations,
shape_results_for_model,
)
from openjarvis.connectors.hybrid_search import SearchHit
class _MockEngine:
@@ -213,3 +222,390 @@ def test_hit_url_unknown_source_returns_empty() -> None:
"""Unsupported sources don't get a guessed URL."""
assert _hit_url("notion", "notion:abc") == ""
assert _hit_url("", "") == ""
def _mk_hit(
*,
source: str = "granola",
document_id: str = "granola:not_abc12345678901",
url: str = "",
title: str = "Sprint Planning",
) -> SearchHit:
"""Tiny SearchHit factory for URL-routing tests."""
return SearchHit(
chunk_id="c1",
document_id=document_id,
chunk_idx=0,
title=title,
content_snippet="...",
source=source,
timestamp="2024-03-15T10:00:00",
participants=["alice@co.com"],
score=0.5,
bm25_score=0.5,
vector_score=0.5,
url=url,
)
def test_build_sources_prefers_stored_url_over_reconstruction() -> None:
"""When the connector stored a URL, the client gets it verbatim.
The doc_id-based reconstruction is a *fallback* — sources like Granola
that supply the URL at ingest time must surface it unchanged.
"""
stored = "https://notes.granola.ai/d/e98b5d85-ff57-46ac-a0ce-849fc68d086f"
sources = build_sources_for_client([_mk_hit(url=stored)])
assert sources[0]["url"] == stored
def test_build_sources_falls_back_to_reconstruction_when_url_missing() -> None:
"""Slack/Gmail still work without a stored URL — reconstructed from doc_id."""
sources = build_sources_for_client(
[
_mk_hit(
source="slack",
document_id="slack:acme:C123:1710500000.000100",
url="",
title="#general",
)
]
)
assert (
sources[0]["url"]
== "https://acme.slack.com/archives/C123/p1710500000000100"
)
def test_hit_url_granola_not_reconstructible() -> None:
"""Granola doc_ids cannot reconstruct a web URL.
The Granola web app uses a UUID that is different from the API
``note_id`` embedded in our doc_id. Citation URLs for Granola must
come from the stored ``SearchHit.url`` (populated at ingest time from
the API's ``web_url`` field). ``_hit_url`` therefore returns empty.
"""
assert _hit_url("granola", "granola:not_abc12345678901") == ""
assert _hit_url("granola", "granola:") == ""
assert _hit_url("granola", "") == ""
# ---------------------------------------------------------------------------
# Sources filter — propagated through _execute_search to HybridSearch.search
# ---------------------------------------------------------------------------
def test_search_with_sources_filter_is_passed_through(
stub_search: MagicMock,
) -> None:
"""A search tool call carrying ``sources=['granola']`` reaches HybridSearch.
Without this propagation, "tell me about my Granola notes" returns Gmail
emails that mention Granola instead of actual meeting notes.
"""
engine = _MockEngine(
responses=[
{
"content": "",
"tool_calls": [
{
"id": "s1",
"name": "search",
"arguments": json.dumps(
{
"query": "recent meetings",
"sources": ["granola"],
}
),
}
],
"usage": {},
},
_text_response("Here are your recent meetings."),
]
)
agent = ResearchAgent(engine, stub_search, model="mock", max_iterations=2)
result = agent.run("tell me about my recent meetings from Granola")
stub_search.search.assert_called_once()
kwargs = stub_search.search.call_args.kwargs
assert kwargs.get("sources") == ["granola"]
# And the loop terminates with the synthesized answer.
assert "Here are your recent meetings." in result.answer
def test_search_sources_coerces_scalar_to_list(stub_search: MagicMock) -> None:
"""If the model sends ``sources='slack'`` (string), wrap it as ``['slack']``."""
engine = _MockEngine(
responses=[
{
"content": "",
"tool_calls": [
{
"id": "s1",
"name": "search",
"arguments": json.dumps(
{"query": "anything", "sources": "slack"}
),
}
],
"usage": {},
},
_text_response("done"),
]
)
agent = ResearchAgent(engine, stub_search, model="mock", max_iterations=2)
agent.run("anything in slack")
kwargs = stub_search.search.call_args.kwargs
assert kwargs.get("sources") == ["slack"]
# ---------------------------------------------------------------------------
# Prompt + tool schema — the planner must see the sources directive
# ---------------------------------------------------------------------------
def test_tool_schema_sources_lists_known_connectors() -> None:
"""The ``sources`` parameter description enumerates the common connector IDs.
Without explicit IDs in the description the planner makes up names like
"Granola" or "Slack workspace" instead of the lowercase connector IDs the
backend filter actually matches against.
"""
sources_prop = SEARCH_TOOL_SPEC["function"]["parameters"]["properties"]["sources"]
desc = sources_prop["description"]
for connector_id in ("granola", "slack", "gmail", "notion"):
assert connector_id in desc
def test_system_prompt_mandates_sources_extraction() -> None:
"""The system prompt has a directive telling the planner to extract sources.
Without it, the model treats "from my Granola notes" as a topical cue
rather than a hard filter, and returns email about Granola instead of
Granola records.
"""
assert "sources=" in SYSTEM_PROMPT
# Synonym mapping for the most-common alias.
assert "granola" in SYSTEM_PROMPT.lower()
assert "meeting notes" in SYSTEM_PROMPT.lower()
# The dynamic placeholder is what's interpolated per-run.
assert "{available_sources}" in SYSTEM_PROMPT
# ---------------------------------------------------------------------------
# Dynamic available_sources — only list what the user actually has connected
# ---------------------------------------------------------------------------
def test_available_sources_override_appears_in_prompt(
stub_search: MagicMock,
) -> None:
"""An explicit override is injected into the system prompt verbatim.
Without this the agent will reference disconnected sources ("I couldn't
find anything in Notion or Apple Notes") even when those connectors
have never been wired up.
"""
engine = _MockEngine(responses=[_text_response("ok")])
agent = ResearchAgent(
engine,
stub_search,
model="mock",
max_iterations=1,
available_sources=["gmail", "slack", "granola"],
)
agent.run("hi")
sys_content = engine.calls[0]["messages"][0].content
# The connected-sources blurb is interpolated verbatim.
assert "gmail, slack, granola" in sys_content
# And the prompt no longer hard-codes a static "Valid IDs include..."
# enumeration of unconnected sources — those that aren't in the
# available_sources list shouldn't appear in any *listing*. (Rule 4a
# still uses Notion as a *narrative example* of how to handle an
# unconnected source, which is intentional.)
assert "obsidian" not in sys_content.lower()
assert "apple_notes" not in sys_content.lower()
assert "gdrive" not in sys_content.lower()
def test_available_sources_fall_back_to_store(stub_search: MagicMock) -> None:
"""When no override is given, the agent queries the KnowledgeStore.
Mirrors the live wiring used by the SSE research router: a HybridSearch
instance that exposes a ``_store`` with ``distinct_sources()``.
"""
fake_store = MagicMock()
fake_store.distinct_sources.return_value = ["granola", "slack"]
stub_search._store = fake_store
engine = _MockEngine(responses=[_text_response("ok")])
agent = ResearchAgent(engine, stub_search, model="mock", max_iterations=1)
agent.run("hi")
fake_store.distinct_sources.assert_called_once()
sys_content = engine.calls[0]["messages"][0].content
assert "granola, slack" in sys_content
def test_available_sources_empty_message_when_nothing_connected(
stub_search: MagicMock,
) -> None:
"""No connected sources surfaces a clear message instead of a stray {}.
A missing placeholder would crash `str.format`; a broken format would
leave a raw ``{available_sources}`` in the prompt. Both are bad UX —
the test pins the friendly fallback string.
"""
stub_search._store = None
engine = _MockEngine(responses=[_text_response("ok")])
agent = ResearchAgent(
engine,
stub_search,
model="mock",
max_iterations=1,
available_sources=[],
)
agent.run("hi")
sys_content = engine.calls[0]["messages"][0].content
assert "no connected sources" in sys_content
assert "{available_sources}" not in sys_content
# ---------------------------------------------------------------------------
# Ref offsetting + citation renumbering
# ---------------------------------------------------------------------------
def test_shape_results_for_model_respects_ref_offset() -> None:
"""Multi-search runs hand the model globally unique refs.
Without offsets, two searches each emit refs 1..20 and the model can't
disambiguate ``[5]`` across calls — which breaks the renumbering pass
that runs over the final synthesis.
"""
hits = [_mk_hit(title=f"t{i}", document_id=f"granola:{i}") for i in range(3)]
shaped = shape_results_for_model(hits, ref_offset=20)
refs = [h["ref"] for h in shaped["hits"]]
assert refs == [21, 22, 23]
def test_build_sources_for_client_respects_ref_offset() -> None:
"""``build_sources_for_client`` agrees with ``shape_results_for_model``."""
hits = [_mk_hit(title=f"t{i}", document_id=f"granola:{i}") for i in range(3)]
sources = build_sources_for_client(hits, ref_offset=10)
assert [s["ref"] for s in sources] == [11, 12, 13]
# And the connector source is carried so the renumbered final list
# can show "Granola • Sprint Planning" style chips.
assert all(s["source"] == "granola" for s in sources)
def test_renumber_citations_first_appearance_order() -> None:
"""Citations are renumbered by first-appearance order in the text."""
text = "First [7]. Then [3] and again [7]. Finally [12]."
ref_to_source = {
3: {"ref": 3, "title": "B"},
7: {"ref": 7, "title": "A"},
12: {"ref": 12, "title": "C"},
}
new_text, sources = renumber_citations(text, ref_to_source)
# Order of appearance: 7 → 1, 3 → 2, 12 → 3 (the repeat of 7 keeps its ref).
assert new_text == "First [1]. Then [2] and again [1]. Finally [3]."
assert [s["ref"] for s in sources] == [1, 2, 3]
assert [s["title"] for s in sources] == ["A", "B", "C"]
def test_renumber_citations_drops_uncited_sources() -> None:
"""Sources the synthesis never cited are excluded from the final list.
The frontend would otherwise show citation chips for hits the model
silently ignored, which clutters the panel and misleads about what
the answer actually relied on.
"""
text = "Only one citation here [5]."
ref_to_source = {
1: {"ref": 1, "title": "uncited-A"},
5: {"ref": 5, "title": "cited"},
9: {"ref": 9, "title": "uncited-B"},
}
new_text, sources = renumber_citations(text, ref_to_source)
assert new_text == "Only one citation here [1]."
assert len(sources) == 1
assert sources[0]["title"] == "cited"
def test_renumber_citations_unknown_ref_left_alone() -> None:
"""A ``[N]`` whose ref isn't in the map is left as-is.
Defensive: a hallucinated citation shouldn't blow up renumbering or
silently disappear — the user sees the broken cite and can ask why.
"""
text = "Real [3], hallucinated [99]."
ref_to_source = {3: {"ref": 3, "title": "Real"}}
new_text, sources = renumber_citations(text, ref_to_source)
assert new_text == "Real [1], hallucinated [99]."
assert [s["title"] for s in sources] == ["Real"]
# ---------------------------------------------------------------------------
# End-to-end: final_answer event carries renumbered text + sources
# ---------------------------------------------------------------------------
def test_final_answer_event_carries_renumbered_sources(
stub_search: MagicMock,
) -> None:
"""The ``final_answer`` event includes the deduped, renumbered sources.
Drives the same renumbering pipeline that's wired into the SSE router's
``done`` frame — when this test passes the frontend will receive a
clean ``[1]..[K]`` numbering aligned with a single sources list.
"""
# Two hits returned by the search. The synthesis cites the second one
# twice and the first one once, in the order [2] [1] [2].
hit_a = _mk_hit(title="A", document_id="granola:a", url="https://a")
hit_b = _mk_hit(title="B", document_id="granola:b", url="https://b")
stub_search.search.return_value = [hit_a, hit_b]
engine = _MockEngine(
responses=[
{
"content": "",
"tool_calls": [
{
"id": "s1",
"name": "search",
"arguments": json.dumps({"query": "topic"}),
}
],
"usage": {},
},
_text_response("First [2], then [1], then [2] again."),
]
)
captured: list[dict] = []
def on_event(ev: dict) -> None:
captured.append(ev)
agent = ResearchAgent(
engine,
stub_search,
model="mock",
max_iterations=2,
on_event=on_event,
available_sources=["granola"],
)
result = agent.run("anything")
# The synthesis is renumbered by first appearance: [2]→[1], [1]→[2].
assert result.answer == "First [1], then [2], then [1] again."
final = next(ev for ev in captured if ev["type"] == "final_answer")
assert final["text"] == "First [1], then [2], then [1] again."
# Two cited sources, in the order they appeared in the synthesis.
assert [s["ref"] for s in final["sources"]] == [1, 2]
assert [s["title"] for s in final["sources"]] == ["B", "A"]
tests/connectors/test_granola.py
+134 -4
View File
@@ -6,6 +6,7 @@ All Granola API calls are mocked; no network access is required.
from __future__ import annotations
import json
import logging
from pathlib import Path
from typing import List
from unittest.mock import patch
@@ -40,6 +41,8 @@ _LIST_RESPONSE = {
"cursor": None,
}
_NOTE_1_WEB_URL = "https://notes.granola.ai/d/e98b5d85-ff57-46ac-a0ce-849fc68d086f"
_NOTE_1 = {
"id": "not_abc12345678901",
"title": "Sprint Planning",
@@ -69,6 +72,7 @@ _NOTE_1 = {
"event_title": "Sprint Planning",
"scheduled_start": "2024-03-15T10:00:00Z",
},
"web_url": _NOTE_1_WEB_URL,
}
_NOTE_2 = {
@@ -81,6 +85,8 @@ _NOTE_2 = {
"summary": {"markdown": "Reviewed new dashboard mockups."},
"transcript": [],
"calendar_event": None,
# Some notes have no associated web_url (e.g. quick notes without a
# calendar event). The connector must tolerate the absence.
}
# ---------------------------------------------------------------------------
@@ -157,22 +163,38 @@ def test_sync_yields_documents(
assert len(docs) == 2
# --- Note 1 ---
# --- Note 1 (has calendar_event + attendees) ---
doc1 = next(d for d in docs if d.doc_id == "granola:not_abc12345678901")
assert doc1.source == "granola"
assert doc1.doc_type == "document"
assert doc1.title == "Sprint Planning"
assert doc1.author == "alice@co.com"
assert "alice@co.com" in doc1.participants
assert "carol@co.com" in doc1.participants
# Participants are lowercased emails (cross-source matching).
assert doc1.participants == ["alice@co.com", "carol@co.com"]
# participants_raw keeps the human-readable names.
assert doc1.participants_raw == ["Alice", "Carol"]
# channel is derived from calendar_event.event_title when present.
assert doc1.channel == "Sprint Planning"
# thread_id namespaces transcript chunks under one note.
assert doc1.thread_id == "not_abc12345678901"
# The connector persists the API-provided web_url verbatim. The web URL
# uses a different UUID than the API note_id so this is the *only* way
# to get a working deep-link to the note.
assert doc1.url == _NOTE_1_WEB_URL
assert "Discussed sprint goals and capacity." in doc1.content
assert "Let's start with the sprint goals." in doc1.content
assert "I think we should focus on auth." in doc1.content
# --- Note 2 ---
# --- Note 2 (no calendar_event → channel falls back to "meeting") ---
doc2 = next(d for d in docs if d.doc_id == "granola:not_def12345678901")
assert doc2.title == "Design Review"
assert doc2.author == "bob@co.com"
assert doc2.participants == ["bob@co.com"]
assert doc2.participants_raw == ["Bob"]
assert doc2.channel == "meeting"
assert doc2.thread_id == "not_def12345678901"
# Note 2's fixture has no web_url; the connector tolerates the absence.
assert doc2.url is None
assert "Reviewed new dashboard mockups." in doc2.content
# Verify the API was called correctly
@@ -245,3 +267,111 @@ def test_registry() -> None:
assert ConnectorRegistry.contains("granola")
cls = ConnectorRegistry.get("granola")
assert cls.connector_id == "granola"
# ---------------------------------------------------------------------------
# Test 9 — sync logs the per-page and total note counts at INFO
# ---------------------------------------------------------------------------
@patch("openjarvis.connectors.granola._granola_api_list_notes")
@patch("openjarvis.connectors.granola._granola_api_get_note")
def test_sync_logs_note_count(
mock_get,
mock_list,
connector,
caplog,
) -> None:
"""sync() emits 'Found N notes' and 'Sync complete' INFO lines.
Matches the Slack connector's per-sync logging shape so operators can
grep server logs for sync activity.
"""
creds_path = Path(connector._credentials_path)
creds_path.parent.mkdir(parents=True, exist_ok=True)
creds_path.write_text(json.dumps({"token": "grl_fake"}), encoding="utf-8")
mock_list.return_value = _LIST_RESPONSE
mock_get.side_effect = [_NOTE_1, _NOTE_2]
with caplog.at_level(logging.INFO, logger="openjarvis.connectors.granola"):
list(connector.sync())
text = caplog.text
assert "Granola: Found 2 notes on this page" in text
assert "Granola: Sync complete, 2 notes total" in text
# ---------------------------------------------------------------------------
# Test 10 — end-to-end: connector → pipeline → KnowledgeStore → HybridSearch
# ---------------------------------------------------------------------------
@patch("openjarvis.connectors.granola._granola_api_list_notes")
@patch("openjarvis.connectors.granola._granola_api_get_note")
def test_end_to_end_ingest_and_search(
mock_get,
mock_list,
connector,
tmp_path: Path,
) -> None:
"""Synced Granola notes are searchable via HybridSearch with v1 fields.
Lexical-only path (no embedder) so this stays a pure unit test — no
Ollama daemon needed. Confirms the v1 contract end-to-end: source,
namespaced thread_id, channel, participants, and that the research-loop
URL builder reconstructs the Granola web deep-link from the doc_id.
"""
from openjarvis.connectors.hybrid_search import HybridSearch # noqa: PLC0415
from openjarvis.connectors.pipeline import IngestionPipeline # noqa: PLC0415
from openjarvis.connectors.store import KnowledgeStore # noqa: PLC0415
creds_path = Path(connector._credentials_path)
creds_path.parent.mkdir(parents=True, exist_ok=True)
creds_path.write_text(json.dumps({"token": "grl_fake"}), encoding="utf-8")
mock_list.return_value = _LIST_RESPONSE
mock_get.side_effect = [_NOTE_1, _NOTE_2]
store = KnowledgeStore(db_path=tmp_path / "granola_e2e.db")
pipeline = IngestionPipeline(store)
chunks_stored = pipeline.ingest(connector.sync())
# Two notes; the chunker may split summary/transcript sections, so
# we just assert at least one chunk per note made it in.
assert chunks_stored >= 2
hybrid = HybridSearch(store)
hits = hybrid.search("sprint goals", limit=5)
assert len(hits) >= 1
target = next(
(h for h in hits if "Sprint Planning" in h.title),
None,
)
assert target is not None
assert target.source == "granola"
assert target.title == "Sprint Planning"
# thread_id is namespaced by the pipeline.
assert target.thread_id == "granola:not_abc12345678901"
assert target.participants == ["alice@co.com", "carol@co.com"]
assert target.document_id == "granola:not_abc12345678901"
# The connector-supplied web_url survives ingest → store → hit and is
# what the research-loop client sees as the citation URL. The doc_id-
# based reconstruction can't recover this because the web UUID is
# different from the API note_id.
assert target.url == _NOTE_1_WEB_URL
from openjarvis.agents.research_loop import ( # noqa: PLC0415
_hit_url,
build_sources_for_client,
)
# _hit_url alone cannot reconstruct a Granola URL — there is no UUID
# in the doc_id. It must return empty, leaving the URL to be sourced
# from the stored ``SearchHit.url``.
assert _hit_url(target.source, target.document_id) == ""
# And the client-facing sources list does end up with the stored URL.
client_sources = build_sources_for_client([target])
assert client_sources[0]["url"] == _NOTE_1_WEB_URL