From 367039073186c86c32b496ab9c578c58150ecb13 Mon Sep 17 00:00:00 2001 From: Oussama Hansal Date: Fri, 4 Sep 2026 09:19:11 -0700 Subject: [PATCH 1/4] Update LlamaIndex deps and make Bedrock batch model dispatch extensible --- .github/scripts/check_llama_index_updates.py | 160 ++++++++++++++++ .github/workflows/llama-index-deps-check.yml | 81 ++++++++ .../byokg_rag/requirements.txt | 2 +- .../indexing/utils/batch_inference_utils.py | 178 ++++++++++++------ .../lexical_graph/requirements.txt | 8 +- .../utils/test_batch_inference_utils.py | 44 ++++- 6 files changed, 411 insertions(+), 62 deletions(-) create mode 100644 .github/scripts/check_llama_index_updates.py create mode 100644 .github/workflows/llama-index-deps-check.yml diff --git a/.github/scripts/check_llama_index_updates.py b/.github/scripts/check_llama_index_updates.py new file mode 100644 index 000000000..09283ca2e --- /dev/null +++ b/.github/scripts/check_llama_index_updates.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Check PyPI for llama-index-* releases newer than our pinned floors. + +Reads the requirements.txt files listed in the REQUIREMENTS_FILES env var +(whitespace-separated), extracts every ``llama-index*`` requirement and its +``>=`` floor, queries PyPI for the latest release of each, and writes a markdown +report of the packages that have a newer version available. + +Outputs (written to $GITHUB_OUTPUT when running under Actions): + has_updates : "true" if at least one package has a newer release, else "false" + +The markdown report is written to $GITHUB_WORKSPACE/llama_index_report.md. + +This is a best-effort nudge: a network or parse error for a single package is +logged and skipped rather than failing the job. +""" + +import json +import os +import re +import sys +import urllib.request +import urllib.error + +try: + from packaging.version import Version, InvalidVersion +except ImportError: # pragma: no cover - packaging ships with pip on the runner + print("packaging not available; install it in the workflow step", file=sys.stderr) + raise + +# Matches e.g. "llama-index-llms-bedrock-converse>=0.15.0". Only >= floors are +# tracked because that is the pinning style used across the toolkit. +REQ_LINE = re.compile(r'^(llama-index[a-z0-9-]*)\s*>=\s*([0-9][0-9A-Za-z.\-]*)\s*$') + +PYPI_URL = 'https://pypi.org/pypi/{pkg}/json' + + +def latest_version(pkg: str) -> str: + """Return the latest (info.version) release of a package on PyPI.""" + with urllib.request.urlopen(PYPI_URL.format(pkg=pkg), timeout=30) as resp: + data = json.load(resp) + return data['info']['version'] + + +def collect_requirements(files): + """Map each llama-index package to its lowest pinned floor and the files it appears in.""" + found = {} + for path in files: + if not path or not os.path.exists(path): + print(f'skipping missing requirements file: {path}', file=sys.stderr) + continue + with open(path, encoding='utf-8') as fh: + for raw in fh: + match = REQ_LINE.match(raw.strip()) + if not match: + continue + pkg, floor = match.group(1), match.group(2) + entry = found.setdefault(pkg, {'floor': floor, 'files': set()}) + entry['files'].add(path) + # Keep the lowest floor seen across files (most conservative). + try: + if Version(floor) < Version(entry['floor']): + entry['floor'] = floor + except InvalidVersion: + pass + return found + + +TEST_COMMANDS = """### Suggested tests to run before merging a bump + +```bash +# lexical-graph unit tests +cd lexical-graph && PYTHONPATH=src python -m pytest tests/ + +# byokg-rag unit tests +cd byokg-rag && PYTHONPATH=src python -m pytest tests/ +``` + +Pay particular attention to the Bedrock / LLM paths, which are most sensitive to +llama-index-llms-bedrock-converse and llama-index-llms-anthropic changes: + +```bash +cd lexical-graph && PYTHONPATH=src python -m pytest \\ + tests/unit/indexing/utils/test_batch_inference_utils.py \\ + tests/unit/indexing/utils/test_batch_inference_aws.py \\ + tests/unit/utils/test_llm_cache.py \\ + tests/unit/utils/test_llm_cache_cross_region.py +``` +""" + + +def build_report(updates): + """Render the markdown issue body for the packages that have newer releases.""" + lines = [ + 'The following LlamaIndex dependencies have newer releases on PyPI than ' + 'the floors pinned in this repo. Dependabot opens the actual bump PRs; ' + 'this issue tracks awareness and the tests to run.', + '', + '| Package | Pinned floor | Latest on PyPI | Requirements file(s) |', + '| --- | --- | --- | --- |', + ] + for pkg in sorted(updates): + info = updates[pkg] + files = '
'.join(sorted(info['files'])) + lines.append(f"| `{pkg}` | {info['floor']} | **{info['latest']}** | {files} |") + lines.extend(['', TEST_COMMANDS]) + return '\n'.join(lines) + + +def set_output(name: str, value: str) -> None: + out = os.environ.get('GITHUB_OUTPUT') + if not out: + return + with open(out, 'a', encoding='utf-8') as fh: + fh.write(f'{name}={value}\n') + + +def main() -> int: + files = os.environ.get('REQUIREMENTS_FILES', '').split() + requirements = collect_requirements(files) + if not requirements: + print('no llama-index requirements found', file=sys.stderr) + set_output('has_updates', 'false') + return 0 + + updates = {} + for pkg, info in sorted(requirements.items()): + try: + latest = latest_version(pkg) + except (urllib.error.URLError, urllib.error.HTTPError, KeyError, ValueError) as exc: + print(f'could not fetch {pkg} from PyPI: {exc}', file=sys.stderr) + continue + try: + is_newer = Version(latest) > Version(info['floor']) + except InvalidVersion: + print(f'could not compare versions for {pkg}: {info["floor"]} vs {latest}', file=sys.stderr) + continue + status = 'newer available' if is_newer else 'up to date' + print(f'{pkg}: floor {info["floor"]} -> latest {latest} ({status})') + if is_newer: + updates[pkg] = {**info, 'latest': latest} + + if not updates: + set_output('has_updates', 'false') + print('all llama-index dependencies are up to date') + return 0 + + report = build_report(updates) + report_path = os.path.join(os.environ.get('GITHUB_WORKSPACE', '.'), 'llama_index_report.md') + with open(report_path, 'w', encoding='utf-8') as fh: + fh.write(report) + set_output('has_updates', 'true') + print(f'{len(updates)} package(s) have newer releases; report written to {report_path}') + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/.github/workflows/llama-index-deps-check.yml b/.github/workflows/llama-index-deps-check.yml new file mode 100644 index 000000000..6548505ad --- /dev/null +++ b/.github/workflows/llama-index-deps-check.yml @@ -0,0 +1,81 @@ +# Monthly check for newer llama-index-* releases than the >= floors pinned in the +# toolkit's requirements.txt files. When newer releases exist it opens (or +# refreshes) a single tracking issue that lists the available upgrades and the +# tests to run before merging a bump. Complements Dependabot, which opens the +# actual version-bump PRs; this workflow adds a curated, test-oriented nudge. +name: LlamaIndex Dependency Freshness + +on: + schedule: + # 08:00 UTC on the 1st of each month. + - cron: "0 8 1 * *" + workflow_dispatch: + +permissions: + contents: read + issues: write + +jobs: + check: + runs-on: ubuntu-latest + timeout-minutes: 10 + + steps: + - uses: actions/checkout@v7 + + - name: Set up Python + uses: actions/setup-python@v7 + with: + python-version: "3.12" + + - name: Install packaging + run: python -m pip install --disable-pip-version-check packaging + + - name: Check PyPI for newer llama-index releases + id: check + env: + REQUIREMENTS_FILES: >- + lexical-graph/src/graphrag_toolkit/lexical_graph/requirements.txt + byokg-rag/src/graphrag_toolkit/byokg_rag/requirements.txt + run: python .github/scripts/check_llama_index_updates.py + + - name: Open or refresh tracking issue + if: steps.check.outputs.has_updates == 'true' + uses: actions/github-script@v9 + with: + script: | + const fs = require('fs'); + const title = 'chore(deps): LlamaIndex dependency updates available'; + const label = 'dependencies'; + const body = fs.readFileSync('llama_index_report.md', 'utf8') + + `\n\n_Last checked: ${new Date().toISOString().slice(0, 10)} ` + + `([run](https://github.com/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}))._`; + + // Reuse a single living issue rather than opening a new one each week. + const existing = await github.rest.issues.listForRepo({ + owner: context.repo.owner, + repo: context.repo.repo, + state: 'open', + labels: label, + per_page: 100, + }); + const match = existing.data.find(i => i.title === title && !i.pull_request); + + if (match) { + await github.rest.issues.update({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: match.number, + body, + }); + console.log(`Updated existing issue #${match.number}`); + } else { + const created = await github.rest.issues.create({ + owner: context.repo.owner, + repo: context.repo.repo, + title, + body, + labels: [label], + }); + console.log(`Created issue #${created.data.number}`); + } diff --git a/byokg-rag/src/graphrag_toolkit/byokg_rag/requirements.txt b/byokg-rag/src/graphrag_toolkit/byokg_rag/requirements.txt index 6b084e13e..516155e6e 100644 --- a/byokg-rag/src/graphrag_toolkit/byokg_rag/requirements.txt +++ b/byokg-rag/src/graphrag_toolkit/byokg_rag/requirements.txt @@ -6,7 +6,7 @@ colorama==0.4.6 faiss-cpu==1.9.0 langchain_huggingface==0.3.1 langchain_aws==0.2.33 -llama-index-embeddings-bedrock>=0.8.2 +llama-index-embeddings-bedrock>=0.9.0 numpy>=1.24.0,<2.0.0 pydantic>=2.8.2 pyyaml==6.0.3 diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py index 1b5b1adca..4029a3b6d 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py @@ -6,7 +6,8 @@ import time import os import json -from typing import Any, List, Dict +from typing import Any, Callable, List, Dict +from dataclasses import dataclass from os import stat, listdir from os.path import isfile, join @@ -65,47 +66,126 @@ def split_nodes(nodes: List[Any], batch_size: int) -> List[List[Any]]: return results -def get_request_body(llm:BedrockConverse, messages:List[ChatMessage], inference_parameters: dict): - - model_id = llm.model - - if 'amazon.nova' in model_id: - converse_messages, system_prompt = messages_to_converse_messages(messages) - request_body = { - 'messages': converse_messages, - 'inferenceConfig': { - 'maxTokens': inference_parameters['max_tokens'], - 'temperature': inference_parameters['temperature'], - } - } - if system_prompt: - request_body['system'] = [{'text': system_prompt}] - return request_body - elif 'anthropic.claude' in model_id: - anthropic_messages, system_prompt = messages_to_anthropic_messages(messages) - request_body = { - 'anthropic_version': inference_parameters.get('anthropic_version', 'bedrock-2023-05-31'), - 'messages': anthropic_messages, - 'max_tokens': inference_parameters['max_tokens'], - 'temperature': inference_parameters['temperature'] +# --- Bedrock batch (InvokeModel JSONL) provider registry ------------------- +# +# Batch inference does not go through the Converse API; each record's +# `modelInput`/`modelOutput` uses the provider-specific InvokeModel schema, so +# each family is described by one BATCH_MODEL_PROVIDERS entry: +# +# * build_request - builds `modelInput` (the request body genuinely differs +# per family, so this stays a small function) +# * output_path - where the generated text lives in a record, walked from +# the record root (so a family can read under 'modelOutput' +# or at the top level) +# * output_mode - 'blocks' joins a [{text: ...}] content list; 'text' +# returns a scalar string node as-is +# +# To add a family, verify its InvokeModel request/response schema against the AWS +# Bedrock docs and add one entry - no edits to get_request_body / +# get_parse_output_text_fn. match_prefixes are substrings tested against +# llm.model; they match through a cross-region inference-profile prefix (e.g. +# 'us.anthropic.claude...' contains 'anthropic.claude') and across model versions +# (Claude Opus 5, Nova 2, Llama 4). + + +def _build_nova_request(messages: List[ChatMessage], params: dict) -> dict: + converse_messages, system_prompt = messages_to_converse_messages(messages) + request_body = { + 'messages': converse_messages, + 'inferenceConfig': { + 'maxTokens': params['max_tokens'], + 'temperature': params['temperature'], } - if system_prompt: - request_body['system'] = system_prompt - return request_body - elif 'meta.llama' in model_id: - converse_messages, system_prompt = messages_to_converse_messages(messages) - request_body = { - 'messages': converse_messages, - 'parameters': { - 'max_new_tokens': inference_parameters['max_tokens'], - 'temperature': inference_parameters['temperature'], - } + } + if system_prompt: + request_body['system'] = [{'text': system_prompt}] + return request_body + + +def _build_claude_request(messages: List[ChatMessage], params: dict) -> dict: + anthropic_messages, system_prompt = messages_to_anthropic_messages(messages) + request_body = { + 'anthropic_version': params.get('anthropic_version', 'bedrock-2023-05-31'), + 'messages': anthropic_messages, + 'max_tokens': params['max_tokens'], + 'temperature': params['temperature'] + } + if system_prompt: + request_body['system'] = system_prompt + return request_body + + +def _build_llama_request(messages: List[ChatMessage], params: dict) -> dict: + converse_messages, system_prompt = messages_to_converse_messages(messages) + return { + 'messages': converse_messages, + 'parameters': { + 'max_new_tokens': params['max_tokens'], + 'temperature': params['temperature'], } - return request_body - else: - raise ValueError(f'Unrecognized model_id: batch extraction for {model_id} is not supported') + } +@dataclass(frozen=True) +class BatchModelProvider: + """A model family's batch (InvokeModel JSONL) request builder and output spec.""" + name: str + match_prefixes: tuple + build_request: Callable[[List[ChatMessage], dict], dict] + output_path: tuple + output_mode: str = 'blocks' + + +BATCH_MODEL_PROVIDERS: List[BatchModelProvider] = [ + BatchModelProvider( + name='amazon.nova', + match_prefixes=('amazon.nova',), + build_request=_build_nova_request, + output_path=('modelOutput', 'output', 'message', 'content'), + ), + BatchModelProvider( + name='anthropic.claude', + match_prefixes=('anthropic.claude',), + build_request=_build_claude_request, + output_path=('modelOutput', 'content'), + ), + BatchModelProvider( + name='meta.llama', + match_prefixes=('meta.llama',), + build_request=_build_llama_request, + output_path=('generation',), + output_mode='text', + ), +] + + +def _resolve_batch_model_provider(model_id: str) -> BatchModelProvider: + for provider in BATCH_MODEL_PROVIDERS: + if any(prefix in model_id for prefix in provider.match_prefixes): + return provider + supported = ', '.join(provider.name for provider in BATCH_MODEL_PROVIDERS) + raise ValueError( + f'Unrecognized model_id: batch extraction for {model_id} is not supported. ' + f'Supported model families: {supported}' + ) + + +def _parse_output_text(json_data: dict, output_path: tuple, output_mode: str) -> str: + """Extract generated text from a batch output record per a provider's output spec.""" + node = json_data + for key in output_path: + if not isinstance(node, dict): + node = None + break + node = node.get(key) + if output_mode == 'text': + return node if isinstance(node, str) else '' + return ''.join(block.get('text', '') for block in (node or [])) + + +def get_request_body(llm:BedrockConverse, messages:List[ChatMessage], inference_parameters: dict): + return _resolve_batch_model_provider(llm.model).build_request(messages, inference_parameters) + def create_inference_inputs_for_messages(llm:BedrockConverse, nodes: List[TextNode], messages_batch: List[List[ChatMessage]], **kwargs) -> List[Dict[str, Any]]: inference_parameters = llm._get_all_kwargs(**kwargs) @@ -238,23 +318,11 @@ def download_output_files(s3_client: Any, bucket_name:str, output_path:str, inpu s3_client.download_file(Bucket=bucket_name, Key=key, Filename=local_file_path) logger.debug(f'Finished downloading {key} to {local_file_path}') -def get_parse_output_text_fn(model_id:str): - if 'amazon.nova' in model_id: - def get_output_text(json_data): - contents = json_data.get('modelOutput', {}).get('output', {}).get('message', {}).get('content', []) - return ''.join([content.get('text', '') for content in contents]) - return get_output_text - elif 'anthropic.claude' in model_id: - def get_output_text(json_data): - contents = json_data.get('modelOutput', {}).get('content', []) - return ''.join([content.get('text', '') for content in contents]) - return get_output_text - elif 'meta.llama' in model_id: - def get_output_text(json_data): - return json_data['generation'] - return get_output_text - else: - raise ValueError(f'Unrecognized model_id: batch extraction for {model_id} is not supported') +def get_parse_output_text_fn(model_id:str): + provider = _resolve_batch_model_provider(model_id) + def parse_output(json_data): + return _parse_output_text(json_data, provider.output_path, provider.output_mode) + return parse_output async def process_batch_output(local_output_directory:str, input_filename:str, llm:LLMCache) -> Dict[str, str]: """Process batch output files and return results.""" diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/requirements.txt b/lexical-graph/src/graphrag_toolkit/lexical_graph/requirements.txt index 3f2b430d7..ef66de4bb 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/requirements.txt +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/requirements.txt @@ -4,10 +4,10 @@ anthropic-bedrock==0.8.0 boto3>=1.40.61 botocore>=1.40.61 json2xml==5.2.0 -llama-index-core>=0.14.23 -llama-index-embeddings-bedrock>=0.8.2 -llama-index-llms-anthropic>=0.11.8 -llama-index-llms-bedrock-converse>=0.14.16 +llama-index-core>=0.14.24 +llama-index-embeddings-bedrock>=0.9.0 +llama-index-llms-anthropic>=0.12.0 +llama-index-llms-bedrock-converse>=0.15.0 lru-dict==1.3.0 # Pinned below 3.10.3: that release's pathsec hardlink check (CWE-59) rejects # uv-hardlinked NLTK data files, breaking sentence splitting. See issue 468. diff --git a/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py b/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py index e225ad44b..c513de269 100644 --- a/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py +++ b/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py @@ -19,10 +19,13 @@ create_inference_inputs_for_messages, create_inference_inputs, get_parse_output_text_fn, + BATCH_MODEL_PROVIDERS, BEDROCK_MIN_BATCH_SIZE, BEDROCK_MAX_BATCH_SIZE ) +CONVERSE_PATCH_TARGET = 'graphrag_toolkit.lexical_graph.indexing.utils.batch_inference_utils.messages_to_converse_messages' + class TestGetFileSize: """Tests for get_file_size_mb function.""" @@ -383,7 +386,7 @@ def test_parse_output_unsupported_model(self): def test_parse_output_empty_content(self): """Verify parsing handles empty content.""" parse_fn = get_parse_output_text_fn('amazon.nova-lite-v1:0') - + json_data = { 'modelOutput': { 'output': { @@ -393,6 +396,43 @@ def test_parse_output_empty_content(self): } } } - + result = parse_fn(json_data) assert result == '' + + +class TestBatchModelProviderRegistry: + """Tests for the extensible provider registry backing batch dispatch.""" + + def test_registry_contains_expected_families(self): + """The registry holds exactly the supported families; adding one is a single entry.""" + names = {provider.name for provider in BATCH_MODEL_PROVIDERS} + assert names == {'amazon.nova', 'anthropic.claude', 'meta.llama'} + + def test_parse_output_cross_region_profile_prefix(self): + """Inference-profile prefixes (us./eu.) still resolve to the right family.""" + claude_fn = get_parse_output_text_fn('us.anthropic.claude-sonnet-4-6') + assert claude_fn({'modelOutput': {'content': [{'text': 'hi'}]}}) == 'hi' + + nova_fn = get_parse_output_text_fn('eu.amazon.nova-pro-v1:0') + assert nova_fn({'modelOutput': {'output': {'message': {'content': [{'text': 'yo'}]}}}}) == 'yo' + + def test_request_body_cross_region_profile_prefix(self): + """get_request_body resolves the family through a cross-region profile prefix.""" + mock_llm = Mock(spec=BedrockConverse) + mock_llm.model = 'us.meta.llama3-3-70b-instruct-v1:0' + + messages = [ChatMessage(role=MessageRole.USER, content="Test")] + + with patch(CONVERSE_PATCH_TARGET) as mock_convert: + mock_convert.return_value = ([{'role': 'user', 'content': [{'text': 'Test'}]}], None) + + body = get_request_body(mock_llm, messages, {'max_tokens': 100, 'temperature': 0.2}) + + assert 'parameters' in body + assert body['parameters']['max_new_tokens'] == 100 + + def test_unsupported_model_error_lists_supported_families(self): + """The unsupported-model error names the families that are supported.""" + with pytest.raises(ValueError, match="Supported model families"): + get_parse_output_text_fn('cohere.command-r-v1:0') From 719b84be2b0fdb43e48bc8b39f57cabe5f013714 Mon Sep 17 00:00:00 2001 From: Oussama Hansal Date: Wed, 9 Sep 2026 08:40:19 -0700 Subject: [PATCH 2/4] adress comments --- .github/workflows/llama-index-deps-check.yml | 1 + .../indexing/utils/batch_inference_utils.py | 23 ++++++++++--- .../utils/test_batch_inference_utils.py | 32 ++++++++++++++----- 3 files changed, 44 insertions(+), 12 deletions(-) diff --git a/.github/workflows/llama-index-deps-check.yml b/.github/workflows/llama-index-deps-check.yml index 6548505ad..30cf3ea0b 100644 --- a/.github/workflows/llama-index-deps-check.yml +++ b/.github/workflows/llama-index-deps-check.yml @@ -37,6 +37,7 @@ jobs: REQUIREMENTS_FILES: >- lexical-graph/src/graphrag_toolkit/lexical_graph/requirements.txt byokg-rag/src/graphrag_toolkit/byokg_rag/requirements.txt + integration-tests/requirements-integ-test.txt run: python .github/scripts/check_llama_index_updates.py - name: Open or refresh tracking issue diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py index 4029a3b6d..5736e8cdc 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py @@ -98,7 +98,10 @@ def _build_nova_request(messages: List[ChatMessage], params: dict) -> dict: } } if system_prompt: - request_body['system'] = [{'text': system_prompt}] + # messages_to_converse_messages already returns the system prompt as a + # list of {'text': ...} blocks, so assign it directly (wrapping it again + # would nest as [{'text': [{'text': ...}]}]). + request_body['system'] = system_prompt return request_body @@ -153,7 +156,7 @@ class BatchModelProvider: name='meta.llama', match_prefixes=('meta.llama',), build_request=_build_llama_request, - output_path=('generation',), + output_path=('modelOutput', 'generation'), output_mode='text', ), ] @@ -171,7 +174,14 @@ def _resolve_batch_model_provider(model_id: str) -> BatchModelProvider: def _parse_output_text(json_data: dict, output_path: tuple, output_mode: str) -> str: - """Extract generated text from a batch output record per a provider's output spec.""" + """Extract generated text from a batch output record per a provider's output spec. + + 'text' mode fails loud if the expected key is absent: a missing scalar means + the output schema is not what we expect, and silently returning '' would let + every record log as a successful-but-empty extraction (masking the failure). + 'blocks' mode stays lenient - an empty/missing content list yields '', which + matches how nova/claude parsing behaved before the registry refactor. + """ node = json_data for key in output_path: if not isinstance(node, dict): @@ -179,7 +189,12 @@ def _parse_output_text(json_data: dict, output_path: tuple, output_mode: str) -> break node = node.get(key) if output_mode == 'text': - return node if isinstance(node, str) else '' + if not isinstance(node, str): + raise ValueError( + f"Expected a string at {'.'.join(output_path)!r} in the batch output " + f"record, got {type(node).__name__}; the model output schema may have changed" + ) + return node return ''.join(block.get('text', '') for block in (node or [])) diff --git a/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py b/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py index c513de269..e0f671413 100644 --- a/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py +++ b/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py @@ -176,10 +176,12 @@ def test_get_request_body_nova_with_system_prompt(self): inference_params = {'max_tokens': 500, 'temperature': 0.5} with patch('graphrag_toolkit.lexical_graph.indexing.utils.batch_inference_utils.messages_to_converse_messages') as mock_convert: - mock_convert.return_value = ([{'role': 'user', 'content': [{'text': 'User message'}]}], 'System prompt') - + # messages_to_converse_messages returns the system prompt as a list of + # {'text': ...} blocks, so the builder must pass it through unwrapped. + mock_convert.return_value = ([{'role': 'user', 'content': [{'text': 'User message'}]}], [{'text': 'System prompt'}]) + request_body = get_request_body(mock_llm, messages, inference_params) - + assert 'system' in request_body assert request_body['system'] == [{'text': 'System prompt'}] @@ -370,14 +372,24 @@ def test_parse_output_claude_model(self): def test_parse_output_llama_model(self): """Verify parsing function works for Llama model output.""" parse_fn = get_parse_output_text_fn('meta.llama3-70b-instruct-v1:0') - + + # Bedrock wraps the provider payload under 'modelOutput', same as nova/claude. json_data = { - 'generation': 'Generated text response' + 'modelOutput': { + 'generation': 'Generated text response' + } } - + result = parse_fn(json_data) assert result == 'Generated text response' - + + def test_parse_output_text_mode_missing_key_raises(self): + """A missing scalar output must fail loud, not silently return ''.""" + parse_fn = get_parse_output_text_fn('meta.llama3-70b-instruct-v1:0') + + with pytest.raises(ValueError, match="model output schema may have changed"): + parse_fn({'modelOutput': {}}) + def test_parse_output_unsupported_model(self): """Verify error raised for unsupported model.""" with pytest.raises(ValueError, match="Unrecognized model_id"): @@ -405,7 +417,11 @@ class TestBatchModelProviderRegistry: """Tests for the extensible provider registry backing batch dispatch.""" def test_registry_contains_expected_families(self): - """The registry holds exactly the supported families; adding one is a single entry.""" + """Exact-set guard: fail if a family is dropped or an unexpected one appears. + + Equality is intentional - it catches an accidental add/remove. Adding a + family is a deliberate change that should update this set too. + """ names = {provider.name for provider in BATCH_MODEL_PROVIDERS} assert names == {'amazon.nova', 'anthropic.claude', 'meta.llama'} From 71a3dbf534f0173b8aa12004767c7680c36c4a1c Mon Sep 17 00:00:00 2001 From: Oussama Hansal Date: Fri, 11 Sep 2026 06:08:09 -0700 Subject: [PATCH 3/4] split into another pr --- .github/scripts/check_llama_index_updates.py | 24 ++- .../indexing/utils/batch_inference_utils.py | 193 +++++------------- .../utils/test_batch_inference_utils.py | 74 +------ 3 files changed, 86 insertions(+), 205 deletions(-) diff --git a/.github/scripts/check_llama_index_updates.py b/.github/scripts/check_llama_index_updates.py index 09283ca2e..63219d345 100644 --- a/.github/scripts/check_llama_index_updates.py +++ b/.github/scripts/check_llama_index_updates.py @@ -38,10 +38,30 @@ def latest_version(pkg: str) -> str: - """Return the latest (info.version) release of a package on PyPI.""" + """Return the latest stable, non-yanked release of a package on PyPI. + + `info.version` can point at a pre-release and does not account for yanked + releases, so walk `releases` instead: skip pre-releases, versions with no + uploaded files, and versions whose files are all yanked, then return the + highest remaining version. Falls back to `info.version` only if that filter + leaves nothing. + """ with urllib.request.urlopen(PYPI_URL.format(pkg=pkg), timeout=30) as resp: data = json.load(resp) - return data['info']['version'] + + candidates = [] + for version, files in data.get('releases', {}).items(): + if not files or all(f.get('yanked') for f in files): + continue + try: + parsed = Version(version) + except InvalidVersion: + continue + if parsed.is_prerelease: + continue + candidates.append(parsed) + + return str(max(candidates)) if candidates else data['info']['version'] def collect_requirements(files): diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py index 5736e8cdc..1b5b1adca 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/batch_inference_utils.py @@ -6,8 +6,7 @@ import time import os import json -from typing import Any, Callable, List, Dict -from dataclasses import dataclass +from typing import Any, List, Dict from os import stat, listdir from os.path import isfile, join @@ -66,140 +65,46 @@ def split_nodes(nodes: List[Any], batch_size: int) -> List[List[Any]]: return results -# --- Bedrock batch (InvokeModel JSONL) provider registry ------------------- -# -# Batch inference does not go through the Converse API; each record's -# `modelInput`/`modelOutput` uses the provider-specific InvokeModel schema, so -# each family is described by one BATCH_MODEL_PROVIDERS entry: -# -# * build_request - builds `modelInput` (the request body genuinely differs -# per family, so this stays a small function) -# * output_path - where the generated text lives in a record, walked from -# the record root (so a family can read under 'modelOutput' -# or at the top level) -# * output_mode - 'blocks' joins a [{text: ...}] content list; 'text' -# returns a scalar string node as-is -# -# To add a family, verify its InvokeModel request/response schema against the AWS -# Bedrock docs and add one entry - no edits to get_request_body / -# get_parse_output_text_fn. match_prefixes are substrings tested against -# llm.model; they match through a cross-region inference-profile prefix (e.g. -# 'us.anthropic.claude...' contains 'anthropic.claude') and across model versions -# (Claude Opus 5, Nova 2, Llama 4). - - -def _build_nova_request(messages: List[ChatMessage], params: dict) -> dict: - converse_messages, system_prompt = messages_to_converse_messages(messages) - request_body = { - 'messages': converse_messages, - 'inferenceConfig': { - 'maxTokens': params['max_tokens'], - 'temperature': params['temperature'], +def get_request_body(llm:BedrockConverse, messages:List[ChatMessage], inference_parameters: dict): + + model_id = llm.model + + if 'amazon.nova' in model_id: + converse_messages, system_prompt = messages_to_converse_messages(messages) + request_body = { + 'messages': converse_messages, + 'inferenceConfig': { + 'maxTokens': inference_parameters['max_tokens'], + 'temperature': inference_parameters['temperature'], + } } - } - if system_prompt: - # messages_to_converse_messages already returns the system prompt as a - # list of {'text': ...} blocks, so assign it directly (wrapping it again - # would nest as [{'text': [{'text': ...}]}]). - request_body['system'] = system_prompt - return request_body - - -def _build_claude_request(messages: List[ChatMessage], params: dict) -> dict: - anthropic_messages, system_prompt = messages_to_anthropic_messages(messages) - request_body = { - 'anthropic_version': params.get('anthropic_version', 'bedrock-2023-05-31'), - 'messages': anthropic_messages, - 'max_tokens': params['max_tokens'], - 'temperature': params['temperature'] - } - if system_prompt: - request_body['system'] = system_prompt - return request_body - - -def _build_llama_request(messages: List[ChatMessage], params: dict) -> dict: - converse_messages, system_prompt = messages_to_converse_messages(messages) - return { - 'messages': converse_messages, - 'parameters': { - 'max_new_tokens': params['max_tokens'], - 'temperature': params['temperature'], + if system_prompt: + request_body['system'] = [{'text': system_prompt}] + return request_body + elif 'anthropic.claude' in model_id: + anthropic_messages, system_prompt = messages_to_anthropic_messages(messages) + request_body = { + 'anthropic_version': inference_parameters.get('anthropic_version', 'bedrock-2023-05-31'), + 'messages': anthropic_messages, + 'max_tokens': inference_parameters['max_tokens'], + 'temperature': inference_parameters['temperature'] } - } - - -@dataclass(frozen=True) -class BatchModelProvider: - """A model family's batch (InvokeModel JSONL) request builder and output spec.""" - name: str - match_prefixes: tuple - build_request: Callable[[List[ChatMessage], dict], dict] - output_path: tuple - output_mode: str = 'blocks' - - -BATCH_MODEL_PROVIDERS: List[BatchModelProvider] = [ - BatchModelProvider( - name='amazon.nova', - match_prefixes=('amazon.nova',), - build_request=_build_nova_request, - output_path=('modelOutput', 'output', 'message', 'content'), - ), - BatchModelProvider( - name='anthropic.claude', - match_prefixes=('anthropic.claude',), - build_request=_build_claude_request, - output_path=('modelOutput', 'content'), - ), - BatchModelProvider( - name='meta.llama', - match_prefixes=('meta.llama',), - build_request=_build_llama_request, - output_path=('modelOutput', 'generation'), - output_mode='text', - ), -] - - -def _resolve_batch_model_provider(model_id: str) -> BatchModelProvider: - for provider in BATCH_MODEL_PROVIDERS: - if any(prefix in model_id for prefix in provider.match_prefixes): - return provider - supported = ', '.join(provider.name for provider in BATCH_MODEL_PROVIDERS) - raise ValueError( - f'Unrecognized model_id: batch extraction for {model_id} is not supported. ' - f'Supported model families: {supported}' - ) - - -def _parse_output_text(json_data: dict, output_path: tuple, output_mode: str) -> str: - """Extract generated text from a batch output record per a provider's output spec. - - 'text' mode fails loud if the expected key is absent: a missing scalar means - the output schema is not what we expect, and silently returning '' would let - every record log as a successful-but-empty extraction (masking the failure). - 'blocks' mode stays lenient - an empty/missing content list yields '', which - matches how nova/claude parsing behaved before the registry refactor. - """ - node = json_data - for key in output_path: - if not isinstance(node, dict): - node = None - break - node = node.get(key) - if output_mode == 'text': - if not isinstance(node, str): - raise ValueError( - f"Expected a string at {'.'.join(output_path)!r} in the batch output " - f"record, got {type(node).__name__}; the model output schema may have changed" - ) - return node - return ''.join(block.get('text', '') for block in (node or [])) - + if system_prompt: + request_body['system'] = system_prompt + return request_body + elif 'meta.llama' in model_id: + converse_messages, system_prompt = messages_to_converse_messages(messages) + request_body = { + 'messages': converse_messages, + 'parameters': { + 'max_new_tokens': inference_parameters['max_tokens'], + 'temperature': inference_parameters['temperature'], + } + } + return request_body + else: + raise ValueError(f'Unrecognized model_id: batch extraction for {model_id} is not supported') -def get_request_body(llm:BedrockConverse, messages:List[ChatMessage], inference_parameters: dict): - return _resolve_batch_model_provider(llm.model).build_request(messages, inference_parameters) def create_inference_inputs_for_messages(llm:BedrockConverse, nodes: List[TextNode], messages_batch: List[List[ChatMessage]], **kwargs) -> List[Dict[str, Any]]: @@ -333,11 +238,23 @@ def download_output_files(s3_client: Any, bucket_name:str, output_path:str, inpu s3_client.download_file(Bucket=bucket_name, Key=key, Filename=local_file_path) logger.debug(f'Finished downloading {key} to {local_file_path}') -def get_parse_output_text_fn(model_id:str): - provider = _resolve_batch_model_provider(model_id) - def parse_output(json_data): - return _parse_output_text(json_data, provider.output_path, provider.output_mode) - return parse_output +def get_parse_output_text_fn(model_id:str): + if 'amazon.nova' in model_id: + def get_output_text(json_data): + contents = json_data.get('modelOutput', {}).get('output', {}).get('message', {}).get('content', []) + return ''.join([content.get('text', '') for content in contents]) + return get_output_text + elif 'anthropic.claude' in model_id: + def get_output_text(json_data): + contents = json_data.get('modelOutput', {}).get('content', []) + return ''.join([content.get('text', '') for content in contents]) + return get_output_text + elif 'meta.llama' in model_id: + def get_output_text(json_data): + return json_data['generation'] + return get_output_text + else: + raise ValueError(f'Unrecognized model_id: batch extraction for {model_id} is not supported') async def process_batch_output(local_output_directory:str, input_filename:str, llm:LLMCache) -> Dict[str, str]: """Process batch output files and return results.""" diff --git a/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py b/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py index e0f671413..e225ad44b 100644 --- a/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py +++ b/lexical-graph/tests/unit/indexing/utils/test_batch_inference_utils.py @@ -19,13 +19,10 @@ create_inference_inputs_for_messages, create_inference_inputs, get_parse_output_text_fn, - BATCH_MODEL_PROVIDERS, BEDROCK_MIN_BATCH_SIZE, BEDROCK_MAX_BATCH_SIZE ) -CONVERSE_PATCH_TARGET = 'graphrag_toolkit.lexical_graph.indexing.utils.batch_inference_utils.messages_to_converse_messages' - class TestGetFileSize: """Tests for get_file_size_mb function.""" @@ -176,12 +173,10 @@ def test_get_request_body_nova_with_system_prompt(self): inference_params = {'max_tokens': 500, 'temperature': 0.5} with patch('graphrag_toolkit.lexical_graph.indexing.utils.batch_inference_utils.messages_to_converse_messages') as mock_convert: - # messages_to_converse_messages returns the system prompt as a list of - # {'text': ...} blocks, so the builder must pass it through unwrapped. - mock_convert.return_value = ([{'role': 'user', 'content': [{'text': 'User message'}]}], [{'text': 'System prompt'}]) - + mock_convert.return_value = ([{'role': 'user', 'content': [{'text': 'User message'}]}], 'System prompt') + request_body = get_request_body(mock_llm, messages, inference_params) - + assert 'system' in request_body assert request_body['system'] == [{'text': 'System prompt'}] @@ -372,24 +367,14 @@ def test_parse_output_claude_model(self): def test_parse_output_llama_model(self): """Verify parsing function works for Llama model output.""" parse_fn = get_parse_output_text_fn('meta.llama3-70b-instruct-v1:0') - - # Bedrock wraps the provider payload under 'modelOutput', same as nova/claude. + json_data = { - 'modelOutput': { - 'generation': 'Generated text response' - } + 'generation': 'Generated text response' } - + result = parse_fn(json_data) assert result == 'Generated text response' - - def test_parse_output_text_mode_missing_key_raises(self): - """A missing scalar output must fail loud, not silently return ''.""" - parse_fn = get_parse_output_text_fn('meta.llama3-70b-instruct-v1:0') - - with pytest.raises(ValueError, match="model output schema may have changed"): - parse_fn({'modelOutput': {}}) - + def test_parse_output_unsupported_model(self): """Verify error raised for unsupported model.""" with pytest.raises(ValueError, match="Unrecognized model_id"): @@ -398,7 +383,7 @@ def test_parse_output_unsupported_model(self): def test_parse_output_empty_content(self): """Verify parsing handles empty content.""" parse_fn = get_parse_output_text_fn('amazon.nova-lite-v1:0') - + json_data = { 'modelOutput': { 'output': { @@ -408,47 +393,6 @@ def test_parse_output_empty_content(self): } } } - + result = parse_fn(json_data) assert result == '' - - -class TestBatchModelProviderRegistry: - """Tests for the extensible provider registry backing batch dispatch.""" - - def test_registry_contains_expected_families(self): - """Exact-set guard: fail if a family is dropped or an unexpected one appears. - - Equality is intentional - it catches an accidental add/remove. Adding a - family is a deliberate change that should update this set too. - """ - names = {provider.name for provider in BATCH_MODEL_PROVIDERS} - assert names == {'amazon.nova', 'anthropic.claude', 'meta.llama'} - - def test_parse_output_cross_region_profile_prefix(self): - """Inference-profile prefixes (us./eu.) still resolve to the right family.""" - claude_fn = get_parse_output_text_fn('us.anthropic.claude-sonnet-4-6') - assert claude_fn({'modelOutput': {'content': [{'text': 'hi'}]}}) == 'hi' - - nova_fn = get_parse_output_text_fn('eu.amazon.nova-pro-v1:0') - assert nova_fn({'modelOutput': {'output': {'message': {'content': [{'text': 'yo'}]}}}}) == 'yo' - - def test_request_body_cross_region_profile_prefix(self): - """get_request_body resolves the family through a cross-region profile prefix.""" - mock_llm = Mock(spec=BedrockConverse) - mock_llm.model = 'us.meta.llama3-3-70b-instruct-v1:0' - - messages = [ChatMessage(role=MessageRole.USER, content="Test")] - - with patch(CONVERSE_PATCH_TARGET) as mock_convert: - mock_convert.return_value = ([{'role': 'user', 'content': [{'text': 'Test'}]}], None) - - body = get_request_body(mock_llm, messages, {'max_tokens': 100, 'temperature': 0.2}) - - assert 'parameters' in body - assert body['parameters']['max_new_tokens'] == 100 - - def test_unsupported_model_error_lists_supported_families(self): - """The unsupported-model error names the families that are supported.""" - with pytest.raises(ValueError, match="Supported model families"): - get_parse_output_text_fn('cohere.command-r-v1:0') From 4aba5631b508b3092794a05dc937b1a23a99263d Mon Sep 17 00:00:00 2001 From: Oussama Hansal Date: Mon, 14 Sep 2026 12:27:55 -0700 Subject: [PATCH 4/4] added a pip entry for /integration-tests to dependabot.yml --- .github/dependabot.yml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 49ecd09e5..9e478d105 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -22,6 +22,17 @@ updates: commit-message: prefix: "chore(deps)" + - package-ecosystem: "pip" + directory: "/integration-tests" + schedule: + interval: "weekly" + day: "monday" + open-pull-requests-limit: 10 + labels: + - "dependencies" + commit-message: + prefix: "chore(deps)" + - package-ecosystem: "npm" directory: "/docs-site" schedule: