diff --git a/src/openjarvis/agents/research_loop.py b/src/openjarvis/agents/research_loop.py index 5831733d..505544c2 100644 --- a/src/openjarvis/agents/research_loop.py +++ b/src/openjarvis/agents/research_loop.py @@ -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": "", "end": ""}}`. 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, diff --git a/src/openjarvis/connectors/granola.py b/src/openjarvis/connectors/granola.py index a2242207..45624933 100644 --- a/src/openjarvis/connectors/granola.py +++ b/src/openjarvis/connectors/granola.py @@ -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.""" diff --git a/src/openjarvis/connectors/hybrid_search.py b/src/openjarvis/connectors/hybrid_search.py index 0465989b..2165db7e 100644 --- a/src/openjarvis/connectors/hybrid_search.py +++ b/src/openjarvis/connectors/hybrid_search.py @@ -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 diff --git a/src/openjarvis/connectors/store.py b/src/openjarvis/connectors/store.py index 2eb7a4d3..7fff7d51 100644 --- a/src/openjarvis/connectors/store.py +++ b/src/openjarvis/connectors/store.py @@ -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: diff --git a/src/openjarvis/server/research_router.py b/src/openjarvis/server/research_router.py index a81c9057..1525324f 100644 --- a/src/openjarvis/server/research_router.py +++ b/src/openjarvis/server/research_router.py @@ -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 diff --git a/tests/agents/test_research_loop.py b/tests/agents/test_research_loop.py index ce66281d..5a377319 100644 --- a/tests/agents/test_research_loop.py +++ b/tests/agents/test_research_loop.py @@ -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"] diff --git a/tests/connectors/test_granola.py b/tests/connectors/test_granola.py index 8bb33118..1c5fd896 100644 --- a/tests/connectors/test_granola.py +++ b/tests/connectors/test_granola.py @@ -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