Draft: feat: add multi-source RAG crawler demo#327
Conversation
Add a cross-source document crawler demo that follows links between platforms (GitHub, web) and indexes discovered documents into a Llama Stack vector store for RAG queries. Demonstrates breadth-first crawling with depth limits, content deduplication via hashing, crawl provenance tracking, and querying across the automatically-discovered knowledge graph. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
amastbau
changed the title
Print file_search results (filename, score, source URL) after the RAG query, matching the pattern from 05_hybrid_search.py. Improve the system instructions to reference the embedded Source: and Crawl path: metadata in each chunk so the LLM cites provenance. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughAdds a new demo script Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes IssuesReview from line-level code quality to architectural design. Security-first. Be direct. Cite CWE/CVE IDs for security findings. Only flag actionable issues; no praise, no summaries.
Addressing the items above will reduce attack surface and improve operational robustness. 🚥 Pre-merge checks | ✅ 2✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Allow users to override the default system prompt used during the RAG query phase. Prints the active prompt so users can see what instructions the model is operating under. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@demos/03_rag/06_multi_source_rag_crawler.py`:
- Around line 309-317: The crawler currently enqueues links returned by
detect_links into self._queue via _QueueItem without validating destination
addresses; add validation before appending: resolve each link.url (and any
redirect target when fetching) to IP(s) and reject addresses in
localhost/127.0.0.1, ::1, and RFC1918/RFC3927 private ranges (e.g., 10.0.0.0/8,
172.16.0.0/12, 192.168.0.0/16, link-local) and non-public reserved ranges, and
ensure redirects do not change host to a private address (reject cross-host
redirects); apply the same check wherever links are enqueued (the detect_links
-> _QueueItem path and the other enqueue sites noted) and log + skip unsafe URLs
instead of enqueueing or fetching them.
- Around line 357-361: The current code reads the entire HTTP response into
memory via resp.text and then slices it, which does not prevent oversized
downloads; change the fetch logic that uses _requests.get(...) and the resp
variable to call _requests.get(..., stream=True), first validate the
Content-Type header (accept only text/* or known safe charsets) and then stream
the body with resp.iter_content(chunk_size) accumulating bytes up to the 15_000
byte cap (stop and abort the request if the limit is exceeded), and return None
for non-text or oversize responses; apply the same streaming + content-type
check/cap fix to the other fetcher block referenced around the 383-387 region so
both places avoid loading full responses into memory.
- Around line 345-357: The code naively rewrites any URL containing "github.com"
using str.replace and may send GITHUB_TOKEN to attacker-controlled hosts; update
the logic around raw_url, headers, and the _requests.get call to (1) parse the
incoming url with urllib.parse (or equivalent) and verify the netloc is exactly
"github.com" or "www.github.com" before deriving a raw.githubusercontent.com URL
(do not use str.replace on the whole string), (2) construct raw_url only after
confirming the origin is GitHub (or when building a README fallback from
identifier, ensure identifier's repo portion is validated), and (3) only attach
the Authorization header (token) when the final raw_url's parsed netloc equals
"raw.githubusercontent.com"; reference the variables raw_url, url, identifier,
headers, token and the request call _requests.get when making these changes.
- Around line 92-95: The LinkType.GITHUB regex is too permissive and captures
repo roots inside non-repo URLs (e.g., /issues, /pull, /releases), causing the
crawler to collapse those links to the repo README; update the LinkType.GITHUB
pattern to only match either a bare repo root (ending at end-of-string or a
trailing slash) or explicit content paths you intend to follow (e.g.,
/(blob|tree)/...), avoiding matching paths like /issues, /pull, /releases, and
then adjust the conversion logic that currently reads the regex groups (the code
around the repo extraction used at lines 130-133) to only convert the match into
a repo identifier when the regex indicates a true repo root or a blob/tree
content path (i.e., when the appropriate capture group is present), otherwise
skip treating it as the repo root so issues/pull/release links are not refetched
as the repo root.
- Around line 505-515: The provenance and enriched chunk currently include raw
URLs (doc.url and entries from doc.trace.path) which may contain sensitive query
strings/fragments; sanitize by stripping query and fragment components before
building provenance and the "Source:" field: for example, canonicalize each URL
from doc.url and each element of doc.trace.path via a URL-parsing routine (used
by _chunk_text caller) that removes the query and fragment, then use those
sanitized values when computing provenance and composing enriched; update the
code paths that construct provenance and enriched (references: provenance,
doc.url, doc.trace.path, enriched) to use the sanitized URLs so no secrets are
persisted or printed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Central YAML (base), Organization UI (inherited)
Review profile: CHILL
Plan: Pro
Run ID: b40e3e19-7b90-4627-9802-deabffdfb2e9
📒 Files selected for processing (2)
demos/03_rag/06_multi_source_rag_crawler.pydemos/03_rag/README.md
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Extend the crawler to support all source types from the original project. Each source has a dedicated fetcher that activates when credentials are configured via environment variables. Seed URL source type is auto-detected. The demo prints which sources are available based on configured credentials. Sources added: - GitLab: fetch files via API with PRIVATE-TOKEN auth - Confluence: fetch pages by ID with Bearer token auth - Jira: fetch issues with description and comments - Google Docs: fetch document text via Google API key Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1. GitHub regex: add end-of-token anchor to prevent matching /issues/, /pull/, /releases/ URLs as repo roots (CWE-918) 2. SSRF guard: add _is_safe_url() that resolves hostnames and rejects private/loopback/link-local/reserved addresses before enqueueing or fetching web URLs (CWE-918) 3. Token leak prevention: validate parsed hostname is exactly "github.com" before rewriting URL, and only attach GITHUB_TOKEN when raw_url hostname is "raw.githubusercontent.com" (CWE-200) 4. Download size limits: add _safe_download() helper that uses stream=True, checks Content-Type and Content-Length headers, and reads at most MAX_DOWNLOAD_BYTES via resp.raw.read() instead of loading full response into memory then slicing (CWE-400) 5. URL redaction: add _redact_url() that strips query strings and fragments before persisting Source: and Crawl path: in indexed chunks, preventing signed tokens or secrets from being stored in the vector DB and echoed by the model (CWE-200) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
demos/03_rag/README.md (1)
83-91: Missing blank line before table.Per MD058, tables should be surrounded by blank lines for consistent rendering.
Fix markdown formatting
**Supported Sources**: + | Source | Auth Env Vars | Seed URL Example |🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@demos/03_rag/README.md` around lines 83 - 91, Add a blank line immediately before the markdown table that starts with the header "| Source | Auth Env Vars | Seed URL Example |" (under the "**Supported Sources**:" section) so the table is separated from the preceding paragraph, and ensure there's also a blank line after the table block to satisfy MD058 and ensure consistent rendering.demos/03_rag/06_multi_source_rag_crawler.py (1)
764-788: Double chunking with inconsistent parameters.This file performs client-side character-based chunking (512 chars, 50 overlap) via
_chunk_text(), then applies server-side token-based chunking (512 tokens, 64 overlap) at index time. Other RAG demos use only server-side chunking withmax_chunk_size_tokens=256, chunk_overlap_tokens=32.The double-chunking creates unpredictable boundaries and the parameter divergence may degrade retrieval quality.
Option: remove client-side chunking, align with other demos
def index_crawled_documents( client: LlamaStackClient, vector_store_id: str, documents: list[CrawledDocument], - chunk_size: int = 512, ) -> int: ... for doc in documents: redacted_path = [*map(_redact_url, doc.trace.path), _redact_url(doc.url)] provenance = " -> ".join(redacted_path) - chunks = _chunk_text(doc.content, chunk_size) - for i, chunk in enumerate(chunks): - enriched = ( - f"Source: {_redact_url(doc.url)}\n" - f"Crawl path: {provenance}\n" - f"Depth: {doc.depth}\n\n" - f"{chunk}" - ) + enriched = ( + f"Source: {_redact_url(doc.url)}\n" + f"Crawl path: {provenance}\n" + f"Depth: {doc.depth}\n\n" + f"{doc.content}" + ) - filename = f"{doc.source_type.value}_{doc.content_hash}_chunk{i}.txt" + filename = f"{doc.source_type.value}_{doc.content_hash}.txt" ... client.vector_stores.files.create( ... chunking_strategy={ "type": "static", - "static": {"max_chunk_size_tokens": 512, "chunk_overlap_tokens": 64}, + "static": {"max_chunk_size_tokens": 256, "chunk_overlap_tokens": 32}, }, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@demos/03_rag/06_multi_source_rag_crawler.py` around lines 764 - 788, The client-side character chunker _chunk_text introduces double-chunking and mismatched parameters; remove or disable its usage and rely on server/token-based chunking with the same settings used by other demos (max_chunk_size_tokens=256 and chunk_overlap_tokens=32). Locate all calls to _chunk_text in this file (e.g., where documents are pre-split before indexing) and either stop calling it (pass full text through) or replace it with a no-op that returns the original text as a single item; ensure any local chunk_size/overlap parameters are removed or updated to match the server-side config so only token-based chunking is applied at index time.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@demos/03_rag/06_multi_source_rag_crawler.py`:
- Around line 513-560: fetch_gitlab currently derives base_url from the input
URL and will send the GITLAB_TOKEN to any host; before using token or building
base_url, validate the request host (parsed.hostname) is an allowed GitLab host
(e.g., "gitlab.com" or matches the hostname of the optional GITLAB_URL env var)
and return None if it is not allowed; update the logic in fetch_gitlab to
perform this hostname check (use variables parsed, base_url, token) so the token
is only sent to trusted GitLab hosts.
---
Nitpick comments:
In `@demos/03_rag/06_multi_source_rag_crawler.py`:
- Around line 764-788: The client-side character chunker _chunk_text introduces
double-chunking and mismatched parameters; remove or disable its usage and rely
on server/token-based chunking with the same settings used by other demos
(max_chunk_size_tokens=256 and chunk_overlap_tokens=32). Locate all calls to
_chunk_text in this file (e.g., where documents are pre-split before indexing)
and either stop calling it (pass full text through) or replace it with a no-op
that returns the original text as a single item; ensure any local
chunk_size/overlap parameters are removed or updated to match the server-side
config so only token-based chunking is applied at index time.
In `@demos/03_rag/README.md`:
- Around line 83-91: Add a blank line immediately before the markdown table that
starts with the header "| Source | Auth Env Vars | Seed URL Example |" (under
the "**Supported Sources**:" section) so the table is separated from the
preceding paragraph, and ensure there's also a blank line after the table block
to satisfy MD058 and ensure consistent rendering.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Central YAML (base), Organization UI (inherited)
Review profile: CHILL
Plan: Pro
Run ID: 56e8772d-e40e-4cb2-9041-2b4b41456d79
📒 Files selected for processing (2)
demos/03_rag/06_multi_source_rag_crawler.pydemos/03_rag/README.md
| def fetch_gitlab( | ||
| identifier: str, url: str, trace: CrawlTrace | ||
| ) -> CrawledDocument | None: | ||
| """Fetch a file from a GitLab repository. | ||
|
|
||
| Requires ``GITLAB_TOKEN`` env var. Optionally set ``GITLAB_URL`` to override | ||
| the base URL (defaults to extracting it from the link URL). | ||
|
|
||
| Example env:: | ||
|
|
||
| GITLAB_URL=https://gitlab.com | ||
| GITLAB_TOKEN=glpat-xxxxxxxxxxxxxxxxxxxx | ||
| """ | ||
| if _requests is None: | ||
| return None | ||
|
|
||
| token = os.environ.get("GITLAB_TOKEN") | ||
| if not token: | ||
| return None | ||
|
|
||
| # Determine GitLab base URL from the matched URL | ||
| parsed = urllib.parse.urlparse(url) | ||
| base_url = os.environ.get("GITLAB_URL", f"{parsed.scheme}://{parsed.netloc}") | ||
|
|
||
| repo = identifier.split(":")[0] | ||
| path = identifier.split(":", 1)[1] if ":" in identifier else "README.md" | ||
| project_encoded = urllib.parse.quote(repo, safe="") | ||
|
|
||
| # Try main branch, then master | ||
| for ref in ("main", "master"): | ||
| file_url = ( | ||
| f"{base_url}/api/v4/projects/{project_encoded}" | ||
| f"/repository/files/{urllib.parse.quote(path, safe='')}/raw" | ||
| ) | ||
| headers = {"PRIVATE-TOKEN": token} | ||
| content = _safe_download(f"{file_url}?ref={ref}", headers) | ||
| if content is not None: | ||
| web_url = f"{base_url}/{repo}/-/blob/{ref}/{path}" | ||
| return CrawledDocument( | ||
| url=web_url, | ||
| title=f"{repo}/{path}", | ||
| content=content, | ||
| source_type=LinkType.GITLAB, | ||
| trace=trace, | ||
| metadata={"repo": repo, "path": path, "ref": ref}, | ||
| ) | ||
|
|
||
| return None |
There was a problem hiding this comment.
Token sent to arbitrary hosts matching GitLab regex pattern (CWE-200).
Unlike fetch_github which validates parsed.hostname == "github.com", this fetcher derives base_url from whatever URL matched the GitLab regex (line 188: gitlab[^/\s]*). A crawled link like https://gitlab.attacker.com/org/repo would send GITLAB_TOKEN to the attacker's server.
Remediation: validate against allowed GitLab hosts
def fetch_gitlab(
identifier: str, url: str, trace: CrawlTrace
) -> CrawledDocument | None:
...
token = os.environ.get("GITLAB_TOKEN")
if not token:
return None
# Determine GitLab base URL from the matched URL
parsed = urllib.parse.urlparse(url)
+ # Only send token to explicitly configured GitLab host
+ allowed_host = urllib.parse.urlparse(os.environ.get("GITLAB_URL", "https://gitlab.com")).hostname
+ if parsed.hostname != allowed_host:
+ return None
base_url = os.environ.get("GITLAB_URL", f"{parsed.scheme}://{parsed.netloc}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@demos/03_rag/06_multi_source_rag_crawler.py` around lines 513 - 560,
fetch_gitlab currently derives base_url from the input URL and will send the
GITLAB_TOKEN to any host; before using token or building base_url, validate the
request host (parsed.hostname) is an allowed GitLab host (e.g., "gitlab.com" or
matches the hostname of the optional GITLAB_URL env var) and return None if it
is not allowed; update the logic in fetch_gitlab to perform this hostname check
(use variables parsed, base_url, token) so the token is only sent to trusted
GitLab hosts.
1 participant
Summary
demos/03_rag/06_multi_source_rag_crawler.py-- a new demo showing how to build a cross-source document crawler that follows links between platforms (GitHub, web pages) and indexes all discovered documents into a Llama Stack vector store for RAG queriesfire.Fire()CLI, shared utils,LlamaStackClient, Responses API withfile_search, and try/finally cleanupMotivation
Existing RAG demos (01-05) index documents from static/inline sources. In real-world scenarios, knowledge is spread across platforms -- a GitHub README links to design docs, which reference issues, which point to other repos. This demo fills that gap by showing how to automatically discover and index connected documents across sources.
The crawler architecture is extensible: new source types (GitLab, Confluence, Jira, etc.) can be added by registering a fetcher function for a new
LinkType, without modifying the core crawl logic.What's included
LinkTypeenum +detect_links()MultiSourceCrawlerCrawlTrace/CrawledDocumentfetch_github_raw/fetch_web_pageindex_crawled_documents()main()Usage
Test plan
--max_depth=0to confirm it only indexes seeds (no link following)--seed_urlspointing to a different public repo🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation