Deploy 698d58f

app/pipeline/nodes/generate.py

@@ -235,6 +235,15 @@ def _format_history(state: "PipelineState") -> str:
     return "Prior conversation (oldest first):\n" + "\n".join(lines) + "\n\n"
 
 
+def _numbered_sources_event(source_refs: list[SourceRef]) -> dict:
+    """
+    The sources the streamed answer's [n] markers count from 1 in, sent before its first token.
+    Retrieval's sources event lists chunks before they are merged by document, and the done
+    event's list is renumbered for the stored answer, so neither matches the streamed words.
+    """
+    return {"type": "sources", "sources": [ref.model_dump() for ref in source_refs]}
+
+
 def _merge_by_source(chunks: list) -> list[dict]:
     """
     Collapse chunks that share the same source_url + section (or source_title
@@ -448,6 +457,7 @@ def make_generate_node(llm_client: LLMClient, gemini_client=None) -> Callable[[P
                     )
                 )
             context_block_enum = "\n\n".join(enum_context_parts)
+            writer(_numbered_sources_event(source_refs))
             prompt_enum = f"Items fetched from database:\n{context_block_enum}\n\nVisitor request: {query}"
             stream = llm_client.complete_with_complexity(
                 prompt=prompt_enum,
@@ -497,6 +507,7 @@ def make_generate_node(llm_client: LLMClient, gemini_client=None) -> Callable[[P
             )
 
         context_block = "\n\n".join(context_parts)
+        writer(_numbered_sources_event(source_refs))
 
         history_prefix = _format_history(state)
         is_criticism = state.get("is_criticism", False)

tests/test_generate_numbered_sources.py

@@ -0,0 +1,44 @@
+from unittest.mock import patch
+
+import pytest
+
+from app.pipeline.nodes.generate import make_generate_node
+
+
+def _chunk(title: str, url: str, text: str) -> dict:
+    return {
+        "text": text,
+        "metadata": {"source_title": title, "source_url": url, "section": "Overview", "source_type": "github_readme", "doc_id": title.lower(), "chunk_index": 0},
+    }
+
+
+class _StreamingLLM:
+    suggests_follow_ups = True
+
+    def __init__(self) -> None:
+        self.prompt = ""
+
+    async def complete_with_complexity(self, prompt, system, stream, complexity):
+        self.prompt = prompt
+        yield "Cortex maps a codebase into a graph [2]."
+
+
+@pytest.mark.asyncio
+async def test_streamed_citations_count_in_the_sources_sent_before_the_answer():
+    # Two chunks of one document share a number in the prompt, so Cortex is [2], not [3].
+    chunks = [
+        _chunk("AimBuddy", "https://github.com/1337Xcode/AimBuddy", "AimBuddy tracks targets on Android in real time."),
+        _chunk("AimBuddy", "https://github.com/1337Xcode/AimBuddy", "It runs a YOLO model on the phone's GPU through NCNN."),
+        _chunk("Cortex", "https://github.com/1337Xcode/Cortex", "Cortex maps a codebase into a SQLite graph for coding agents."),
+    ]
+    events: list[dict] = []
+    llm = _StreamingLLM()
+
+    with patch("app.pipeline.nodes.generate.get_stream_writer", return_value=events.append):
+        await make_generate_node(llm)({"query": "What is Cortex?", "reranked_chunks": chunks, "query_complexity": "simple"})
+
+    kinds = [event["type"] for event in events]
+    [sources] = [event["sources"] for event in events if event["type"] == "sources"]
+    assert [source["title"] for source in sources] == ["AimBuddy", "Cortex"]
+    assert "[2] Cortex" in llm.prompt
+    assert kinds.index("sources") < kinds.index("token")

tests/test_site_profile.py

@@ -0,0 +1,38 @@
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[2]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from ingestion.site_profile import profile_documents
+
+LLMS_TXT = """# Darshan Chheda
+
+> London-based software engineer.
+
+## Availability
+Open to software engineering roles in the UK
+
+## Links
+- [Website](https://www.darshanchheda.com): Canonical portfolio and profile
+- [Email](mailto:me@darshanchheda.com): Direct contact
+- [GitHub](https://github.com/1337XCode)
+- [LinkedIn](https://linkedin.com/in/darshchheda)
+
+## Pages
+- [Photos](https://www.darshanchheda.com/photos): Photographs.
+"""
+
+
+def test_the_page_is_one_profile_and_each_off_site_profile_its_own_document() -> None:
+    documents = profile_documents(LLMS_TXT)
+
+    assert [(d["title"], d["url"], d["slug"]) for d in documents] == [
+        ("darshanchheda.com", "https://www.darshanchheda.com/", "site-profile"),
+        ("GitHub", "https://github.com/1337XCode", "profile-github"),
+        ("LinkedIn", "https://linkedin.com/in/darshchheda", "profile-linkedin"),
+    ]
+    assert "Open to software engineering roles in the UK" in documents[0]["clean_content"]
+    assert documents[2]["clean_content"].endswith("https://linkedin.com/in/darshchheda")
+    assert len(documents[2]["clean_content"].split()) >= 8