From 1fcc112563f29791966ae6408b6a285f7fef35ca Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Fri, 18 Sep 2026 10:34:24 -0700 Subject: [PATCH 01/23] Initial outline --- .github/workflows/propose_release.yml | 28 ++ .github/workflows/publish_release.yml | 19 + .github/workflows/publish_test_build.yml | 18 + .github/workflows/unit_tests.yml | 33 ++ .gitignore | 377 +++++++++++++++++++ CHANGELOG.md | 4 + README.md | 71 ++++ brainforge_inference_utils/__init__.py | 48 +++ brainforge_inference_utils/config.py | 73 ++++ brainforge_inference_utils/pipeline.py | 149 ++++++++ brainforge_inference_utils/steps/__init__.py | 34 ++ brainforge_inference_utils/steps/base.py | 128 +++++++ brainforge_inference_utils/transaction.py | 39 ++ brainforge_inference_utils/version.py | 5 + requirements/requirements.txt | 2 + setup.py | 88 +++++ 16 files changed, 1116 insertions(+) create mode 100644 .github/workflows/propose_release.yml create mode 100644 .github/workflows/publish_release.yml create mode 100644 .github/workflows/publish_test_build.yml create mode 100644 .github/workflows/unit_tests.yml create mode 100644 .gitignore create mode 100644 CHANGELOG.md create mode 100644 README.md create mode 100644 brainforge_inference_utils/__init__.py create mode 100644 brainforge_inference_utils/config.py create mode 100644 brainforge_inference_utils/pipeline.py create mode 100644 brainforge_inference_utils/steps/__init__.py create mode 100644 brainforge_inference_utils/steps/base.py create mode 100644 brainforge_inference_utils/transaction.py create mode 100644 brainforge_inference_utils/version.py create mode 100644 requirements/requirements.txt create mode 100644 setup.py diff --git a/.github/workflows/propose_release.yml b/.github/workflows/propose_release.yml new file mode 100644 index 0000000..49448c3 --- /dev/null +++ b/.github/workflows/propose_release.yml @@ -0,0 +1,28 @@ +name: Propose Stable Release +on: + workflow_dispatch: + inputs: + release_type: + type: choice + description: Release Type + options: + - patch + - minor + - major +jobs: + update_version: + uses: neongeckocom/.github/.github/workflows/propose_semver_release.yml@master + with: + branch: dev + release_type: ${{ inputs.release_type }} + update_changelog: True + version_file: "brainforge_inference_utils/version.py" + pull_changes: + uses: neongeckocom/.github/.github/workflows/pull_master.yml@master + needs: update_version + with: + pr_reviewer: neonreviewers + pr_assignee: ${{ github.actor }} + pr_draft: false + pr_title: ${{ needs.update_version.outputs.version }} + pr_body: ${{ needs.update_version.outputs.changelog }} diff --git a/.github/workflows/publish_release.yml b/.github/workflows/publish_release.yml new file mode 100644 index 0000000..4901b63 --- /dev/null +++ b/.github/workflows/publish_release.yml @@ -0,0 +1,19 @@ +name: Publish GitHub Release +on: + push: + branches: + - master + +jobs: + tag_release: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Get Version + run: | + VERSION=$(python setup.py --version) + echo "VERSION=${VERSION}" >> $GITHUB_ENV + - uses: ncipollo/release-action@v1 + with: + token: ${{secrets.GITHUB_TOKEN}} + tag: ${{env.VERSION}} diff --git a/.github/workflows/publish_test_build.yml b/.github/workflows/publish_test_build.yml new file mode 100644 index 0000000..3df6b1b --- /dev/null +++ b/.github/workflows/publish_test_build.yml @@ -0,0 +1,18 @@ +# This workflow will generate a distribution and upload it to PyPI + +name: Publish Alpha Build +on: + push: + branches: + - dev + paths-ignore: + - 'brainforge_inference_utils/version.py' + +jobs: + publish_alpha_release: + uses: neongeckocom/.github/.github/workflows/publish_alpha_release.yml@master + secrets: inherit + with: + publish_pypi: false + version_file: "brainforge_inference_utils/version.py" + publish_prerelease: true diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml new file mode 100644 index 0000000..e65b252 --- /dev/null +++ b/.github/workflows/unit_tests.yml @@ -0,0 +1,33 @@ +name: Run Unit Tests +on: + pull_request: + workflow_dispatch: + +jobs: + python_build_tests: + uses: neongeckocom/.github/.github/workflows/python_build_tests.yml@master + unit_tests: + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12"] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install --no-cache-dir .[test] + - name: Unit tests + run: pytest tests -v --junitxml=tests/unit-test-results.xml + - name: Upload test results + if: always() + uses: actions/upload-artifact@v4 + with: + name: unit-test-results-${{ matrix.python-version }} + path: tests/unit-test-results.xml diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..65b543a --- /dev/null +++ b/.gitignore @@ -0,0 +1,377 @@ +## Ignore Visual Studio temporary files, build results, and +## files generated by popular Visual Studio add-ons. +## +## Get latest from https://github.com/github/gitignore/blob/master/VisualStudio.gitignore + +# User-specific files +*.rsuser +*.suo +*.user +*.userosscache +*.sln.docstates + +# User-specific files (MonoDevelop/Xamarin Studio) +*.userprefs + +# Mono auto generated files +mono_crash.* + +# Build results +[Dd]ebug/ +[Dd]ebugPublic/ +[Rr]elease/ +[Rr]eleases/ +x64/ +x86/ +[Ww][Ii][Nn]32/ +[Aa][Rr][Mm]/ +[Aa][Rr][Mm]64/ +bld/ +[Bb]in/ +[Oo]bj/ +[Oo]ut/ +[Ll]og/ +[Ll]ogs/ + +# Visual Studio 2015/2017 cache/options directory +.vs/ +# Uncomment if you have tasks that create the project's static files in wwwroot +#wwwroot/ + +# Visual Studio 2017 auto generated files +Generated\ Files/ + +# MSTest test Results +[Tt]est[Rr]esult*/ +[Bb]uild[Ll]og.* + +# NUnit +*.VisualState.xml +TestResult.xml +nunit-*.xml + +# Build Results of an ATL Project +[Dd]ebugPS/ +[Rr]eleasePS/ +dlldata.c + +# Benchmark Results +BenchmarkDotNet.Artifacts/ + +# .NET Core +project.lock.json +project.fragment.lock.json +artifacts/ + +# ASP.NET Scaffolding +ScaffoldingReadMe.txt + +# StyleCop +StyleCopReport.xml + +# Files built by Visual Studio +*_i.c +*_p.c +*_h.h +*.ilk +*.meta +*.obj +*.iobj +*.pch +*.pdb +*.ipdb +*.pgc +*.pgd +*.rsp +*.sbr +*.tlb +*.tli +*.tlh +*.tmp +*.tmp_proj +*_wpftmp.csproj +*.log +*.vspscc +*.vssscc +.builds +*.pidb +*.svclog +*.scc + +# Chutzpah Test files +_Chutzpah* + +# Visual C++ cache files +ipch/ +*.aps +*.ncb +*.opendb +*.opensdf +*.sdf +*.cachefile +*.VC.db +*.VC.VC.opendb + +# Visual Studio profiler +*.psess +*.vsp +*.vspx +*.sap + +# Visual Studio Trace Files +*.e2e + +# TFS 2012 Local Workspace +$tf/ + +# Guidance Automation Toolkit +*.gpState + +# ReSharper is a .NET coding add-in +_ReSharper*/ +*.[Rr]e[Ss]harper +*.DotSettings.user + +# TeamCity is a build add-in +_TeamCity* + +# DotCover is a Code Coverage Tool +*.dotCover + +# AxoCover is a Code Coverage Tool +.axoCover/* +!.axoCover/settings.json + +# Coverlet is a free, cross platform Code Coverage Tool +coverage*.json +coverage*.xml +coverage*.info + +# Visual Studio code coverage results +*.coverage +*.coveragexml + +# NCrunch +_NCrunch_* +.*crunch*.local.xml +nCrunchTemp_* + +# MightyMoose +*.mm.* +AutoTest.Net/ + +# Web workbench (sass) +.sass-cache/ + +# Installshield output folder +[Ee]xpress/ + +# DocProject is a documentation generator add-in +DocProject/buildhelp/ +DocProject/Help/*.HxT +DocProject/Help/*.HxC +DocProject/Help/*.hhc +DocProject/Help/*.hhk +DocProject/Help/*.hhp +DocProject/Help/Html2 +DocProject/Help/html + +# Click-Once directory +publish/ + +# Publish Web Output +*.[Pp]ublish.xml +*.azurePubxml +# Note: Comment the next line if you want to checkin your web deploy settings, +# but database connection strings (with potential passwords) will be unencrypted +*.pubxml +*.publishproj + +# Microsoft Azure Web App publish settings. Comment the next line if you want to +# checkin your Azure Web App publish settings, but sensitive information contained +# in these scripts will be unencrypted +PublishScripts/ + +# NuGet Packages +*.nupkg +# NuGet Symbol Packages +*.snupkg +# The packages folder can be ignored because of Package Restore +**/[Pp]ackages/* +# except build/, which is used as an MSBuild target. +!**/[Pp]ackages/build/ +# Uncomment if necessary however generally it will be regenerated when needed +#!**/[Pp]ackages/repositories.config +# NuGet v3's project.json files produces more ignorable files +*.nuget.props +*.nuget.targets + +# Microsoft Azure Build Output +csx/ +*.build.csdef + +# Microsoft Azure Emulator +ecf/ +rcf/ + +# Windows Store app package directories and files +AppPackages/ +BundleArtifacts/ +Package.StoreAssociation.xml +_pkginfo.txt +*.appx +*.appxbundle +*.appxupload + +# Visual Studio cache files +# files ending in .cache can be ignored +*.[Cc]ache +# but keep track of directories ending in .cache +!?*.[Cc]ache/ + +# Others +ClientBin/ +~$* +*~ +*.dbmdl +*.dbproj.schemaview +*.jfm +*.pfx +*.publishsettings +orleans.codegen.cs + +# Including strong name files can present a security risk +# (https://github.com/github/gitignore/pull/2483#issue-259490424) +#*.snk + +# Since there are multiple workflows, uncomment next line to ignore bower_components +# (https://github.com/github/gitignore/pull/1529#issuecomment-104372622) +#bower_components/ + +# RIA/Silverlight projects +Generated_Code/ + +# Backup & report files from converting an old project file +# to a newer Visual Studio version. Backup files are not needed, +# because we have git ;-) +_UpgradeReport_Files/ +Backup*/ +UpgradeLog*.XML +UpgradeLog*.htm +ServiceFabricBackup/ +*.rptproj.bak + +# SQL Server files +*.mdf +*.ldf +*.ndf + +# Business Intelligence projects +*.rdl.data +*.bim.layout +*.bim_*.settings +*.rptproj.rsuser +*- [Bb]ackup.rdl +*- [Bb]ackup ([0-9]).rdl +*- [Bb]ackup ([0-9][0-9]).rdl + +# Microsoft Fakes +FakesAssemblies/ + +# GhostDoc plugin setting file +*.GhostDoc.xml + +# Node.js Tools for Visual Studio +.ntvs_analysis.dat +node_modules/ + +# Visual Studio 6 build log +*.plg + +# Visual Studio 6 workspace options file +*.opt + +# Visual Studio 6 auto-generated workspace file (contains which files were open etc.) +*.vbw + +# Visual Studio LightSwitch build output +**/*.HTMLClient/GeneratedArtifacts +**/*.DesktopClient/GeneratedArtifacts +**/*.DesktopClient/ModelManifest.xml +**/*.Server/GeneratedArtifacts +**/*.Server/ModelManifest.xml +_Pvt_Extensions + +# Paket dependency manager +.paket/paket.exe +paket-files/ + +# FAKE - F# Make +.fake/ + +# CodeRush personal settings +.cr/personal + +# Python Tools for Visual Studio (PTVS) +__pycache__/ +*.pyc + +# Cake - Uncomment if you are using it +# tools/** +# !tools/packages.config + +# Tabs Studio +*.tss + +# Telerik's JustMock configuration file +*.jmconfig + +# BizTalk build output +*.btp.cs +*.btm.cs +*.odx.cs +*.xsd.cs + +# OpenCover UI analysis results +OpenCover/ + +# Azure Stream Analytics local run output +ASALocalRun/ + +# MSBuild Binary and Structured Log +*.binlog + +# NVidia Nsight GPU debugger configuration file +*.nvuser + +# MFractors (Xamarin productivity tool) working folder +.mfractor/ + +# Local History for Visual Studio +.localhistory/ + +# BeatPulse healthcheck temp database +healthchecksdb + +# Backup folder for Package Reference Convert tool in Visual Studio 2017 +MigrationBackup/ + +# Ionide (cross platform F# VS Code tools) working folder +.ionide/ + +# Fody - auto-generated XML schema +FodyWeavers.xsd + +# Python files +*.pyc +__pycache__/ +*.egg-info/ +.python-version +.venv/ +venv/ +dist/ +build/ +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +*.py[cod] diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..36d431e --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,4 @@ +# Changelog + +All notable changes to this project are tracked in GitHub releases. This file +is maintained automatically by the release workflow. diff --git a/README.md b/README.md new file mode 100644 index 0000000..b5a83d0 --- /dev/null +++ b/README.md @@ -0,0 +1,71 @@ +# BrainForge Inference Utilities + +Pluggable pipelines for shaping BrainForge LLM requests and responses. + +This package factors the request/response orchestration currently embedded in +`brainforge-webapp` into a standalone, framework-agnostic library. Concrete +behavior (conversation history handling, RAG/MCP context injection, persona +validation, citations, response sanitizing, safety guardrails, ...) is +implemented as **steps** that plug into one of two pipelines: + +- **Input pipeline** — runs before generation and shapes the request: selecting + a persona, windowing/compacting history, retrieving RAG or MCP context, etc. +- **Response pipeline** — runs after generation and shapes the response: + inserting citations, stripping artifacts, applying guardrails, etc. + +Both pipelines operate on a `UserTransaction`, which carries the request and +response as the existing `neon-data-models` models (`LLMRequest` / `LLMResponse`) +used across the BrainForge stack. + +## Status + +Early scaffolding, deliberately minimal. This cut defines only the pipelines and +the step base classes they run; it performs no LLM calls. Generation, a step +registry, and richer per-request state are intentionally deferred. + +## Concepts + +| Concept | Purpose | +| --- | --- | +| `UserTransaction` | The request/response pair a pipeline operates on (`request: LLMRequest`, `response: LLMResponse`). | +| `InputStep` | Shapes `transaction.request` before generation. | +| `ResponseStep` | Annotates/rewrites `transaction.response` after generation. | +| `InputPipeline` / `ResponsePipeline` | Run an ordered list of steps over a transaction. | +| `PipelineConfig` / `StepConfig` | Declarative, ordered step lists (see below). | +| `steps_by_id` | Static `id -> Step class` mapping (mirrors `collectors_by_id` in `brainforge_dataset_utils`); concrete steps are registered here as they land. | +| `build_pipelines(config)` | Build both pipelines from config, resolving step ids via `steps_by_id`. | + +## Usage sketch + +```python +input_pipeline, response_pipeline = build_pipelines(pipeline_config) + +await input_pipeline.run(transaction) # shape request +# ... caller performs generation, sets transaction.response ... +await response_pipeline.run(transaction) # shape response +``` + +Generation sits between the two pipelines and is the caller's responsibility; +this package does not perform or depend on any particular LLM transport. + +## Design notes + +- **Steps are independently defensive.** There is no dependency graph between + steps. A step that needs something it does not find (e.g. a citation step with + no retrieved context) reports itself inapplicable via `applicable()` and is + skipped with a logged reason, rather than failing the pipeline. +- **Ordering is explicit.** Step order is exactly the order declared in + configuration; the pipeline does not reorder. +- **Format contracts live together.** Steps that share an output/parse format + (e.g. a citation formatter and the sanitizer that strips it) belong in the + same module and should derive from a shared constant so they cannot drift. + +## Deferred (stashed for later) + +- A step **registry** (in-process registration + entry-point discovery under + `brainforge.inference.steps`). For now, steps are resolved from the static + `steps_by_id` mapping. +- A **services** adapter for LLM completion/streaming/tool calls and token + counting. +- Richer per-request **state** (context blocks, retrieved documents, execution + trace) beyond the request/response on `UserTransaction`. diff --git a/brainforge_inference_utils/__init__.py b/brainforge_inference_utils/__init__.py new file mode 100644 index 0000000..e38f3bd --- /dev/null +++ b/brainforge_inference_utils/__init__.py @@ -0,0 +1,48 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +"""BrainForge inference pipeline utilities.""" + +from neon_data_models.models.api.llm import ( + LLMPersona, + LLMRequest, + LLMResponse, +) + +from brainforge_inference_utils.version import __version__ + +from brainforge_inference_utils.config import PipelineConfig, StepConfig +from brainforge_inference_utils.pipeline import ( + InputPipeline, + ResponsePipeline, + build_pipelines, +) +from brainforge_inference_utils.steps import steps_by_id +from brainforge_inference_utils.steps.base import ( + InputStep, + ResponseStep, + Step, +) +from brainforge_inference_utils.transaction import UserTransaction + +__all__ = [ + "__version__", + # Request/response models (re-exported from neon-data-models) + "LLMRequest", + "LLMResponse", + "LLMPersona", + # Transaction + "UserTransaction", + # Steps + "Step", + "InputStep", + "ResponseStep", + "steps_by_id", + # Pipelines + "InputPipeline", + "ResponsePipeline", + "build_pipelines", + # Config + "PipelineConfig", + "StepConfig", +] diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py new file mode 100644 index 0000000..b273ebe --- /dev/null +++ b/brainforge_inference_utils/config.py @@ -0,0 +1,73 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +""" +Pipeline configuration schema. + +Configuration is two explicitly ordered lists of step entries: `input` and +`response`. Order is authoritative; the pipeline runs steps in exactly the +order given and does not reorder. Each entry names a step `id` (resolved via the +`steps_by_id` mapping) and optionally carries step-specific `config`. + +Example (YAML):: + + pipeline: + input: + - persona # shorthand for {id: persona} + - {id: ambient_context, config: {timeout: 5.0}} + - {id: rag, config: {top_k: 10}} + - {id: history_window, config: {max_messages: 5}} + - {id: sanitize} + response: + - {id: reasoning_normalize} + - {id: citations} + +A host typically loads its own config and hands the `pipeline` section here:: + + pipeline_config = PipelineConfig.model_validate(config["pipeline"]) +""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class StepConfig(BaseModel): + """ + A single step entry within a pipeline. + + :param id: Step id to look up in `steps_by_id` and instantiate. + :param enabled: When false, the entry is skipped at build time (a + convenient way to turn a step off without deleting its config). + :param config: Step-specific configuration passed to the step constructor. + """ + + model_config = ConfigDict(extra="forbid") + + id: str + enabled: bool = True + config: Dict[str, Any] = Field(default_factory=dict) + + @model_validator(mode="before") + @classmethod + def _coerce_shorthand(cls, data: Any) -> Any: + """Allow a bare string entry as shorthand for `{id: }`.""" + if isinstance(data, str): + return {"id": data} + return data + + +class PipelineConfig(BaseModel): + """ + Ordered step lists for the input and response pipelines. + + :param input: Ordered `InputStep` entries, run before generation. + :param response: Ordered `ResponseStep` entries, run after generation. + """ + + model_config = ConfigDict(extra="forbid") + + input: List[StepConfig] = Field(default_factory=list) + response: List[StepConfig] = Field(default_factory=list) diff --git a/brainforge_inference_utils/pipeline.py b/brainforge_inference_utils/pipeline.py new file mode 100644 index 0000000..738118e --- /dev/null +++ b/brainforge_inference_utils/pipeline.py @@ -0,0 +1,149 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +""" +The input and response pipelines. + +- `InputPipeline` runs ordered `InputStep`s over a `UserTransaction` to shape + its request before generation. +- `ResponsePipeline` runs ordered `ResponseStep`s over a `UserTransaction` to + annotate or rewrite its response after generation. + +Generation itself is out of scope here: a caller runs the input pipeline, +performs generation separately to populate `transaction.response`, then runs the +response pipeline. Steps that are inapplicable to a given transaction are +skipped (with a logged reason) rather than failing the pipeline. +""" + +from __future__ import annotations + +import logging +from typing import List, Mapping, Optional, Sequence, Tuple, Type + +from brainforge_inference_utils.config import PipelineConfig +from brainforge_inference_utils.steps import steps_by_id +from brainforge_inference_utils.steps.base import ( + InputStep, + ResponseStep, + Step, +) +from brainforge_inference_utils.transaction import UserTransaction + +LOG = logging.getLogger(__name__) + + +class InputPipeline: + """ + Runs `InputStep`s in order over a transaction's request. + + :param steps: Ordered input steps. + """ + + def __init__(self, steps: Sequence[InputStep]): + self.steps: List[InputStep] = list(steps) + + async def run(self, transaction: UserTransaction) -> UserTransaction: + """ + Run every applicable input step, mutating the transaction in place. + + :param transaction: The transaction whose request to shape. + :returns: The same `transaction`. + """ + for step in self.steps: + reason = step.applicable(transaction) + if reason is not None: + LOG.debug("Skipping input step %r: %s", step.id, reason) + continue + await step.process(transaction) + return transaction + + +class ResponsePipeline: + """ + Runs `ResponseStep`s in order over a transaction's response. + + :param steps: Ordered response steps. + """ + + def __init__(self, steps: Sequence[ResponseStep]): + self.steps: List[ResponseStep] = list(steps) + + async def run(self, transaction: UserTransaction) -> UserTransaction: + """ + Run every applicable response step, threading the (possibly replaced) + response back onto the transaction. + + :param transaction: The transaction whose response to shape. + :returns: The same `transaction`. + """ + for step in self.steps: + reason = step.applicable(transaction) + if reason is not None: + LOG.debug("Skipping response step %r: %s", step.id, reason) + continue + revised = await step.process(transaction) + if revised is not None: + transaction.response = revised + return transaction + + +def _resolve_step(stage_id: str, + steps: Mapping[str, Type[Step]]) -> Type[Step]: + try: + return steps[stage_id] + except KeyError: + raise KeyError( + f"No step registered with id {stage_id!r}. " + f"Known steps: {sorted(steps)}") + + +def build_pipelines( + config: PipelineConfig, + steps: Optional[Mapping[str, Type[Step]]] = None, +) -> Tuple[InputPipeline, ResponsePipeline]: + """ + Instantiate both pipelines from configuration. + + Step classes are looked up by id in `steps` (defaulting to the static + `steps_by_id` mapping) and validated against the pipeline they are declared + in, so a response step placed in the input list (or vice versa) fails fast + at build time. + + :param config: Parsed pipeline configuration. + :param steps: Optional id -> `Step` class mapping. Defaults to + `brainforge_inference_utils.steps.steps_by_id`. + :returns: `(input_pipeline, response_pipeline)`. + :raises KeyError: If a configured step id is not present in `steps`. + :raises TypeError: If a configured step resolves to the wrong step type. + """ + steps = steps_by_id if steps is None else steps + + input_steps: List[InputStep] = [] + for entry in config.input: + if not entry.enabled: + LOG.info("Input step %r disabled by config", entry.id) + continue + cls = _resolve_step(entry.id, steps) + if not (isinstance(cls, type) and issubclass(cls, InputStep)): + raise TypeError( + f"Step {entry.id!r} is configured as an input step but " + f"{getattr(cls, '__name__', cls)!r} is not an InputStep") + input_steps.append(cls(entry.config)) + + response_steps: List[ResponseStep] = [] + for entry in config.response: + if not entry.enabled: + LOG.info("Response step %r disabled by config", entry.id) + continue + cls = _resolve_step(entry.id, steps) + if not (isinstance(cls, type) and issubclass(cls, ResponseStep)): + raise TypeError( + f"Step {entry.id!r} is configured as a response step but " + f"{getattr(cls, '__name__', cls)!r} is not a ResponseStep") + response_steps.append(cls(entry.config)) + + LOG.info("Built input pipeline: %s", + [s.id for s in input_steps] or "(empty)") + LOG.info("Built response pipeline: %s", + [s.id for s in response_steps] or "(empty)") + return InputPipeline(input_steps), ResponsePipeline(response_steps) diff --git a/brainforge_inference_utils/steps/__init__.py b/brainforge_inference_utils/steps/__init__.py new file mode 100644 index 0000000..a5bb189 --- /dev/null +++ b/brainforge_inference_utils/steps/__init__.py @@ -0,0 +1,34 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +""" +Step base classes, first-party step implementations, and their id mapping. + +Concrete steps are intentionally not shipped yet. New steps subclass +`InputStep` or `ResponseStep`, set a unique `id`, and are added to +`steps_by_id` below (mirroring `collectors_by_id` in `brainforge_dataset_utils`) +so pipelines can resolve them by id. +""" + +from __future__ import annotations + +from typing import Dict, Type + +from brainforge_inference_utils.steps.base import ( + InputStep, + ResponseStep, + Step, +) + +# Static id -> Step class mapping. Populate as concrete steps are implemented, +# e.g.: +# from brainforge_inference_utils.steps.persona import PersonaStep +# from brainforge_inference_utils.steps.citations import CitationFormatStep +# +# steps_by_id = { +# PersonaStep.id: PersonaStep, +# CitationFormatStep.id: CitationFormatStep, +# } +steps_by_id: Dict[str, Type[Step]] = {} + +__all__ = ["Step", "InputStep", "ResponseStep", "steps_by_id"] diff --git a/brainforge_inference_utils/steps/base.py b/brainforge_inference_utils/steps/base.py new file mode 100644 index 0000000..33f8291 --- /dev/null +++ b/brainforge_inference_utils/steps/base.py @@ -0,0 +1,128 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +""" +Step base classes. + +Two step types plug into the two pipelines: + +- `InputStep` runs before generation. It shapes the request on the + `UserTransaction` (persona selection, history windowing/compaction, RAG/MCP + context, ...), mutating it in place. +- `ResponseStep` runs after generation. It has the complete response on the + `UserTransaction` and may return a revised response (or mutate it in place) to + annotate or rewrite it (citations, sanitizing, guardrails, ...). + +Both share two conventions: + +1. **Defensive, not dependent.** Steps declare no relationships to one another. + A step that requires something it does not find reports itself inapplicable + from `applicable()` and is skipped (with a logged reason) rather than failing + the pipeline. This keeps configuration order the only ordering contract and + lets steps be added, removed, or substituted freely. + +2. **Shared object, narrow surface.** Data that must pass between steps travels + on the `UserTransaction`, not through a step-to-step interface. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import ClassVar, Optional + +from neon_data_models.models.api.llm import LLMResponse + +from brainforge_inference_utils.transaction import UserTransaction + + +class Step(ABC): + """ + Common base for all pipeline steps. + + A step is constructed once, at pipeline build time, with its resolved + configuration. The same instance is reused across transactions, so steps + should keep per-request state on the `UserTransaction`, not on `self`. + + :cvar id: Unique step identifier used in configuration. Subclasses must set + this to a stable, non-empty string. + """ + + id: ClassVar[str] = "" + + def __init__(self, config: Optional[dict] = None): + """ + :param config: Step-specific configuration from the pipeline config + entry's `config` mapping. Stored as `self.config`. + """ + self.config: dict = dict(config or {}) + + def __repr__(self) -> str: + return f"<{type(self).__name__} id={self.id!r}>" + + +class InputStep(Step, ABC): + """ + A step that shapes the request before generation. + + Implementations mutate `transaction.request` in place (selecting a persona, + trimming or compacting history, injecting retrieved context, adjusting + generation params, ...). + """ + + def applicable(self, transaction: UserTransaction) -> Optional[str]: + """ + Decide whether this step should run for the given transaction. + + :param transaction: The current transaction. + :returns: `None` to run the step, or a short human-readable reason to + skip it (logged). The default runs unconditionally; override to add + per-request gating. + """ + return None + + @abstractmethod + async def process(self, transaction: UserTransaction) -> None: + """ + Apply this step's effect to the request. + + Only called when `applicable(transaction)` returned `None`. + Implementations mutate `transaction` in place and return nothing. + + :param transaction: The transaction whose `request` to shape. + """ + raise NotImplementedError + + +class ResponseStep(Step, ABC): + """ + A step that shapes the generated response. + + Implementations read the response from `transaction.response` and either + mutate it in place or return a new `LLMResponse` to replace it. Returning + `None` leaves the current response unchanged. + """ + + def applicable(self, transaction: UserTransaction) -> Optional[str]: + """ + Decide whether this step should run for the given transaction. + + :param transaction: The current transaction (with `response` populated). + :returns: `None` to run the step, or a short human-readable reason to + skip it (logged). The default runs unconditionally; override to add + per-request gating (e.g. skip when there is nothing to annotate). + """ + return None + + @abstractmethod + async def process(self, + transaction: UserTransaction) -> Optional[LLMResponse]: + """ + Annotate or rewrite the response. + + Only called when `applicable(transaction)` returned `None`. + + :param transaction: The transaction whose `response` to shape. + :returns: A revised `LLMResponse` to replace the current one, or `None` + to keep it (including when the step mutated the response in place). + """ + raise NotImplementedError diff --git a/brainforge_inference_utils/transaction.py b/brainforge_inference_utils/transaction.py new file mode 100644 index 0000000..81fe7bf --- /dev/null +++ b/brainforge_inference_utils/transaction.py @@ -0,0 +1,39 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +""" +The object a pipeline operates on. + +`UserTransaction` represents one user request/response cycle. Input steps shape +its `request`; generation (performed elsewhere) populates its `response`; +response steps annotate or rewrite that `response`. + +The request and response are the existing `neon-data-models` models used across +the BrainForge stack (HANA, `brainforge-llm-service`), so this package does not +redefine their shape. +""" + +from __future__ import annotations + +from typing import Optional + +from pydantic import BaseModel, Field + +from neon_data_models.models.api.llm import LLMRequest, LLMResponse + + +class UserTransaction(BaseModel): + """ + One user request/response cycle as it flows through the pipelines. + + :param request: The request being shaped by the input pipeline. + :param response: The generated response, populated after generation and + shaped by the response pipeline. `None` until generation completes. + """ + + request: LLMRequest = Field( + description="The LLM request shaped by the input pipeline") + response: Optional[LLMResponse] = Field( + default=None, + description="The generated response shaped by the response pipeline; " + "None until generation completes") diff --git a/brainforge_inference_utils/version.py b/brainforge_inference_utils/version.py new file mode 100644 index 0000000..439820c --- /dev/null +++ b/brainforge_inference_utils/version.py @@ -0,0 +1,5 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. + +__version__ = "0.0.1a0" diff --git a/requirements/requirements.txt b/requirements/requirements.txt new file mode 100644 index 0000000..e9b0d14 --- /dev/null +++ b/requirements/requirements.txt @@ -0,0 +1,2 @@ +pydantic~=2.0 +neon-data-models~=0.0.2,>=0.0.3a8 diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..b3dae2e --- /dev/null +++ b/setup.py @@ -0,0 +1,88 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. + +from os import getenv, path + +from setuptools import find_packages, setup + +BASE_PATH = path.abspath(path.dirname(__file__)) + + +def get_requirements(requirements_filename: str): + requirements_file = path.join( + BASE_PATH, "requirements", requirements_filename + ) + with open(requirements_file, "r", encoding="utf-8") as r: + requirements = r.readlines() + requirements = [ + r.strip() + for r in requirements + if r.strip() and not r.strip().startswith("#") + ] + + for i in range(0, len(requirements)): + r = requirements[i] + if "@" in r: + parts = [ + p.lower() if p.strip().startswith("git+http") else p + for p in r.split("@") + ] + r = "@".join(parts) + if getenv("GITHUB_TOKEN"): + if "github.com" in r: + requirements[i] = r.replace( + "github.com", f"{getenv('GITHUB_TOKEN')}@github.com" + ) + return requirements + + +with open(path.join(BASE_PATH, "README.md"), "r") as f: + long_description = f.read() + +with open( + path.join(BASE_PATH, "brainforge_inference_utils", "version.py"), + "r", + encoding="utf-8", +) as v: + for line in v.readlines(): + if line.startswith("__version__"): + if '"' in line: + version = line.split('"')[1] + else: + version = line.split("'")[1] + + +setup( + name="brainforge-inference-utils", + version=version, + description="Pluggable input/response pipelines for BrainForge LLM inference", + long_description=long_description, + long_description_content_type="text/markdown", + url="https://github.com/NeonGeckoCom/brainforge-inference-utils", + author="NeonGecko", + author_email="developers@neon.ai", + license="", + packages=find_packages(), + include_package_data=True, + install_requires=get_requirements("requirements.txt"), + extras_require={ + "test": get_requirements("test.txt"), + }, + zip_safe=True, + classifiers=[ + "Intended Audience :: Developers", + "Programming Language :: Python :: 3", + ], + # Reserved for a future step registry: first-party and third-party steps + # will be advertised under this entry point group so pipelines can resolve + # them by id. Example: + # entry_points={ + # "brainforge.inference.steps": [ + # "persona=brainforge_inference_utils.steps.persona:PersonaStep", + # ], + # }, + entry_points={ + "brainforge.inference.steps": [], + }, +) From 6d9dbc49373a19fac38fa63399bd387f4f875c56 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Fri, 18 Sep 2026 10:45:58 -0700 Subject: [PATCH 02/23] Implement tests --- requirements/test.txt | 1 + tests/helpers.py | 106 ++++++++++++++++++++++ tests/test_config.py | 100 +++++++++++++++++++++ tests/test_pipeline.py | 183 ++++++++++++++++++++++++++++++++++++++ tests/test_steps.py | 47 ++++++++++ tests/test_transaction.py | 29 ++++++ 6 files changed, 466 insertions(+) create mode 100644 requirements/test.txt create mode 100644 tests/helpers.py create mode 100644 tests/test_config.py create mode 100644 tests/test_pipeline.py create mode 100644 tests/test_steps.py create mode 100644 tests/test_transaction.py diff --git a/requirements/test.txt b/requirements/test.txt new file mode 100644 index 0000000..0c5b054 --- /dev/null +++ b/requirements/test.txt @@ -0,0 +1 @@ +pytest~=8.0 diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..083ffea --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,106 @@ +from typing import Optional + +from neon_data_models.models.api.llm import ( + LLMPersona, + LLMRequest, + LLMResponse, +) + +from brainforge_inference_utils.steps.base import InputStep, ResponseStep +from brainforge_inference_utils.transaction import UserTransaction + + +def make_request(**overrides) -> LLMRequest: + values = dict( + query="hello", + history=[("user", "hi"), ("llm", "hey")], + persona=LLMPersona(name="vanilla"), + model="demo@1", + ) + values.update(overrides) + return LLMRequest(**values) + + +def make_response(text: str = "ok", request: Optional[LLMRequest] = None, + **overrides) -> LLMResponse: + request = request or make_request() + values = dict( + response=text, + history=list(request.history) + [("user", request.query), ("llm", text)], + ) + values.update(overrides) + return LLMResponse(**values) + + +def make_transaction(response: Optional[LLMResponse] = None, + **request_overrides) -> UserTransaction: + request = make_request(**request_overrides) + return UserTransaction(request=request, response=response) + + +class SetMaxHistoryStep(InputStep): + """Input step that writes `config["max_history"]` onto the request.""" + + id = "set_max_history" + + async def process(self, transaction: UserTransaction) -> None: + transaction.request.max_history = self.config["max_history"] + + +class SkipUnlessQueryStep(InputStep): + """Input step skipped unless `request.query` matches `config["query"]`.""" + + id = "skip_unless_query" + + def applicable(self, transaction: UserTransaction) -> Optional[str]: + expected = self.config.get("query") + if transaction.request.query != expected: + return f"query is not {expected!r}" + return None + + async def process(self, transaction: UserTransaction) -> None: + transaction.request.max_tokens = self.config.get("max_tokens", 256) + + +class AppendQueryStep(InputStep): + """Input step that appends `config["suffix"]` to `request.query`.""" + + id = "append_query" + + async def process(self, transaction: UserTransaction) -> None: + transaction.request.query += self.config.get("suffix", "") + + +class PrefixResponseStep(ResponseStep): + """Response step that returns a copy with `config["prefix"]` prepended.""" + + id = "prefix_response" + + def applicable(self, transaction: UserTransaction) -> Optional[str]: + if transaction.response is None: + return "no response" + return None + + async def process(self, transaction: UserTransaction) -> Optional[LLMResponse]: + current = transaction.response + prefix = self.config.get("prefix", "") + return current.model_copy(update={"response": prefix + current.response}) + + +class MutateResponseInPlaceStep(ResponseStep): + """Response step that mutates `response.response` and returns None.""" + + id = "mutate_in_place" + + async def process(self, transaction: UserTransaction) -> Optional[LLMResponse]: + transaction.response.response += self.config.get("suffix", "!") + return None + + +TEST_STEPS = { + SetMaxHistoryStep.id: SetMaxHistoryStep, + SkipUnlessQueryStep.id: SkipUnlessQueryStep, + AppendQueryStep.id: AppendQueryStep, + PrefixResponseStep.id: PrefixResponseStep, + MutateResponseInPlaceStep.id: MutateResponseInPlaceStep, +} diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..802bca5 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,100 @@ +import pytest +from pydantic import ValidationError + +from brainforge_inference_utils.config import PipelineConfig, StepConfig + + +class TestStepConfig: + def test_shorthand_string_becomes_id(self): + parsed = StepConfig.model_validate("persona") + assert parsed.id == "persona" + assert parsed.enabled is True + assert parsed.config == {} + + def test_explicit_mapping(self): + parsed = StepConfig.model_validate({ + "id": "rag", + "enabled": False, + "config": {"top_k": 10}, + }) + assert parsed.id == "rag" + assert parsed.enabled is False + assert parsed.config == {"top_k": 10} + + def test_missing_id_is_invalid(self): + with pytest.raises(ValidationError): + StepConfig.model_validate({"config": {"top_k": 1}}) + + def test_extra_fields_are_forbidden(self): + with pytest.raises(ValidationError): + StepConfig.model_validate({"id": "persona", "unknown": True}) + + def test_constructor_copies_config(self): + payload = {"id": "rag", "config": {"top_k": 3}} + parsed = StepConfig.model_validate(payload) + payload["config"]["top_k"] = 99 + assert parsed.config["top_k"] == 3 + + +class TestPipelineConfig: + def test_defaults_to_empty_lists(self): + parsed = PipelineConfig() + assert parsed.input == [] + assert parsed.response == [] + + def test_mixed_shorthand_and_mapping_entries(self): + parsed = PipelineConfig.model_validate({ + "input": [ + "persona", + {"id": "rag", "config": {"top_k": 10}}, + ], + "response": [ + {"id": "citations"}, + ], + }) + assert [entry.id for entry in parsed.input] == ["persona", "rag"] + assert parsed.input[1].config == {"top_k": 10} + assert parsed.response[0].id == "citations" + + def test_duplicate_step_ids_are_allowed_with_distinct_configs(self): + parsed = PipelineConfig.model_validate({ + "input": [ + {"id": "append_query", "config": {"suffix": " A"}}, + {"id": "append_query", "config": {"suffix": " B"}}, + ], + "response": [ + {"id": "prefix_response", "config": {"prefix": "["}}, + {"id": "prefix_response", "config": {"prefix": "<"}}, + ], + }) + assert [entry.id for entry in parsed.input] == [ + "append_query", + "append_query", + ] + assert [entry.config for entry in parsed.input] == [ + {"suffix": " A"}, + {"suffix": " B"}, + ] + assert [entry.config["prefix"] for entry in parsed.response] == [ + "[", + "<", + ] + + def test_duplicate_identical_entries_are_allowed(self): + parsed = PipelineConfig.model_validate({ + "input": ["sanitize", "sanitize"], + }) + assert [entry.id for entry in parsed.input] == ["sanitize", "sanitize"] + + def test_empty_lists_are_valid(self): + parsed = PipelineConfig.model_validate({"input": [], "response": []}) + assert parsed.input == [] + assert parsed.response == [] + + def test_extra_fields_are_forbidden(self): + with pytest.raises(ValidationError): + PipelineConfig.model_validate({"input": [], "chunk": []}) + + def test_invalid_entry_type_is_rejected(self): + with pytest.raises(ValidationError): + PipelineConfig.model_validate({"input": [123]}) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..df856c3 --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,183 @@ +import unittest + +import pytest + +from brainforge_inference_utils.config import PipelineConfig +from brainforge_inference_utils.pipeline import ( + InputPipeline, + ResponsePipeline, + build_pipelines, +) +from brainforge_inference_utils.steps import steps_by_id +from helpers import ( + AppendQueryStep, + MutateResponseInPlaceStep, + PrefixResponseStep, + SetMaxHistoryStep, + SkipUnlessQueryStep, + TEST_STEPS, + make_response, + make_transaction, +) + + +class TestBuildPipelines: + def test_empty_config_builds_empty_pipelines(self): + input_pipeline, response_pipeline = build_pipelines( + PipelineConfig(), TEST_STEPS) + assert input_pipeline.steps == [] + assert response_pipeline.steps == [] + + def test_default_mapping_is_empty(self): + assert steps_by_id == {} + with pytest.raises(KeyError, match="Known steps: \\[\\]"): + build_pipelines(PipelineConfig.model_validate({ + "input": ["append_query"], + })) + + def test_unknown_id_lists_known_steps(self): + with pytest.raises(KeyError, match="append_query"): + build_pipelines( + PipelineConfig.model_validate({"input": ["nope"]}), + TEST_STEPS, + ) + + def test_disabled_steps_are_omitted(self): + input_pipeline, response_pipeline = build_pipelines( + PipelineConfig.model_validate({ + "input": [ + {"id": "append_query", "enabled": False, + "config": {"suffix": " skipped"}}, + {"id": "append_query", "config": {"suffix": " kept"}}, + ], + "response": [ + {"id": "prefix_response", "enabled": False, + "config": {"prefix": "X"}}, + ], + }), + TEST_STEPS, + ) + assert len(input_pipeline.steps) == 1 + assert input_pipeline.steps[0].config == {"suffix": " kept"} + assert response_pipeline.steps == [] + + def test_response_step_in_input_list_raises(self): + with pytest.raises(TypeError, match="not an InputStep"): + build_pipelines( + PipelineConfig.model_validate({"input": ["prefix_response"]}), + TEST_STEPS, + ) + + def test_input_step_in_response_list_raises(self): + with pytest.raises(TypeError, match="not a ResponseStep"): + build_pipelines( + PipelineConfig.model_validate( + {"response": ["append_query"]}), + TEST_STEPS, + ) + + def test_duplicate_ids_instantiate_independent_instances(self): + input_pipeline, response_pipeline = build_pipelines( + PipelineConfig.model_validate({ + "input": [ + {"id": "append_query", "config": {"suffix": " A"}}, + {"id": "append_query", "config": {"suffix": " B"}}, + ], + "response": [ + {"id": "prefix_response", "config": {"prefix": "["}}, + {"id": "prefix_response", "config": {"prefix": "<"}}, + ], + }), + TEST_STEPS, + ) + assert len(input_pipeline.steps) == 2 + assert input_pipeline.steps[0] is not input_pipeline.steps[1] + assert [type(step) for step in input_pipeline.steps] == [ + AppendQueryStep, + AppendQueryStep, + ] + assert [step.config for step in input_pipeline.steps] == [ + {"suffix": " A"}, + {"suffix": " B"}, + ] + assert [step.config["prefix"] for step in response_pipeline.steps] == [ + "[", + "<", + ] + + +class TestPipelineRun(unittest.IsolatedAsyncioTestCase): + async def test_input_steps_run_in_declared_order(self): + pipeline = InputPipeline([ + AppendQueryStep({"suffix": " A"}), + AppendQueryStep({"suffix": " B"}), + ]) + transaction = make_transaction(query="hello") + result = await pipeline.run(transaction) + assert result is transaction + assert transaction.request.query == "hello A B" + + async def test_inapplicable_input_step_is_skipped(self): + pipeline = InputPipeline([ + SkipUnlessQueryStep({"query": "other", "max_tokens": 64}), + SetMaxHistoryStep({"max_history": 7}), + ]) + transaction = make_transaction(query="hello", max_history=2) + await pipeline.run(transaction) + assert transaction.request.max_tokens == 512 + assert transaction.request.max_history == 7 + + async def test_response_steps_thread_returned_revisions(self): + pipeline = ResponsePipeline([ + PrefixResponseStep({"prefix": "["}), + PrefixResponseStep({"prefix": "<"}), + ]) + transaction = make_transaction() + transaction.response = make_response("ok", transaction.request) + await pipeline.run(transaction) + assert transaction.response.response == "<[ok" + + async def test_response_step_none_keeps_in_place_mutation(self): + pipeline = ResponsePipeline([ + MutateResponseInPlaceStep({"suffix": "!"}), + PrefixResponseStep({"prefix": ">"}), + ]) + transaction = make_transaction() + transaction.response = make_response("ok", transaction.request) + await pipeline.run(transaction) + assert transaction.response.response == ">ok!" + + async def test_inapplicable_response_step_is_skipped(self): + pipeline = ResponsePipeline([ + PrefixResponseStep({"prefix": "["}), + ]) + transaction = make_transaction() + assert transaction.response is None + await pipeline.run(transaction) + assert transaction.response is None + + async def test_build_then_run_duplicate_configured_steps(self): + input_pipeline, response_pipeline = build_pipelines( + PipelineConfig.model_validate({ + "input": [ + {"id": "append_query", "config": {"suffix": " A"}}, + {"id": "skip_unless_query", + "config": {"query": "nope", "max_tokens": 64}}, + {"id": "append_query", "config": {"suffix": " B"}}, + ], + "response": [ + {"id": "prefix_response", "config": {"prefix": "["}}, + {"id": "mutate_in_place", "config": {"suffix": "!"}}, + {"id": "prefix_response", "config": {"prefix": "<"}}, + ], + }), + TEST_STEPS, + ) + transaction = make_transaction(query="hello") + await input_pipeline.run(transaction) + assert transaction.request.query == "hello A B" + assert transaction.request.max_tokens == 512 + + transaction.response = make_response("ok", transaction.request) + await response_pipeline.run(transaction) + assert transaction.response.response == "<[ok!" diff --git a/tests/test_steps.py b/tests/test_steps.py new file mode 100644 index 0000000..75da546 --- /dev/null +++ b/tests/test_steps.py @@ -0,0 +1,47 @@ +from abc import ABC + +import pytest + +from brainforge_inference_utils.steps.base import InputStep, ResponseStep +from helpers import ( + AppendQueryStep, + make_transaction, +) + + +class TestStepBase: + def test_config_defaults_to_empty_dict(self): + step = AppendQueryStep() + assert step.config == {} + + def test_none_config_defaults_to_empty_dict(self): + step = AppendQueryStep(None) + assert step.config == {} + + def test_constructor_copies_config(self): + payload = {"suffix": "!"} + step = AppendQueryStep(payload) + payload["suffix"] = "?" + assert step.config["suffix"] == "!" + + def test_repr_includes_class_and_id(self): + step = AppendQueryStep() + assert repr(step) == "" + + def test_default_applicable_runs_unconditionally(self): + class DefaultResponse(ResponseStep): + id = "default_response" + + async def process(self, transaction): + return None + + assert AppendQueryStep().applicable(make_transaction()) is None + assert DefaultResponse().applicable(make_transaction()) is None + + def test_abstract_bases_cannot_be_instantiated(self): + assert issubclass(InputStep, ABC) + assert issubclass(ResponseStep, ABC) + with pytest.raises(TypeError): + InputStep() + with pytest.raises(TypeError): + ResponseStep() diff --git a/tests/test_transaction.py b/tests/test_transaction.py new file mode 100644 index 0000000..992f288 --- /dev/null +++ b/tests/test_transaction.py @@ -0,0 +1,29 @@ +from pydantic import ValidationError + +from brainforge_inference_utils.transaction import UserTransaction +from helpers import make_request, make_response + + +class TestUserTransaction: + def test_response_defaults_to_none(self): + transaction = UserTransaction(request=make_request()) + assert transaction.response is None + + def test_request_and_response_round_trip(self): + request = make_request(query="ping") + response = make_response("pong", request=request) + transaction = UserTransaction(request=request, response=response) + assert transaction.request.query == "ping" + assert transaction.response.response == "pong" + + def test_missing_request_is_invalid(self): + try: + UserTransaction() + except ValidationError: + return + raise AssertionError("expected ValidationError for missing request") + + def test_request_can_be_mutated_in_place(self): + transaction = UserTransaction(request=make_request(max_history=2)) + transaction.request.max_history = 5 + assert transaction.request.max_history == 5 From 52628b330fa5c6ac0ff54136f501c30918dd5d9e Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Fri, 18 Sep 2026 11:19:24 -0700 Subject: [PATCH 03/23] Manual updates to README content --- README.md | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index b5a83d0..1863ad5 100644 --- a/README.md +++ b/README.md @@ -2,10 +2,9 @@ Pluggable pipelines for shaping BrainForge LLM requests and responses. -This package factors the request/response orchestration currently embedded in -`brainforge-webapp` into a standalone, framework-agnostic library. Concrete +This package defines LLM request/response orchestration. Concrete behavior (conversation history handling, RAG/MCP context injection, persona -validation, citations, response sanitizing, safety guardrails, ...) is +validation, citations, response sanitizing, safety guardrails, ...) can be implemented as **steps** that plug into one of two pipelines: - **Input pipeline** — runs before generation and shapes the request: selecting @@ -17,12 +16,6 @@ Both pipelines operate on a `UserTransaction`, which carries the request and response as the existing `neon-data-models` models (`LLMRequest` / `LLMResponse`) used across the BrainForge stack. -## Status - -Early scaffolding, deliberately minimal. This cut defines only the pipelines and -the step base classes they run; it performs no LLM calls. Generation, a step -registry, and richer per-request state are intentionally deferred. - ## Concepts | Concept | Purpose | @@ -32,7 +25,6 @@ registry, and richer per-request state are intentionally deferred. | `ResponseStep` | Annotates/rewrites `transaction.response` after generation. | | `InputPipeline` / `ResponsePipeline` | Run an ordered list of steps over a transaction. | | `PipelineConfig` / `StepConfig` | Declarative, ordered step lists (see below). | -| `steps_by_id` | Static `id -> Step class` mapping (mirrors `collectors_by_id` in `brainforge_dataset_utils`); concrete steps are registered here as they land. | | `build_pipelines(config)` | Build both pipelines from config, resolving step ids via `steps_by_id`. | ## Usage sketch @@ -60,7 +52,7 @@ this package does not perform or depend on any particular LLM transport. (e.g. a citation formatter and the sanitizer that strips it) belong in the same module and should derive from a shared constant so they cannot drift. -## Deferred (stashed for later) +## Deferred - A step **registry** (in-process registration + entry-point discovery under `brainforge.inference.steps`). For now, steps are resolved from the static From 7816ca43aa0c9a48fad98b2091ebea3b5c179cd9 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Mon, 21 Sep 2026 18:06:31 -0700 Subject: [PATCH 04/23] Implement Client with contract for input/response handling as separate pipelines --- README.md | 57 +++++++++----- brainforge_inference_utils/__init__.py | 2 + brainforge_inference_utils/client.py | 58 +++++++++++++++ brainforge_inference_utils/steps/base.py | 6 +- brainforge_inference_utils/transaction.py | 17 ++++- tests/test_client.py | 90 +++++++++++++++++++++++ tests/test_pipeline.py | 3 +- tests/test_transaction.py | 15 ++++ 8 files changed, 224 insertions(+), 24 deletions(-) create mode 100644 brainforge_inference_utils/client.py create mode 100644 tests/test_client.py diff --git a/README.md b/README.md index 1863ad5..fa729f0 100644 --- a/README.md +++ b/README.md @@ -2,25 +2,32 @@ Pluggable pipelines for shaping BrainForge LLM requests and responses. -This package defines LLM request/response orchestration. Concrete -behavior (conversation history handling, RAG/MCP context injection, persona -validation, citations, response sanitizing, safety guardrails, ...) can be -implemented as **steps** that plug into one of two pipelines: +This package defines request/response orchestration that attaches to an +existing inference workflow. Concrete behavior (conversation history handling, +RAG/MCP context injection, persona validation, citations, response sanitizing, +safety guardrails, ...) can be implemented as **steps** that plug into one of +two pipelines: - **Input pipeline** — runs before generation and shapes the request: selecting a persona, windowing/compacting history, retrieving RAG or MCP context, etc. - **Response pipeline** — runs after generation and shapes the response: inserting citations, stripping artifacts, applying guardrails, etc. -Both pipelines operate on a `UserTransaction`, which carries the request and -response as the existing `neon-data-models` models (`LLMRequest` / `LLMResponse`) -used across the BrainForge stack. +This package does not call an LLM. The host performs generation between the +two pipelines and keeps the `UserTransaction` returned by the input pipeline +until the response is ready. + +Both pipelines operate on a `UserTransaction`, which carries the request, +response, and two per-side context dicts (`input_context` and +`output_context`). Requests and responses use the existing +`neon-data-models` models (`LLMRequest` / `LLMResponse`). ## Concepts | Concept | Purpose | | --- | --- | -| `UserTransaction` | The request/response pair a pipeline operates on (`request: LLMRequest`, `response: LLMResponse`). | +| `UserTransaction` | The request/response pair plus `input_context` / `output_context` dicts. The host holds this object across generation. | +| `BrainForgeClient` | Builds pipelines from config and exposes `run_input_pipeline` / `run_output_pipeline`. | | `InputStep` | Shapes `transaction.request` before generation. | | `ResponseStep` | Annotates/rewrites `transaction.response` after generation. | | `InputPipeline` / `ResponsePipeline` | Run an ordered list of steps over a transaction. | @@ -30,15 +37,19 @@ used across the BrainForge stack. ## Usage sketch ```python -input_pipeline, response_pipeline = build_pipelines(pipeline_config) +client = BrainForgeClient(pipeline_config) -await input_pipeline.run(transaction) # shape request -# ... caller performs generation, sets transaction.response ... -await response_pipeline.run(transaction) # shape response +transaction = await client.run_input_pipeline(llm_request) +# host performs generation using transaction.request +transaction = await client.run_output_pipeline(transaction, llm_response) ``` -Generation sits between the two pipelines and is the caller's responsibility; -this package does not perform or depend on any particular LLM transport. +Input steps write `transaction.input_context[step.id]`. Response steps write +`transaction.output_context[step.id]` and may *read* the input dict. Those +namespaces are independent, so the same step id may exist on both sides +without colliding. Context is available across generation only because the +same `UserTransaction` is passed back in; this package does not store it +between calls. ## Design notes @@ -48,6 +59,17 @@ this package does not perform or depend on any particular LLM transport. skipped with a logged reason, rather than failing the pipeline. - **Ordering is explicit.** Step order is exactly the order declared in configuration; the pipeline does not reorder. +- **Steps may be repeated.** The same step id may appear more than once in a + pipeline, including with different `config`. Each appearance is a separate + instance and runs in declared order. The pipeline does not merge or + de-duplicate entries. +- **Steps own their context.** Each side has its own dict keyed by step id. + A repeated step on the *same* side shares that key; a step id on the input + side does not clash with the same id on the output side. Each step is + responsible for how it stores and reads its own data across one or more + calls (for example appending to a list instead of overwriting a scalar). + Response steps that consume input-step data should read + `transaction.input_context` and skip when nothing usable is present. - **Format contracts live together.** Steps that share an output/parse format (e.g. a citation formatter and the sanitizer that strips it) belong in the same module and should derive from a shared constant so they cannot drift. @@ -57,7 +79,6 @@ this package does not perform or depend on any particular LLM transport. - A step **registry** (in-process registration + entry-point discovery under `brainforge.inference.steps`). For now, steps are resolved from the static `steps_by_id` mapping. -- A **services** adapter for LLM completion/streaming/tool calls and token - counting. -- Richer per-request **state** (context blocks, retrieved documents, execution - trace) beyond the request/response on `UserTransaction`. +- A **services** adapter for streaming, tool calls, and token counting. +- Richer persistent state or execution tracing beyond the transient + `input_context` / `output_context` dicts on `UserTransaction`. diff --git a/brainforge_inference_utils/__init__.py b/brainforge_inference_utils/__init__.py index e38f3bd..06e9bde 100644 --- a/brainforge_inference_utils/__init__.py +++ b/brainforge_inference_utils/__init__.py @@ -11,6 +11,7 @@ from brainforge_inference_utils.version import __version__ +from brainforge_inference_utils.client import BrainForgeClient from brainforge_inference_utils.config import PipelineConfig, StepConfig from brainforge_inference_utils.pipeline import ( InputPipeline, @@ -27,6 +28,7 @@ __all__ = [ "__version__", + "BrainForgeClient", # Request/response models (re-exported from neon-data-models) "LLMRequest", "LLMResponse", diff --git a/brainforge_inference_utils/client.py b/brainforge_inference_utils/client.py new file mode 100644 index 0000000..68ffb81 --- /dev/null +++ b/brainforge_inference_utils/client.py @@ -0,0 +1,58 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +"""Configured runner for the input and response pipelines. + +`BrainForgeClient` does not call an LLM. The host runs inference between +`run_input_pipeline` and `run_output_pipeline`, and is responsible for keeping +the returned `UserTransaction` until the response is ready. +""" + +from __future__ import annotations + +from typing import Mapping, Optional, Type + +from neon_data_models.models.api.llm import LLMRequest, LLMResponse + +from brainforge_inference_utils.config import PipelineConfig +from brainforge_inference_utils.pipeline import build_pipelines +from brainforge_inference_utils.steps.base import Step +from brainforge_inference_utils.transaction import UserTransaction + + +class BrainForgeClient: + """Build pipelines from config and run them around host-owned inference.""" + + def __init__( + self, + config: PipelineConfig, + steps: Optional[Mapping[str, Type[Step]]] = None, + ): + self.config = config + self.input_pipeline, self.response_pipeline = build_pipelines( + config, steps) + + async def run_input_pipeline( + self, request: LLMRequest + ) -> UserTransaction: + """ + Shape `request` and collect step context on a new transaction. + + The caller must hold the returned transaction until + `run_output_pipeline` (or drop it if the response pipeline is unused). + """ + transaction = UserTransaction(request=request) + return await self.input_pipeline.run(transaction) + + async def run_output_pipeline( + self, transaction: UserTransaction, response: LLMResponse + ) -> UserTransaction: + """ + Attach `response` to `transaction` and run the response pipeline. + + Input-step context is on `transaction.input_context`; response steps + write to `transaction.output_context`. The caller must pass through + the transaction from `run_input_pipeline`. + """ + transaction.response = response + return await self.response_pipeline.run(transaction) diff --git a/brainforge_inference_utils/steps/base.py b/brainforge_inference_utils/steps/base.py index 33f8291..62ade3b 100644 --- a/brainforge_inference_utils/steps/base.py +++ b/brainforge_inference_utils/steps/base.py @@ -21,8 +21,10 @@ the pipeline. This keeps configuration order the only ordering contract and lets steps be added, removed, or substituted freely. -2. **Shared object, narrow surface.** Data that must pass between steps travels - on the `UserTransaction`, not through a step-to-step interface. +2. **Shared object, separate namespaces.** Data that must pass between steps + travels on the `UserTransaction`, not through a step-to-step interface. + Input steps use `transaction.input_context`; response steps use + `transaction.output_context`. Step ids are unique only within one side. """ from __future__ import annotations diff --git a/brainforge_inference_utils/transaction.py b/brainforge_inference_utils/transaction.py index 81fe7bf..8543539 100644 --- a/brainforge_inference_utils/transaction.py +++ b/brainforge_inference_utils/transaction.py @@ -5,8 +5,9 @@ The object a pipeline operates on. `UserTransaction` represents one user request/response cycle. Input steps shape -its `request`; generation (performed elsewhere) populates its `response`; -response steps annotate or rewrite that `response`. +its `request` and write `input_context`; generation (performed elsewhere) +populates its `response`; response steps annotate that `response` and write +`output_context`. The request and response are the existing `neon-data-models` models used across the BrainForge stack (HANA, `brainforge-llm-service`), so this package does not @@ -15,7 +16,7 @@ from __future__ import annotations -from typing import Optional +from typing import Any, Dict, Optional from pydantic import BaseModel, Field @@ -29,6 +30,10 @@ class UserTransaction(BaseModel): :param request: The request being shaped by the input pipeline. :param response: The generated response, populated after generation and shaped by the response pipeline. `None` until generation completes. + :param input_context: Transient data written by input steps, keyed by + input step id. + :param output_context: Transient data written by response steps, keyed by + response step id. """ request: LLMRequest = Field( @@ -37,3 +42,9 @@ class UserTransaction(BaseModel): default=None, description="The generated response shaped by the response pipeline; " "None until generation completes") + input_context: Dict[str, Any] = Field( + default_factory=dict, + description="Input-step context keyed by input step id") + output_context: Dict[str, Any] = Field( + default_factory=dict, + description="Response-step context keyed by response step id") diff --git a/tests/test_client.py b/tests/test_client.py new file mode 100644 index 0000000..b0e2b84 --- /dev/null +++ b/tests/test_client.py @@ -0,0 +1,90 @@ +import unittest +from typing import Optional + +from neon_data_models.models.api.llm import LLMResponse + +from brainforge_inference_utils.client import BrainForgeClient +from brainforge_inference_utils.config import PipelineConfig +from brainforge_inference_utils.steps.base import InputStep, ResponseStep +from brainforge_inference_utils.transaction import UserTransaction +from helpers import ( + PrefixResponseStep, + make_request, + make_response, +) + + +class RecordInputStep(InputStep): + id = "record_input" + + async def process(self, transaction: UserTransaction) -> None: + transaction.input_context[self.id] = self.config["value"] + + +class RecordOutputStep(ResponseStep): + id = "record_output" + + async def process( + self, transaction: UserTransaction + ) -> Optional[LLMResponse]: + transaction.output_context[self.id] = self.config["value"] + return None + + +class TestBrainForgeClient(unittest.IsolatedAsyncioTestCase): + def make_client(self): + config = PipelineConfig.model_validate({ + "input": [{ + "id": "record_input", + "config": {"value": {"documents": ["one"]}}, + }], + "response": [ + { + "id": "prefix_response", + "config": {"prefix": "processed: "}, + }, + { + "id": "record_output", + "config": {"value": {"citations": [1]}}, + }, + ], + }) + return BrainForgeClient( + config, + { + RecordInputStep.id: RecordInputStep, + RecordOutputStep.id: RecordOutputStep, + PrefixResponseStep.id: PrefixResponseStep, + }, + ) + + async def test_run_input_pipeline_returns_transaction_with_context(self): + client = self.make_client() + request = make_request() + transaction = await client.run_input_pipeline(request) + + self.assertIs(transaction.request, request) + self.assertIsNone(transaction.response) + self.assertEqual( + transaction.input_context, + {"record_input": {"documents": ["one"]}}, + ) + self.assertEqual(transaction.output_context, {}) + + async def test_run_output_pipeline_keeps_separate_step_namespaces(self): + client = self.make_client() + transaction = await client.run_input_pipeline(make_request()) + response = make_response("answer", transaction.request) + + result = await client.run_output_pipeline(transaction, response) + + self.assertIs(result, transaction) + self.assertEqual(transaction.response.response, "processed: answer") + self.assertEqual( + transaction.input_context, + {"record_input": {"documents": ["one"]}}, + ) + self.assertEqual( + transaction.output_context, + {"record_output": {"citations": [1]}}, + ) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index df856c3..79809a3 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -64,7 +64,8 @@ def test_disabled_steps_are_omitted(self): def test_response_step_in_input_list_raises(self): with pytest.raises(TypeError, match="not an InputStep"): build_pipelines( - PipelineConfig.model_validate({"input": ["prefix_response"]}), + PipelineConfig.model_validate({ + "input": ["prefix_response"]}), TEST_STEPS, ) diff --git a/tests/test_transaction.py b/tests/test_transaction.py index 992f288..bc2f994 100644 --- a/tests/test_transaction.py +++ b/tests/test_transaction.py @@ -27,3 +27,18 @@ def test_request_can_be_mutated_in_place(self): transaction = UserTransaction(request=make_request(max_history=2)) transaction.request.max_history = 5 assert transaction.request.max_history == 5 + + def test_pipeline_context_defaults_to_independent_dicts(self): + first = UserTransaction(request=make_request()) + second = UserTransaction(request=make_request()) + first.input_context["rag"] = {"documents": ["document"]} + first.output_context["citations"] = {"markers": [1]} + assert second.input_context == {} + assert second.output_context == {} + + def test_input_and_output_step_ids_do_not_share_a_namespace(self): + transaction = UserTransaction(request=make_request()) + transaction.input_context["shared_id"] = {"from": "input"} + transaction.output_context["shared_id"] = {"from": "output"} + assert transaction.input_context["shared_id"] == {"from": "input"} + assert transaction.output_context["shared_id"] == {"from": "output"} From ef4425034c5aefcfa418882da1173b1386739c3d Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Tue, 22 Sep 2026 10:31:23 -0700 Subject: [PATCH 05/23] Add `request_context` to enable client-provided context for input pipeline Necessary for functionality like time/weather which need location data not otherwise available --- brainforge_inference_utils/transaction.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/brainforge_inference_utils/transaction.py b/brainforge_inference_utils/transaction.py index 8543539..dc5b80f 100644 --- a/brainforge_inference_utils/transaction.py +++ b/brainforge_inference_utils/transaction.py @@ -42,6 +42,9 @@ class UserTransaction(BaseModel): default=None, description="The generated response shaped by the response pipeline; " "None until generation completes") + request_context: Dict[str, Any] = Field( + default_factory=dict, + description="Client-provided context for the request for input steps") input_context: Dict[str, Any] = Field( default_factory=dict, description="Input-step context keyed by input step id") From f575bc0795d8d2ce06508c89e231f7f16204a0d5 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Tue, 22 Sep 2026 15:01:53 -0700 Subject: [PATCH 06/23] Add pipeline step for time/weather context --- brainforge_inference_utils/config.py | 6 +- brainforge_inference_utils/mcp.py | 66 +++++++ brainforge_inference_utils/steps/__init__.py | 22 +-- .../steps/ambient_context.py | 173 ++++++++++++++++ requirements/requirements.txt | 2 + tests/test_ambient_context.py | 186 ++++++++++++++++++ tests/test_mcp.py | 34 ++++ tests/test_pipeline.py | 12 +- 8 files changed, 486 insertions(+), 15 deletions(-) create mode 100644 brainforge_inference_utils/mcp.py create mode 100644 brainforge_inference_utils/steps/ambient_context.py create mode 100644 tests/test_ambient_context.py create mode 100644 tests/test_mcp.py diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py index b273ebe..1b85113 100644 --- a/brainforge_inference_utils/config.py +++ b/brainforge_inference_utils/config.py @@ -14,7 +14,11 @@ pipeline: input: - persona # shorthand for {id: persona} - - {id: ambient_context, config: {timeout: 5.0}} + - id: ambient_context + config: + mcp_url: https://mcp.neonaiservices2.com/mcp + default_location: Renton, WA, USA + timeout: 5.0 - {id: rag, config: {top_k: 10}} - {id: history_window, config: {max_messages: 5}} - {id: sanitize} diff --git a/brainforge_inference_utils/mcp.py b/brainforge_inference_utils/mcp.py new file mode 100644 index 0000000..6b30940 --- /dev/null +++ b/brainforge_inference_utils/mcp.py @@ -0,0 +1,66 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +"""Generic MCP client helpers shared by pipeline steps.""" + +from __future__ import annotations + +import json +from contextlib import asynccontextmanager +from typing import Any, AsyncIterator + +from fastmcp import Client as MCPClient +from fastmcp.client.transports import StreamableHttpTransport + + +@asynccontextmanager +async def mcp_client(mcp_url: str) -> AsyncIterator[MCPClient]: + """ + Open a short-lived MCP client session for one or more tool calls. + + :param mcp_url: MCP server URL from step configuration. + """ + transport = StreamableHttpTransport(mcp_url) + async with MCPClient(transport) as client: + yield client + + +def mcp_tool_inner_payload(result: Any) -> dict: + """ + Server tool payload from tools/call structuredContent envelope. + + Newer mcp CallToolResult stores structuredContent as a model field; older + stacks put it only in model_extra. Fall back to JSON in text content. + """ + envelope = getattr(result, "structuredContent", None) + if not isinstance(envelope, dict): + envelope = (getattr(result, "model_extra", None) or {}).get( + "structuredContent") or {} + inner = envelope.get("result") or {} + payload = inner.get("result") + if isinstance(payload, dict): + return payload + content = getattr(result, "content", None) or [] + if not content: + return {} + try: + parsed = json.loads(content[0].text) + except (json.JSONDecodeError, TypeError, AttributeError, IndexError): + return {} + inner = parsed.get("result") or {} + payload = inner.get("result") + return payload if isinstance(payload, dict) else {} + + +def mcp_tool_content_parts(tool_result: Any) -> list: + """ + Normalize tools/call handling across fastmcp and raw protocol objects. + """ + if isinstance(tool_result, list): + return tool_result + content = getattr(tool_result, "content", None) + if isinstance(content, list): + return content + raise TypeError( + f"Unexpected MCP tool result {type(tool_result)!r}; " + "expected list or .content list") diff --git a/brainforge_inference_utils/steps/__init__.py b/brainforge_inference_utils/steps/__init__.py index a5bb189..9e2f133 100644 --- a/brainforge_inference_utils/steps/__init__.py +++ b/brainforge_inference_utils/steps/__init__.py @@ -14,21 +14,21 @@ from typing import Dict, Type +from brainforge_inference_utils.steps.ambient_context import AmbientContextStep from brainforge_inference_utils.steps.base import ( InputStep, ResponseStep, Step, ) -# Static id -> Step class mapping. Populate as concrete steps are implemented, -# e.g.: -# from brainforge_inference_utils.steps.persona import PersonaStep -# from brainforge_inference_utils.steps.citations import CitationFormatStep -# -# steps_by_id = { -# PersonaStep.id: PersonaStep, -# CitationFormatStep.id: CitationFormatStep, -# } -steps_by_id: Dict[str, Type[Step]] = {} +steps_by_id: Dict[str, Type[Step]] = { + AmbientContextStep.id: AmbientContextStep, +} -__all__ = ["Step", "InputStep", "ResponseStep", "steps_by_id"] +__all__ = [ + "Step", + "InputStep", + "ResponseStep", + "AmbientContextStep", + "steps_by_id", +] diff --git a/brainforge_inference_utils/steps/ambient_context.py b/brainforge_inference_utils/steps/ambient_context.py new file mode 100644 index 0000000..f1034a4 --- /dev/null +++ b/brainforge_inference_utils/steps/ambient_context.py @@ -0,0 +1,173 @@ +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +""" +Inject ambient time and weather context as `information` messages. + +Mirrors the time/weather injection in `brainforge-webapp`: resolve the user's +location (from `transaction.request_context` or step config), call MCP tools, +and prepend two `information` tuples to `request.history`. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from typing import Any, Dict, List, Optional, Tuple + +import requests +from neon_data_models.types import LlmMessageRole + +from brainforge_inference_utils.mcp import mcp_client, mcp_tool_inner_payload +from brainforge_inference_utils.steps.base import InputStep +from brainforge_inference_utils.transaction import UserTransaction + +LOG = logging.getLogger(__name__) + + +class AmbientContextStep(InputStep): + """ + Prepend time and weather `information` messages to `request.history`. + + Required configuration: + + - `mcp_url`: MCP server URL. + - `default_location`: Fallback when no location can be resolved. + - `timeout`: Seconds to wait for MCP calls. + + `transaction.request_context` may include: + + - `location`: Pre-resolved location description (preferred when set). + - `client_ip`: Client IP used to resolve location via ip-api.com. + """ + + id = "ambient_context" + + def __init__(self, config: Optional[dict] = None) -> None: + resolved = dict(config or {}) + resolved.setdefault("mcp_url", "https://mcp.neonaiservices2.com/mcp") + resolved.setdefault("default_location", "Renton, WA, USA") + resolved.setdefault("timeout", 5.0) + super().__init__(resolved) + + def applicable(self, transaction: UserTransaction) -> Optional[str]: + for key in ("mcp_url", "timeout"): + if self.config.get(key) in (None, ""): + return f"{key} not configured" + return None + + async def process(self, transaction: UserTransaction) -> None: + location = self._resolve_location(transaction) + mcp_url = self.config["mcp_url"] + timeout = float(self.config["timeout"]) + + step_context: Dict[str, Any] = {"location": location} + try: + time_str, weather_str = await self._fetch_time_and_weather( + mcp_url, location, timeout=timeout) + except Exception as exc: + LOG.warning("Skipping ambient context for %r: %s", location, exc) + step_context["error"] = str(exc) + transaction.input_context[self.id] = step_context + return + + messages = self._build_information_messages( + location, time_str, weather_str) + transaction.request.history = ( + list(messages) + list(transaction.request.history)) + + step_context.update({ + "time": time_str, + "weather": weather_str, + "messages": [content for _, content in messages], + }) + transaction.input_context[self.id] = step_context + + def _resolve_location(self, transaction: UserTransaction) -> str: + ctx = transaction.request_context + location = ctx.get("location") + if location: + return str(location) + + client_ip = ctx.get("client_ip") + if client_ip: + return self._get_location_from_ip( + str(client_ip), self.config["default_location"]) + + return self.config["default_location"] + + @staticmethod + def _get_location_from_ip(ip_addr: str, default_location: str) -> str: + """Resolve a human-readable location string from a client IP address.""" + location = None + if ip_addr: + try: + response = requests.get( + f"http://ip-api.com/json/{ip_addr}", + timeout=5, + ).json() + location = ", ".join( + part + for part in ( + response.get("city"), + response.get("region"), + response.get("country"), + ) + if part + ) + except Exception: + location = None + return location or default_location + + async def _fetch_time_and_weather( + self, + mcp_url: str, + location: str, + timeout: float, + ) -> tuple[str, str]: + async def _fetch() -> tuple[str, str]: + async with mcp_client(mcp_url) as client: + time_result = await client.call_tool( + "get_formatted_time", + {"location": location}, + ) + time_str = json.dumps(mcp_tool_inner_payload(time_result)) + + units = "imperial" if "USA" in location else "metric" + weather_result = await client.call_tool_mcp( + "get_weather_forecast", + {"location": location, "unit": units}, + ) + weather_str = self._format_weather_payload( + mcp_tool_inner_payload(weather_result)) + return time_str, weather_str + + return await asyncio.wait_for(_fetch(), timeout=timeout) + + @staticmethod + def _format_weather_payload(weather_dict: dict) -> str: + """Trim weather forecast fields the same way brainforge-webapp does.""" + weather_dict = dict(weather_dict) + weather_dict.pop("minutely", None) + if "hourly" in weather_dict: + weather_dict["hourly"] = weather_dict["hourly"][:24] + if "daily" in weather_dict: + weather_dict["daily"] = weather_dict["daily"][:4] + return json.dumps(weather_dict) + + @staticmethod + def _build_information_messages( + location: str, + time_str: str, + weather_str: str, + ) -> List[Tuple[LlmMessageRole, str]]: + return [ + ("information", + f"The current time and date is: {time_str}.\n" + f"The user's location is {location}."), + ("information", + "Weather forcast reported in Imperial units for USA locations, " + f"otherwise Metric:\n{weather_str}\n---"), + ] + diff --git a/requirements/requirements.txt b/requirements/requirements.txt index e9b0d14..2cf9e1c 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -1,2 +1,4 @@ pydantic~=2.0 neon-data-models~=0.0.2,>=0.0.3a8 +fastmcp~=2.3 +requests diff --git a/tests/test_ambient_context.py b/tests/test_ambient_context.py new file mode 100644 index 0000000..3b80354 --- /dev/null +++ b/tests/test_ambient_context.py @@ -0,0 +1,186 @@ +import json +import unittest +from unittest.mock import AsyncMock, patch + +from brainforge_inference_utils.steps.ambient_context import AmbientContextStep +from brainforge_inference_utils.transaction import UserTransaction +from helpers import make_request + + +STEP_CONFIG = { + "mcp_url": "http://localhost:8000/mcp", + "default_location": "Renton, WA, USA", + "timeout": 5.0, +} + + +class TestFormatWeatherPayload(unittest.TestCase): + def test_trims_forecast_fields(self): + payload = { + "minutely": [{"time": 1}], + "hourly": list(range(30)), + "daily": list(range(10)), + } + formatted = AmbientContextStep._format_weather_payload(payload) + trimmed = json.loads(formatted) + self.assertNotIn("minutely", trimmed) + self.assertEqual(len(trimmed["hourly"]), 24) + self.assertEqual(len(trimmed["daily"]), 4) + + +class TestAmbientContextStep(unittest.IsolatedAsyncioTestCase): + async def test_prepends_information_to_request_history(self): + step = AmbientContextStep(STEP_CONFIG) + transaction = UserTransaction( + request=make_request(history=[("user", "hi"), ("llm", "hey")]), + request_context={"location": "Seattle, WA, USA"}, + ) + + with patch.object( + step, + "_fetch_time_and_weather", + new=AsyncMock(return_value=('{"now": "Monday"}', '{"temp": 60}')), + ) as fetch: + await step.process(transaction) + + fetch.assert_awaited_once_with( + STEP_CONFIG["mcp_url"], "Seattle, WA, USA", timeout=5.0) + self.assertEqual(len(transaction.request.history), 4) + self.assertEqual(transaction.request.history[0][0], "information") + self.assertIn("Monday", transaction.request.history[0][1]) + self.assertIn("Seattle, WA, USA", transaction.request.history[0][1]) + self.assertEqual(transaction.request.history[1][0], "information") + self.assertIn("Weather forcast", transaction.request.history[1][1]) + self.assertIn('{"temp": 60}', transaction.request.history[1][1]) + self.assertEqual(transaction.request.history[2:], [ + ("user", "hi"), + ("llm", "hey"), + ]) + self.assertEqual( + transaction.input_context["ambient_context"]["location"], + "Seattle, WA, USA", + ) + + async def test_resolves_location_from_client_ip(self): + step = AmbientContextStep(STEP_CONFIG) + transaction = UserTransaction( + request=make_request(), + request_context={"client_ip": "203.0.113.10"}, + ) + + with patch.object( + step, + "_get_location_from_ip", + return_value="Portland, OR, USA", + ) as get_location, patch.object( + step, + "_fetch_time_and_weather", + new=AsyncMock(return_value=('{"now": "Tuesday"}', '{"temp": 55}')), + ) as fetch: + await step.process(transaction) + + get_location.assert_called_once_with( + "203.0.113.10", STEP_CONFIG["default_location"]) + fetch.assert_awaited_once_with( + STEP_CONFIG["mcp_url"], "Portland, OR, USA", timeout=5.0) + + async def test_uses_default_location_when_request_context_empty(self): + step = AmbientContextStep(STEP_CONFIG) + transaction = UserTransaction(request=make_request()) + + with patch.object( + step, + "_fetch_time_and_weather", + new=AsyncMock(return_value=('{"now": "Wednesday"}', '{"temp": 50}')), + ) as fetch: + await step.process(transaction) + + fetch.assert_awaited_once_with( + STEP_CONFIG["mcp_url"], + STEP_CONFIG["default_location"], + timeout=5.0, + ) + + async def test_uses_configured_mcp_url(self): + config = { + **STEP_CONFIG, + "mcp_url": "https://mcp.example.com/mcp", + "timeout": 2.5, + } + step = AmbientContextStep(config) + transaction = UserTransaction( + request=make_request(), + request_context={"location": "Austin, TX, USA"}, + ) + + with patch.object( + step, + "_fetch_time_and_weather", + new=AsyncMock(return_value=('{"now": "Thursday"}', '{"temp": 70}')), + ) as fetch: + await step.process(transaction) + + fetch.assert_awaited_once_with( + "https://mcp.example.com/mcp", "Austin, TX, USA", timeout=2.5) + + async def test_mcp_failure_leaves_history_unchanged(self): + step = AmbientContextStep(STEP_CONFIG) + original_history = [("user", "hi"), ("llm", "hey")] + transaction = UserTransaction( + request=make_request(history=list(original_history)), + request_context={"location": "Denver, CO, USA"}, + ) + + with patch.object( + step, + "_fetch_time_and_weather", + new=AsyncMock(side_effect=TimeoutError("timed out")), + ): + await step.process(transaction) + + self.assertEqual(transaction.request.history, original_history) + self.assertEqual( + transaction.input_context["ambient_context"], + {"location": "Denver, CO, USA", "error": "timed out"}, + ) + + async def test_prefers_request_context_location_over_client_ip(self): + step = AmbientContextStep(STEP_CONFIG) + transaction = UserTransaction( + request=make_request(), + request_context={ + "location": "Chicago, IL, USA", + "client_ip": "203.0.113.10", + }, + ) + + with patch.object( + step, + "_get_location_from_ip", + ) as get_location, patch.object( + step, + "_fetch_time_and_weather", + new=AsyncMock(return_value=('{"now": "Friday"}', '{"temp": 40}')), + ) as fetch: + await step.process(transaction) + + get_location.assert_not_called() + fetch.assert_awaited_once_with( + STEP_CONFIG["mcp_url"], "Chicago, IL, USA", timeout=5.0) + + def test_skips_when_timeout_missing(self): + step = AmbientContextStep({ + "mcp_url": "http://localhost:8000/mcp", + "default_location": "Renton, WA, USA", + "timeout": "", + }) + reason = step.applicable(UserTransaction(request=make_request())) + self.assertEqual(reason, "timeout not configured") + + def test_uses_default_location_from_step_config(self): + step = AmbientContextStep({"mcp_url": "http://localhost:8000/mcp"}) + transaction = UserTransaction(request=make_request()) + self.assertEqual( + step._resolve_location(transaction), + "Renton, WA, USA", + ) diff --git a/tests/test_mcp.py b/tests/test_mcp.py new file mode 100644 index 0000000..366282e --- /dev/null +++ b/tests/test_mcp.py @@ -0,0 +1,34 @@ +import json +import unittest +from types import SimpleNamespace + +from brainforge_inference_utils.mcp import ( + mcp_tool_content_parts, + mcp_tool_inner_payload, +) + + +class TestMcpToolHelpers(unittest.TestCase): + def test_inner_payload_from_structured_content(self): + result = SimpleNamespace( + structuredContent={"result": {"result": {"now": "Monday"}}}, + content=[], + ) + self.assertEqual(mcp_tool_inner_payload(result), {"now": "Monday"}) + + def test_inner_payload_from_text_content(self): + payload = {"result": {"result": {"temp": 60}}} + result = SimpleNamespace( + structuredContent=None, + content=[SimpleNamespace(text=json.dumps(payload))], + ) + self.assertEqual(mcp_tool_inner_payload(result), {"temp": 60}) + + def test_content_parts_accepts_list(self): + parts = [{"type": "text", "text": "hello"}] + self.assertEqual(mcp_tool_content_parts(parts), parts) + + def test_content_parts_accepts_result_object(self): + parts = [{"type": "text", "text": "hello"}] + result = SimpleNamespace(content=parts) + self.assertEqual(mcp_tool_content_parts(result), parts) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 79809a3..9f51793 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -28,9 +28,15 @@ def test_empty_config_builds_empty_pipelines(self): assert input_pipeline.steps == [] assert response_pipeline.steps == [] - def test_default_mapping_is_empty(self): - assert steps_by_id == {} - with pytest.raises(KeyError, match="Known steps: \\[\\]"): + def test_default_mapping_includes_first_party_steps(self): + assert "ambient_context" in steps_by_id + input_pipeline, _ = build_pipelines(PipelineConfig.model_validate({ + "input": ["ambient_context"], + })) + assert len(input_pipeline.steps) == 1 + assert input_pipeline.steps[0].id == "ambient_context" + + with pytest.raises(KeyError, match="append_query"): build_pipelines(PipelineConfig.model_validate({ "input": ["append_query"], })) From 2cd1af6e7a09ed3333e3a316ce26b12a93e967a4 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Tue, 22 Sep 2026 17:35:23 -0700 Subject: [PATCH 07/23] Update license notices to AGPL --- brainforge_inference_utils/__init__.py | 17 ++++++++++++++--- brainforge_inference_utils/client.py | 17 ++++++++++++++--- brainforge_inference_utils/config.py | 17 ++++++++++++++--- brainforge_inference_utils/mcp.py | 17 ++++++++++++++--- brainforge_inference_utils/pipeline.py | 17 ++++++++++++++--- brainforge_inference_utils/steps/__init__.py | 17 ++++++++++++++--- .../steps/ambient_context.py | 17 ++++++++++++++--- brainforge_inference_utils/steps/base.py | 17 ++++++++++++++--- brainforge_inference_utils/transaction.py | 17 ++++++++++++++--- brainforge_inference_utils/version.py | 17 ++++++++++++++--- setup.py | 19 +++++++++++++++---- tests/helpers.py | 15 +++++++++++++++ tests/test_ambient_context.py | 15 +++++++++++++++ tests/test_client.py | 15 +++++++++++++++ tests/test_config.py | 15 +++++++++++++++ tests/test_mcp.py | 15 +++++++++++++++ tests/test_pipeline.py | 15 +++++++++++++++ tests/test_steps.py | 15 +++++++++++++++ tests/test_transaction.py | 15 +++++++++++++++ 19 files changed, 275 insertions(+), 34 deletions(-) diff --git a/brainforge_inference_utils/__init__.py b/brainforge_inference_utils/__init__.py index 06e9bde..d56c642 100644 --- a/brainforge_inference_utils/__init__.py +++ b/brainforge_inference_utils/__init__.py @@ -1,6 +1,17 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . """BrainForge inference pipeline utilities.""" from neon_data_models.models.api.llm import ( diff --git a/brainforge_inference_utils/client.py b/brainforge_inference_utils/client.py index 68ffb81..1a6362f 100644 --- a/brainforge_inference_utils/client.py +++ b/brainforge_inference_utils/client.py @@ -1,6 +1,17 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . """Configured runner for the input and response pipelines. `BrainForgeClient` does not call an LLM. The host runs inference between diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py index 1b85113..b05297a 100644 --- a/brainforge_inference_utils/config.py +++ b/brainforge_inference_utils/config.py @@ -1,6 +1,17 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . """ Pipeline configuration schema. diff --git a/brainforge_inference_utils/mcp.py b/brainforge_inference_utils/mcp.py index 6b30940..4109cec 100644 --- a/brainforge_inference_utils/mcp.py +++ b/brainforge_inference_utils/mcp.py @@ -1,6 +1,17 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . """Generic MCP client helpers shared by pipeline steps.""" from __future__ import annotations diff --git a/brainforge_inference_utils/pipeline.py b/brainforge_inference_utils/pipeline.py index 738118e..5cbaf00 100644 --- a/brainforge_inference_utils/pipeline.py +++ b/brainforge_inference_utils/pipeline.py @@ -1,6 +1,17 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . """ The input and response pipelines. diff --git a/brainforge_inference_utils/steps/__init__.py b/brainforge_inference_utils/steps/__init__.py index 9e2f133..6a876cf 100644 --- a/brainforge_inference_utils/steps/__init__.py +++ b/brainforge_inference_utils/steps/__init__.py @@ -1,6 +1,17 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . """ Step base classes, first-party step implementations, and their id mapping. diff --git a/brainforge_inference_utils/steps/ambient_context.py b/brainforge_inference_utils/steps/ambient_context.py index f1034a4..7b8cca4 100644 --- a/brainforge_inference_utils/steps/ambient_context.py +++ b/brainforge_inference_utils/steps/ambient_context.py @@ -1,6 +1,17 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . """ Inject ambient time and weather context as `information` messages. diff --git a/brainforge_inference_utils/steps/base.py b/brainforge_inference_utils/steps/base.py index 62ade3b..7b43d7e 100644 --- a/brainforge_inference_utils/steps/base.py +++ b/brainforge_inference_utils/steps/base.py @@ -1,6 +1,17 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . """ Step base classes. diff --git a/brainforge_inference_utils/transaction.py b/brainforge_inference_utils/transaction.py index dc5b80f..8404650 100644 --- a/brainforge_inference_utils/transaction.py +++ b/brainforge_inference_utils/transaction.py @@ -1,6 +1,17 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . """ The object a pipeline operates on. diff --git a/brainforge_inference_utils/version.py b/brainforge_inference_utils/version.py index 439820c..b355adb 100644 --- a/brainforge_inference_utils/version.py +++ b/brainforge_inference_utils/version.py @@ -1,5 +1,16 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . __version__ = "0.0.1a0" diff --git a/setup.py b/setup.py index b3dae2e..183c464 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,17 @@ -# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework -# All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . from os import getenv, path @@ -62,7 +73,7 @@ def get_requirements(requirements_filename: str): url="https://github.com/NeonGeckoCom/brainforge-inference-utils", author="NeonGecko", author_email="developers@neon.ai", - license="", + license="AGPL-3.0-only", packages=find_packages(), include_package_data=True, install_requires=get_requirements("requirements.txt"), diff --git a/tests/helpers.py b/tests/helpers.py index 083ffea..d2bf12e 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -1,3 +1,18 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + from typing import Optional from neon_data_models.models.api.llm import ( diff --git a/tests/test_ambient_context.py b/tests/test_ambient_context.py index 3b80354..8c86a62 100644 --- a/tests/test_ambient_context.py +++ b/tests/test_ambient_context.py @@ -1,3 +1,18 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + import json import unittest from unittest.mock import AsyncMock, patch diff --git a/tests/test_client.py b/tests/test_client.py index b0e2b84..4c7ea2d 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,3 +1,18 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + import unittest from typing import Optional diff --git a/tests/test_config.py b/tests/test_config.py index 802bca5..6cc5e14 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,3 +1,18 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + import pytest from pydantic import ValidationError diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 366282e..26bcc74 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,3 +1,18 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + import json import unittest from types import SimpleNamespace diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 9f51793..8b40e8f 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -1,3 +1,18 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + import unittest import pytest diff --git a/tests/test_steps.py b/tests/test_steps.py index 75da546..91a00b6 100644 --- a/tests/test_steps.py +++ b/tests/test_steps.py @@ -1,3 +1,18 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + from abc import ABC import pytest diff --git a/tests/test_transaction.py b/tests/test_transaction.py index bc2f994..64f3c5f 100644 --- a/tests/test_transaction.py +++ b/tests/test_transaction.py @@ -1,3 +1,18 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + from pydantic import ValidationError from brainforge_inference_utils.transaction import UserTransaction From c891e93dc92d44fe94d77ebf0c03d57d96097c18 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 23 Sep 2026 17:26:24 -0700 Subject: [PATCH 08/23] Advertise optional params for Steps --- brainforge_inference_utils/__init__.py | 2 + brainforge_inference_utils/client.py | 42 ++++++++++++- .../steps/ambient_context.py | 12 +++- brainforge_inference_utils/steps/base.py | 47 +++++++++++++- tests/test_ambient_context.py | 6 ++ tests/test_client.py | 61 ++++++++++++++++++- tests/test_steps.py | 3 + 7 files changed, 166 insertions(+), 7 deletions(-) diff --git a/brainforge_inference_utils/__init__.py b/brainforge_inference_utils/__init__.py index d56c642..104ec1f 100644 --- a/brainforge_inference_utils/__init__.py +++ b/brainforge_inference_utils/__init__.py @@ -32,6 +32,7 @@ from brainforge_inference_utils.steps import steps_by_id from brainforge_inference_utils.steps.base import ( InputStep, + RequestContextOption, ResponseStep, Step, ) @@ -50,6 +51,7 @@ "Step", "InputStep", "ResponseStep", + "RequestContextOption", "steps_by_id", # Pipelines "InputPipeline", diff --git a/brainforge_inference_utils/client.py b/brainforge_inference_utils/client.py index 1a6362f..1d7d4c4 100644 --- a/brainforge_inference_utils/client.py +++ b/brainforge_inference_utils/client.py @@ -21,13 +21,13 @@ from __future__ import annotations -from typing import Mapping, Optional, Type +from typing import Dict, Mapping, Optional, Type from neon_data_models.models.api.llm import LLMRequest, LLMResponse from brainforge_inference_utils.config import PipelineConfig from brainforge_inference_utils.pipeline import build_pipelines -from brainforge_inference_utils.steps.base import Step +from brainforge_inference_utils.steps.base import RequestContextOption, Step from brainforge_inference_utils.transaction import UserTransaction @@ -43,6 +43,23 @@ def __init__( self.input_pipeline, self.response_pipeline = build_pipelines( config, steps) + async def supported_request_context( + self, timeout: float = 10.0 + ) -> Dict[str, RequestContextOption]: + """ + Collect `request_context` keys and allowed values from the pipeline. + + The host uses this to populate UI; it does not call MCP itself. + Duplicate keys across steps are merged: enumerated values are unioned + in first-seen order, and if any declaration is unconstrained + (`values is None`) the merged option is unconstrained. + """ + merged: Dict[str, RequestContextOption] = {} + for step in (*self.input_pipeline.steps, *self.response_pipeline.steps): + for option in await step.supported_request_context(timeout=timeout): + _merge_request_context_option(merged, option) + return merged + async def run_input_pipeline( self, request: LLMRequest ) -> UserTransaction: @@ -67,3 +84,24 @@ async def run_output_pipeline( """ transaction.response = response return await self.response_pipeline.run(transaction) + + +def _merge_request_context_option( + merged: Dict[str, RequestContextOption], + option: RequestContextOption, +) -> None: + prior = merged.get(option.key) + if prior is None: + merged[option.key] = option + return + description = prior.description or option.description + if prior.values is None or option.values is None: + merged[option.key] = RequestContextOption( + option.key, None, description) + return + values = list(prior.values) + for value in option.values: + if value not in values: + values.append(value) + merged[option.key] = RequestContextOption( + option.key, values, description) diff --git a/brainforge_inference_utils/steps/ambient_context.py b/brainforge_inference_utils/steps/ambient_context.py index 7b8cca4..399b103 100644 --- a/brainforge_inference_utils/steps/ambient_context.py +++ b/brainforge_inference_utils/steps/ambient_context.py @@ -31,7 +31,7 @@ from neon_data_models.types import LlmMessageRole from brainforge_inference_utils.mcp import mcp_client, mcp_tool_inner_payload -from brainforge_inference_utils.steps.base import InputStep +from brainforge_inference_utils.steps.base import InputStep, RequestContextOption from brainforge_inference_utils.transaction import UserTransaction LOG = logging.getLogger(__name__) @@ -54,6 +54,16 @@ class AmbientContextStep(InputStep): """ id = "ambient_context" + request_context_options = ( + RequestContextOption( + "location", + description="Pre-resolved location description (preferred when set).", + ), + RequestContextOption( + "client_ip", + description="Client IP used to resolve location via ip-api.com.", + ), + ) def __init__(self, config: Optional[dict] = None) -> None: resolved = dict(config or {}) diff --git a/brainforge_inference_utils/steps/base.py b/brainforge_inference_utils/steps/base.py index 7b43d7e..b629788 100644 --- a/brainforge_inference_utils/steps/base.py +++ b/brainforge_inference_utils/steps/base.py @@ -18,8 +18,8 @@ Two step types plug into the two pipelines: - `InputStep` runs before generation. It shapes the request on the - `UserTransaction` (persona selection, history windowing/compaction, RAG/MCP - context, ...), mutating it in place. + `UserTransaction` (RAG/MCP context, ambient context, request history + validation, ...), mutating it in place. - `ResponseStep` runs after generation. It has the complete response on the `UserTransaction` and may return a revised response (or mutate it in place) to annotate or rewrite it (citations, sanitizing, guardrails, ...). @@ -41,13 +41,31 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import ClassVar, Optional +from dataclasses import dataclass +from typing import Any, ClassVar, List, Optional, Sequence from neon_data_models.models.api.llm import LLMResponse from brainforge_inference_utils.transaction import UserTransaction +@dataclass(frozen=True) +class RequestContextOption: + """ + One `UserTransaction.request_context` key a step reads. + + Hosts collect these from the built pipeline to present the keys and, when + known, the values a client may send. `values is None` means the key is + supported but unconstrained (any value of the implied type). An empty + `values` list means the key is supported but no choices are available + right now (for example an empty live catalog). + """ + + key: str + values: Optional[Sequence[Any]] = None + description: str = "" + + class Step(ABC): """ Common base for all pipeline steps. @@ -58,9 +76,13 @@ class Step(ABC): :cvar id: Unique step identifier used in configuration. Subclasses must set this to a stable, non-empty string. + :cvar request_context_options: Static `request_context` keys this step + reads. Override `supported_request_context` when allowed values must be + fetched live (for example a provider catalog). """ id: ClassVar[str] = "" + request_context_options: ClassVar[Sequence[RequestContextOption]] = () def __init__(self, config: Optional[dict] = None): """ @@ -72,6 +94,20 @@ def __init__(self, config: Optional[dict] = None): def __repr__(self) -> str: return f"<{type(self).__name__} id={self.id!r}>" + async def supported_request_context( + self, timeout: float = 10.0 + ) -> List[RequestContextOption]: + """ + Return the `request_context` keys this step reads, with allowed values. + + The default copies `request_context_options`. Override when values must + be fetched live (for example a provider catalog). `timeout` is unused + for static declarations. + + :param timeout: Seconds to wait for a live catalog, when overridden. + """ + return list(self.request_context_options) + class InputStep(Step, ABC): """ @@ -80,6 +116,11 @@ class InputStep(Step, ABC): Implementations mutate `transaction.request` in place (selecting a persona, trimming or compacting history, injecting retrieved context, adjusting generation params, ...). + + Steps that read `transaction.request_context` should declare those keys on + `request_context_options` (static choices) or override + `supported_request_context` (live catalogs). Hosts collect the result via + `BrainForgeClient.supported_request_context`. """ def applicable(self, transaction: UserTransaction) -> Optional[str]: diff --git a/tests/test_ambient_context.py b/tests/test_ambient_context.py index 8c86a62..b5f9033 100644 --- a/tests/test_ambient_context.py +++ b/tests/test_ambient_context.py @@ -199,3 +199,9 @@ def test_uses_default_location_from_step_config(self): step._resolve_location(transaction), "Renton, WA, USA", ) + + async def test_supported_request_context_is_unconstrained(self): + step = AmbientContextStep(STEP_CONFIG) + options = await step.supported_request_context() + self.assertEqual([option.key for option in options], ["location", "client_ip"]) + self.assertTrue(all(option.values is None for option in options)) diff --git a/tests/test_client.py b/tests/test_client.py index 4c7ea2d..25f955c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -20,7 +20,11 @@ from brainforge_inference_utils.client import BrainForgeClient from brainforge_inference_utils.config import PipelineConfig -from brainforge_inference_utils.steps.base import InputStep, ResponseStep +from brainforge_inference_utils.steps.base import ( + InputStep, + RequestContextOption, + ResponseStep, +) from brainforge_inference_utils.transaction import UserTransaction from helpers import ( PrefixResponseStep, @@ -36,6 +40,35 @@ async def process(self, transaction: UserTransaction) -> None: transaction.input_context[self.id] = self.config["value"] +class CatalogStep(InputStep): + id = "catalog" + + async def process(self, transaction: UserTransaction) -> None: + return None + + async def supported_request_context(self, timeout: float = 10.0): + return [ + RequestContextOption( + "rag_providers", + list(self.config.get("providers") or []), + ), + RequestContextOption("location", ["Seattle"]), + ] + + +class FlagStep(ResponseStep): + id = "flag" + request_context_options = ( + RequestContextOption("citations_enabled", (True, False)), + RequestContextOption("location"), + ) + + async def process( + self, transaction: UserTransaction + ) -> Optional[LLMResponse]: + return None + + class RecordOutputStep(ResponseStep): id = "record_output" @@ -103,3 +136,29 @@ async def test_run_output_pipeline_keeps_separate_step_namespaces(self): transaction.output_context, {"record_output": {"citations": [1]}}, ) + + async def test_supported_request_context_merges_pipeline_steps(self): + config = PipelineConfig.model_validate({ + "input": [ + {"id": "catalog", "config": {"providers": ["news", "docs"]}}, + {"id": "record_input", "config": {"value": {}}}, + {"id": "catalog", "config": {"providers": ["docs", "web"]}}, + ], + "response": [{"id": "flag"}], + }) + client = BrainForgeClient( + config, + { + CatalogStep.id: CatalogStep, + RecordInputStep.id: RecordInputStep, + FlagStep.id: FlagStep, + }, + ) + + options = await client.supported_request_context() + self.assertEqual( + list(options["rag_providers"].values), ["news", "docs", "web"]) + self.assertEqual( + list(options["citations_enabled"].values), [True, False]) + self.assertIsNone(options["location"].values) + self.assertNotIn("record_input", options) diff --git a/tests/test_steps.py b/tests/test_steps.py index 91a00b6..a10f7fd 100644 --- a/tests/test_steps.py +++ b/tests/test_steps.py @@ -43,6 +43,9 @@ def test_repr_includes_class_and_id(self): step = AppendQueryStep() assert repr(step) == "" + def test_default_request_context_options_are_empty(self): + assert list(AppendQueryStep().request_context_options) == [] + def test_default_applicable_runs_unconditionally(self): class DefaultResponse(ResponseStep): id = "default_response" From 1a9649e4aa1d95eae50f45aca1a1b64e4e5c468d Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 23 Sep 2026 17:27:21 -0700 Subject: [PATCH 09/23] Add parser for history pruning from brainforge-webapp --- .../steps/request_history.py | 177 +++++++++++++++++ tests/test_request_history.py | 185 ++++++++++++++++++ 2 files changed, 362 insertions(+) create mode 100644 brainforge_inference_utils/steps/request_history.py create mode 100644 tests/test_request_history.py diff --git a/brainforge_inference_utils/steps/request_history.py b/brainforge_inference_utils/steps/request_history.py new file mode 100644 index 0000000..58d3ad6 --- /dev/null +++ b/brainforge_inference_utils/steps/request_history.py @@ -0,0 +1,177 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +""" +Guarantee `LLMRequest.history` is valid, including truncation. + +`LLMRequest` contract: history excludes the system prompt (on `persona`) and +the incoming query (on `request.query`). Roles are `information` / +`instruction` / `tool` context, then complete `user`/`llm` pairs. + +This step is meant to run last on the input pipeline so earlier steps can +prepend context freely. It repairs rather than skipping: unknown or empty +messages are dropped, prior citation markup is stripped, `assistant` becomes +`llm`, a trailing user turn that duplicates `query` is removed, then +conversational turns are truncated to `max_history` complete pairs. +""" + +from __future__ import annotations + +import re +from typing import Any, List, Optional, Sequence, Tuple + +from neon_data_models.types import LlmMessageRole + +from brainforge_inference_utils.steps.base import InputStep, RequestContextOption +from brainforge_inference_utils.transaction import UserTransaction + +# Matches `[[n](url)]` written by `ResponseParser.apply_citation_markup`. +_CITATION_MARKUP = re.compile(r"\[\[\d+\]\([^\)]*\)\]") + +_CONTEXT_ROLES = {"information", "instruction", "tool"} +_TURN_ROLES = {"user", "llm", "assistant"} +HistoryItem = Tuple[LlmMessageRole, str] + + +def remove_citation_strings(text: str) -> str: + """Remove citation strings like `[[1](https://example)]` from text.""" + return _CITATION_MARKUP.sub("", text) + + +class RequestHistoryStep(InputStep): + """ + Normalize, truncate, and order `request.history` into a valid LLM history. + + Optional configuration: + + - `max_history`: User/llm pairs to keep, excluding `request.query`. + Defaults to `transaction.request.max_history`. + - `max_messages`: Alternate cap that includes `request.query` + (webapp `history[-5:]` after appending the query). If set, pair budget + is `(max_messages - 1) // 2`. + + `transaction.request_context` may include: + + - `conversational`: When false, drop all prior turns (query only). + Defaults to true. + """ + + id = "request_history" + request_context_options = ( + RequestContextOption( + "conversational", + (True, False), + "When false, drop all prior turns (query only). Defaults to true.", + ), + ) + + def applicable(self, transaction: UserTransaction) -> Optional[str]: + return None + + async def process(self, transaction: UserTransaction) -> None: + conversational = bool( + transaction.request_context.get("conversational", True)) + max_history = self._pair_budget(transaction, conversational) + + cleaned = self._clean_messages( + transaction.request.history, transaction.request.query) + context_messages, turn_messages = self._split_history(cleaned) + windowed_turns = self._last_complete_pairs(turn_messages, max_history) + dropped_turns = len(turn_messages) - len(windowed_turns) + + transaction.request.history = list(context_messages) + list(windowed_turns) + transaction.request.max_history = max_history + transaction.input_context[self.id] = { + "conversational": conversational, + "max_history": max_history, + "kept_turns": len(windowed_turns), + "dropped_turns": dropped_turns, + } + + def _pair_budget(self, transaction: UserTransaction, + conversational: bool) -> int: + if not conversational: + return 0 + if self.config.get("max_messages") not in (None, ""): + max_messages = int(self.config["max_messages"]) + return max(0, (max_messages - 1) // 2) + if self.config.get("max_history") not in (None, ""): + return max(0, int(self.config["max_history"])) + return max(0, int(transaction.request.max_history)) + + @staticmethod + def _clean_messages( + history: Sequence[Tuple[Any, Any]], + query: str, + ) -> List[HistoryItem]: + cleaned: List[HistoryItem] = [] + for role, content in history: + if not isinstance(content, str): + continue + content = remove_citation_strings(content).strip() + if not content: + continue + if role == "assistant": + role = "llm" + if role == "system": + # System prompt belongs on `request.persona`, not history. + continue + if role not in _CONTEXT_ROLES and role not in _TURN_ROLES: + continue + cleaned.append((role, content)) + query_text = remove_citation_strings(query).strip() + if ( + cleaned + and cleaned[-1][0] == "user" + and cleaned[-1][1] == query_text + ): + cleaned.pop() + return cleaned + + @staticmethod + def _split_history( + history: Sequence[HistoryItem], + ) -> Tuple[List[HistoryItem], List[HistoryItem]]: + context_messages: List[HistoryItem] = [] + turn_messages: List[HistoryItem] = [] + for role, content in history: + if role in _TURN_ROLES: + turn_messages.append((role, content)) + else: + context_messages.append((role, content)) + return context_messages, turn_messages + + @staticmethod + def _last_complete_pairs( + turns: Sequence[HistoryItem], + max_history: int, + ) -> List[HistoryItem]: + if max_history <= 0 or not turns: + return [] + pairs: List[Tuple[HistoryItem, HistoryItem]] = [] + idx = len(turns) - 1 + while idx >= 1 and len(pairs) < max_history: + later = turns[idx] + earlier = turns[idx - 1] + if later[0] == "llm" and earlier[0] == "user": + pairs.append((earlier, later)) + idx -= 2 + else: + idx -= 1 + pairs.reverse() + windowed: List[HistoryItem] = [] + for earlier, later in pairs: + windowed.append(earlier) + windowed.append(later) + return windowed diff --git a/tests/test_request_history.py b/tests/test_request_history.py new file mode 100644 index 0000000..ecf8d21 --- /dev/null +++ b/tests/test_request_history.py @@ -0,0 +1,185 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import unittest + +from brainforge_inference_utils.steps.request_history import RequestHistoryStep +from brainforge_inference_utils.transaction import UserTransaction +from helpers import make_request + + +def _turns(*labels: str): + history = [] + for label in labels: + role, _, text = label.partition(":") + history.append((role, text)) + return history + + +class TestRequestHistoryStep(unittest.IsolatedAsyncioTestCase): + async def test_keeps_last_complete_pairs(self): + step = RequestHistoryStep({"max_messages": 5}) + history = _turns( + "user:u1", "llm:a1", "user:u2", "llm:a2", "user:u3", "llm:a3") + transaction = UserTransaction(request=make_request(history=history)) + + await step.process(transaction) + + self.assertEqual(transaction.request.history, [ + ("user", "u2"), + ("llm", "a2"), + ("user", "u3"), + ("llm", "a3"), + ]) + self.assertEqual(transaction.request.max_history, 2) + self.assertEqual( + transaction.input_context["request_history"]["dropped_turns"], 2) + self.assertEqual( + transaction.input_context["request_history"]["kept_turns"], 4) + + async def test_preserves_context_messages_outside_the_window(self): + step = RequestHistoryStep({"max_history": 2}) + history = _turns( + "information:time", + "information:weather", + "user:u1", "llm:a1", + "user:u2", "llm:a2", + "user:u3", "llm:a3", + ) + transaction = UserTransaction(request=make_request(history=history)) + + await step.process(transaction) + + self.assertEqual(transaction.request.history[0], ("information", "time")) + self.assertEqual( + transaction.request.history[1], ("information", "weather")) + self.assertEqual(transaction.request.history[2:], [ + ("user", "u2"), + ("llm", "a2"), + ("user", "u3"), + ("llm", "a3"), + ]) + + async def test_non_conversational_drops_prior_turns(self): + step = RequestHistoryStep({"max_history": 2}) + history = _turns( + "information:rag", + "user:u1", "llm:a1", "user:u2", "llm:a2", + ) + transaction = UserTransaction( + request=make_request(history=history), + request_context={"conversational": False}, + ) + + await step.process(transaction) + + self.assertEqual(transaction.request.history, [("information", "rag")]) + self.assertEqual(transaction.request.max_history, 0) + ctx = transaction.input_context["request_history"] + self.assertFalse(ctx["conversational"]) + self.assertEqual(ctx["dropped_turns"], 4) + + async def test_short_history_is_unchanged(self): + step = RequestHistoryStep({"max_history": 2}) + history = _turns("user:u1", "llm:a1") + transaction = UserTransaction(request=make_request(history=history)) + + await step.process(transaction) + + self.assertEqual(transaction.request.history, history) + self.assertEqual( + transaction.input_context["request_history"]["dropped_turns"], 0) + + async def test_drops_system_empty_and_duplicate_query(self): + step = RequestHistoryStep({"max_history": 2}) + transaction = UserTransaction( + request=make_request( + query="hello", + history=[ + ("information", "ctx"), + ("user", ""), + ("user", "u1"), + ("llm", "a1"), + ("user", "hello"), + ], + ), + ) + # `system` is not an LLMRequest history role; earlier steps or hosts + # may still append it by mutating the list. + transaction.request.history.insert(0, ("system", "you are a bot")) + + await step.process(transaction) + + self.assertEqual(transaction.request.history, [ + ("information", "ctx"), + ("user", "u1"), + ("llm", "a1"), + ]) + + async def test_drops_unpaired_turns_to_keep_complete_pairs(self): + step = RequestHistoryStep({"max_history": 2}) + history = _turns("llm:orphan", "user:u1", "llm:a1", "user:dangling") + transaction = UserTransaction(request=make_request(history=history)) + + await step.process(transaction) + + self.assertEqual(transaction.request.history, [ + ("user", "u1"), + ("llm", "a1"), + ]) + + async def test_uses_request_max_history_by_default(self): + step = RequestHistoryStep() + history = _turns( + "user:u1", "llm:a1", "user:u2", "llm:a2", "user:u3", "llm:a3") + transaction = UserTransaction( + request=make_request(history=history, max_history=1), + ) + + await step.process(transaction) + + self.assertEqual(transaction.request.history, [ + ("user", "u3"), + ("llm", "a3"), + ]) + self.assertEqual(transaction.request.max_history, 1) + + def test_always_applicable(self): + step = RequestHistoryStep() + self.assertIsNone( + step.applicable(UserTransaction(request=make_request()))) + + async def test_strips_citation_markup_from_history(self): + step = RequestHistoryStep({"max_history": 2}) + transaction = UserTransaction( + request=make_request(history=[ + ("user", "what about that?"), + ("llm", "See [[1](https://example.com/a)] for details."), + ]), + ) + + await step.process(transaction) + + self.assertEqual(transaction.request.history, [ + ("user", "what about that?"), + ("llm", "See for details."), + ]) + + async def test_supported_request_context_lists_conversational(self): + step = RequestHistoryStep() + options = await step.supported_request_context() + self.assertEqual(len(options), 1) + self.assertEqual(options[0].key, "conversational") + self.assertEqual(list(options[0].values), [True, False]) From bfa6d17c1bbe45ec15f6e341b13f33806e34fff4 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 23 Sep 2026 17:27:50 -0700 Subject: [PATCH 10/23] Add RAG and Citations from brainforge-webapp with updated config and documentation --- README.md | 26 +- brainforge_inference_utils/config.py | 12 +- brainforge_inference_utils/steps/__init__.py | 17 +- brainforge_inference_utils/steps/rag.py | 388 +++++++++++++++++++ tests/test_pipeline.py | 19 +- tests/test_rag.py | 129 ++++++ tests/test_response_parser.py | 141 +++++++ 7 files changed, 707 insertions(+), 25 deletions(-) create mode 100644 brainforge_inference_utils/steps/rag.py create mode 100644 tests/test_rag.py create mode 100644 tests/test_response_parser.py diff --git a/README.md b/README.md index fa729f0..aa4f026 100644 --- a/README.md +++ b/README.md @@ -3,19 +3,21 @@ Pluggable pipelines for shaping BrainForge LLM requests and responses. This package defines request/response orchestration that attaches to an -existing inference workflow. Concrete behavior (conversation history handling, -RAG/MCP context injection, persona validation, citations, response sanitizing, +existing inference workflow. Concrete behavior (ambient context, RAG/MCP +injection, request history validation, citations, response sanitizing, safety guardrails, ...) can be implemented as **steps** that plug into one of two pipelines: -- **Input pipeline** — runs before generation and shapes the request: selecting - a persona, windowing/compacting history, retrieving RAG or MCP context, etc. +- **Input pipeline** — runs before generation and shapes the request: + injecting ambient time/weather, retrieving RAG for host-selected + providers, then validating and truncating history last. - **Response pipeline** — runs after generation and shapes the response: - inserting citations, stripping artifacts, applying guardrails, etc. + citation rewrite via `ResponseParser`, sanitizing, guardrails, etc. -This package does not call an LLM. The host performs generation between the -two pipelines and keeps the `UserTransaction` returned by the input pipeline -until the response is ready. +Primary generation stays with the host, between the two pipelines. The host +keeps the `UserTransaction` from the input pipeline until the response is +ready. A response step may call a configured OpenAI-compatible endpoint +(for example citation rewrite) without the host owning that second request. Both pipelines operate on a `UserTransaction`, which carries the request, response, and two per-side context dicts (`input_context` and @@ -27,9 +29,10 @@ response, and two per-side context dicts (`input_context` and | Concept | Purpose | | --- | --- | | `UserTransaction` | The request/response pair plus `input_context` / `output_context` dicts. The host holds this object across generation. | -| `BrainForgeClient` | Builds pipelines from config and exposes `run_input_pipeline` / `run_output_pipeline`. | -| `InputStep` | Shapes `transaction.request` before generation. | +| `BrainForgeClient` | Builds pipelines from config and exposes `run_input_pipeline` / `run_output_pipeline` / `supported_request_context`. | +| `InputStep` | Shapes `transaction.request` before generation. Declares supported `request_context` keys/values. | | `ResponseStep` | Annotates/rewrites `transaction.response` after generation. | +| `RequestContextOption` | One `request_context` key a step reads, with optional allowed values for hosts to present. | | `InputPipeline` / `ResponsePipeline` | Run an ordered list of steps over a transaction. | | `PipelineConfig` / `StepConfig` | Declarative, ordered step lists (see below). | | `build_pipelines(config)` | Build both pipelines from config, resolving step ids via `steps_by_id`. | @@ -39,6 +42,9 @@ response, and two per-side context dicts (`input_context` and ```python client = BrainForgeClient(pipeline_config) +options = await client.supported_request_context() +# e.g. options["rag_providers"].values, options["conversational"].values + transaction = await client.run_input_pipeline(llm_request) # host performs generation using transaction.request transaction = await client.run_output_pipeline(transaction, llm_response) diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py index b05297a..90cbd70 100644 --- a/brainforge_inference_utils/config.py +++ b/brainforge_inference_utils/config.py @@ -24,18 +24,18 @@ pipeline: input: - - persona # shorthand for {id: persona} - id: ambient_context config: mcp_url: https://mcp.neonaiservices2.com/mcp default_location: Renton, WA, USA timeout: 5.0 - - {id: rag, config: {top_k: 10}} - - {id: history_window, config: {max_messages: 5}} - - {id: sanitize} + - {id: rag, config: {mcp_url: https://mcp.neonaiservices2.com/mcp, top_k: 10}} + - request_history # last: valid, truncated history response: - - {id: reasoning_normalize} - - {id: citations} + - id: citations + config: + mcp_url: https://mcp.neonaiservices2.com/mcp + openai_url: https://hana.example/brainforge/openai/chat/completions A host typically loads its own config and hands the `pipeline` section here:: diff --git a/brainforge_inference_utils/steps/__init__.py b/brainforge_inference_utils/steps/__init__.py index 6a876cf..7eb0ff9 100644 --- a/brainforge_inference_utils/steps/__init__.py +++ b/brainforge_inference_utils/steps/__init__.py @@ -15,10 +15,9 @@ """ Step base classes, first-party step implementations, and their id mapping. -Concrete steps are intentionally not shipped yet. New steps subclass -`InputStep` or `ResponseStep`, set a unique `id`, and are added to -`steps_by_id` below (mirroring `collectors_by_id` in `brainforge_dataset_utils`) -so pipelines can resolve them by id. +New steps subclass `InputStep` or `ResponseStep`, set a unique `id`, and are +added to `steps_by_id` below (mirroring `collectors_by_id` in +`brainforge_dataset_utils`) so pipelines can resolve them by id. """ from __future__ import annotations @@ -28,18 +27,28 @@ from brainforge_inference_utils.steps.ambient_context import AmbientContextStep from brainforge_inference_utils.steps.base import ( InputStep, + RequestContextOption, ResponseStep, Step, ) +from brainforge_inference_utils.steps.rag import RagStep, ResponseParser +from brainforge_inference_utils.steps.request_history import RequestHistoryStep steps_by_id: Dict[str, Type[Step]] = { AmbientContextStep.id: AmbientContextStep, + RagStep.id: RagStep, + ResponseParser.id: ResponseParser, + RequestHistoryStep.id: RequestHistoryStep, } __all__ = [ "Step", "InputStep", "ResponseStep", + "RequestContextOption", "AmbientContextStep", + "RagStep", + "ResponseParser", + "RequestHistoryStep", "steps_by_id", ] diff --git a/brainforge_inference_utils/steps/rag.py b/brainforge_inference_utils/steps/rag.py new file mode 100644 index 0000000..370005e --- /dev/null +++ b/brainforge_inference_utils/steps/rag.py @@ -0,0 +1,388 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +""" +RAG retrieval and citation rewrite. + +`RagStep` lists providers and injects retrieved documents on the input +pipeline. `ResponseParser` (id `citations`) runs on the response pipeline: +it filters those documents, builds a citation completion via MCP, POSTs it +to a configured OpenAI-compatible endpoint, and writes the annotated text +back onto the same `UserTransaction`. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from typing import Any, Dict, List, Optional, Tuple + +import requests +from neon_data_models.models.api.llm import LLMResponse + +from brainforge_inference_utils.mcp import mcp_client, mcp_tool_content_parts +from brainforge_inference_utils.steps.base import ( + InputStep, + RequestContextOption, + ResponseStep, +) +from brainforge_inference_utils.transaction import UserTransaction + +LOG = logging.getLogger(__name__) + + +class RagStep(InputStep): + """ + Look up selected RAG providers and prepend retrieved context. + + Required configuration: + + - `mcp_url`: MCP server URL. + - `top_k`: Default document cap for `lookup_rag`. + - `timeout`: Seconds to wait for MCP calls. + + `transaction.request_context` may include: + + - `rag_providers`: Provider ids selected by the host (required to run). + - `rag_top_k`: Optional per-request override of configured `top_k`. + """ + + id = "rag" + + def __init__(self, config: Optional[dict] = None) -> None: + resolved = dict(config or {}) + resolved.setdefault("mcp_url", "https://mcp.neonaiservices2.com/mcp") + resolved.setdefault("top_k", 10) + resolved.setdefault("timeout", 10.0) + super().__init__(resolved) + + async def supported_request_context( + self, timeout: float = 10.0 + ) -> List[RequestContextOption]: + """Advertise selected RAG providers (live catalog) and optional top_k.""" + providers = await self._available_providers(timeout=timeout) + return [ + RequestContextOption( + "rag_providers", + providers, + "Provider ids selected by the host (required to run).", + ), + RequestContextOption( + "rag_top_k", + description="Optional per-request override of configured top_k.", + ), + ] + + def applicable(self, transaction: UserTransaction) -> Optional[str]: + if self.config.get("mcp_url") in (None, ""): + return "mcp_url not configured" + if self.config.get("top_k") in (None, ""): + return "top_k not configured" + providers = transaction.request_context.get("rag_providers") or [] + if not providers: + return "no rag providers selected" + return None + + async def _available_providers(self, timeout: float = 10.0) -> List[str]: + """Return the live provider catalog from the configured MCP server.""" + timeout = float(timeout if timeout is not None else self.config["timeout"]) + + async def _fetch() -> List[str]: + async with mcp_client(self.config["mcp_url"]) as client: + result = await client.call_tool_mcp("available_contexts", {}) + if getattr(result, "isError", False): + err_text = ( + result.content[0].text + if getattr(result, "content", None) + else "available_contexts failed" + ) + raise RuntimeError(err_text) + return [ + part.text + for part in mcp_tool_content_parts(result) + if getattr(part, "text", None) + ] + + return await asyncio.wait_for(_fetch(), timeout=timeout) + + async def process(self, transaction: UserTransaction) -> None: + providers = [str(p) for p in ( + transaction.request_context.get("rag_providers") or [])] + top_k = int( + transaction.request_context.get("rag_top_k") + if transaction.request_context.get("rag_top_k") not in (None, "") + else self.config["top_k"] + ) + timeout = float(self.config["timeout"]) + step_context: Dict[str, Any] = { + "providers": providers, + "top_k": top_k, + "documents": [], + } + try: + documents = await self._lookup_rag( + transaction.request.query, providers, top_k, timeout) + except Exception as exc: + LOG.warning("Skipping RAG lookup for %s: %s", providers, exc) + step_context["error"] = str(exc) + transaction.input_context[self.id] = step_context + return + + step_context["documents"] = documents + if documents: + transaction.request.history = ( + [self._information_message(documents)] + + list(transaction.request.history) + ) + transaction.input_context[self.id] = step_context + + async def _lookup_rag( + self, + query: str, + providers: List[str], + top_k: int, + timeout: float, + ) -> List[dict]: + async def _fetch() -> List[dict]: + async with mcp_client(self.config["mcp_url"]) as client: + result = await client.call_tool( + "lookup_rag", + { + "message": query, + "top_k": top_k, + "context_providers": providers, + }, + ) + documents = [] + for part in mcp_tool_content_parts(result): + documents.append(json.loads(part.text)) + return documents + + return await asyncio.wait_for(_fetch(), timeout=timeout) + + @staticmethod + def _information_message(documents: List[dict]) -> tuple: + context_str = "\n".join( + f"{doc.get('metadata', {}).get('file_name', doc.get('metadata', {}).get('citation_uri'))}: " + f"{doc.get('text', '')}\n---" + for doc in documents + ) + return ("information", context_str) + + +def _citation_key(document: dict) -> str: + meta = document.get("metadata") or {} + return str( + meta.get("citation") + or meta.get("citation_uri") + or meta.get("file_name") + or "" + ) + + +def apply_citation_markup( + response: str, documents: List[dict] +) -> Tuple[str, List[str], str, str]: + """ + Map raw `[n]` citation tags onto unique documents and markdown links. + + :returns: rewritten text, unique citation labels, compact credits, + formatted credits for a UI box. + """ + unique_contexts = list({_citation_key(doc) for doc in documents}) + raw_to_unique_idx = { + documents.index(doc) + 1: unique_contexts.index(_citation_key(doc)) + 1 + for doc in documents + } + rewritten = response.replace("[0]", "") + for raw_idx, unique_idx in raw_to_unique_idx.items(): + citation_uri = (documents[raw_idx - 1].get("metadata") or {}).get( + "citation_uri", "") + citation_string = ( + f"[[0{unique_idx}]({citation_uri})]" if citation_uri + else f"[{unique_idx}]" + ) + if f"`[{raw_idx}]`" in rewritten: + rewritten = rewritten.replace(f"`[{raw_idx}]`", citation_string) + else: + rewritten = rewritten.replace(f"[{raw_idx}]", citation_string) + rewritten = rewritten.replace("[[0", "[[") + credits = ", ".join( + f"[{unique_contexts.index(label) + 1}] {label}" + for label in unique_contexts + ) + formatted_credits = "\n".join( + f"[{unique_contexts.index(label) + 1}] {label}" + for label in unique_contexts + ) + return rewritten, unique_contexts, credits, formatted_credits + + +class ResponseParser(ResponseStep): + """ + Rewrite the generated response with citations using RAG input-context. + + The whole MCP filter → citation completion → markup path runs here so the + host does not handle a second request/response pair. + + Required configuration: + + - `mcp_url`: MCP server URL (same server as `RagStep`). + - `openai_url`: OpenAI-compatible chat completions URL (or API base). + - `timeout`: Seconds to wait for MCP and the completion call. + + Optional configuration: + + - `openai_api_key`: Bearer token for `openai_url`. + + `transaction.request_context` may include: + + - `citations_enabled`: When false, the step is skipped. + """ + + id = "citations" + request_context_options = ( + RequestContextOption( + "citations_enabled", + (True, False), + "When false, the step is skipped.", + ), + ) + + def __init__(self, config: Optional[dict] = None) -> None: + resolved = dict(config or {}) + resolved.setdefault("mcp_url", "https://mcp.neonaiservices2.com/mcp") + resolved.setdefault("timeout", 30.0) + super().__init__(resolved) + + def applicable(self, transaction: UserTransaction) -> Optional[str]: + if transaction.response is None: + return "no response" + if self.config.get("mcp_url") in (None, ""): + return "mcp_url not configured" + if self.config.get("openai_url") in (None, ""): + return "openai_url not configured" + if transaction.request_context.get("citations_enabled") is False: + return "citations disabled" + documents = (transaction.input_context.get("rag") or {}).get( + "documents") or [] + if not documents: + return "no rag documents" + return None + + async def process( + self, transaction: UserTransaction + ) -> Optional[LLMResponse]: + documents = list( + (transaction.input_context.get("rag") or {}).get("documents") or []) + original = transaction.response.response + timeout = float(self.config["timeout"]) + try: + filtered = await self._filter_sources(original, documents, timeout) + if not filtered: + transaction.output_context[self.id] = { + "documents": [], + "credits": "", + "formatted_credits": "", + } + return None + payload = await self._build_citation_request( + original, filtered, timeout) + raw_cited = await self._complete( + payload, transaction.request.model, timeout) + rewritten, unique, credits, formatted = apply_citation_markup( + raw_cited, filtered) + except Exception as exc: + LOG.warning("Skipping citation rewrite: %s", exc) + transaction.output_context[self.id] = { + "documents": documents, + "error": str(exc), + } + return None + + transaction.output_context[self.id] = { + "documents": filtered, + "unique_citations": unique, + "credits": credits, + "formatted_credits": formatted, + } + history = list(transaction.response.history) + if history and history[-1][0] == "llm": + history[-1] = ("llm", rewritten) + return transaction.response.model_copy( + update={"response": rewritten, "history": history}) + + async def _filter_sources( + self, response: str, documents: List[dict], timeout: float + ) -> List[dict]: + async def _fetch() -> List[dict]: + async with mcp_client(self.config["mcp_url"]) as client: + result = await client.call_tool( + "filter_citation_sources", + { + "generated_response": response, + "rag_documents": documents, + }, + ) + return [ + json.loads(part.text) + for part in mcp_tool_content_parts(result) + ] + + return await asyncio.wait_for(_fetch(), timeout=timeout) + + async def _build_citation_request( + self, response: str, documents: List[dict], timeout: float + ) -> dict: + async def _fetch() -> dict: + async with mcp_client(self.config["mcp_url"]) as client: + result = await client.call_tool( + "build_citation_request", + { + "original_response": response, + "rag_documents": documents, + }, + ) + parts = mcp_tool_content_parts(result) + return json.loads(parts[0].text) + + return await asyncio.wait_for(_fetch(), timeout=timeout) + + async def _complete( + self, payload: dict, model: str, timeout: float + ) -> str: + request = dict(payload) + request.setdefault("model", model) + request["stream"] = False + url = self._completions_url() + headers = {"Content-Type": "application/json"} + api_key = self.config.get("openai_api_key") + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + def _post() -> dict: + response = requests.post( + url, json=request, headers=headers, timeout=timeout) + response.raise_for_status() + return response.json() + + completion = await asyncio.to_thread(_post) + return completion["choices"][0]["message"]["content"] + + def _completions_url(self) -> str: + url = str(self.config["openai_url"]).rstrip("/") + if url.endswith("/chat/completions") or url.endswith("/completions"): + return url + return f"{url}/chat/completions" diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 8b40e8f..c3d0a94 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -45,11 +45,20 @@ def test_empty_config_builds_empty_pipelines(self): def test_default_mapping_includes_first_party_steps(self): assert "ambient_context" in steps_by_id - input_pipeline, _ = build_pipelines(PipelineConfig.model_validate({ - "input": ["ambient_context"], - })) - assert len(input_pipeline.steps) == 1 - assert input_pipeline.steps[0].id == "ambient_context" + assert "rag" in steps_by_id + assert "citations" in steps_by_id + assert "request_history" in steps_by_id + input_pipeline, response_pipeline = build_pipelines( + PipelineConfig.model_validate({ + "input": ["ambient_context", "rag", "request_history"], + "response": ["citations"], + })) + assert [step.id for step in input_pipeline.steps] == [ + "ambient_context", + "rag", + "request_history", + ] + assert [step.id for step in response_pipeline.steps] == ["citations"] with pytest.raises(KeyError, match="append_query"): build_pipelines(PipelineConfig.model_validate({ diff --git a/tests/test_rag.py b/tests/test_rag.py new file mode 100644 index 0000000..c022060 --- /dev/null +++ b/tests/test_rag.py @@ -0,0 +1,129 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import unittest +from unittest.mock import AsyncMock, patch + +from brainforge_inference_utils.steps.rag import RagStep +from brainforge_inference_utils.transaction import UserTransaction +from helpers import make_request + + +STEP_CONFIG = { + "mcp_url": "http://localhost:8000/mcp", + "top_k": 10, + "timeout": 10.0, +} + +DOCUMENTS = [ + { + "text": "chunk one", + "metadata": {"file_name": "doc-a.md", "citation_uri": "https://a"}, + }, + { + "text": "chunk two", + "metadata": {"citation_uri": "https://b"}, + }, +] + + +class TestRagStep(unittest.IsolatedAsyncioTestCase): + async def test_injects_information_and_stores_documents(self): + step = RagStep(STEP_CONFIG) + transaction = UserTransaction( + request=make_request(query="what?", history=[("user", "hi")]), + request_context={"rag_providers": ["news"]}, + ) + + with patch.object( + step, "_lookup_rag", new=AsyncMock(return_value=DOCUMENTS) + ) as lookup: + await step.process(transaction) + + lookup.assert_awaited_once_with( + "what?", ["news"], 10, STEP_CONFIG["timeout"]) + self.assertEqual(transaction.request.history[0][0], "information") + self.assertIn("doc-a.md: chunk one", transaction.request.history[0][1]) + self.assertIn("https://b: chunk two", transaction.request.history[0][1]) + self.assertEqual(transaction.request.history[1], ("user", "hi")) + self.assertEqual( + transaction.input_context["rag"]["documents"], DOCUMENTS) + self.assertEqual( + transaction.input_context["rag"]["providers"], ["news"]) + + async def test_uses_request_top_k_override(self): + step = RagStep(STEP_CONFIG) + transaction = UserTransaction( + request=make_request(), + request_context={"rag_providers": ["docs"], "rag_top_k": 3}, + ) + + with patch.object( + step, "_lookup_rag", new=AsyncMock(return_value=[]) + ) as lookup: + await step.process(transaction) + + lookup.assert_awaited_once_with( + transaction.request.query, ["docs"], 3, STEP_CONFIG["timeout"]) + self.assertEqual(transaction.request.history, make_request().history) + self.assertEqual(transaction.input_context["rag"]["documents"], []) + + async def test_lookup_failure_leaves_history_unchanged(self): + step = RagStep(STEP_CONFIG) + original = [("user", "hi"), ("llm", "hey")] + transaction = UserTransaction( + request=make_request(history=list(original)), + request_context={"rag_providers": ["news"]}, + ) + + with patch.object( + step, + "_lookup_rag", + new=AsyncMock(side_effect=TimeoutError("timed out")), + ): + await step.process(transaction) + + self.assertEqual(transaction.request.history, original) + self.assertEqual( + transaction.input_context["rag"]["error"], "timed out") + + def test_skips_when_no_providers_selected(self): + step = RagStep(STEP_CONFIG) + reason = step.applicable(UserTransaction(request=make_request())) + self.assertEqual(reason, "no rag providers selected") + + async def test_supported_request_context_uses_live_catalog(self): + step = RagStep(STEP_CONFIG) + with patch.object( + step, + "_available_providers", + new=AsyncMock(return_value=["news", "docs"]), + ) as catalog: + options = { + option.key: option + for option in await step.supported_request_context(timeout=2.0) + } + + catalog.assert_awaited_once_with(timeout=2.0) + self.assertEqual(list(options["rag_providers"].values), ["news", "docs"]) + self.assertIsNone(options["rag_top_k"].values) + + def test_skips_when_mcp_url_missing(self): + step = RagStep({"mcp_url": "", "top_k": 10}) + transaction = UserTransaction( + request=make_request(), + request_context={"rag_providers": ["news"]}, + ) + self.assertEqual(step.applicable(transaction), "mcp_url not configured") diff --git a/tests/test_response_parser.py b/tests/test_response_parser.py new file mode 100644 index 0000000..2cb9b7a --- /dev/null +++ b/tests/test_response_parser.py @@ -0,0 +1,141 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import unittest +from unittest.mock import AsyncMock, patch + +from brainforge_inference_utils.steps.rag import ( + ResponseParser, + apply_citation_markup, +) +from brainforge_inference_utils.transaction import UserTransaction +from helpers import make_request, make_response + + +STEP_CONFIG = { + "mcp_url": "http://localhost:8000/mcp", + "openai_url": "https://hana.example/brainforge/openai/chat/completions", + "openai_api_key": "token", + "timeout": 30.0, +} + +DOCUMENTS = [ + {"text": "one", "metadata": {"citation": "Doc A", "citation_uri": "https://a"}}, + {"text": "two", "metadata": {"citation": "Doc A", "citation_uri": "https://a"}}, + {"text": "three", "metadata": {"citation": "Doc B", "citation_uri": "https://b"}}, +] + + +class TestApplyCitationMarkup(unittest.TestCase): + def test_rewrites_raw_indexes_to_unique_links(self): + rewritten, unique, credits, formatted = apply_citation_markup( + "See [1] and [3] and [0].", DOCUMENTS) + self.assertNotIn("[0]", rewritten) + self.assertIn("](https://", rewritten) + self.assertEqual(len(unique), 2) + self.assertIn("Doc A", credits) + self.assertIn("\n", formatted) + + +class TestResponseParser(unittest.IsolatedAsyncioTestCase): + def _transaction(self, text="See the docs.", **request_context): + request = make_request() + transaction = UserTransaction( + request=request, + response=make_response(text, request), + request_context=request_context, + input_context={"rag": {"documents": DOCUMENTS}}, + ) + return transaction + + async def test_rewrites_response_via_openai_endpoint(self): + step = ResponseParser(STEP_CONFIG) + transaction = self._transaction() + + with patch.object( + step, "_filter_sources", new=AsyncMock(return_value=DOCUMENTS) + ), patch.object( + step, "_build_citation_request", + new=AsyncMock(return_value={"messages": []}), + ) as build, patch.object( + step, "_complete", + new=AsyncMock(return_value="See [1] and [3]."), + ) as complete: + revised = await step.process(transaction) + + build.assert_awaited_once() + complete.assert_awaited_once_with( + {"messages": []}, transaction.request.model, 30.0) + self.assertIsNotNone(revised) + self.assertNotEqual(revised.response, "See the docs.") + self.assertEqual(revised.history[-1][0], "llm") + self.assertEqual(revised.history[-1][1], revised.response) + self.assertIn("Doc A", transaction.output_context["citations"]["credits"]) + + async def test_skips_without_rag_documents(self): + step = ResponseParser(STEP_CONFIG) + transaction = UserTransaction( + request=make_request(), + response=make_response(), + ) + self.assertEqual(step.applicable(transaction), "no rag documents") + + async def test_skips_when_openai_url_missing(self): + step = ResponseParser({"mcp_url": "http://localhost:8000/mcp"}) + transaction = self._transaction() + self.assertEqual(step.applicable(transaction), "openai_url not configured") + + async def test_skips_when_citations_disabled(self): + step = ResponseParser(STEP_CONFIG) + transaction = self._transaction(citations_enabled=False) + self.assertEqual(step.applicable(transaction), "citations disabled") + + async def test_completion_failure_leaves_response(self): + step = ResponseParser(STEP_CONFIG) + transaction = self._transaction() + original = transaction.response.response + + with patch.object( + step, "_filter_sources", new=AsyncMock(return_value=DOCUMENTS) + ), patch.object( + step, "_build_citation_request", + new=AsyncMock(return_value={"messages": []}), + ), patch.object( + step, "_complete", + new=AsyncMock(side_effect=RuntimeError("nope")), + ): + revised = await step.process(transaction) + + self.assertIsNone(revised) + self.assertEqual(transaction.response.response, original) + self.assertEqual( + transaction.output_context["citations"]["error"], "nope") + + def test_appends_chat_completions_to_base_url(self): + step = ResponseParser({ + **STEP_CONFIG, + "openai_url": "https://api.openai.com/v1", + }) + self.assertEqual( + step._completions_url(), + "https://api.openai.com/v1/chat/completions", + ) + + async def test_supported_request_context_lists_citations_enabled(self): + step = ResponseParser(STEP_CONFIG) + options = await step.supported_request_context() + self.assertEqual(len(options), 1) + self.assertEqual(options[0].key, "citations_enabled") + self.assertEqual(list(options[0].values), [True, False]) From 3e6f1a837e72c75bf2c38046e5b416c8abf2b1b6 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 13:49:53 -0700 Subject: [PATCH 11/23] Refactor "Input" and "Output" to consistent "Request" and "Response" --- README.md | 36 ++++++------ brainforge_inference_utils/__init__.py | 8 +-- brainforge_inference_utils/client.py | 22 +++---- brainforge_inference_utils/config.py | 10 ++-- brainforge_inference_utils/pipeline.py | 48 +++++++-------- brainforge_inference_utils/steps/__init__.py | 6 +- .../steps/ambient_context.py | 8 +-- brainforge_inference_utils/steps/base.py | 8 +-- brainforge_inference_utils/steps/rag.py | 22 +++---- .../steps/request_history.py | 8 +-- brainforge_inference_utils/transaction.py | 24 ++++---- setup.py | 2 +- tests/helpers.py | 14 ++--- tests/test_ambient_context.py | 4 +- tests/test_client.py | 58 +++++++++---------- tests/test_config.py | 26 ++++----- tests/test_pipeline.py | 58 +++++++++---------- tests/test_rag.py | 8 +-- tests/test_request_history.py | 8 +-- tests/test_response_parser.py | 6 +- tests/test_steps.py | 6 +- tests/test_transaction.py | 18 +++--- 22 files changed, 204 insertions(+), 204 deletions(-) diff --git a/README.md b/README.md index aa4f026..e014cfa 100644 --- a/README.md +++ b/README.md @@ -8,32 +8,32 @@ injection, request history validation, citations, response sanitizing, safety guardrails, ...) can be implemented as **steps** that plug into one of two pipelines: -- **Input pipeline** — runs before generation and shapes the request: +- **Request pipeline** — runs before generation and shapes the request: injecting ambient time/weather, retrieving RAG for host-selected providers, then validating and truncating history last. - **Response pipeline** — runs after generation and shapes the response: citation rewrite via `ResponseParser`, sanitizing, guardrails, etc. Primary generation stays with the host, between the two pipelines. The host -keeps the `UserTransaction` from the input pipeline until the response is +keeps the `UserTransaction` from the request pipeline until the response is ready. A response step may call a configured OpenAI-compatible endpoint (for example citation rewrite) without the host owning that second request. Both pipelines operate on a `UserTransaction`, which carries the request, -response, and two per-side context dicts (`input_context` and -`output_context`). Requests and responses use the existing +response, and two per-side context dicts (`request_step_context` and +`response_step_context`). Requests and responses use the existing `neon-data-models` models (`LLMRequest` / `LLMResponse`). ## Concepts | Concept | Purpose | | --- | --- | -| `UserTransaction` | The request/response pair plus `input_context` / `output_context` dicts. The host holds this object across generation. | -| `BrainForgeClient` | Builds pipelines from config and exposes `run_input_pipeline` / `run_output_pipeline` / `supported_request_context`. | -| `InputStep` | Shapes `transaction.request` before generation. Declares supported `request_context` keys/values. | +| `UserTransaction` | The request/response pair plus `request_step_context` / `response_step_context` dicts. The host holds this object across generation. | +| `BrainForgeClient` | Builds pipelines from config and exposes `run_request_pipeline` / `run_response_pipeline` / `supported_request_context`. | +| `RequestStep` | Shapes `transaction.request` before generation. Declares supported `request_context` keys/values. | | `ResponseStep` | Annotates/rewrites `transaction.response` after generation. | | `RequestContextOption` | One `request_context` key a step reads, with optional allowed values for hosts to present. | -| `InputPipeline` / `ResponsePipeline` | Run an ordered list of steps over a transaction. | +| `RequestPipeline` / `ResponsePipeline` | Run an ordered list of steps over a transaction. | | `PipelineConfig` / `StepConfig` | Declarative, ordered step lists (see below). | | `build_pipelines(config)` | Build both pipelines from config, resolving step ids via `steps_by_id`. | @@ -45,13 +45,13 @@ client = BrainForgeClient(pipeline_config) options = await client.supported_request_context() # e.g. options["rag_providers"].values, options["conversational"].values -transaction = await client.run_input_pipeline(llm_request) +transaction = await client.run_request_pipeline(llm_request) # host performs generation using transaction.request -transaction = await client.run_output_pipeline(transaction, llm_response) +transaction = await client.run_response_pipeline(transaction, llm_response) ``` -Input steps write `transaction.input_context[step.id]`. Response steps write -`transaction.output_context[step.id]` and may *read* the input dict. Those +Request steps write `transaction.request_step_context[step.id]`. Response steps write +`transaction.response_step_context[step.id]` and may *read* the request-step context. Those namespaces are independent, so the same step id may exist on both sides without colliding. Context is available across generation only because the same `UserTransaction` is passed back in; this package does not store it @@ -70,13 +70,13 @@ between calls. instance and runs in declared order. The pipeline does not merge or de-duplicate entries. - **Steps own their context.** Each side has its own dict keyed by step id. - A repeated step on the *same* side shares that key; a step id on the input - side does not clash with the same id on the output side. Each step is + A repeated step on the *same* side shares that key; a step id on the request + side does not clash with the same id on the response side. Each step is responsible for how it stores and reads its own data across one or more calls (for example appending to a list instead of overwriting a scalar). - Response steps that consume input-step data should read - `transaction.input_context` and skip when nothing usable is present. -- **Format contracts live together.** Steps that share an output/parse format + Response steps that consume request-step data should read + `transaction.request_step_context` and skip when nothing usable is present. +- **Format contracts live together.** Steps that share a parse format (e.g. a citation formatter and the sanitizer that strips it) belong in the same module and should derive from a shared constant so they cannot drift. @@ -87,4 +87,4 @@ between calls. `steps_by_id` mapping. - A **services** adapter for streaming, tool calls, and token counting. - Richer persistent state or execution tracing beyond the transient - `input_context` / `output_context` dicts on `UserTransaction`. + `request_step_context` / `response_step_context` dicts on `UserTransaction`. diff --git a/brainforge_inference_utils/__init__.py b/brainforge_inference_utils/__init__.py index 104ec1f..4e14254 100644 --- a/brainforge_inference_utils/__init__.py +++ b/brainforge_inference_utils/__init__.py @@ -25,13 +25,13 @@ from brainforge_inference_utils.client import BrainForgeClient from brainforge_inference_utils.config import PipelineConfig, StepConfig from brainforge_inference_utils.pipeline import ( - InputPipeline, + RequestPipeline, ResponsePipeline, build_pipelines, ) from brainforge_inference_utils.steps import steps_by_id from brainforge_inference_utils.steps.base import ( - InputStep, + RequestStep, RequestContextOption, ResponseStep, Step, @@ -49,12 +49,12 @@ "UserTransaction", # Steps "Step", - "InputStep", + "RequestStep", "ResponseStep", "RequestContextOption", "steps_by_id", # Pipelines - "InputPipeline", + "RequestPipeline", "ResponsePipeline", "build_pipelines", # Config diff --git a/brainforge_inference_utils/client.py b/brainforge_inference_utils/client.py index 1d7d4c4..7faa00b 100644 --- a/brainforge_inference_utils/client.py +++ b/brainforge_inference_utils/client.py @@ -12,10 +12,10 @@ # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -"""Configured runner for the input and response pipelines. +"""Configured runner for the request and response pipelines. `BrainForgeClient` does not call an LLM. The host runs inference between -`run_input_pipeline` and `run_output_pipeline`, and is responsible for keeping +`run_request_pipeline` and `run_response_pipeline`, and is responsible for keeping the returned `UserTransaction` until the response is ready. """ @@ -40,7 +40,7 @@ def __init__( steps: Optional[Mapping[str, Type[Step]]] = None, ): self.config = config - self.input_pipeline, self.response_pipeline = build_pipelines( + self.request_pipeline, self.response_pipeline = build_pipelines( config, steps) async def supported_request_context( @@ -55,32 +55,32 @@ async def supported_request_context( (`values is None`) the merged option is unconstrained. """ merged: Dict[str, RequestContextOption] = {} - for step in (*self.input_pipeline.steps, *self.response_pipeline.steps): + for step in (*self.request_pipeline.steps, *self.response_pipeline.steps): for option in await step.supported_request_context(timeout=timeout): _merge_request_context_option(merged, option) return merged - async def run_input_pipeline( + async def run_request_pipeline( self, request: LLMRequest ) -> UserTransaction: """ Shape `request` and collect step context on a new transaction. The caller must hold the returned transaction until - `run_output_pipeline` (or drop it if the response pipeline is unused). + `run_response_pipeline` (or drop it if the response pipeline is unused). """ transaction = UserTransaction(request=request) - return await self.input_pipeline.run(transaction) + return await self.request_pipeline.run(transaction) - async def run_output_pipeline( + async def run_response_pipeline( self, transaction: UserTransaction, response: LLMResponse ) -> UserTransaction: """ Attach `response` to `transaction` and run the response pipeline. - Input-step context is on `transaction.input_context`; response steps - write to `transaction.output_context`. The caller must pass through - the transaction from `run_input_pipeline`. + Request-step context is on `transaction.request_step_context`; response steps + write to `transaction.response_step_context`. The caller must pass through + the transaction from `run_request_pipeline`. """ transaction.response = response return await self.response_pipeline.run(transaction) diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py index 90cbd70..50f0359 100644 --- a/brainforge_inference_utils/config.py +++ b/brainforge_inference_utils/config.py @@ -15,7 +15,7 @@ """ Pipeline configuration schema. -Configuration is two explicitly ordered lists of step entries: `input` and +Configuration is two explicitly ordered lists of step entries: `request` and `response`. Order is authoritative; the pipeline runs steps in exactly the order given and does not reorder. Each entry names a step `id` (resolved via the `steps_by_id` mapping) and optionally carries step-specific `config`. @@ -23,7 +23,7 @@ Example (YAML):: pipeline: - input: + request: - id: ambient_context config: mcp_url: https://mcp.neonaiservices2.com/mcp @@ -76,13 +76,13 @@ def _coerce_shorthand(cls, data: Any) -> Any: class PipelineConfig(BaseModel): """ - Ordered step lists for the input and response pipelines. + Ordered step lists for the request and response pipelines. - :param input: Ordered `InputStep` entries, run before generation. + :param request: Ordered `RequestStep` entries, run before generation. :param response: Ordered `ResponseStep` entries, run after generation. """ model_config = ConfigDict(extra="forbid") - input: List[StepConfig] = Field(default_factory=list) + request: List[StepConfig] = Field(default_factory=list) response: List[StepConfig] = Field(default_factory=list) diff --git a/brainforge_inference_utils/pipeline.py b/brainforge_inference_utils/pipeline.py index 5cbaf00..35f1cdc 100644 --- a/brainforge_inference_utils/pipeline.py +++ b/brainforge_inference_utils/pipeline.py @@ -13,14 +13,14 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . """ -The input and response pipelines. +The request and response pipelines. -- `InputPipeline` runs ordered `InputStep`s over a `UserTransaction` to shape +- `RequestPipeline` runs ordered `RequestStep`s over a `UserTransaction` to shape its request before generation. - `ResponsePipeline` runs ordered `ResponseStep`s over a `UserTransaction` to annotate or rewrite its response after generation. -Generation itself is out of scope here: a caller runs the input pipeline, +Generation itself is out of scope here: a caller runs the request pipeline, performs generation separately to populate `transaction.response`, then runs the response pipeline. Steps that are inapplicable to a given transaction are skipped (with a logged reason) rather than failing the pipeline. @@ -34,7 +34,7 @@ from brainforge_inference_utils.config import PipelineConfig from brainforge_inference_utils.steps import steps_by_id from brainforge_inference_utils.steps.base import ( - InputStep, + RequestStep, ResponseStep, Step, ) @@ -43,19 +43,19 @@ LOG = logging.getLogger(__name__) -class InputPipeline: +class RequestPipeline: """ - Runs `InputStep`s in order over a transaction's request. + Runs `RequestStep`s in order over a transaction's request. - :param steps: Ordered input steps. + :param steps: Ordered request steps. """ - def __init__(self, steps: Sequence[InputStep]): - self.steps: List[InputStep] = list(steps) + def __init__(self, steps: Sequence[RequestStep]): + self.steps: List[RequestStep] = list(steps) async def run(self, transaction: UserTransaction) -> UserTransaction: """ - Run every applicable input step, mutating the transaction in place. + Run every applicable request step, mutating the transaction in place. :param transaction: The transaction whose request to shape. :returns: The same `transaction`. @@ -63,7 +63,7 @@ async def run(self, transaction: UserTransaction) -> UserTransaction: for step in self.steps: reason = step.applicable(transaction) if reason is not None: - LOG.debug("Skipping input step %r: %s", step.id, reason) + LOG.debug("Skipping request step %r: %s", step.id, reason) continue await step.process(transaction) return transaction @@ -111,35 +111,35 @@ def _resolve_step(stage_id: str, def build_pipelines( config: PipelineConfig, steps: Optional[Mapping[str, Type[Step]]] = None, -) -> Tuple[InputPipeline, ResponsePipeline]: +) -> Tuple[RequestPipeline, ResponsePipeline]: """ Instantiate both pipelines from configuration. Step classes are looked up by id in `steps` (defaulting to the static `steps_by_id` mapping) and validated against the pipeline they are declared - in, so a response step placed in the input list (or vice versa) fails fast + in, so a response step placed in the request list (or vice versa) fails fast at build time. :param config: Parsed pipeline configuration. :param steps: Optional id -> `Step` class mapping. Defaults to `brainforge_inference_utils.steps.steps_by_id`. - :returns: `(input_pipeline, response_pipeline)`. + :returns: `(request_pipeline, response_pipeline)`. :raises KeyError: If a configured step id is not present in `steps`. :raises TypeError: If a configured step resolves to the wrong step type. """ steps = steps_by_id if steps is None else steps - input_steps: List[InputStep] = [] - for entry in config.input: + request_steps: List[RequestStep] = [] + for entry in config.request: if not entry.enabled: - LOG.info("Input step %r disabled by config", entry.id) + LOG.info("Request step %r disabled by config", entry.id) continue cls = _resolve_step(entry.id, steps) - if not (isinstance(cls, type) and issubclass(cls, InputStep)): + if not (isinstance(cls, type) and issubclass(cls, RequestStep)): raise TypeError( - f"Step {entry.id!r} is configured as an input step but " - f"{getattr(cls, '__name__', cls)!r} is not an InputStep") - input_steps.append(cls(entry.config)) + f"Step {entry.id!r} is configured as a request step but " + f"{getattr(cls, '__name__', cls)!r} is not a RequestStep") + request_steps.append(cls(entry.config)) response_steps: List[ResponseStep] = [] for entry in config.response: @@ -153,8 +153,8 @@ def build_pipelines( f"{getattr(cls, '__name__', cls)!r} is not a ResponseStep") response_steps.append(cls(entry.config)) - LOG.info("Built input pipeline: %s", - [s.id for s in input_steps] or "(empty)") + LOG.info("Built request pipeline: %s", + [s.id for s in request_steps] or "(empty)") LOG.info("Built response pipeline: %s", [s.id for s in response_steps] or "(empty)") - return InputPipeline(input_steps), ResponsePipeline(response_steps) + return RequestPipeline(request_steps), ResponsePipeline(response_steps) diff --git a/brainforge_inference_utils/steps/__init__.py b/brainforge_inference_utils/steps/__init__.py index 7eb0ff9..d9b0b36 100644 --- a/brainforge_inference_utils/steps/__init__.py +++ b/brainforge_inference_utils/steps/__init__.py @@ -15,7 +15,7 @@ """ Step base classes, first-party step implementations, and their id mapping. -New steps subclass `InputStep` or `ResponseStep`, set a unique `id`, and are +New steps subclass `RequestStep` or `ResponseStep`, set a unique `id`, and are added to `steps_by_id` below (mirroring `collectors_by_id` in `brainforge_dataset_utils`) so pipelines can resolve them by id. """ @@ -26,7 +26,7 @@ from brainforge_inference_utils.steps.ambient_context import AmbientContextStep from brainforge_inference_utils.steps.base import ( - InputStep, + RequestStep, RequestContextOption, ResponseStep, Step, @@ -43,7 +43,7 @@ __all__ = [ "Step", - "InputStep", + "RequestStep", "ResponseStep", "RequestContextOption", "AmbientContextStep", diff --git a/brainforge_inference_utils/steps/ambient_context.py b/brainforge_inference_utils/steps/ambient_context.py index 399b103..960e422 100644 --- a/brainforge_inference_utils/steps/ambient_context.py +++ b/brainforge_inference_utils/steps/ambient_context.py @@ -31,13 +31,13 @@ from neon_data_models.types import LlmMessageRole from brainforge_inference_utils.mcp import mcp_client, mcp_tool_inner_payload -from brainforge_inference_utils.steps.base import InputStep, RequestContextOption +from brainforge_inference_utils.steps.base import RequestStep, RequestContextOption from brainforge_inference_utils.transaction import UserTransaction LOG = logging.getLogger(__name__) -class AmbientContextStep(InputStep): +class AmbientContextStep(RequestStep): """ Prepend time and weather `information` messages to `request.history`. @@ -90,7 +90,7 @@ async def process(self, transaction: UserTransaction) -> None: except Exception as exc: LOG.warning("Skipping ambient context for %r: %s", location, exc) step_context["error"] = str(exc) - transaction.input_context[self.id] = step_context + transaction.request_step_context[self.id] = step_context return messages = self._build_information_messages( @@ -103,7 +103,7 @@ async def process(self, transaction: UserTransaction) -> None: "weather": weather_str, "messages": [content for _, content in messages], }) - transaction.input_context[self.id] = step_context + transaction.request_step_context[self.id] = step_context def _resolve_location(self, transaction: UserTransaction) -> str: ctx = transaction.request_context diff --git a/brainforge_inference_utils/steps/base.py b/brainforge_inference_utils/steps/base.py index b629788..5c83c48 100644 --- a/brainforge_inference_utils/steps/base.py +++ b/brainforge_inference_utils/steps/base.py @@ -17,7 +17,7 @@ Two step types plug into the two pipelines: -- `InputStep` runs before generation. It shapes the request on the +- `RequestStep` runs before generation. It shapes the request on the `UserTransaction` (RAG/MCP context, ambient context, request history validation, ...), mutating it in place. - `ResponseStep` runs after generation. It has the complete response on the @@ -34,8 +34,8 @@ 2. **Shared object, separate namespaces.** Data that must pass between steps travels on the `UserTransaction`, not through a step-to-step interface. - Input steps use `transaction.input_context`; response steps use - `transaction.output_context`. Step ids are unique only within one side. + Request steps use `transaction.request_step_context`; response steps use + `transaction.response_step_context`. Step ids are unique only within one side. """ from __future__ import annotations @@ -109,7 +109,7 @@ async def supported_request_context( return list(self.request_context_options) -class InputStep(Step, ABC): +class RequestStep(Step, ABC): """ A step that shapes the request before generation. diff --git a/brainforge_inference_utils/steps/rag.py b/brainforge_inference_utils/steps/rag.py index 370005e..81adddd 100644 --- a/brainforge_inference_utils/steps/rag.py +++ b/brainforge_inference_utils/steps/rag.py @@ -15,7 +15,7 @@ """ RAG retrieval and citation rewrite. -`RagStep` lists providers and injects retrieved documents on the input +`RagStep` lists providers and injects retrieved documents on the request pipeline. `ResponseParser` (id `citations`) runs on the response pipeline: it filters those documents, builds a citation completion via MCP, POSTs it to a configured OpenAI-compatible endpoint, and writes the annotated text @@ -34,7 +34,7 @@ from brainforge_inference_utils.mcp import mcp_client, mcp_tool_content_parts from brainforge_inference_utils.steps.base import ( - InputStep, + RequestStep, RequestContextOption, ResponseStep, ) @@ -43,7 +43,7 @@ LOG = logging.getLogger(__name__) -class RagStep(InputStep): +class RagStep(RequestStep): """ Look up selected RAG providers and prepend retrieved context. @@ -137,7 +137,7 @@ async def process(self, transaction: UserTransaction) -> None: except Exception as exc: LOG.warning("Skipping RAG lookup for %s: %s", providers, exc) step_context["error"] = str(exc) - transaction.input_context[self.id] = step_context + transaction.request_step_context[self.id] = step_context return step_context["documents"] = documents @@ -146,7 +146,7 @@ async def process(self, transaction: UserTransaction) -> None: [self._information_message(documents)] + list(transaction.request.history) ) - transaction.input_context[self.id] = step_context + transaction.request_step_context[self.id] = step_context async def _lookup_rag( self, @@ -232,7 +232,7 @@ def apply_citation_markup( class ResponseParser(ResponseStep): """ - Rewrite the generated response with citations using RAG input-context. + Rewrite the generated response with citations using RAG request-step context. The whole MCP filter → citation completion → markup path runs here so the host does not handle a second request/response pair. @@ -276,7 +276,7 @@ def applicable(self, transaction: UserTransaction) -> Optional[str]: return "openai_url not configured" if transaction.request_context.get("citations_enabled") is False: return "citations disabled" - documents = (transaction.input_context.get("rag") or {}).get( + documents = (transaction.request_step_context.get("rag") or {}).get( "documents") or [] if not documents: return "no rag documents" @@ -286,13 +286,13 @@ async def process( self, transaction: UserTransaction ) -> Optional[LLMResponse]: documents = list( - (transaction.input_context.get("rag") or {}).get("documents") or []) + (transaction.request_step_context.get("rag") or {}).get("documents") or []) original = transaction.response.response timeout = float(self.config["timeout"]) try: filtered = await self._filter_sources(original, documents, timeout) if not filtered: - transaction.output_context[self.id] = { + transaction.response_step_context[self.id] = { "documents": [], "credits": "", "formatted_credits": "", @@ -306,13 +306,13 @@ async def process( raw_cited, filtered) except Exception as exc: LOG.warning("Skipping citation rewrite: %s", exc) - transaction.output_context[self.id] = { + transaction.response_step_context[self.id] = { "documents": documents, "error": str(exc), } return None - transaction.output_context[self.id] = { + transaction.response_step_context[self.id] = { "documents": filtered, "unique_citations": unique, "credits": credits, diff --git a/brainforge_inference_utils/steps/request_history.py b/brainforge_inference_utils/steps/request_history.py index 58d3ad6..ee14e15 100644 --- a/brainforge_inference_utils/steps/request_history.py +++ b/brainforge_inference_utils/steps/request_history.py @@ -19,7 +19,7 @@ the incoming query (on `request.query`). Roles are `information` / `instruction` / `tool` context, then complete `user`/`llm` pairs. -This step is meant to run last on the input pipeline so earlier steps can +This step is meant to run last on the request pipeline so earlier steps can prepend context freely. It repairs rather than skipping: unknown or empty messages are dropped, prior citation markup is stripped, `assistant` becomes `llm`, a trailing user turn that duplicates `query` is removed, then @@ -33,7 +33,7 @@ from neon_data_models.types import LlmMessageRole -from brainforge_inference_utils.steps.base import InputStep, RequestContextOption +from brainforge_inference_utils.steps.base import RequestStep, RequestContextOption from brainforge_inference_utils.transaction import UserTransaction # Matches `[[n](url)]` written by `ResponseParser.apply_citation_markup`. @@ -49,7 +49,7 @@ def remove_citation_strings(text: str) -> str: return _CITATION_MARKUP.sub("", text) -class RequestHistoryStep(InputStep): +class RequestHistoryStep(RequestStep): """ Normalize, truncate, and order `request.history` into a valid LLM history. @@ -92,7 +92,7 @@ async def process(self, transaction: UserTransaction) -> None: transaction.request.history = list(context_messages) + list(windowed_turns) transaction.request.max_history = max_history - transaction.input_context[self.id] = { + transaction.request_step_context[self.id] = { "conversational": conversational, "max_history": max_history, "kept_turns": len(windowed_turns), diff --git a/brainforge_inference_utils/transaction.py b/brainforge_inference_utils/transaction.py index 8404650..bb5bc39 100644 --- a/brainforge_inference_utils/transaction.py +++ b/brainforge_inference_utils/transaction.py @@ -15,10 +15,10 @@ """ The object a pipeline operates on. -`UserTransaction` represents one user request/response cycle. Input steps shape -its `request` and write `input_context`; generation (performed elsewhere) +`UserTransaction` represents one user request/response cycle. Request steps shape +its `request` and write `request_step_context`; generation (performed elsewhere) populates its `response`; response steps annotate that `response` and write -`output_context`. +`response_step_context`. The request and response are the existing `neon-data-models` models used across the BrainForge stack (HANA, `brainforge-llm-service`), so this package does not @@ -38,27 +38,27 @@ class UserTransaction(BaseModel): """ One user request/response cycle as it flows through the pipelines. - :param request: The request being shaped by the input pipeline. + :param request: The request being shaped by the request pipeline. :param response: The generated response, populated after generation and shaped by the response pipeline. `None` until generation completes. - :param input_context: Transient data written by input steps, keyed by - input step id. - :param output_context: Transient data written by response steps, keyed by + :param request_step_context: Transient data written by request steps, keyed by + request step id. + :param response_step_context: Transient data written by response steps, keyed by response step id. """ request: LLMRequest = Field( - description="The LLM request shaped by the input pipeline") + description="The LLM request shaped by the request pipeline") response: Optional[LLMResponse] = Field( default=None, description="The generated response shaped by the response pipeline; " "None until generation completes") request_context: Dict[str, Any] = Field( default_factory=dict, - description="Client-provided context for the request for input steps") - input_context: Dict[str, Any] = Field( + description="Client-provided context read by request steps") + request_step_context: Dict[str, Any] = Field( default_factory=dict, - description="Input-step context keyed by input step id") - output_context: Dict[str, Any] = Field( + description="Request-step context keyed by request step id") + response_step_context: Dict[str, Any] = Field( default_factory=dict, description="Response-step context keyed by response step id") diff --git a/setup.py b/setup.py index 183c464..d6fb170 100644 --- a/setup.py +++ b/setup.py @@ -67,7 +67,7 @@ def get_requirements(requirements_filename: str): setup( name="brainforge-inference-utils", version=version, - description="Pluggable input/response pipelines for BrainForge LLM inference", + description="Pluggable request/response pipelines for BrainForge LLM inference", long_description=long_description, long_description_content_type="text/markdown", url="https://github.com/NeonGeckoCom/brainforge-inference-utils", diff --git a/tests/helpers.py b/tests/helpers.py index d2bf12e..3f195a3 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -21,7 +21,7 @@ LLMResponse, ) -from brainforge_inference_utils.steps.base import InputStep, ResponseStep +from brainforge_inference_utils.steps.base import RequestStep, ResponseStep from brainforge_inference_utils.transaction import UserTransaction @@ -53,8 +53,8 @@ def make_transaction(response: Optional[LLMResponse] = None, return UserTransaction(request=request, response=response) -class SetMaxHistoryStep(InputStep): - """Input step that writes `config["max_history"]` onto the request.""" +class SetMaxHistoryStep(RequestStep): + """Request step that writes `config["max_history"]` onto the request.""" id = "set_max_history" @@ -62,8 +62,8 @@ async def process(self, transaction: UserTransaction) -> None: transaction.request.max_history = self.config["max_history"] -class SkipUnlessQueryStep(InputStep): - """Input step skipped unless `request.query` matches `config["query"]`.""" +class SkipUnlessQueryStep(RequestStep): + """Request step skipped unless `request.query` matches `config["query"]`.""" id = "skip_unless_query" @@ -77,8 +77,8 @@ async def process(self, transaction: UserTransaction) -> None: transaction.request.max_tokens = self.config.get("max_tokens", 256) -class AppendQueryStep(InputStep): - """Input step that appends `config["suffix"]` to `request.query`.""" +class AppendQueryStep(RequestStep): + """Request step that appends `config["suffix"]` to `request.query`.""" id = "append_query" diff --git a/tests/test_ambient_context.py b/tests/test_ambient_context.py index b5f9033..164e4af 100644 --- a/tests/test_ambient_context.py +++ b/tests/test_ambient_context.py @@ -72,7 +72,7 @@ async def test_prepends_information_to_request_history(self): ("llm", "hey"), ]) self.assertEqual( - transaction.input_context["ambient_context"]["location"], + transaction.request_step_context["ambient_context"]["location"], "Seattle, WA, USA", ) @@ -155,7 +155,7 @@ async def test_mcp_failure_leaves_history_unchanged(self): self.assertEqual(transaction.request.history, original_history) self.assertEqual( - transaction.input_context["ambient_context"], + transaction.request_step_context["ambient_context"], {"location": "Denver, CO, USA", "error": "timed out"}, ) diff --git a/tests/test_client.py b/tests/test_client.py index 25f955c..fdfa22c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -21,7 +21,7 @@ from brainforge_inference_utils.client import BrainForgeClient from brainforge_inference_utils.config import PipelineConfig from brainforge_inference_utils.steps.base import ( - InputStep, + RequestStep, RequestContextOption, ResponseStep, ) @@ -33,14 +33,14 @@ ) -class RecordInputStep(InputStep): - id = "record_input" +class RecordRequestStep(RequestStep): + id = "record_request" async def process(self, transaction: UserTransaction) -> None: - transaction.input_context[self.id] = self.config["value"] + transaction.request_step_context[self.id] = self.config["value"] -class CatalogStep(InputStep): +class CatalogStep(RequestStep): id = "catalog" async def process(self, transaction: UserTransaction) -> None: @@ -69,21 +69,21 @@ async def process( return None -class RecordOutputStep(ResponseStep): - id = "record_output" +class RecordResponseStep(ResponseStep): + id = "record_response" async def process( self, transaction: UserTransaction ) -> Optional[LLMResponse]: - transaction.output_context[self.id] = self.config["value"] + transaction.response_step_context[self.id] = self.config["value"] return None class TestBrainForgeClient(unittest.IsolatedAsyncioTestCase): def make_client(self): config = PipelineConfig.model_validate({ - "input": [{ - "id": "record_input", + "request": [{ + "id": "record_request", "config": {"value": {"documents": ["one"]}}, }], "response": [ @@ -92,7 +92,7 @@ def make_client(self): "config": {"prefix": "processed: "}, }, { - "id": "record_output", + "id": "record_response", "config": {"value": {"citations": [1]}}, }, ], @@ -100,48 +100,48 @@ def make_client(self): return BrainForgeClient( config, { - RecordInputStep.id: RecordInputStep, - RecordOutputStep.id: RecordOutputStep, + RecordRequestStep.id: RecordRequestStep, + RecordResponseStep.id: RecordResponseStep, PrefixResponseStep.id: PrefixResponseStep, }, ) - async def test_run_input_pipeline_returns_transaction_with_context(self): + async def test_run_request_pipeline_returns_transaction_with_context(self): client = self.make_client() request = make_request() - transaction = await client.run_input_pipeline(request) + transaction = await client.run_request_pipeline(request) self.assertIs(transaction.request, request) self.assertIsNone(transaction.response) self.assertEqual( - transaction.input_context, - {"record_input": {"documents": ["one"]}}, + transaction.request_step_context, + {"record_request": {"documents": ["one"]}}, ) - self.assertEqual(transaction.output_context, {}) + self.assertEqual(transaction.response_step_context, {}) - async def test_run_output_pipeline_keeps_separate_step_namespaces(self): + async def test_run_response_pipeline_keeps_separate_step_namespaces(self): client = self.make_client() - transaction = await client.run_input_pipeline(make_request()) + transaction = await client.run_request_pipeline(make_request()) response = make_response("answer", transaction.request) - result = await client.run_output_pipeline(transaction, response) + result = await client.run_response_pipeline(transaction, response) self.assertIs(result, transaction) self.assertEqual(transaction.response.response, "processed: answer") self.assertEqual( - transaction.input_context, - {"record_input": {"documents": ["one"]}}, + transaction.request_step_context, + {"record_request": {"documents": ["one"]}}, ) self.assertEqual( - transaction.output_context, - {"record_output": {"citations": [1]}}, + transaction.response_step_context, + {"record_response": {"citations": [1]}}, ) async def test_supported_request_context_merges_pipeline_steps(self): config = PipelineConfig.model_validate({ - "input": [ + "request": [ {"id": "catalog", "config": {"providers": ["news", "docs"]}}, - {"id": "record_input", "config": {"value": {}}}, + {"id": "record_request", "config": {"value": {}}}, {"id": "catalog", "config": {"providers": ["docs", "web"]}}, ], "response": [{"id": "flag"}], @@ -150,7 +150,7 @@ async def test_supported_request_context_merges_pipeline_steps(self): config, { CatalogStep.id: CatalogStep, - RecordInputStep.id: RecordInputStep, + RecordRequestStep.id: RecordRequestStep, FlagStep.id: FlagStep, }, ) @@ -161,4 +161,4 @@ async def test_supported_request_context_merges_pipeline_steps(self): self.assertEqual( list(options["citations_enabled"].values), [True, False]) self.assertIsNone(options["location"].values) - self.assertNotIn("record_input", options) + self.assertNotIn("record_request", options) diff --git a/tests/test_config.py b/tests/test_config.py index 6cc5e14..92ad0af 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -54,12 +54,12 @@ def test_constructor_copies_config(self): class TestPipelineConfig: def test_defaults_to_empty_lists(self): parsed = PipelineConfig() - assert parsed.input == [] + assert parsed.request == [] assert parsed.response == [] def test_mixed_shorthand_and_mapping_entries(self): parsed = PipelineConfig.model_validate({ - "input": [ + "request": [ "persona", {"id": "rag", "config": {"top_k": 10}}, ], @@ -67,13 +67,13 @@ def test_mixed_shorthand_and_mapping_entries(self): {"id": "citations"}, ], }) - assert [entry.id for entry in parsed.input] == ["persona", "rag"] - assert parsed.input[1].config == {"top_k": 10} + assert [entry.id for entry in parsed.request] == ["persona", "rag"] + assert parsed.request[1].config == {"top_k": 10} assert parsed.response[0].id == "citations" def test_duplicate_step_ids_are_allowed_with_distinct_configs(self): parsed = PipelineConfig.model_validate({ - "input": [ + "request": [ {"id": "append_query", "config": {"suffix": " A"}}, {"id": "append_query", "config": {"suffix": " B"}}, ], @@ -82,11 +82,11 @@ def test_duplicate_step_ids_are_allowed_with_distinct_configs(self): {"id": "prefix_response", "config": {"prefix": "<"}}, ], }) - assert [entry.id for entry in parsed.input] == [ + assert [entry.id for entry in parsed.request] == [ "append_query", "append_query", ] - assert [entry.config for entry in parsed.input] == [ + assert [entry.config for entry in parsed.request] == [ {"suffix": " A"}, {"suffix": " B"}, ] @@ -97,19 +97,19 @@ def test_duplicate_step_ids_are_allowed_with_distinct_configs(self): def test_duplicate_identical_entries_are_allowed(self): parsed = PipelineConfig.model_validate({ - "input": ["sanitize", "sanitize"], + "request": ["sanitize", "sanitize"], }) - assert [entry.id for entry in parsed.input] == ["sanitize", "sanitize"] + assert [entry.id for entry in parsed.request] == ["sanitize", "sanitize"] def test_empty_lists_are_valid(self): - parsed = PipelineConfig.model_validate({"input": [], "response": []}) - assert parsed.input == [] + parsed = PipelineConfig.model_validate({"request": [], "response": []}) + assert parsed.request == [] assert parsed.response == [] def test_extra_fields_are_forbidden(self): with pytest.raises(ValidationError): - PipelineConfig.model_validate({"input": [], "chunk": []}) + PipelineConfig.model_validate({"request": [], "chunk": []}) def test_invalid_entry_type_is_rejected(self): with pytest.raises(ValidationError): - PipelineConfig.model_validate({"input": [123]}) + PipelineConfig.model_validate({"request": [123]}) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index c3d0a94..82e45ca 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -19,7 +19,7 @@ from brainforge_inference_utils.config import PipelineConfig from brainforge_inference_utils.pipeline import ( - InputPipeline, + RequestPipeline, ResponsePipeline, build_pipelines, ) @@ -38,9 +38,9 @@ class TestBuildPipelines: def test_empty_config_builds_empty_pipelines(self): - input_pipeline, response_pipeline = build_pipelines( + request_pipeline, response_pipeline = build_pipelines( PipelineConfig(), TEST_STEPS) - assert input_pipeline.steps == [] + assert request_pipeline.steps == [] assert response_pipeline.steps == [] def test_default_mapping_includes_first_party_steps(self): @@ -48,12 +48,12 @@ def test_default_mapping_includes_first_party_steps(self): assert "rag" in steps_by_id assert "citations" in steps_by_id assert "request_history" in steps_by_id - input_pipeline, response_pipeline = build_pipelines( + request_pipeline, response_pipeline = build_pipelines( PipelineConfig.model_validate({ - "input": ["ambient_context", "rag", "request_history"], + "request": ["ambient_context", "rag", "request_history"], "response": ["citations"], })) - assert [step.id for step in input_pipeline.steps] == [ + assert [step.id for step in request_pipeline.steps] == [ "ambient_context", "rag", "request_history", @@ -62,20 +62,20 @@ def test_default_mapping_includes_first_party_steps(self): with pytest.raises(KeyError, match="append_query"): build_pipelines(PipelineConfig.model_validate({ - "input": ["append_query"], + "request": ["append_query"], })) def test_unknown_id_lists_known_steps(self): with pytest.raises(KeyError, match="append_query"): build_pipelines( - PipelineConfig.model_validate({"input": ["nope"]}), + PipelineConfig.model_validate({"request": ["nope"]}), TEST_STEPS, ) def test_disabled_steps_are_omitted(self): - input_pipeline, response_pipeline = build_pipelines( + request_pipeline, response_pipeline = build_pipelines( PipelineConfig.model_validate({ - "input": [ + "request": [ {"id": "append_query", "enabled": False, "config": {"suffix": " skipped"}}, {"id": "append_query", "config": {"suffix": " kept"}}, @@ -87,19 +87,19 @@ def test_disabled_steps_are_omitted(self): }), TEST_STEPS, ) - assert len(input_pipeline.steps) == 1 - assert input_pipeline.steps[0].config == {"suffix": " kept"} + assert len(request_pipeline.steps) == 1 + assert request_pipeline.steps[0].config == {"suffix": " kept"} assert response_pipeline.steps == [] - def test_response_step_in_input_list_raises(self): - with pytest.raises(TypeError, match="not an InputStep"): + def test_response_step_in_request_list_raises(self): + with pytest.raises(TypeError, match="not a RequestStep"): build_pipelines( PipelineConfig.model_validate({ - "input": ["prefix_response"]}), + "request": ["prefix_response"]}), TEST_STEPS, ) - def test_input_step_in_response_list_raises(self): + def test_request_step_in_response_list_raises(self): with pytest.raises(TypeError, match="not a ResponseStep"): build_pipelines( PipelineConfig.model_validate( @@ -108,9 +108,9 @@ def test_input_step_in_response_list_raises(self): ) def test_duplicate_ids_instantiate_independent_instances(self): - input_pipeline, response_pipeline = build_pipelines( + request_pipeline, response_pipeline = build_pipelines( PipelineConfig.model_validate({ - "input": [ + "request": [ {"id": "append_query", "config": {"suffix": " A"}}, {"id": "append_query", "config": {"suffix": " B"}}, ], @@ -121,13 +121,13 @@ def test_duplicate_ids_instantiate_independent_instances(self): }), TEST_STEPS, ) - assert len(input_pipeline.steps) == 2 - assert input_pipeline.steps[0] is not input_pipeline.steps[1] - assert [type(step) for step in input_pipeline.steps] == [ + assert len(request_pipeline.steps) == 2 + assert request_pipeline.steps[0] is not request_pipeline.steps[1] + assert [type(step) for step in request_pipeline.steps] == [ AppendQueryStep, AppendQueryStep, ] - assert [step.config for step in input_pipeline.steps] == [ + assert [step.config for step in request_pipeline.steps] == [ {"suffix": " A"}, {"suffix": " B"}, ] @@ -138,8 +138,8 @@ def test_duplicate_ids_instantiate_independent_instances(self): class TestPipelineRun(unittest.IsolatedAsyncioTestCase): - async def test_input_steps_run_in_declared_order(self): - pipeline = InputPipeline([ + async def test_request_steps_run_in_declared_order(self): + pipeline = RequestPipeline([ AppendQueryStep({"suffix": " A"}), AppendQueryStep({"suffix": " B"}), ]) @@ -148,8 +148,8 @@ async def test_input_steps_run_in_declared_order(self): assert result is transaction assert transaction.request.query == "hello A B" - async def test_inapplicable_input_step_is_skipped(self): - pipeline = InputPipeline([ + async def test_inapplicable_request_step_is_skipped(self): + pipeline = RequestPipeline([ SkipUnlessQueryStep({"query": "other", "max_tokens": 64}), SetMaxHistoryStep({"max_history": 7}), ]) @@ -188,9 +188,9 @@ async def test_inapplicable_response_step_is_skipped(self): assert transaction.response is None async def test_build_then_run_duplicate_configured_steps(self): - input_pipeline, response_pipeline = build_pipelines( + request_pipeline, response_pipeline = build_pipelines( PipelineConfig.model_validate({ - "input": [ + "request": [ {"id": "append_query", "config": {"suffix": " A"}}, {"id": "skip_unless_query", "config": {"query": "nope", "max_tokens": 64}}, @@ -205,7 +205,7 @@ async def test_build_then_run_duplicate_configured_steps(self): TEST_STEPS, ) transaction = make_transaction(query="hello") - await input_pipeline.run(transaction) + await request_pipeline.run(transaction) assert transaction.request.query == "hello A B" assert transaction.request.max_tokens == 512 diff --git a/tests/test_rag.py b/tests/test_rag.py index c022060..18bc0ea 100644 --- a/tests/test_rag.py +++ b/tests/test_rag.py @@ -59,9 +59,9 @@ async def test_injects_information_and_stores_documents(self): self.assertIn("https://b: chunk two", transaction.request.history[0][1]) self.assertEqual(transaction.request.history[1], ("user", "hi")) self.assertEqual( - transaction.input_context["rag"]["documents"], DOCUMENTS) + transaction.request_step_context["rag"]["documents"], DOCUMENTS) self.assertEqual( - transaction.input_context["rag"]["providers"], ["news"]) + transaction.request_step_context["rag"]["providers"], ["news"]) async def test_uses_request_top_k_override(self): step = RagStep(STEP_CONFIG) @@ -78,7 +78,7 @@ async def test_uses_request_top_k_override(self): lookup.assert_awaited_once_with( transaction.request.query, ["docs"], 3, STEP_CONFIG["timeout"]) self.assertEqual(transaction.request.history, make_request().history) - self.assertEqual(transaction.input_context["rag"]["documents"], []) + self.assertEqual(transaction.request_step_context["rag"]["documents"], []) async def test_lookup_failure_leaves_history_unchanged(self): step = RagStep(STEP_CONFIG) @@ -97,7 +97,7 @@ async def test_lookup_failure_leaves_history_unchanged(self): self.assertEqual(transaction.request.history, original) self.assertEqual( - transaction.input_context["rag"]["error"], "timed out") + transaction.request_step_context["rag"]["error"], "timed out") def test_skips_when_no_providers_selected(self): step = RagStep(STEP_CONFIG) diff --git a/tests/test_request_history.py b/tests/test_request_history.py index ecf8d21..9475a3d 100644 --- a/tests/test_request_history.py +++ b/tests/test_request_history.py @@ -45,9 +45,9 @@ async def test_keeps_last_complete_pairs(self): ]) self.assertEqual(transaction.request.max_history, 2) self.assertEqual( - transaction.input_context["request_history"]["dropped_turns"], 2) + transaction.request_step_context["request_history"]["dropped_turns"], 2) self.assertEqual( - transaction.input_context["request_history"]["kept_turns"], 4) + transaction.request_step_context["request_history"]["kept_turns"], 4) async def test_preserves_context_messages_outside_the_window(self): step = RequestHistoryStep({"max_history": 2}) @@ -87,7 +87,7 @@ async def test_non_conversational_drops_prior_turns(self): self.assertEqual(transaction.request.history, [("information", "rag")]) self.assertEqual(transaction.request.max_history, 0) - ctx = transaction.input_context["request_history"] + ctx = transaction.request_step_context["request_history"] self.assertFalse(ctx["conversational"]) self.assertEqual(ctx["dropped_turns"], 4) @@ -100,7 +100,7 @@ async def test_short_history_is_unchanged(self): self.assertEqual(transaction.request.history, history) self.assertEqual( - transaction.input_context["request_history"]["dropped_turns"], 0) + transaction.request_step_context["request_history"]["dropped_turns"], 0) async def test_drops_system_empty_and_duplicate_query(self): step = RequestHistoryStep({"max_history": 2}) diff --git a/tests/test_response_parser.py b/tests/test_response_parser.py index 2cb9b7a..5d32e69 100644 --- a/tests/test_response_parser.py +++ b/tests/test_response_parser.py @@ -56,7 +56,7 @@ def _transaction(self, text="See the docs.", **request_context): request=request, response=make_response(text, request), request_context=request_context, - input_context={"rag": {"documents": DOCUMENTS}}, + request_step_context={"rag": {"documents": DOCUMENTS}}, ) return transaction @@ -82,7 +82,7 @@ async def test_rewrites_response_via_openai_endpoint(self): self.assertNotEqual(revised.response, "See the docs.") self.assertEqual(revised.history[-1][0], "llm") self.assertEqual(revised.history[-1][1], revised.response) - self.assertIn("Doc A", transaction.output_context["citations"]["credits"]) + self.assertIn("Doc A", transaction.response_step_context["citations"]["credits"]) async def test_skips_without_rag_documents(self): step = ResponseParser(STEP_CONFIG) @@ -121,7 +121,7 @@ async def test_completion_failure_leaves_response(self): self.assertIsNone(revised) self.assertEqual(transaction.response.response, original) self.assertEqual( - transaction.output_context["citations"]["error"], "nope") + transaction.response_step_context["citations"]["error"], "nope") def test_appends_chat_completions_to_base_url(self): step = ResponseParser({ diff --git a/tests/test_steps.py b/tests/test_steps.py index a10f7fd..6d23128 100644 --- a/tests/test_steps.py +++ b/tests/test_steps.py @@ -17,7 +17,7 @@ import pytest -from brainforge_inference_utils.steps.base import InputStep, ResponseStep +from brainforge_inference_utils.steps.base import RequestStep, ResponseStep from helpers import ( AppendQueryStep, make_transaction, @@ -57,9 +57,9 @@ async def process(self, transaction): assert DefaultResponse().applicable(make_transaction()) is None def test_abstract_bases_cannot_be_instantiated(self): - assert issubclass(InputStep, ABC) + assert issubclass(RequestStep, ABC) assert issubclass(ResponseStep, ABC) with pytest.raises(TypeError): - InputStep() + RequestStep() with pytest.raises(TypeError): ResponseStep() diff --git a/tests/test_transaction.py b/tests/test_transaction.py index 64f3c5f..34f5932 100644 --- a/tests/test_transaction.py +++ b/tests/test_transaction.py @@ -46,14 +46,14 @@ def test_request_can_be_mutated_in_place(self): def test_pipeline_context_defaults_to_independent_dicts(self): first = UserTransaction(request=make_request()) second = UserTransaction(request=make_request()) - first.input_context["rag"] = {"documents": ["document"]} - first.output_context["citations"] = {"markers": [1]} - assert second.input_context == {} - assert second.output_context == {} + first.request_step_context["rag"] = {"documents": ["document"]} + first.response_step_context["citations"] = {"markers": [1]} + assert second.request_step_context == {} + assert second.response_step_context == {} - def test_input_and_output_step_ids_do_not_share_a_namespace(self): + def test_request_and_response_step_ids_do_not_share_a_namespace(self): transaction = UserTransaction(request=make_request()) - transaction.input_context["shared_id"] = {"from": "input"} - transaction.output_context["shared_id"] = {"from": "output"} - assert transaction.input_context["shared_id"] == {"from": "input"} - assert transaction.output_context["shared_id"] == {"from": "output"} + transaction.request_step_context["shared_id"] = {"from": "request"} + transaction.response_step_context["shared_id"] = {"from": "response"} + assert transaction.request_step_context["shared_id"] == {"from": "request"} + assert transaction.response_step_context["shared_id"] == {"from": "response"} From 5fade98ca9887f77bb3e19f5ddc381dfcee32513 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 14:01:12 -0700 Subject: [PATCH 12/23] Refactor to use plugins for request and response pipelines --- README.md | 52 ++++- brainforge_inference_utils/__init__.py | 6 + brainforge_inference_utils/config.py | 8 +- brainforge_inference_utils/pipeline.py | 34 +++- brainforge_inference_utils/registry.py | 124 ++++++++++++ brainforge_inference_utils/steps/__init__.py | 7 +- brainforge_inference_utils/steps/base.py | 3 +- setup.py | 16 +- tests/test_registry.py | 190 +++++++++++++++++++ 9 files changed, 411 insertions(+), 29 deletions(-) create mode 100644 brainforge_inference_utils/registry.py create mode 100644 tests/test_registry.py diff --git a/README.md b/README.md index e014cfa..09fb3ab 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ response, and two per-side context dicts (`request_step_context` and | `RequestContextOption` | One `request_context` key a step reads, with optional allowed values for hosts to present. | | `RequestPipeline` / `ResponsePipeline` | Run an ordered list of steps over a transaction. | | `PipelineConfig` / `StepConfig` | Declarative, ordered step lists (see below). | -| `build_pipelines(config)` | Build both pipelines from config, resolving step ids via `steps_by_id`. | +| `build_pipelines(config)` | Build both pipelines from config. Bundled ids come from `steps_by_id`; plugin ids come from entry points. | ## Usage sketch @@ -80,11 +80,55 @@ between calls. (e.g. a citation formatter and the sanitizer that strips it) belong in the same module and should derive from a shared constant so they cannot drift. +## Step plugins + +Installed packages advertise steps with setuptools entry points. Request and +response steps use different groups: + +| Pipeline | Entry point group | +| --- | --- | +| Request | `brainforge.inference.request_steps` | +| Response | `brainforge.inference.response_steps` | + +Both groups use the step id `author.step_name`: + +- `author` is the plugin developer's slug (`acme`, `neon`). +- `step_name` names that one step (`guardrails`, `rag`). + +Each part is lowercase letters, digits, and underscores, and starts with a +letter. The class `id`, the config `id`, and the context-dict key are that +same string. + +The same `author.step_name` may be registered once in each group. Each +pipeline reads only its own group, so the request class and the response class +stay distinct. Two packages registering the same name in one group is an +error. + +```python +entry_points={ + "brainforge.inference.request_steps": [ + "acme.guardrails=acme_guardrails:GuardrailRequestStep", + ], + "brainforge.inference.response_steps": [ + "acme.guardrails=acme_guardrails:GuardrailResponseStep", + ], +} +``` + +```yaml +pipeline: + request: + - id: acme.guardrails + response: + - id: acme.guardrails +``` + +Bundled steps (`ambient_context`, `rag`, `request_history`, `citations`) are +still resolved from `steps_by_id` in this package. They do not use +`author.step_name` yet. + ## Deferred -- A step **registry** (in-process registration + entry-point discovery under - `brainforge.inference.steps`). For now, steps are resolved from the static - `steps_by_id` mapping. - A **services** adapter for streaming, tool calls, and token counting. - Richer persistent state or execution tracing beyond the transient `request_step_context` / `response_step_context` dicts on `UserTransaction`. diff --git a/brainforge_inference_utils/__init__.py b/brainforge_inference_utils/__init__.py index 4e14254..9892978 100644 --- a/brainforge_inference_utils/__init__.py +++ b/brainforge_inference_utils/__init__.py @@ -29,6 +29,10 @@ ResponsePipeline, build_pipelines, ) +from brainforge_inference_utils.registry import ( + REQUEST_STEP_ENTRY_POINT, + RESPONSE_STEP_ENTRY_POINT, +) from brainforge_inference_utils.steps import steps_by_id from brainforge_inference_utils.steps.base import ( RequestStep, @@ -57,6 +61,8 @@ "RequestPipeline", "ResponsePipeline", "build_pipelines", + "REQUEST_STEP_ENTRY_POINT", + "RESPONSE_STEP_ENTRY_POINT", # Config "PipelineConfig", "StepConfig", diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py index 50f0359..11a213e 100644 --- a/brainforge_inference_utils/config.py +++ b/brainforge_inference_utils/config.py @@ -17,8 +17,9 @@ Configuration is two explicitly ordered lists of step entries: `request` and `response`. Order is authoritative; the pipeline runs steps in exactly the -order given and does not reorder. Each entry names a step `id` (resolved via the -`steps_by_id` mapping) and optionally carries step-specific `config`. +order given and does not reorder. Each entry names a step `id` (a bundled id +from `steps_by_id`, or a plugin id `author.step_name`) and optionally carries +step-specific `config`. Example (YAML):: @@ -53,7 +54,8 @@ class StepConfig(BaseModel): """ A single step entry within a pipeline. - :param id: Step id to look up in `steps_by_id` and instantiate. + :param id: Step id to resolve and instantiate. Bundled steps use their + current ids. Plugin steps use `author.step_name`. :param enabled: When false, the entry is skipped at build time (a convenient way to turn a step off without deleting its config). :param config: Step-specific configuration passed to the step constructor. diff --git a/brainforge_inference_utils/pipeline.py b/brainforge_inference_utils/pipeline.py index 35f1cdc..d3a3e6a 100644 --- a/brainforge_inference_utils/pipeline.py +++ b/brainforge_inference_utils/pipeline.py @@ -32,6 +32,11 @@ from typing import List, Mapping, Optional, Sequence, Tuple, Type from brainforge_inference_utils.config import PipelineConfig +from brainforge_inference_utils.registry import ( + REQUEST_STEP_ENTRY_POINT, + RESPONSE_STEP_ENTRY_POINT, + merge_builtin_and_entry_points, +) from brainforge_inference_utils.steps import steps_by_id from brainforge_inference_utils.steps.base import ( RequestStep, @@ -115,26 +120,35 @@ def build_pipelines( """ Instantiate both pipelines from configuration. - Step classes are looked up by id in `steps` (defaulting to the static - `steps_by_id` mapping) and validated against the pipeline they are declared - in, so a response step placed in the request list (or vice versa) fails fast - at build time. + Step classes are looked up by id in `steps` when that mapping is passed. + Otherwise each side uses `steps_by_id` plus the entry point group for that + pipeline (`brainforge.inference.request_steps` or + `brainforge.inference.response_steps`). A response step placed in the + request list (or vice versa) fails fast at build time. :param config: Parsed pipeline configuration. - :param steps: Optional id -> `Step` class mapping. Defaults to - `brainforge_inference_utils.steps.steps_by_id`. + :param steps: Optional id -> `Step` class mapping. When omitted, bundled + steps are combined with installed entry points. :returns: `(request_pipeline, response_pipeline)`. - :raises KeyError: If a configured step id is not present in `steps`. + :raises KeyError: If a configured step id is not present in the mapping + used for that side. :raises TypeError: If a configured step resolves to the wrong step type. """ - steps = steps_by_id if steps is None else steps + if steps is None: + request_step_classes = merge_builtin_and_entry_points( + steps_by_id, REQUEST_STEP_ENTRY_POINT, RequestStep) + response_step_classes = merge_builtin_and_entry_points( + steps_by_id, RESPONSE_STEP_ENTRY_POINT, ResponseStep) + else: + request_step_classes = steps + response_step_classes = steps request_steps: List[RequestStep] = [] for entry in config.request: if not entry.enabled: LOG.info("Request step %r disabled by config", entry.id) continue - cls = _resolve_step(entry.id, steps) + cls = _resolve_step(entry.id, request_step_classes) if not (isinstance(cls, type) and issubclass(cls, RequestStep)): raise TypeError( f"Step {entry.id!r} is configured as a request step but " @@ -146,7 +160,7 @@ def build_pipelines( if not entry.enabled: LOG.info("Response step %r disabled by config", entry.id) continue - cls = _resolve_step(entry.id, steps) + cls = _resolve_step(entry.id, response_step_classes) if not (isinstance(cls, type) and issubclass(cls, ResponseStep)): raise TypeError( f"Step {entry.id!r} is configured as a response step but " diff --git a/brainforge_inference_utils/registry.py b/brainforge_inference_utils/registry.py new file mode 100644 index 0000000..ae4af31 --- /dev/null +++ b/brainforge_inference_utils/registry.py @@ -0,0 +1,124 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . +""" +Discovery of plugin steps via setuptools entry points. + +Request and response steps are registered in separate groups, so the same +`author.step_name` id can exist on both sides. Bundled steps shipped in this +package stay in `steps_by_id` and are merged with whichever group is loaded. +""" + +from __future__ import annotations + +import logging +import re +from importlib.metadata import entry_points +from typing import Dict, Mapping, Type + +from brainforge_inference_utils.steps.base import Step + +LOG = logging.getLogger(__name__) + +REQUEST_STEP_ENTRY_POINT = "brainforge.inference.request_steps" +RESPONSE_STEP_ENTRY_POINT = "brainforge.inference.response_steps" + +# author.step_name — one developer slug, one step name. +STEP_ID_PATTERN = re.compile(r"^[a-z][a-z0-9_]*\.[a-z][a-z0-9_]*$") + + +def _distribution_name(entry_point) -> str: + dist = getattr(entry_point, "dist", None) + return getattr(dist, "name", None) or "unknown" + + +def load_entry_point_steps( + group: str, + expected_type: Type[Step], +) -> Dict[str, Type[Step]]: + """ + Load step classes advertised under `group`. + + Each entry point name must be `author.step_name` and must equal the + loaded class's `id`. Names must be unique within the group. The same name + may still be registered under the other pipeline's group. + + :param group: Entry point group to read. + :param expected_type: `RequestStep` or `ResponseStep`. Loaded classes must + subclass it. + :returns: Entry point name -> step class. + :raises ValueError: If a name is malformed, duplicated, or differs from + the class `id`. + :raises TypeError: If a loaded object is not a subclass of `expected_type`. + """ + seen: Dict[str, str] = {} + loaded: Dict[str, Type[Step]] = {} + for entry_point in entry_points(group=group): + dist_name = _distribution_name(entry_point) + if not STEP_ID_PATTERN.fullmatch(entry_point.name): + raise ValueError( + f"Entry point {entry_point.name!r} in {group} from " + f"{dist_name} must use the form author.step_name") + prior = seen.get(entry_point.name) + if prior is not None: + raise ValueError( + f"Multiple packages register {group} step " + f"{entry_point.name!r}: {prior} and {dist_name}") + seen[entry_point.name] = dist_name + + try: + step_cls = entry_point.load() + except Exception as exc: + raise ImportError( + f"Failed to load step {entry_point.name!r} from {dist_name}" + ) from exc + if not (isinstance(step_cls, type) and issubclass(step_cls, expected_type)): + raise TypeError( + f"Entry point {entry_point.name!r} in {group} from {dist_name} " + f"loaded {getattr(step_cls, '__name__', step_cls)!r}, which is " + f"not a {expected_type.__name__}") + if step_cls.id != entry_point.name: + raise ValueError( + f"Step {step_cls.__name__} from {dist_name} has id " + f"{step_cls.id!r}, which does not match entry point " + f"{entry_point.name!r}") + loaded[entry_point.name] = step_cls + LOG.debug("Loaded %s step %r from %s", + expected_type.__name__, entry_point.name, dist_name) + return loaded + + +def merge_builtin_and_entry_points( + builtin: Mapping[str, Type[Step]], + group: str, + expected_type: Type[Step], +) -> Dict[str, Type[Step]]: + """ + Combine bundled steps with classes discovered from `group`. + + :param builtin: In-package id -> step class mapping (`steps_by_id`). + :param group: Entry point group for this pipeline side. + :param expected_type: Step base class that group must provide. + :returns: A new mapping. `builtin` is copied unchanged, then plugin steps + are added. + :raises ValueError: If a plugin id is already present in `builtin`. + """ + resolved = dict(builtin) + for step_id, step_cls in load_entry_point_steps(group, expected_type).items(): + if step_id in resolved: + raise ValueError( + f"Step id {step_id!r} from entry point group {group} conflicts " + f"with a built-in step") + resolved[step_id] = step_cls + return resolved diff --git a/brainforge_inference_utils/steps/__init__.py b/brainforge_inference_utils/steps/__init__.py index d9b0b36..aae42ba 100644 --- a/brainforge_inference_utils/steps/__init__.py +++ b/brainforge_inference_utils/steps/__init__.py @@ -15,9 +15,10 @@ """ Step base classes, first-party step implementations, and their id mapping. -New steps subclass `RequestStep` or `ResponseStep`, set a unique `id`, and are -added to `steps_by_id` below (mirroring `collectors_by_id` in -`brainforge_dataset_utils`) so pipelines can resolve them by id. +Bundled steps subclass `RequestStep` or `ResponseStep`, set a unique `id`, and +are added to `steps_by_id` below so pipelines can resolve them by id. Plugin +steps are discovered from entry points instead of this mapping. Their ids use +`author.step_name`. """ from __future__ import annotations diff --git a/brainforge_inference_utils/steps/base.py b/brainforge_inference_utils/steps/base.py index 5c83c48..9924e83 100644 --- a/brainforge_inference_utils/steps/base.py +++ b/brainforge_inference_utils/steps/base.py @@ -75,7 +75,8 @@ class Step(ABC): should keep per-request state on the `UserTransaction`, not on `self`. :cvar id: Unique step identifier used in configuration. Subclasses must set - this to a stable, non-empty string. + this to a stable, non-empty string. Plugin steps use `author.step_name` + and that string must match the entry point name. :cvar request_context_options: Static `request_context` keys this step reads. Override `supported_request_context` when allowed values must be fetched live (for example a provider catalog). diff --git a/setup.py b/setup.py index d6fb170..20d7fed 100644 --- a/setup.py +++ b/setup.py @@ -85,15 +85,15 @@ def get_requirements(requirements_filename: str): "Intended Audience :: Developers", "Programming Language :: Python :: 3", ], - # Reserved for a future step registry: first-party and third-party steps - # will be advertised under this entry point group so pipelines can resolve - # them by id. Example: + # Plugin steps register under brainforge.inference.request_steps and + # brainforge.inference.response_steps. Entry point names are + # author.step_name. Bundled steps stay in steps_by_id for now. # entry_points={ - # "brainforge.inference.steps": [ - # "persona=brainforge_inference_utils.steps.persona:PersonaStep", + # "brainforge.inference.request_steps": [ + # "acme.guardrails=acme_guardrails:GuardrailRequestStep", + # ], + # "brainforge.inference.response_steps": [ + # "acme.guardrails=acme_guardrails:GuardrailResponseStep", # ], # }, - entry_points={ - "brainforge.inference.steps": [], - }, ) diff --git a/tests/test_registry.py b/tests/test_registry.py new file mode 100644 index 0000000..a384f73 --- /dev/null +++ b/tests/test_registry.py @@ -0,0 +1,190 @@ +# Copyright (C) 2008-2026 Neongecko.com Inc. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +from types import SimpleNamespace + +import pytest + +from brainforge_inference_utils.config import PipelineConfig +from brainforge_inference_utils.pipeline import build_pipelines +from brainforge_inference_utils.registry import ( + REQUEST_STEP_ENTRY_POINT, + RESPONSE_STEP_ENTRY_POINT, + load_entry_point_steps, +) +from brainforge_inference_utils.steps.base import RequestStep, ResponseStep +from helpers import TEST_STEPS + + +class _EntryPoint: + def __init__(self, name, cls, dist="acme-plugin"): + self.name = name + self._cls = cls + self.dist = SimpleNamespace(name=dist) + + def load(self): + return self._cls + + +class AcmeRequestGuard(RequestStep): + id = "acme.guard" + + async def process(self, transaction): + return None + + +class AcmeResponseGuard(ResponseStep): + id = "acme.guard" + + async def process(self, transaction): + return None + + +class MismatchedId(RequestStep): + id = "acme.other" + + async def process(self, transaction): + return None + + +def _entry_points_for(request_points=(), response_points=()): + def fake_entry_points(group): + if group == REQUEST_STEP_ENTRY_POINT: + return list(request_points) + if group == RESPONSE_STEP_ENTRY_POINT: + return list(response_points) + return [] + + return fake_entry_points + + +def test_same_author_step_name_loads_from_both_groups(monkeypatch): + monkeypatch.setattr( + "brainforge_inference_utils.registry.entry_points", + _entry_points_for( + [_EntryPoint("acme.guard", AcmeRequestGuard)], + [_EntryPoint("acme.guard", AcmeResponseGuard, dist="acme-plugin")], + ), + ) + + request_steps = load_entry_point_steps( + REQUEST_STEP_ENTRY_POINT, RequestStep) + response_steps = load_entry_point_steps( + RESPONSE_STEP_ENTRY_POINT, ResponseStep) + + assert request_steps["acme.guard"] is AcmeRequestGuard + assert response_steps["acme.guard"] is AcmeResponseGuard + + +def test_build_pipelines_resolves_plugins_and_builtin_steps(monkeypatch): + monkeypatch.setattr( + "brainforge_inference_utils.registry.entry_points", + _entry_points_for( + [_EntryPoint("acme.guard", AcmeRequestGuard)], + [_EntryPoint("acme.guard", AcmeResponseGuard)], + ), + ) + + request_pipeline, response_pipeline = build_pipelines( + PipelineConfig.model_validate({ + "request": ["ambient_context", "acme.guard"], + "response": ["citations", "acme.guard"], + })) + + assert [step.id for step in request_pipeline.steps] == [ + "ambient_context", + "acme.guard", + ] + assert [step.id for step in response_pipeline.steps] == [ + "citations", + "acme.guard", + ] + assert type(request_pipeline.steps[1]) is AcmeRequestGuard + assert type(response_pipeline.steps[1]) is AcmeResponseGuard + + +def test_explicit_step_map_skips_entry_points(monkeypatch): + def fail_entry_points(group): + raise AssertionError(f"entry points consulted for {group}") + + monkeypatch.setattr( + "brainforge_inference_utils.registry.entry_points", + fail_entry_points, + ) + request_pipeline, response_pipeline = build_pipelines( + PipelineConfig.model_validate({ + "request": ["append_query"], + "response": ["prefix_response"], + }), + TEST_STEPS, + ) + assert [step.id for step in request_pipeline.steps] == ["append_query"] + assert [step.id for step in response_pipeline.steps] == ["prefix_response"] + + +def test_entry_point_name_must_be_author_step_name(monkeypatch): + monkeypatch.setattr( + "brainforge_inference_utils.registry.entry_points", + _entry_points_for([_EntryPoint("guardrails", AcmeRequestGuard)]), + ) + with pytest.raises(ValueError, match="author.step_name"): + load_entry_point_steps(REQUEST_STEP_ENTRY_POINT, RequestStep) + + +def test_duplicate_name_in_one_group_names_both_packages(monkeypatch): + monkeypatch.setattr( + "brainforge_inference_utils.registry.entry_points", + _entry_points_for([ + _EntryPoint("acme.guard", AcmeRequestGuard, dist="acme-a"), + _EntryPoint("acme.guard", AcmeRequestGuard, dist="acme-b"), + ]), + ) + with pytest.raises(ValueError, match="acme-a and acme-b"): + load_entry_point_steps(REQUEST_STEP_ENTRY_POINT, RequestStep) + + +def test_class_id_must_match_entry_point_name(monkeypatch): + monkeypatch.setattr( + "brainforge_inference_utils.registry.entry_points", + _entry_points_for([_EntryPoint("acme.guard", MismatchedId)]), + ) + with pytest.raises(ValueError, match="does not match entry point"): + load_entry_point_steps(REQUEST_STEP_ENTRY_POINT, RequestStep) + + +def test_response_class_in_request_group_raises(monkeypatch): + monkeypatch.setattr( + "brainforge_inference_utils.registry.entry_points", + _entry_points_for([_EntryPoint("acme.guard", AcmeResponseGuard)]), + ) + with pytest.raises(TypeError, match="not a RequestStep"): + load_entry_point_steps(REQUEST_STEP_ENTRY_POINT, RequestStep) + + +def test_plugin_id_that_collides_with_builtin_raises(monkeypatch): + monkeypatch.setattr( + "brainforge_inference_utils.pipeline.steps_by_id", + {"acme.guard": AcmeRequestGuard}, + ) + monkeypatch.setattr( + "brainforge_inference_utils.registry.entry_points", + _entry_points_for([ + _EntryPoint("acme.guard", AcmeRequestGuard, dist="acme-plugin"), + ]), + ) + with pytest.raises(ValueError, match="conflicts with a built-in step"): + build_pipelines(PipelineConfig.model_validate({ + "request": ["acme.guard"], + })) From f4f2a9ea72c5ce62b0972afc0d45897a41be1e86 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 14:22:01 -0700 Subject: [PATCH 13/23] Refactor default plugins out of this package --- README.md | 12 +- brainforge_inference_utils/config.py | 8 +- brainforge_inference_utils/steps/__init__.py | 23 +- .../steps/ambient_context.py | 194 --------- brainforge_inference_utils/steps/rag.py | 388 ------------------ .../steps/request_history.py | 177 -------- tests/test_ambient_context.py | 207 ---------- tests/test_pipeline.py | 19 +- tests/test_rag.py | 129 ------ tests/test_registry.py | 20 +- tests/test_request_history.py | 185 --------- tests/test_response_parser.py | 141 ------- 12 files changed, 27 insertions(+), 1476 deletions(-) delete mode 100644 brainforge_inference_utils/steps/ambient_context.py delete mode 100644 brainforge_inference_utils/steps/rag.py delete mode 100644 brainforge_inference_utils/steps/request_history.py delete mode 100644 tests/test_ambient_context.py delete mode 100644 tests/test_rag.py delete mode 100644 tests/test_request_history.py delete mode 100644 tests/test_response_parser.py diff --git a/README.md b/README.md index 09fb3ab..33aeb4b 100644 --- a/README.md +++ b/README.md @@ -123,9 +123,15 @@ pipeline: - id: acme.guardrails ``` -Bundled steps (`ambient_context`, `rag`, `request_history`, `citations`) are -still resolved from `steps_by_id` in this package. They do not use -`author.step_name` yet. +First-party steps ship in their own packages: + +| Package | Steps | +| --- | --- | +| `neon-inference-plugin-locale` | `neon.locale` | +| `neon-inference-plugin-rag` | `neon.rag`, `neon.citations` | +| `neon-inference-plugin-history` | `neon.history` | + +This package does not register those steps. Install the plugin to use the id. ## Deferred diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py index 11a213e..85cf1f8 100644 --- a/brainforge_inference_utils/config.py +++ b/brainforge_inference_utils/config.py @@ -25,15 +25,15 @@ pipeline: request: - - id: ambient_context + - id: neon.locale config: mcp_url: https://mcp.neonaiservices2.com/mcp default_location: Renton, WA, USA timeout: 5.0 - - {id: rag, config: {mcp_url: https://mcp.neonaiservices2.com/mcp, top_k: 10}} - - request_history # last: valid, truncated history + - {id: neon.rag, config: {mcp_url: https://mcp.neonaiservices2.com/mcp, top_k: 10}} + - neon.history # last: valid, truncated history response: - - id: citations + - id: neon.citations config: mcp_url: https://mcp.neonaiservices2.com/mcp openai_url: https://hana.example/brainforge/openai/chat/completions diff --git a/brainforge_inference_utils/steps/__init__.py b/brainforge_inference_utils/steps/__init__.py index aae42ba..3ede871 100644 --- a/brainforge_inference_utils/steps/__init__.py +++ b/brainforge_inference_utils/steps/__init__.py @@ -13,43 +13,30 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . """ -Step base classes, first-party step implementations, and their id mapping. +Step base classes. -Bundled steps subclass `RequestStep` or `ResponseStep`, set a unique `id`, and -are added to `steps_by_id` below so pipelines can resolve them by id. Plugin -steps are discovered from entry points instead of this mapping. Their ids use -`author.step_name`. +Concrete steps ship in separate plugins and are discovered from entry points. +`steps_by_id` remains for any steps still bundled with this package. Plugin +step ids use `author.step_name`. """ from __future__ import annotations from typing import Dict, Type -from brainforge_inference_utils.steps.ambient_context import AmbientContextStep from brainforge_inference_utils.steps.base import ( RequestStep, RequestContextOption, ResponseStep, Step, ) -from brainforge_inference_utils.steps.rag import RagStep, ResponseParser -from brainforge_inference_utils.steps.request_history import RequestHistoryStep -steps_by_id: Dict[str, Type[Step]] = { - AmbientContextStep.id: AmbientContextStep, - RagStep.id: RagStep, - ResponseParser.id: ResponseParser, - RequestHistoryStep.id: RequestHistoryStep, -} +steps_by_id: Dict[str, Type[Step]] = {} __all__ = [ "Step", "RequestStep", "ResponseStep", "RequestContextOption", - "AmbientContextStep", - "RagStep", - "ResponseParser", - "RequestHistoryStep", "steps_by_id", ] diff --git a/brainforge_inference_utils/steps/ambient_context.py b/brainforge_inference_utils/steps/ambient_context.py deleted file mode 100644 index 960e422..0000000 --- a/brainforge_inference_utils/steps/ambient_context.py +++ /dev/null @@ -1,194 +0,0 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -""" -Inject ambient time and weather context as `information` messages. - -Mirrors the time/weather injection in `brainforge-webapp`: resolve the user's -location (from `transaction.request_context` or step config), call MCP tools, -and prepend two `information` tuples to `request.history`. -""" - -from __future__ import annotations - -import asyncio -import json -import logging -from typing import Any, Dict, List, Optional, Tuple - -import requests -from neon_data_models.types import LlmMessageRole - -from brainforge_inference_utils.mcp import mcp_client, mcp_tool_inner_payload -from brainforge_inference_utils.steps.base import RequestStep, RequestContextOption -from brainforge_inference_utils.transaction import UserTransaction - -LOG = logging.getLogger(__name__) - - -class AmbientContextStep(RequestStep): - """ - Prepend time and weather `information` messages to `request.history`. - - Required configuration: - - - `mcp_url`: MCP server URL. - - `default_location`: Fallback when no location can be resolved. - - `timeout`: Seconds to wait for MCP calls. - - `transaction.request_context` may include: - - - `location`: Pre-resolved location description (preferred when set). - - `client_ip`: Client IP used to resolve location via ip-api.com. - """ - - id = "ambient_context" - request_context_options = ( - RequestContextOption( - "location", - description="Pre-resolved location description (preferred when set).", - ), - RequestContextOption( - "client_ip", - description="Client IP used to resolve location via ip-api.com.", - ), - ) - - def __init__(self, config: Optional[dict] = None) -> None: - resolved = dict(config or {}) - resolved.setdefault("mcp_url", "https://mcp.neonaiservices2.com/mcp") - resolved.setdefault("default_location", "Renton, WA, USA") - resolved.setdefault("timeout", 5.0) - super().__init__(resolved) - - def applicable(self, transaction: UserTransaction) -> Optional[str]: - for key in ("mcp_url", "timeout"): - if self.config.get(key) in (None, ""): - return f"{key} not configured" - return None - - async def process(self, transaction: UserTransaction) -> None: - location = self._resolve_location(transaction) - mcp_url = self.config["mcp_url"] - timeout = float(self.config["timeout"]) - - step_context: Dict[str, Any] = {"location": location} - try: - time_str, weather_str = await self._fetch_time_and_weather( - mcp_url, location, timeout=timeout) - except Exception as exc: - LOG.warning("Skipping ambient context for %r: %s", location, exc) - step_context["error"] = str(exc) - transaction.request_step_context[self.id] = step_context - return - - messages = self._build_information_messages( - location, time_str, weather_str) - transaction.request.history = ( - list(messages) + list(transaction.request.history)) - - step_context.update({ - "time": time_str, - "weather": weather_str, - "messages": [content for _, content in messages], - }) - transaction.request_step_context[self.id] = step_context - - def _resolve_location(self, transaction: UserTransaction) -> str: - ctx = transaction.request_context - location = ctx.get("location") - if location: - return str(location) - - client_ip = ctx.get("client_ip") - if client_ip: - return self._get_location_from_ip( - str(client_ip), self.config["default_location"]) - - return self.config["default_location"] - - @staticmethod - def _get_location_from_ip(ip_addr: str, default_location: str) -> str: - """Resolve a human-readable location string from a client IP address.""" - location = None - if ip_addr: - try: - response = requests.get( - f"http://ip-api.com/json/{ip_addr}", - timeout=5, - ).json() - location = ", ".join( - part - for part in ( - response.get("city"), - response.get("region"), - response.get("country"), - ) - if part - ) - except Exception: - location = None - return location or default_location - - async def _fetch_time_and_weather( - self, - mcp_url: str, - location: str, - timeout: float, - ) -> tuple[str, str]: - async def _fetch() -> tuple[str, str]: - async with mcp_client(mcp_url) as client: - time_result = await client.call_tool( - "get_formatted_time", - {"location": location}, - ) - time_str = json.dumps(mcp_tool_inner_payload(time_result)) - - units = "imperial" if "USA" in location else "metric" - weather_result = await client.call_tool_mcp( - "get_weather_forecast", - {"location": location, "unit": units}, - ) - weather_str = self._format_weather_payload( - mcp_tool_inner_payload(weather_result)) - return time_str, weather_str - - return await asyncio.wait_for(_fetch(), timeout=timeout) - - @staticmethod - def _format_weather_payload(weather_dict: dict) -> str: - """Trim weather forecast fields the same way brainforge-webapp does.""" - weather_dict = dict(weather_dict) - weather_dict.pop("minutely", None) - if "hourly" in weather_dict: - weather_dict["hourly"] = weather_dict["hourly"][:24] - if "daily" in weather_dict: - weather_dict["daily"] = weather_dict["daily"][:4] - return json.dumps(weather_dict) - - @staticmethod - def _build_information_messages( - location: str, - time_str: str, - weather_str: str, - ) -> List[Tuple[LlmMessageRole, str]]: - return [ - ("information", - f"The current time and date is: {time_str}.\n" - f"The user's location is {location}."), - ("information", - "Weather forcast reported in Imperial units for USA locations, " - f"otherwise Metric:\n{weather_str}\n---"), - ] - diff --git a/brainforge_inference_utils/steps/rag.py b/brainforge_inference_utils/steps/rag.py deleted file mode 100644 index 81adddd..0000000 --- a/brainforge_inference_utils/steps/rag.py +++ /dev/null @@ -1,388 +0,0 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -""" -RAG retrieval and citation rewrite. - -`RagStep` lists providers and injects retrieved documents on the request -pipeline. `ResponseParser` (id `citations`) runs on the response pipeline: -it filters those documents, builds a citation completion via MCP, POSTs it -to a configured OpenAI-compatible endpoint, and writes the annotated text -back onto the same `UserTransaction`. -""" - -from __future__ import annotations - -import asyncio -import json -import logging -from typing import Any, Dict, List, Optional, Tuple - -import requests -from neon_data_models.models.api.llm import LLMResponse - -from brainforge_inference_utils.mcp import mcp_client, mcp_tool_content_parts -from brainforge_inference_utils.steps.base import ( - RequestStep, - RequestContextOption, - ResponseStep, -) -from brainforge_inference_utils.transaction import UserTransaction - -LOG = logging.getLogger(__name__) - - -class RagStep(RequestStep): - """ - Look up selected RAG providers and prepend retrieved context. - - Required configuration: - - - `mcp_url`: MCP server URL. - - `top_k`: Default document cap for `lookup_rag`. - - `timeout`: Seconds to wait for MCP calls. - - `transaction.request_context` may include: - - - `rag_providers`: Provider ids selected by the host (required to run). - - `rag_top_k`: Optional per-request override of configured `top_k`. - """ - - id = "rag" - - def __init__(self, config: Optional[dict] = None) -> None: - resolved = dict(config or {}) - resolved.setdefault("mcp_url", "https://mcp.neonaiservices2.com/mcp") - resolved.setdefault("top_k", 10) - resolved.setdefault("timeout", 10.0) - super().__init__(resolved) - - async def supported_request_context( - self, timeout: float = 10.0 - ) -> List[RequestContextOption]: - """Advertise selected RAG providers (live catalog) and optional top_k.""" - providers = await self._available_providers(timeout=timeout) - return [ - RequestContextOption( - "rag_providers", - providers, - "Provider ids selected by the host (required to run).", - ), - RequestContextOption( - "rag_top_k", - description="Optional per-request override of configured top_k.", - ), - ] - - def applicable(self, transaction: UserTransaction) -> Optional[str]: - if self.config.get("mcp_url") in (None, ""): - return "mcp_url not configured" - if self.config.get("top_k") in (None, ""): - return "top_k not configured" - providers = transaction.request_context.get("rag_providers") or [] - if not providers: - return "no rag providers selected" - return None - - async def _available_providers(self, timeout: float = 10.0) -> List[str]: - """Return the live provider catalog from the configured MCP server.""" - timeout = float(timeout if timeout is not None else self.config["timeout"]) - - async def _fetch() -> List[str]: - async with mcp_client(self.config["mcp_url"]) as client: - result = await client.call_tool_mcp("available_contexts", {}) - if getattr(result, "isError", False): - err_text = ( - result.content[0].text - if getattr(result, "content", None) - else "available_contexts failed" - ) - raise RuntimeError(err_text) - return [ - part.text - for part in mcp_tool_content_parts(result) - if getattr(part, "text", None) - ] - - return await asyncio.wait_for(_fetch(), timeout=timeout) - - async def process(self, transaction: UserTransaction) -> None: - providers = [str(p) for p in ( - transaction.request_context.get("rag_providers") or [])] - top_k = int( - transaction.request_context.get("rag_top_k") - if transaction.request_context.get("rag_top_k") not in (None, "") - else self.config["top_k"] - ) - timeout = float(self.config["timeout"]) - step_context: Dict[str, Any] = { - "providers": providers, - "top_k": top_k, - "documents": [], - } - try: - documents = await self._lookup_rag( - transaction.request.query, providers, top_k, timeout) - except Exception as exc: - LOG.warning("Skipping RAG lookup for %s: %s", providers, exc) - step_context["error"] = str(exc) - transaction.request_step_context[self.id] = step_context - return - - step_context["documents"] = documents - if documents: - transaction.request.history = ( - [self._information_message(documents)] - + list(transaction.request.history) - ) - transaction.request_step_context[self.id] = step_context - - async def _lookup_rag( - self, - query: str, - providers: List[str], - top_k: int, - timeout: float, - ) -> List[dict]: - async def _fetch() -> List[dict]: - async with mcp_client(self.config["mcp_url"]) as client: - result = await client.call_tool( - "lookup_rag", - { - "message": query, - "top_k": top_k, - "context_providers": providers, - }, - ) - documents = [] - for part in mcp_tool_content_parts(result): - documents.append(json.loads(part.text)) - return documents - - return await asyncio.wait_for(_fetch(), timeout=timeout) - - @staticmethod - def _information_message(documents: List[dict]) -> tuple: - context_str = "\n".join( - f"{doc.get('metadata', {}).get('file_name', doc.get('metadata', {}).get('citation_uri'))}: " - f"{doc.get('text', '')}\n---" - for doc in documents - ) - return ("information", context_str) - - -def _citation_key(document: dict) -> str: - meta = document.get("metadata") or {} - return str( - meta.get("citation") - or meta.get("citation_uri") - or meta.get("file_name") - or "" - ) - - -def apply_citation_markup( - response: str, documents: List[dict] -) -> Tuple[str, List[str], str, str]: - """ - Map raw `[n]` citation tags onto unique documents and markdown links. - - :returns: rewritten text, unique citation labels, compact credits, - formatted credits for a UI box. - """ - unique_contexts = list({_citation_key(doc) for doc in documents}) - raw_to_unique_idx = { - documents.index(doc) + 1: unique_contexts.index(_citation_key(doc)) + 1 - for doc in documents - } - rewritten = response.replace("[0]", "") - for raw_idx, unique_idx in raw_to_unique_idx.items(): - citation_uri = (documents[raw_idx - 1].get("metadata") or {}).get( - "citation_uri", "") - citation_string = ( - f"[[0{unique_idx}]({citation_uri})]" if citation_uri - else f"[{unique_idx}]" - ) - if f"`[{raw_idx}]`" in rewritten: - rewritten = rewritten.replace(f"`[{raw_idx}]`", citation_string) - else: - rewritten = rewritten.replace(f"[{raw_idx}]", citation_string) - rewritten = rewritten.replace("[[0", "[[") - credits = ", ".join( - f"[{unique_contexts.index(label) + 1}] {label}" - for label in unique_contexts - ) - formatted_credits = "\n".join( - f"[{unique_contexts.index(label) + 1}] {label}" - for label in unique_contexts - ) - return rewritten, unique_contexts, credits, formatted_credits - - -class ResponseParser(ResponseStep): - """ - Rewrite the generated response with citations using RAG request-step context. - - The whole MCP filter → citation completion → markup path runs here so the - host does not handle a second request/response pair. - - Required configuration: - - - `mcp_url`: MCP server URL (same server as `RagStep`). - - `openai_url`: OpenAI-compatible chat completions URL (or API base). - - `timeout`: Seconds to wait for MCP and the completion call. - - Optional configuration: - - - `openai_api_key`: Bearer token for `openai_url`. - - `transaction.request_context` may include: - - - `citations_enabled`: When false, the step is skipped. - """ - - id = "citations" - request_context_options = ( - RequestContextOption( - "citations_enabled", - (True, False), - "When false, the step is skipped.", - ), - ) - - def __init__(self, config: Optional[dict] = None) -> None: - resolved = dict(config or {}) - resolved.setdefault("mcp_url", "https://mcp.neonaiservices2.com/mcp") - resolved.setdefault("timeout", 30.0) - super().__init__(resolved) - - def applicable(self, transaction: UserTransaction) -> Optional[str]: - if transaction.response is None: - return "no response" - if self.config.get("mcp_url") in (None, ""): - return "mcp_url not configured" - if self.config.get("openai_url") in (None, ""): - return "openai_url not configured" - if transaction.request_context.get("citations_enabled") is False: - return "citations disabled" - documents = (transaction.request_step_context.get("rag") or {}).get( - "documents") or [] - if not documents: - return "no rag documents" - return None - - async def process( - self, transaction: UserTransaction - ) -> Optional[LLMResponse]: - documents = list( - (transaction.request_step_context.get("rag") or {}).get("documents") or []) - original = transaction.response.response - timeout = float(self.config["timeout"]) - try: - filtered = await self._filter_sources(original, documents, timeout) - if not filtered: - transaction.response_step_context[self.id] = { - "documents": [], - "credits": "", - "formatted_credits": "", - } - return None - payload = await self._build_citation_request( - original, filtered, timeout) - raw_cited = await self._complete( - payload, transaction.request.model, timeout) - rewritten, unique, credits, formatted = apply_citation_markup( - raw_cited, filtered) - except Exception as exc: - LOG.warning("Skipping citation rewrite: %s", exc) - transaction.response_step_context[self.id] = { - "documents": documents, - "error": str(exc), - } - return None - - transaction.response_step_context[self.id] = { - "documents": filtered, - "unique_citations": unique, - "credits": credits, - "formatted_credits": formatted, - } - history = list(transaction.response.history) - if history and history[-1][0] == "llm": - history[-1] = ("llm", rewritten) - return transaction.response.model_copy( - update={"response": rewritten, "history": history}) - - async def _filter_sources( - self, response: str, documents: List[dict], timeout: float - ) -> List[dict]: - async def _fetch() -> List[dict]: - async with mcp_client(self.config["mcp_url"]) as client: - result = await client.call_tool( - "filter_citation_sources", - { - "generated_response": response, - "rag_documents": documents, - }, - ) - return [ - json.loads(part.text) - for part in mcp_tool_content_parts(result) - ] - - return await asyncio.wait_for(_fetch(), timeout=timeout) - - async def _build_citation_request( - self, response: str, documents: List[dict], timeout: float - ) -> dict: - async def _fetch() -> dict: - async with mcp_client(self.config["mcp_url"]) as client: - result = await client.call_tool( - "build_citation_request", - { - "original_response": response, - "rag_documents": documents, - }, - ) - parts = mcp_tool_content_parts(result) - return json.loads(parts[0].text) - - return await asyncio.wait_for(_fetch(), timeout=timeout) - - async def _complete( - self, payload: dict, model: str, timeout: float - ) -> str: - request = dict(payload) - request.setdefault("model", model) - request["stream"] = False - url = self._completions_url() - headers = {"Content-Type": "application/json"} - api_key = self.config.get("openai_api_key") - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - - def _post() -> dict: - response = requests.post( - url, json=request, headers=headers, timeout=timeout) - response.raise_for_status() - return response.json() - - completion = await asyncio.to_thread(_post) - return completion["choices"][0]["message"]["content"] - - def _completions_url(self) -> str: - url = str(self.config["openai_url"]).rstrip("/") - if url.endswith("/chat/completions") or url.endswith("/completions"): - return url - return f"{url}/chat/completions" diff --git a/brainforge_inference_utils/steps/request_history.py b/brainforge_inference_utils/steps/request_history.py deleted file mode 100644 index ee14e15..0000000 --- a/brainforge_inference_utils/steps/request_history.py +++ /dev/null @@ -1,177 +0,0 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -""" -Guarantee `LLMRequest.history` is valid, including truncation. - -`LLMRequest` contract: history excludes the system prompt (on `persona`) and -the incoming query (on `request.query`). Roles are `information` / -`instruction` / `tool` context, then complete `user`/`llm` pairs. - -This step is meant to run last on the request pipeline so earlier steps can -prepend context freely. It repairs rather than skipping: unknown or empty -messages are dropped, prior citation markup is stripped, `assistant` becomes -`llm`, a trailing user turn that duplicates `query` is removed, then -conversational turns are truncated to `max_history` complete pairs. -""" - -from __future__ import annotations - -import re -from typing import Any, List, Optional, Sequence, Tuple - -from neon_data_models.types import LlmMessageRole - -from brainforge_inference_utils.steps.base import RequestStep, RequestContextOption -from brainforge_inference_utils.transaction import UserTransaction - -# Matches `[[n](url)]` written by `ResponseParser.apply_citation_markup`. -_CITATION_MARKUP = re.compile(r"\[\[\d+\]\([^\)]*\)\]") - -_CONTEXT_ROLES = {"information", "instruction", "tool"} -_TURN_ROLES = {"user", "llm", "assistant"} -HistoryItem = Tuple[LlmMessageRole, str] - - -def remove_citation_strings(text: str) -> str: - """Remove citation strings like `[[1](https://example)]` from text.""" - return _CITATION_MARKUP.sub("", text) - - -class RequestHistoryStep(RequestStep): - """ - Normalize, truncate, and order `request.history` into a valid LLM history. - - Optional configuration: - - - `max_history`: User/llm pairs to keep, excluding `request.query`. - Defaults to `transaction.request.max_history`. - - `max_messages`: Alternate cap that includes `request.query` - (webapp `history[-5:]` after appending the query). If set, pair budget - is `(max_messages - 1) // 2`. - - `transaction.request_context` may include: - - - `conversational`: When false, drop all prior turns (query only). - Defaults to true. - """ - - id = "request_history" - request_context_options = ( - RequestContextOption( - "conversational", - (True, False), - "When false, drop all prior turns (query only). Defaults to true.", - ), - ) - - def applicable(self, transaction: UserTransaction) -> Optional[str]: - return None - - async def process(self, transaction: UserTransaction) -> None: - conversational = bool( - transaction.request_context.get("conversational", True)) - max_history = self._pair_budget(transaction, conversational) - - cleaned = self._clean_messages( - transaction.request.history, transaction.request.query) - context_messages, turn_messages = self._split_history(cleaned) - windowed_turns = self._last_complete_pairs(turn_messages, max_history) - dropped_turns = len(turn_messages) - len(windowed_turns) - - transaction.request.history = list(context_messages) + list(windowed_turns) - transaction.request.max_history = max_history - transaction.request_step_context[self.id] = { - "conversational": conversational, - "max_history": max_history, - "kept_turns": len(windowed_turns), - "dropped_turns": dropped_turns, - } - - def _pair_budget(self, transaction: UserTransaction, - conversational: bool) -> int: - if not conversational: - return 0 - if self.config.get("max_messages") not in (None, ""): - max_messages = int(self.config["max_messages"]) - return max(0, (max_messages - 1) // 2) - if self.config.get("max_history") not in (None, ""): - return max(0, int(self.config["max_history"])) - return max(0, int(transaction.request.max_history)) - - @staticmethod - def _clean_messages( - history: Sequence[Tuple[Any, Any]], - query: str, - ) -> List[HistoryItem]: - cleaned: List[HistoryItem] = [] - for role, content in history: - if not isinstance(content, str): - continue - content = remove_citation_strings(content).strip() - if not content: - continue - if role == "assistant": - role = "llm" - if role == "system": - # System prompt belongs on `request.persona`, not history. - continue - if role not in _CONTEXT_ROLES and role not in _TURN_ROLES: - continue - cleaned.append((role, content)) - query_text = remove_citation_strings(query).strip() - if ( - cleaned - and cleaned[-1][0] == "user" - and cleaned[-1][1] == query_text - ): - cleaned.pop() - return cleaned - - @staticmethod - def _split_history( - history: Sequence[HistoryItem], - ) -> Tuple[List[HistoryItem], List[HistoryItem]]: - context_messages: List[HistoryItem] = [] - turn_messages: List[HistoryItem] = [] - for role, content in history: - if role in _TURN_ROLES: - turn_messages.append((role, content)) - else: - context_messages.append((role, content)) - return context_messages, turn_messages - - @staticmethod - def _last_complete_pairs( - turns: Sequence[HistoryItem], - max_history: int, - ) -> List[HistoryItem]: - if max_history <= 0 or not turns: - return [] - pairs: List[Tuple[HistoryItem, HistoryItem]] = [] - idx = len(turns) - 1 - while idx >= 1 and len(pairs) < max_history: - later = turns[idx] - earlier = turns[idx - 1] - if later[0] == "llm" and earlier[0] == "user": - pairs.append((earlier, later)) - idx -= 2 - else: - idx -= 1 - pairs.reverse() - windowed: List[HistoryItem] = [] - for earlier, later in pairs: - windowed.append(earlier) - windowed.append(later) - return windowed diff --git a/tests/test_ambient_context.py b/tests/test_ambient_context.py deleted file mode 100644 index 164e4af..0000000 --- a/tests/test_ambient_context.py +++ /dev/null @@ -1,207 +0,0 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import json -import unittest -from unittest.mock import AsyncMock, patch - -from brainforge_inference_utils.steps.ambient_context import AmbientContextStep -from brainforge_inference_utils.transaction import UserTransaction -from helpers import make_request - - -STEP_CONFIG = { - "mcp_url": "http://localhost:8000/mcp", - "default_location": "Renton, WA, USA", - "timeout": 5.0, -} - - -class TestFormatWeatherPayload(unittest.TestCase): - def test_trims_forecast_fields(self): - payload = { - "minutely": [{"time": 1}], - "hourly": list(range(30)), - "daily": list(range(10)), - } - formatted = AmbientContextStep._format_weather_payload(payload) - trimmed = json.loads(formatted) - self.assertNotIn("minutely", trimmed) - self.assertEqual(len(trimmed["hourly"]), 24) - self.assertEqual(len(trimmed["daily"]), 4) - - -class TestAmbientContextStep(unittest.IsolatedAsyncioTestCase): - async def test_prepends_information_to_request_history(self): - step = AmbientContextStep(STEP_CONFIG) - transaction = UserTransaction( - request=make_request(history=[("user", "hi"), ("llm", "hey")]), - request_context={"location": "Seattle, WA, USA"}, - ) - - with patch.object( - step, - "_fetch_time_and_weather", - new=AsyncMock(return_value=('{"now": "Monday"}', '{"temp": 60}')), - ) as fetch: - await step.process(transaction) - - fetch.assert_awaited_once_with( - STEP_CONFIG["mcp_url"], "Seattle, WA, USA", timeout=5.0) - self.assertEqual(len(transaction.request.history), 4) - self.assertEqual(transaction.request.history[0][0], "information") - self.assertIn("Monday", transaction.request.history[0][1]) - self.assertIn("Seattle, WA, USA", transaction.request.history[0][1]) - self.assertEqual(transaction.request.history[1][0], "information") - self.assertIn("Weather forcast", transaction.request.history[1][1]) - self.assertIn('{"temp": 60}', transaction.request.history[1][1]) - self.assertEqual(transaction.request.history[2:], [ - ("user", "hi"), - ("llm", "hey"), - ]) - self.assertEqual( - transaction.request_step_context["ambient_context"]["location"], - "Seattle, WA, USA", - ) - - async def test_resolves_location_from_client_ip(self): - step = AmbientContextStep(STEP_CONFIG) - transaction = UserTransaction( - request=make_request(), - request_context={"client_ip": "203.0.113.10"}, - ) - - with patch.object( - step, - "_get_location_from_ip", - return_value="Portland, OR, USA", - ) as get_location, patch.object( - step, - "_fetch_time_and_weather", - new=AsyncMock(return_value=('{"now": "Tuesday"}', '{"temp": 55}')), - ) as fetch: - await step.process(transaction) - - get_location.assert_called_once_with( - "203.0.113.10", STEP_CONFIG["default_location"]) - fetch.assert_awaited_once_with( - STEP_CONFIG["mcp_url"], "Portland, OR, USA", timeout=5.0) - - async def test_uses_default_location_when_request_context_empty(self): - step = AmbientContextStep(STEP_CONFIG) - transaction = UserTransaction(request=make_request()) - - with patch.object( - step, - "_fetch_time_and_weather", - new=AsyncMock(return_value=('{"now": "Wednesday"}', '{"temp": 50}')), - ) as fetch: - await step.process(transaction) - - fetch.assert_awaited_once_with( - STEP_CONFIG["mcp_url"], - STEP_CONFIG["default_location"], - timeout=5.0, - ) - - async def test_uses_configured_mcp_url(self): - config = { - **STEP_CONFIG, - "mcp_url": "https://mcp.example.com/mcp", - "timeout": 2.5, - } - step = AmbientContextStep(config) - transaction = UserTransaction( - request=make_request(), - request_context={"location": "Austin, TX, USA"}, - ) - - with patch.object( - step, - "_fetch_time_and_weather", - new=AsyncMock(return_value=('{"now": "Thursday"}', '{"temp": 70}')), - ) as fetch: - await step.process(transaction) - - fetch.assert_awaited_once_with( - "https://mcp.example.com/mcp", "Austin, TX, USA", timeout=2.5) - - async def test_mcp_failure_leaves_history_unchanged(self): - step = AmbientContextStep(STEP_CONFIG) - original_history = [("user", "hi"), ("llm", "hey")] - transaction = UserTransaction( - request=make_request(history=list(original_history)), - request_context={"location": "Denver, CO, USA"}, - ) - - with patch.object( - step, - "_fetch_time_and_weather", - new=AsyncMock(side_effect=TimeoutError("timed out")), - ): - await step.process(transaction) - - self.assertEqual(transaction.request.history, original_history) - self.assertEqual( - transaction.request_step_context["ambient_context"], - {"location": "Denver, CO, USA", "error": "timed out"}, - ) - - async def test_prefers_request_context_location_over_client_ip(self): - step = AmbientContextStep(STEP_CONFIG) - transaction = UserTransaction( - request=make_request(), - request_context={ - "location": "Chicago, IL, USA", - "client_ip": "203.0.113.10", - }, - ) - - with patch.object( - step, - "_get_location_from_ip", - ) as get_location, patch.object( - step, - "_fetch_time_and_weather", - new=AsyncMock(return_value=('{"now": "Friday"}', '{"temp": 40}')), - ) as fetch: - await step.process(transaction) - - get_location.assert_not_called() - fetch.assert_awaited_once_with( - STEP_CONFIG["mcp_url"], "Chicago, IL, USA", timeout=5.0) - - def test_skips_when_timeout_missing(self): - step = AmbientContextStep({ - "mcp_url": "http://localhost:8000/mcp", - "default_location": "Renton, WA, USA", - "timeout": "", - }) - reason = step.applicable(UserTransaction(request=make_request())) - self.assertEqual(reason, "timeout not configured") - - def test_uses_default_location_from_step_config(self): - step = AmbientContextStep({"mcp_url": "http://localhost:8000/mcp"}) - transaction = UserTransaction(request=make_request()) - self.assertEqual( - step._resolve_location(transaction), - "Renton, WA, USA", - ) - - async def test_supported_request_context_is_unconstrained(self): - step = AmbientContextStep(STEP_CONFIG) - options = await step.supported_request_context() - self.assertEqual([option.key for option in options], ["location", "client_ip"]) - self.assertTrue(all(option.values is None for option in options)) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 82e45ca..91b2a33 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -43,23 +43,8 @@ def test_empty_config_builds_empty_pipelines(self): assert request_pipeline.steps == [] assert response_pipeline.steps == [] - def test_default_mapping_includes_first_party_steps(self): - assert "ambient_context" in steps_by_id - assert "rag" in steps_by_id - assert "citations" in steps_by_id - assert "request_history" in steps_by_id - request_pipeline, response_pipeline = build_pipelines( - PipelineConfig.model_validate({ - "request": ["ambient_context", "rag", "request_history"], - "response": ["citations"], - })) - assert [step.id for step in request_pipeline.steps] == [ - "ambient_context", - "rag", - "request_history", - ] - assert [step.id for step in response_pipeline.steps] == ["citations"] - + def test_default_registry_has_no_bundled_steps(self): + assert steps_by_id == {} with pytest.raises(KeyError, match="append_query"): build_pipelines(PipelineConfig.model_validate({ "request": ["append_query"], diff --git a/tests/test_rag.py b/tests/test_rag.py deleted file mode 100644 index 18bc0ea..0000000 --- a/tests/test_rag.py +++ /dev/null @@ -1,129 +0,0 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import unittest -from unittest.mock import AsyncMock, patch - -from brainforge_inference_utils.steps.rag import RagStep -from brainforge_inference_utils.transaction import UserTransaction -from helpers import make_request - - -STEP_CONFIG = { - "mcp_url": "http://localhost:8000/mcp", - "top_k": 10, - "timeout": 10.0, -} - -DOCUMENTS = [ - { - "text": "chunk one", - "metadata": {"file_name": "doc-a.md", "citation_uri": "https://a"}, - }, - { - "text": "chunk two", - "metadata": {"citation_uri": "https://b"}, - }, -] - - -class TestRagStep(unittest.IsolatedAsyncioTestCase): - async def test_injects_information_and_stores_documents(self): - step = RagStep(STEP_CONFIG) - transaction = UserTransaction( - request=make_request(query="what?", history=[("user", "hi")]), - request_context={"rag_providers": ["news"]}, - ) - - with patch.object( - step, "_lookup_rag", new=AsyncMock(return_value=DOCUMENTS) - ) as lookup: - await step.process(transaction) - - lookup.assert_awaited_once_with( - "what?", ["news"], 10, STEP_CONFIG["timeout"]) - self.assertEqual(transaction.request.history[0][0], "information") - self.assertIn("doc-a.md: chunk one", transaction.request.history[0][1]) - self.assertIn("https://b: chunk two", transaction.request.history[0][1]) - self.assertEqual(transaction.request.history[1], ("user", "hi")) - self.assertEqual( - transaction.request_step_context["rag"]["documents"], DOCUMENTS) - self.assertEqual( - transaction.request_step_context["rag"]["providers"], ["news"]) - - async def test_uses_request_top_k_override(self): - step = RagStep(STEP_CONFIG) - transaction = UserTransaction( - request=make_request(), - request_context={"rag_providers": ["docs"], "rag_top_k": 3}, - ) - - with patch.object( - step, "_lookup_rag", new=AsyncMock(return_value=[]) - ) as lookup: - await step.process(transaction) - - lookup.assert_awaited_once_with( - transaction.request.query, ["docs"], 3, STEP_CONFIG["timeout"]) - self.assertEqual(transaction.request.history, make_request().history) - self.assertEqual(transaction.request_step_context["rag"]["documents"], []) - - async def test_lookup_failure_leaves_history_unchanged(self): - step = RagStep(STEP_CONFIG) - original = [("user", "hi"), ("llm", "hey")] - transaction = UserTransaction( - request=make_request(history=list(original)), - request_context={"rag_providers": ["news"]}, - ) - - with patch.object( - step, - "_lookup_rag", - new=AsyncMock(side_effect=TimeoutError("timed out")), - ): - await step.process(transaction) - - self.assertEqual(transaction.request.history, original) - self.assertEqual( - transaction.request_step_context["rag"]["error"], "timed out") - - def test_skips_when_no_providers_selected(self): - step = RagStep(STEP_CONFIG) - reason = step.applicable(UserTransaction(request=make_request())) - self.assertEqual(reason, "no rag providers selected") - - async def test_supported_request_context_uses_live_catalog(self): - step = RagStep(STEP_CONFIG) - with patch.object( - step, - "_available_providers", - new=AsyncMock(return_value=["news", "docs"]), - ) as catalog: - options = { - option.key: option - for option in await step.supported_request_context(timeout=2.0) - } - - catalog.assert_awaited_once_with(timeout=2.0) - self.assertEqual(list(options["rag_providers"].values), ["news", "docs"]) - self.assertIsNone(options["rag_top_k"].values) - - def test_skips_when_mcp_url_missing(self): - step = RagStep({"mcp_url": "", "top_k": 10}) - transaction = UserTransaction( - request=make_request(), - request_context={"rag_providers": ["news"]}, - ) - self.assertEqual(step.applicable(transaction), "mcp_url not configured") diff --git a/tests/test_registry.py b/tests/test_registry.py index a384f73..ec7c308 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -88,7 +88,7 @@ def test_same_author_step_name_loads_from_both_groups(monkeypatch): assert response_steps["acme.guard"] is AcmeResponseGuard -def test_build_pipelines_resolves_plugins_and_builtin_steps(monkeypatch): +def test_build_pipelines_resolves_entry_point_steps(monkeypatch): monkeypatch.setattr( "brainforge_inference_utils.registry.entry_points", _entry_points_for( @@ -99,20 +99,14 @@ def test_build_pipelines_resolves_plugins_and_builtin_steps(monkeypatch): request_pipeline, response_pipeline = build_pipelines( PipelineConfig.model_validate({ - "request": ["ambient_context", "acme.guard"], - "response": ["citations", "acme.guard"], + "request": ["acme.guard"], + "response": ["acme.guard"], })) - assert [step.id for step in request_pipeline.steps] == [ - "ambient_context", - "acme.guard", - ] - assert [step.id for step in response_pipeline.steps] == [ - "citations", - "acme.guard", - ] - assert type(request_pipeline.steps[1]) is AcmeRequestGuard - assert type(response_pipeline.steps[1]) is AcmeResponseGuard + assert [step.id for step in request_pipeline.steps] == ["acme.guard"] + assert [step.id for step in response_pipeline.steps] == ["acme.guard"] + assert type(request_pipeline.steps[0]) is AcmeRequestGuard + assert type(response_pipeline.steps[0]) is AcmeResponseGuard def test_explicit_step_map_skips_entry_points(monkeypatch): diff --git a/tests/test_request_history.py b/tests/test_request_history.py deleted file mode 100644 index 9475a3d..0000000 --- a/tests/test_request_history.py +++ /dev/null @@ -1,185 +0,0 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import unittest - -from brainforge_inference_utils.steps.request_history import RequestHistoryStep -from brainforge_inference_utils.transaction import UserTransaction -from helpers import make_request - - -def _turns(*labels: str): - history = [] - for label in labels: - role, _, text = label.partition(":") - history.append((role, text)) - return history - - -class TestRequestHistoryStep(unittest.IsolatedAsyncioTestCase): - async def test_keeps_last_complete_pairs(self): - step = RequestHistoryStep({"max_messages": 5}) - history = _turns( - "user:u1", "llm:a1", "user:u2", "llm:a2", "user:u3", "llm:a3") - transaction = UserTransaction(request=make_request(history=history)) - - await step.process(transaction) - - self.assertEqual(transaction.request.history, [ - ("user", "u2"), - ("llm", "a2"), - ("user", "u3"), - ("llm", "a3"), - ]) - self.assertEqual(transaction.request.max_history, 2) - self.assertEqual( - transaction.request_step_context["request_history"]["dropped_turns"], 2) - self.assertEqual( - transaction.request_step_context["request_history"]["kept_turns"], 4) - - async def test_preserves_context_messages_outside_the_window(self): - step = RequestHistoryStep({"max_history": 2}) - history = _turns( - "information:time", - "information:weather", - "user:u1", "llm:a1", - "user:u2", "llm:a2", - "user:u3", "llm:a3", - ) - transaction = UserTransaction(request=make_request(history=history)) - - await step.process(transaction) - - self.assertEqual(transaction.request.history[0], ("information", "time")) - self.assertEqual( - transaction.request.history[1], ("information", "weather")) - self.assertEqual(transaction.request.history[2:], [ - ("user", "u2"), - ("llm", "a2"), - ("user", "u3"), - ("llm", "a3"), - ]) - - async def test_non_conversational_drops_prior_turns(self): - step = RequestHistoryStep({"max_history": 2}) - history = _turns( - "information:rag", - "user:u1", "llm:a1", "user:u2", "llm:a2", - ) - transaction = UserTransaction( - request=make_request(history=history), - request_context={"conversational": False}, - ) - - await step.process(transaction) - - self.assertEqual(transaction.request.history, [("information", "rag")]) - self.assertEqual(transaction.request.max_history, 0) - ctx = transaction.request_step_context["request_history"] - self.assertFalse(ctx["conversational"]) - self.assertEqual(ctx["dropped_turns"], 4) - - async def test_short_history_is_unchanged(self): - step = RequestHistoryStep({"max_history": 2}) - history = _turns("user:u1", "llm:a1") - transaction = UserTransaction(request=make_request(history=history)) - - await step.process(transaction) - - self.assertEqual(transaction.request.history, history) - self.assertEqual( - transaction.request_step_context["request_history"]["dropped_turns"], 0) - - async def test_drops_system_empty_and_duplicate_query(self): - step = RequestHistoryStep({"max_history": 2}) - transaction = UserTransaction( - request=make_request( - query="hello", - history=[ - ("information", "ctx"), - ("user", ""), - ("user", "u1"), - ("llm", "a1"), - ("user", "hello"), - ], - ), - ) - # `system` is not an LLMRequest history role; earlier steps or hosts - # may still append it by mutating the list. - transaction.request.history.insert(0, ("system", "you are a bot")) - - await step.process(transaction) - - self.assertEqual(transaction.request.history, [ - ("information", "ctx"), - ("user", "u1"), - ("llm", "a1"), - ]) - - async def test_drops_unpaired_turns_to_keep_complete_pairs(self): - step = RequestHistoryStep({"max_history": 2}) - history = _turns("llm:orphan", "user:u1", "llm:a1", "user:dangling") - transaction = UserTransaction(request=make_request(history=history)) - - await step.process(transaction) - - self.assertEqual(transaction.request.history, [ - ("user", "u1"), - ("llm", "a1"), - ]) - - async def test_uses_request_max_history_by_default(self): - step = RequestHistoryStep() - history = _turns( - "user:u1", "llm:a1", "user:u2", "llm:a2", "user:u3", "llm:a3") - transaction = UserTransaction( - request=make_request(history=history, max_history=1), - ) - - await step.process(transaction) - - self.assertEqual(transaction.request.history, [ - ("user", "u3"), - ("llm", "a3"), - ]) - self.assertEqual(transaction.request.max_history, 1) - - def test_always_applicable(self): - step = RequestHistoryStep() - self.assertIsNone( - step.applicable(UserTransaction(request=make_request()))) - - async def test_strips_citation_markup_from_history(self): - step = RequestHistoryStep({"max_history": 2}) - transaction = UserTransaction( - request=make_request(history=[ - ("user", "what about that?"), - ("llm", "See [[1](https://example.com/a)] for details."), - ]), - ) - - await step.process(transaction) - - self.assertEqual(transaction.request.history, [ - ("user", "what about that?"), - ("llm", "See for details."), - ]) - - async def test_supported_request_context_lists_conversational(self): - step = RequestHistoryStep() - options = await step.supported_request_context() - self.assertEqual(len(options), 1) - self.assertEqual(options[0].key, "conversational") - self.assertEqual(list(options[0].values), [True, False]) diff --git a/tests/test_response_parser.py b/tests/test_response_parser.py deleted file mode 100644 index 5d32e69..0000000 --- a/tests/test_response_parser.py +++ /dev/null @@ -1,141 +0,0 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - -import unittest -from unittest.mock import AsyncMock, patch - -from brainforge_inference_utils.steps.rag import ( - ResponseParser, - apply_citation_markup, -) -from brainforge_inference_utils.transaction import UserTransaction -from helpers import make_request, make_response - - -STEP_CONFIG = { - "mcp_url": "http://localhost:8000/mcp", - "openai_url": "https://hana.example/brainforge/openai/chat/completions", - "openai_api_key": "token", - "timeout": 30.0, -} - -DOCUMENTS = [ - {"text": "one", "metadata": {"citation": "Doc A", "citation_uri": "https://a"}}, - {"text": "two", "metadata": {"citation": "Doc A", "citation_uri": "https://a"}}, - {"text": "three", "metadata": {"citation": "Doc B", "citation_uri": "https://b"}}, -] - - -class TestApplyCitationMarkup(unittest.TestCase): - def test_rewrites_raw_indexes_to_unique_links(self): - rewritten, unique, credits, formatted = apply_citation_markup( - "See [1] and [3] and [0].", DOCUMENTS) - self.assertNotIn("[0]", rewritten) - self.assertIn("](https://", rewritten) - self.assertEqual(len(unique), 2) - self.assertIn("Doc A", credits) - self.assertIn("\n", formatted) - - -class TestResponseParser(unittest.IsolatedAsyncioTestCase): - def _transaction(self, text="See the docs.", **request_context): - request = make_request() - transaction = UserTransaction( - request=request, - response=make_response(text, request), - request_context=request_context, - request_step_context={"rag": {"documents": DOCUMENTS}}, - ) - return transaction - - async def test_rewrites_response_via_openai_endpoint(self): - step = ResponseParser(STEP_CONFIG) - transaction = self._transaction() - - with patch.object( - step, "_filter_sources", new=AsyncMock(return_value=DOCUMENTS) - ), patch.object( - step, "_build_citation_request", - new=AsyncMock(return_value={"messages": []}), - ) as build, patch.object( - step, "_complete", - new=AsyncMock(return_value="See [1] and [3]."), - ) as complete: - revised = await step.process(transaction) - - build.assert_awaited_once() - complete.assert_awaited_once_with( - {"messages": []}, transaction.request.model, 30.0) - self.assertIsNotNone(revised) - self.assertNotEqual(revised.response, "See the docs.") - self.assertEqual(revised.history[-1][0], "llm") - self.assertEqual(revised.history[-1][1], revised.response) - self.assertIn("Doc A", transaction.response_step_context["citations"]["credits"]) - - async def test_skips_without_rag_documents(self): - step = ResponseParser(STEP_CONFIG) - transaction = UserTransaction( - request=make_request(), - response=make_response(), - ) - self.assertEqual(step.applicable(transaction), "no rag documents") - - async def test_skips_when_openai_url_missing(self): - step = ResponseParser({"mcp_url": "http://localhost:8000/mcp"}) - transaction = self._transaction() - self.assertEqual(step.applicable(transaction), "openai_url not configured") - - async def test_skips_when_citations_disabled(self): - step = ResponseParser(STEP_CONFIG) - transaction = self._transaction(citations_enabled=False) - self.assertEqual(step.applicable(transaction), "citations disabled") - - async def test_completion_failure_leaves_response(self): - step = ResponseParser(STEP_CONFIG) - transaction = self._transaction() - original = transaction.response.response - - with patch.object( - step, "_filter_sources", new=AsyncMock(return_value=DOCUMENTS) - ), patch.object( - step, "_build_citation_request", - new=AsyncMock(return_value={"messages": []}), - ), patch.object( - step, "_complete", - new=AsyncMock(side_effect=RuntimeError("nope")), - ): - revised = await step.process(transaction) - - self.assertIsNone(revised) - self.assertEqual(transaction.response.response, original) - self.assertEqual( - transaction.response_step_context["citations"]["error"], "nope") - - def test_appends_chat_completions_to_base_url(self): - step = ResponseParser({ - **STEP_CONFIG, - "openai_url": "https://api.openai.com/v1", - }) - self.assertEqual( - step._completions_url(), - "https://api.openai.com/v1/chat/completions", - ) - - async def test_supported_request_context_lists_citations_enabled(self): - step = ResponseParser(STEP_CONFIG) - options = await step.supported_request_context() - self.assertEqual(len(options), 1) - self.assertEqual(options[0].key, "citations_enabled") - self.assertEqual(list(options[0].values), [True, False]) From eed1e4d4559421f00a9ca06548b95da301b3d2e2 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 14:30:55 -0700 Subject: [PATCH 14/23] Replace GNU Affero license with BSD-3 --- LICENSE.md | 681 +------------------ brainforge_inference_utils/__init__.py | 43 +- brainforge_inference_utils/client.py | 43 +- brainforge_inference_utils/config.py | 43 +- brainforge_inference_utils/mcp.py | 43 +- brainforge_inference_utils/pipeline.py | 43 +- brainforge_inference_utils/registry.py | 43 +- brainforge_inference_utils/steps/__init__.py | 43 +- brainforge_inference_utils/steps/base.py | 43 +- brainforge_inference_utils/transaction.py | 43 +- brainforge_inference_utils/version.py | 43 +- setup.py | 43 +- tests/helpers.py | 44 +- tests/test_client.py | 43 +- tests/test_config.py | 43 +- tests/test_mcp.py | 43 +- tests/test_pipeline.py | 43 +- tests/test_registry.py | 43 +- tests/test_steps.py | 43 +- tests/test_transaction.py | 43 +- 20 files changed, 572 insertions(+), 927 deletions(-) diff --git a/LICENSE.md b/LICENSE.md index c6f01c6..52b3834 100644 --- a/LICENSE.md +++ b/LICENSE.md @@ -1,660 +1,21 @@ -# GNU AFFERO GENERAL PUBLIC LICENSE - -Version 3, 19 November 2007 - -Copyright (C) 2007 Free Software Foundation, Inc. - - -Everyone is permitted to copy and distribute verbatim copies of this -license document, but changing it is not allowed. - -## Preamble - -The GNU Affero General Public License is a free, copyleft license for -software and other kinds of works, specifically designed to ensure -cooperation with the community in the case of network server software. - -The licenses for most software and other practical works are designed -to take away your freedom to share and change the works. By contrast, -our General Public Licenses are intended to guarantee your freedom to -share and change all versions of a program--to make sure it remains -free software for all its users. - -When we speak of free software, we are referring to freedom, not -price. Our General Public Licenses are designed to make sure that you -have the freedom to distribute copies of free software (and charge for -them if you wish), that you receive source code or can get it if you -want it, that you can change the software or use pieces of it in new -free programs, and that you know you can do these things. - -Developers that use our General Public Licenses protect your rights -with two steps: (1) assert copyright on the software, and (2) offer -you this License which gives you legal permission to copy, distribute -and/or modify the software. - -A secondary benefit of defending all users' freedom is that -improvements made in alternate versions of the program, if they -receive widespread use, become available for other developers to -incorporate. Many developers of free software are heartened and -encouraged by the resulting cooperation. However, in the case of -software used on network servers, this result may fail to come about. -The GNU General Public License permits making a modified version and -letting the public access it on a server without ever releasing its -source code to the public. - -The GNU Affero General Public License is designed specifically to -ensure that, in such cases, the modified source code becomes available -to the community. It requires the operator of a network server to -provide the source code of the modified version running there to the -users of that server. Therefore, public use of a modified version, on -a publicly accessible server, gives the public access to the source -code of the modified version. - -An older license, called the Affero General Public License and -published by Affero, was designed to accomplish similar goals. This is -a different license, not a version of the Affero GPL, but Affero has -released a new version of the Affero GPL which permits relicensing -under this license. - -The precise terms and conditions for copying, distribution and -modification follow. - -## TERMS AND CONDITIONS - -### 0. Definitions. - -"This License" refers to version 3 of the GNU Affero General Public -License. - -"Copyright" also means copyright-like laws that apply to other kinds -of works, such as semiconductor masks. - -"The Program" refers to any copyrightable work licensed under this -License. Each licensee is addressed as "you". "Licensees" and -"recipients" may be individuals or organizations. - -To "modify" a work means to copy from or adapt all or part of the work -in a fashion requiring copyright permission, other than the making of -an exact copy. The resulting work is called a "modified version" of -the earlier work or a work "based on" the earlier work. - -A "covered work" means either the unmodified Program or a work based -on the Program. - -To "propagate" a work means to do anything with it that, without -permission, would make you directly or secondarily liable for -infringement under applicable copyright law, except executing it on a -computer or modifying a private copy. Propagation includes copying, -distribution (with or without modification), making available to the -public, and in some countries other activities as well. - -To "convey" a work means any kind of propagation that enables other -parties to make or receive copies. Mere interaction with a user -through a computer network, with no transfer of a copy, is not -conveying. - -An interactive user interface displays "Appropriate Legal Notices" to -the extent that it includes a convenient and prominently visible -feature that (1) displays an appropriate copyright notice, and (2) -tells the user that there is no warranty for the work (except to the -extent that warranties are provided), that licensees may convey the -work under this License, and how to view a copy of this License. If -the interface presents a list of user commands or options, such as a -menu, a prominent item in the list meets this criterion. - -### 1. Source Code. - -The "source code" for a work means the preferred form of the work for -making modifications to it. "Object code" means any non-source form of -a work. - -A "Standard Interface" means an interface that either is an official -standard defined by a recognized standards body, or, in the case of -interfaces specified for a particular programming language, one that -is widely used among developers working in that language. - -The "System Libraries" of an executable work include anything, other -than the work as a whole, that (a) is included in the normal form of -packaging a Major Component, but which is not part of that Major -Component, and (b) serves only to enable use of the work with that -Major Component, or to implement a Standard Interface for which an -implementation is available to the public in source code form. A -"Major Component", in this context, means a major essential component -(kernel, window system, and so on) of the specific operating system -(if any) on which the executable work runs, or a compiler used to -produce the work, or an object code interpreter used to run it. - -The "Corresponding Source" for a work in object code form means all -the source code needed to generate, install, and (for an executable -work) run the object code and to modify the work, including scripts to -control those activities. However, it does not include the work's -System Libraries, or general-purpose tools or generally available free -programs which are used unmodified in performing those activities but -which are not part of the work. For example, Corresponding Source -includes interface definition files associated with source files for -the work, and the source code for shared libraries and dynamically -linked subprograms that the work is specifically designed to require, -such as by intimate data communication or control flow between those -subprograms and other parts of the work. - -The Corresponding Source need not include anything that users can -regenerate automatically from other parts of the Corresponding Source. - -The Corresponding Source for a work in source code form is that same -work. - -### 2. Basic Permissions. - -All rights granted under this License are granted for the term of -copyright on the Program, and are irrevocable provided the stated -conditions are met. This License explicitly affirms your unlimited -permission to run the unmodified Program. The output from running a -covered work is covered by this License only if the output, given its -content, constitutes a covered work. This License acknowledges your -rights of fair use or other equivalent, as provided by copyright law. - -You may make, run and propagate covered works that you do not convey, -without conditions so long as your license otherwise remains in force. -You may convey covered works to others for the sole purpose of having -them make modifications exclusively for you, or provide you with -facilities for running those works, provided that you comply with the -terms of this License in conveying all material for which you do not -control copyright. Those thus making or running the covered works for -you must do so exclusively on your behalf, under your direction and -control, on terms that prohibit them from making any copies of your -copyrighted material outside their relationship with you. - -Conveying under any other circumstances is permitted solely under the -conditions stated below. Sublicensing is not allowed; section 10 makes -it unnecessary. - -### 3. Protecting Users' Legal Rights From Anti-Circumvention Law. - -No covered work shall be deemed part of an effective technological -measure under any applicable law fulfilling obligations under article -11 of the WIPO copyright treaty adopted on 20 December 1996, or -similar laws prohibiting or restricting circumvention of such -measures. - -When you convey a covered work, you waive any legal power to forbid -circumvention of technological measures to the extent such -circumvention is effected by exercising rights under this License with -respect to the covered work, and you disclaim any intention to limit -operation or modification of the work as a means of enforcing, against -the work's users, your or third parties' legal rights to forbid -circumvention of technological measures. - -### 4. Conveying Verbatim Copies. - -You may convey verbatim copies of the Program's source code as you -receive it, in any medium, provided that you conspicuously and -appropriately publish on each copy an appropriate copyright notice; -keep intact all notices stating that this License and any -non-permissive terms added in accord with section 7 apply to the code; -keep intact all notices of the absence of any warranty; and give all -recipients a copy of this License along with the Program. - -You may charge any price or no price for each copy that you convey, -and you may offer support or warranty protection for a fee. - -### 5. Conveying Modified Source Versions. - -You may convey a work based on the Program, or the modifications to -produce it from the Program, in the form of source code under the -terms of section 4, provided that you also meet all of these -conditions: - -- a) The work must carry prominent notices stating that you modified - it, and giving a relevant date. -- b) The work must carry prominent notices stating that it is - released under this License and any conditions added under - section 7. This requirement modifies the requirement in section 4 - to "keep intact all notices". -- c) You must license the entire work, as a whole, under this - License to anyone who comes into possession of a copy. This - License will therefore apply, along with any applicable section 7 - additional terms, to the whole of the work, and all its parts, - regardless of how they are packaged. This License gives no - permission to license the work in any other way, but it does not - invalidate such permission if you have separately received it. -- d) If the work has interactive user interfaces, each must display - Appropriate Legal Notices; however, if the Program has interactive - interfaces that do not display Appropriate Legal Notices, your - work need not make them do so. - -A compilation of a covered work with other separate and independent -works, which are not by their nature extensions of the covered work, -and which are not combined with it such as to form a larger program, -in or on a volume of a storage or distribution medium, is called an -"aggregate" if the compilation and its resulting copyright are not -used to limit the access or legal rights of the compilation's users -beyond what the individual works permit. Inclusion of a covered work -in an aggregate does not cause this License to apply to the other -parts of the aggregate. - -### 6. Conveying Non-Source Forms. - -You may convey a covered work in object code form under the terms of -sections 4 and 5, provided that you also convey the machine-readable -Corresponding Source under the terms of this License, in one of these -ways: - -- a) Convey the object code in, or embodied in, a physical product - (including a physical distribution medium), accompanied by the - Corresponding Source fixed on a durable physical medium - customarily used for software interchange. -- b) Convey the object code in, or embodied in, a physical product - (including a physical distribution medium), accompanied by a - written offer, valid for at least three years and valid for as - long as you offer spare parts or customer support for that product - model, to give anyone who possesses the object code either (1) a - copy of the Corresponding Source for all the software in the - product that is covered by this License, on a durable physical - medium customarily used for software interchange, for a price no - more than your reasonable cost of physically performing this - conveying of source, or (2) access to copy the Corresponding - Source from a network server at no charge. -- c) Convey individual copies of the object code with a copy of the - written offer to provide the Corresponding Source. This - alternative is allowed only occasionally and noncommercially, and - only if you received the object code with such an offer, in accord - with subsection 6b. -- d) Convey the object code by offering access from a designated - place (gratis or for a charge), and offer equivalent access to the - Corresponding Source in the same way through the same place at no - further charge. You need not require recipients to copy the - Corresponding Source along with the object code. If the place to - copy the object code is a network server, the Corresponding Source - may be on a different server (operated by you or a third party) - that supports equivalent copying facilities, provided you maintain - clear directions next to the object code saying where to find the - Corresponding Source. Regardless of what server hosts the - Corresponding Source, you remain obligated to ensure that it is - available for as long as needed to satisfy these requirements. -- e) Convey the object code using peer-to-peer transmission, - provided you inform other peers where the object code and - Corresponding Source of the work are being offered to the general - public at no charge under subsection 6d. - -A separable portion of the object code, whose source code is excluded -from the Corresponding Source as a System Library, need not be -included in conveying the object code work. - -A "User Product" is either (1) a "consumer product", which means any -tangible personal property which is normally used for personal, -family, or household purposes, or (2) anything designed or sold for -incorporation into a dwelling. In determining whether a product is a -consumer product, doubtful cases shall be resolved in favor of -coverage. For a particular product received by a particular user, -"normally used" refers to a typical or common use of that class of -product, regardless of the status of the particular user or of the way -in which the particular user actually uses, or expects or is expected -to use, the product. A product is a consumer product regardless of -whether the product has substantial commercial, industrial or -non-consumer uses, unless such uses represent the only significant -mode of use of the product. - -"Installation Information" for a User Product means any methods, -procedures, authorization keys, or other information required to -install and execute modified versions of a covered work in that User -Product from a modified version of its Corresponding Source. The -information must suffice to ensure that the continued functioning of -the modified object code is in no case prevented or interfered with -solely because modification has been made. - -If you convey an object code work under this section in, or with, or -specifically for use in, a User Product, and the conveying occurs as -part of a transaction in which the right of possession and use of the -User Product is transferred to the recipient in perpetuity or for a -fixed term (regardless of how the transaction is characterized), the -Corresponding Source conveyed under this section must be accompanied -by the Installation Information. But this requirement does not apply -if neither you nor any third party retains the ability to install -modified object code on the User Product (for example, the work has -been installed in ROM). - -The requirement to provide Installation Information does not include a -requirement to continue to provide support service, warranty, or -updates for a work that has been modified or installed by the -recipient, or for the User Product in which it has been modified or -installed. Access to a network may be denied when the modification -itself materially and adversely affects the operation of the network -or violates the rules and protocols for communication across the -network. - -Corresponding Source conveyed, and Installation Information provided, -in accord with this section must be in a format that is publicly -documented (and with an implementation available to the public in -source code form), and must require no special password or key for -unpacking, reading or copying. - -### 7. Additional Terms. - -"Additional permissions" are terms that supplement the terms of this -License by making exceptions from one or more of its conditions. -Additional permissions that are applicable to the entire Program shall -be treated as though they were included in this License, to the extent -that they are valid under applicable law. If additional permissions -apply only to part of the Program, that part may be used separately -under those permissions, but the entire Program remains governed by -this License without regard to the additional permissions. - -When you convey a copy of a covered work, you may at your option -remove any additional permissions from that copy, or from any part of -it. (Additional permissions may be written to require their own -removal in certain cases when you modify the work.) You may place -additional permissions on material, added by you to a covered work, -for which you have or can give appropriate copyright permission. - -Notwithstanding any other provision of this License, for material you -add to a covered work, you may (if authorized by the copyright holders -of that material) supplement the terms of this License with terms: - -- a) Disclaiming warranty or limiting liability differently from the - terms of sections 15 and 16 of this License; or -- b) Requiring preservation of specified reasonable legal notices or - author attributions in that material or in the Appropriate Legal - Notices displayed by works containing it; or -- c) Prohibiting misrepresentation of the origin of that material, - or requiring that modified versions of such material be marked in - reasonable ways as different from the original version; or -- d) Limiting the use for publicity purposes of names of licensors - or authors of the material; or -- e) Declining to grant rights under trademark law for use of some - trade names, trademarks, or service marks; or -- f) Requiring indemnification of licensors and authors of that - material by anyone who conveys the material (or modified versions - of it) with contractual assumptions of liability to the recipient, - for any liability that these contractual assumptions directly - impose on those licensors and authors. - -All other non-permissive additional terms are considered "further -restrictions" within the meaning of section 10. If the Program as you -received it, or any part of it, contains a notice stating that it is -governed by this License along with a term that is a further -restriction, you may remove that term. If a license document contains -a further restriction but permits relicensing or conveying under this -License, you may add to a covered work material governed by the terms -of that license document, provided that the further restriction does -not survive such relicensing or conveying. - -If you add terms to a covered work in accord with this section, you -must place, in the relevant source files, a statement of the -additional terms that apply to those files, or a notice indicating -where to find the applicable terms. - -Additional terms, permissive or non-permissive, may be stated in the -form of a separately written license, or stated as exceptions; the -above requirements apply either way. - -### 8. Termination. - -You may not propagate or modify a covered work except as expressly -provided under this License. Any attempt otherwise to propagate or -modify it is void, and will automatically terminate your rights under -this License (including any patent licenses granted under the third -paragraph of section 11). - -However, if you cease all violation of this License, then your license -from a particular copyright holder is reinstated (a) provisionally, -unless and until the copyright holder explicitly and finally -terminates your license, and (b) permanently, if the copyright holder -fails to notify you of the violation by some reasonable means prior to -60 days after the cessation. - -Moreover, your license from a particular copyright holder is -reinstated permanently if the copyright holder notifies you of the -violation by some reasonable means, this is the first time you have -received notice of violation of this License (for any work) from that -copyright holder, and you cure the violation prior to 30 days after -your receipt of the notice. - -Termination of your rights under this section does not terminate the -licenses of parties who have received copies or rights from you under -this License. If your rights have been terminated and not permanently -reinstated, you do not qualify to receive new licenses for the same -material under section 10. - -### 9. Acceptance Not Required for Having Copies. - -You are not required to accept this License in order to receive or run -a copy of the Program. Ancillary propagation of a covered work -occurring solely as a consequence of using peer-to-peer transmission -to receive a copy likewise does not require acceptance. However, -nothing other than this License grants you permission to propagate or -modify any covered work. These actions infringe copyright if you do -not accept this License. Therefore, by modifying or propagating a -covered work, you indicate your acceptance of this License to do so. - -### 10. Automatic Licensing of Downstream Recipients. - -Each time you convey a covered work, the recipient automatically -receives a license from the original licensors, to run, modify and -propagate that work, subject to this License. You are not responsible -for enforcing compliance by third parties with this License. - -An "entity transaction" is a transaction transferring control of an -organization, or substantially all assets of one, or subdividing an -organization, or merging organizations. If propagation of a covered -work results from an entity transaction, each party to that -transaction who receives a copy of the work also receives whatever -licenses to the work the party's predecessor in interest had or could -give under the previous paragraph, plus a right to possession of the -Corresponding Source of the work from the predecessor in interest, if -the predecessor has it or can get it with reasonable efforts. - -You may not impose any further restrictions on the exercise of the -rights granted or affirmed under this License. For example, you may -not impose a license fee, royalty, or other charge for exercise of -rights granted under this License, and you may not initiate litigation -(including a cross-claim or counterclaim in a lawsuit) alleging that -any patent claim is infringed by making, using, selling, offering for -sale, or importing the Program or any portion of it. - -### 11. Patents. - -A "contributor" is a copyright holder who authorizes use under this -License of the Program or a work on which the Program is based. The -work thus licensed is called the contributor's "contributor version". - -A contributor's "essential patent claims" are all patent claims owned -or controlled by the contributor, whether already acquired or -hereafter acquired, that would be infringed by some manner, permitted -by this License, of making, using, or selling its contributor version, -but do not include claims that would be infringed only as a -consequence of further modification of the contributor version. For -purposes of this definition, "control" includes the right to grant -patent sublicenses in a manner consistent with the requirements of -this License. - -Each contributor grants you a non-exclusive, worldwide, royalty-free -patent license under the contributor's essential patent claims, to -make, use, sell, offer for sale, import and otherwise run, modify and -propagate the contents of its contributor version. - -In the following three paragraphs, a "patent license" is any express -agreement or commitment, however denominated, not to enforce a patent -(such as an express permission to practice a patent or covenant not to -sue for patent infringement). To "grant" such a patent license to a -party means to make such an agreement or commitment not to enforce a -patent against the party. - -If you convey a covered work, knowingly relying on a patent license, -and the Corresponding Source of the work is not available for anyone -to copy, free of charge and under the terms of this License, through a -publicly available network server or other readily accessible means, -then you must either (1) cause the Corresponding Source to be so -available, or (2) arrange to deprive yourself of the benefit of the -patent license for this particular work, or (3) arrange, in a manner -consistent with the requirements of this License, to extend the patent -license to downstream recipients. "Knowingly relying" means you have -actual knowledge that, but for the patent license, your conveying the -covered work in a country, or your recipient's use of the covered work -in a country, would infringe one or more identifiable patents in that -country that you have reason to believe are valid. - -If, pursuant to or in connection with a single transaction or -arrangement, you convey, or propagate by procuring conveyance of, a -covered work, and grant a patent license to some of the parties -receiving the covered work authorizing them to use, propagate, modify -or convey a specific copy of the covered work, then the patent license -you grant is automatically extended to all recipients of the covered -work and works based on it. - -A patent license is "discriminatory" if it does not include within the -scope of its coverage, prohibits the exercise of, or is conditioned on -the non-exercise of one or more of the rights that are specifically -granted under this License. You may not convey a covered work if you -are a party to an arrangement with a third party that is in the -business of distributing software, under which you make payment to the -third party based on the extent of your activity of conveying the -work, and under which the third party grants, to any of the parties -who would receive the covered work from you, a discriminatory patent -license (a) in connection with copies of the covered work conveyed by -you (or copies made from those copies), or (b) primarily for and in -connection with specific products or compilations that contain the -covered work, unless you entered into that arrangement, or that patent -license was granted, prior to 28 March 2007. - -Nothing in this License shall be construed as excluding or limiting -any implied license or other defenses to infringement that may -otherwise be available to you under applicable patent law. - -### 12. No Surrender of Others' Freedom. - -If conditions are imposed on you (whether by court order, agreement or -otherwise) that contradict the conditions of this License, they do not -excuse you from the conditions of this License. If you cannot convey a -covered work so as to satisfy simultaneously your obligations under -this License and any other pertinent obligations, then as a -consequence you may not convey it at all. For example, if you agree to -terms that obligate you to collect a royalty for further conveying -from those to whom you convey the Program, the only way you could -satisfy both those terms and this License would be to refrain entirely -from conveying the Program. - -### 13. Remote Network Interaction; Use with the GNU General Public License. - -Notwithstanding any other provision of this License, if you modify the -Program, your modified version must prominently offer all users -interacting with it remotely through a computer network (if your -version supports such interaction) an opportunity to receive the -Corresponding Source of your version by providing access to the -Corresponding Source from a network server at no charge, through some -standard or customary means of facilitating copying of software. This -Corresponding Source shall include the Corresponding Source for any -work covered by version 3 of the GNU General Public License that is -incorporated pursuant to the following paragraph. - -Notwithstanding any other provision of this License, you have -permission to link or combine any covered work with a work licensed -under version 3 of the GNU General Public License into a single -combined work, and to convey the resulting work. The terms of this -License will continue to apply to the part which is the covered work, -but the work with which it is combined will remain governed by version -3 of the GNU General Public License. - -### 14. Revised Versions of this License. - -The Free Software Foundation may publish revised and/or new versions -of the GNU Affero General Public License from time to time. Such new -versions will be similar in spirit to the present version, but may -differ in detail to address new problems or concerns. - -Each version is given a distinguishing version number. If the Program -specifies that a certain numbered version of the GNU Affero General -Public License "or any later version" applies to it, you have the -option of following the terms and conditions either of that numbered -version or of any later version published by the Free Software -Foundation. If the Program does not specify a version number of the -GNU Affero General Public License, you may choose any version ever -published by the Free Software Foundation. - -If the Program specifies that a proxy can decide which future versions -of the GNU Affero General Public License can be used, that proxy's -public statement of acceptance of a version permanently authorizes you -to choose that version for the Program. - -Later license versions may give you additional or different -permissions. However, no additional obligations are imposed on any -author or copyright holder as a result of your choosing to follow a -later version. - -### 15. Disclaimer of Warranty. - -THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY -APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT -HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT -WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT -LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR -A PARTICULAR PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND -PERFORMANCE OF THE PROGRAM IS WITH YOU. SHOULD THE PROGRAM PROVE -DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING, REPAIR OR -CORRECTION. - -### 16. Limitation of Liability. - -IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING -WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR -CONVEYS THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, -INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES -ARISING OUT OF THE USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT -NOT LIMITED TO LOSS OF DATA OR DATA BEING RENDERED INACCURATE OR -LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A FAILURE OF THE PROGRAM -TO OPERATE WITH ANY OTHER PROGRAMS), EVEN IF SUCH HOLDER OR OTHER -PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. - -### 17. Interpretation of Sections 15 and 16. - -If the disclaimer of warranty and limitation of liability provided -above cannot be given local legal effect according to their terms, -reviewing courts shall apply local law that most closely approximates -an absolute waiver of all civil liability in connection with the -Program, unless a warranty or assumption of liability accompanies a -copy of the Program in return for a fee. - -END OF TERMS AND CONDITIONS - -## How to Apply These Terms to Your New Programs - -If you develop a new program, and you want it to be of the greatest -possible use to the public, the best way to achieve this is to make it -free software which everyone can redistribute and change under these -terms. - -To do so, attach the following notices to the program. It is safest to -attach them to the start of each source file to most effectively state -the exclusion of warranty; and each file should have at least the -"copyright" line and a pointer to where the full notice is found. - - - Copyright (C) - - This program is free software: you can redistribute it and/or modify - it under the terms of the GNU Affero General Public License as - published by the Free Software Foundation, either version 3 of the - License, or (at your option) any later version. - - This program is distributed in the hope that it will be useful, - but WITHOUT ANY WARRANTY; without even the implied warranty of - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - GNU Affero General Public License for more details. - - You should have received a copy of the GNU Affero General Public License - along with this program. If not, see . - -Also add information on how to contact you by electronic and paper -mail. - -If your software can interact with users remotely through a computer -network, you should also make sure that it provides a way for users to -get its source. For example, if your program is a web application, its -interface could display a "Source" link that leads users to an archive -of the code. There are many ways you could offer source, and different -solutions will be better for different programs; see section 13 for -the specific requirements. - -You should also get your employer (if you work as a programmer) or -school, if any, to sign a "copyright disclaimer" for the program, if -necessary. For more information on this, and how to apply and follow -the GNU AGPL, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Development System +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2026 Neongecko.com Inc. +# BSD-3 License + +Redistribution and use in source and binary forms, with or without modification, are permitted provided that the +following conditions are met: +1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following +disclaimer. +2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following +disclaimer in the documentation and/or other materials provided with the distribution. +3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products +derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF +THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/brainforge_inference_utils/__init__.py b/brainforge_inference_utils/__init__.py index 9892978..7418a0b 100644 --- a/brainforge_inference_utils/__init__.py +++ b/brainforge_inference_utils/__init__.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """BrainForge inference pipeline utilities.""" from neon_data_models.models.api.llm import ( diff --git a/brainforge_inference_utils/client.py b/brainforge_inference_utils/client.py index 7faa00b..af77b8d 100644 --- a/brainforge_inference_utils/client.py +++ b/brainforge_inference_utils/client.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """Configured runner for the request and response pipelines. `BrainForgeClient` does not call an LLM. The host runs inference between diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py index 85cf1f8..1304679 100644 --- a/brainforge_inference_utils/config.py +++ b/brainforge_inference_utils/config.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """ Pipeline configuration schema. diff --git a/brainforge_inference_utils/mcp.py b/brainforge_inference_utils/mcp.py index 4109cec..d0b6272 100644 --- a/brainforge_inference_utils/mcp.py +++ b/brainforge_inference_utils/mcp.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """Generic MCP client helpers shared by pipeline steps.""" from __future__ import annotations diff --git a/brainforge_inference_utils/pipeline.py b/brainforge_inference_utils/pipeline.py index d3a3e6a..7c2282c 100644 --- a/brainforge_inference_utils/pipeline.py +++ b/brainforge_inference_utils/pipeline.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """ The request and response pipelines. diff --git a/brainforge_inference_utils/registry.py b/brainforge_inference_utils/registry.py index ae4af31..41ab51b 100644 --- a/brainforge_inference_utils/registry.py +++ b/brainforge_inference_utils/registry.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """ Discovery of plugin steps via setuptools entry points. diff --git a/brainforge_inference_utils/steps/__init__.py b/brainforge_inference_utils/steps/__init__.py index 3ede871..19d1fc3 100644 --- a/brainforge_inference_utils/steps/__init__.py +++ b/brainforge_inference_utils/steps/__init__.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """ Step base classes. diff --git a/brainforge_inference_utils/steps/base.py b/brainforge_inference_utils/steps/base.py index 9924e83..23f2661 100644 --- a/brainforge_inference_utils/steps/base.py +++ b/brainforge_inference_utils/steps/base.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """ Step base classes. diff --git a/brainforge_inference_utils/transaction.py b/brainforge_inference_utils/transaction.py index bb5bc39..fdab320 100644 --- a/brainforge_inference_utils/transaction.py +++ b/brainforge_inference_utils/transaction.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. """ The object a pipeline operates on. diff --git a/brainforge_inference_utils/version.py b/brainforge_inference_utils/version.py index b355adb..fcc63ed 100644 --- a/brainforge_inference_utils/version.py +++ b/brainforge_inference_utils/version.py @@ -1,16 +1,31 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. __version__ = "0.0.1a0" diff --git a/setup.py b/setup.py index 20d7fed..5ca10d0 100644 --- a/setup.py +++ b/setup.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. from os import getenv, path diff --git a/tests/helpers.py b/tests/helpers.py index 3f195a3..9721c5e 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -1,18 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. from typing import Optional from neon_data_models.models.api.llm import ( diff --git a/tests/test_client.py b/tests/test_client.py index fdfa22c..0c2771d 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import unittest from typing import Optional diff --git a/tests/test_config.py b/tests/test_config.py index 92ad0af..38b066f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import pytest from pydantic import ValidationError diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 26bcc74..1db1b59 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import json import unittest diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 91b2a33..0f2973e 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import unittest diff --git a/tests/test_registry.py b/tests/test_registry.py index ec7c308..596fb7a 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. from types import SimpleNamespace diff --git a/tests/test_steps.py b/tests/test_steps.py index 6d23128..70887f8 100644 --- a/tests/test_steps.py +++ b/tests/test_steps.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. from abc import ABC diff --git a/tests/test_transaction.py b/tests/test_transaction.py index 34f5932..6011d4f 100644 --- a/tests/test_transaction.py +++ b/tests/test_transaction.py @@ -1,17 +1,32 @@ -# Copyright (C) 2008-2026 Neongecko.com Inc. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as -# published by the Free Software Foundation, either version 3 of the -# License, or (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +# NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework +# All trademark and other rights reserved by their respective owners +# Copyright 2008-2025 Neongecko.com Inc. +# Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, +# Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, +# Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, +# Regina Bloomstine, Casimiro Ferreira, Andrii Pernatii, Kirill Hrymailo +# BSD-3 License +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# 1. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. from pydantic import ValidationError From ae8ca886502a3457d7f823f8e2fd479f19d53fd3 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 15:30:45 -0700 Subject: [PATCH 15/23] Fix default plugin names to specify `brainforge` instead of `neon` --- README.md | 8 ++++---- brainforge_inference_utils/config.py | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 33aeb4b..cee2823 100644 --- a/README.md +++ b/README.md @@ -92,7 +92,7 @@ response steps use different groups: Both groups use the step id `author.step_name`: -- `author` is the plugin developer's slug (`acme`, `neon`). +- `author` is the plugin developer's slug (`acme`, `brainforge`). - `step_name` names that one step (`guardrails`, `rag`). Each part is lowercase letters, digits, and underscores, and starts with a @@ -127,9 +127,9 @@ First-party steps ship in their own packages: | Package | Steps | | --- | --- | -| `neon-inference-plugin-locale` | `neon.locale` | -| `neon-inference-plugin-rag` | `neon.rag`, `neon.citations` | -| `neon-inference-plugin-history` | `neon.history` | +| `brainforge-inference-plugin-locale` | `brainforge.locale` | +| `brainforge-inference-plugin-rag` | `brainforge.rag`, `brainforge.citations` | +| `brainforge-inference-plugin-history` | `brainforge.history` | This package does not register those steps. Install the plugin to use the id. diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py index 1304679..9830743 100644 --- a/brainforge_inference_utils/config.py +++ b/brainforge_inference_utils/config.py @@ -40,15 +40,15 @@ pipeline: request: - - id: neon.locale + - id: brainforge.locale config: mcp_url: https://mcp.neonaiservices2.com/mcp default_location: Renton, WA, USA timeout: 5.0 - - {id: neon.rag, config: {mcp_url: https://mcp.neonaiservices2.com/mcp, top_k: 10}} - - neon.history # last: valid, truncated history + - {id: brainforge.rag, config: {mcp_url: https://mcp.neonaiservices2.com/mcp, top_k: 10}} + - brainforge.history # last: valid, truncated history response: - - id: neon.citations + - id: brainforge.citations config: mcp_url: https://mcp.neonaiservices2.com/mcp openai_url: https://hana.example/brainforge/openai/chat/completions From 274b0afc3a32c33ace0bee9a92a925e7231265e4 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 15:48:56 -0700 Subject: [PATCH 16/23] Update GitHub actions for PyPI publication and maintained Python versions --- .github/workflows/publish_release.yml | 19 ++++++------------- .github/workflows/publish_test_build.yml | 4 +--- .github/workflows/unit_tests.yml | 2 +- 3 files changed, 8 insertions(+), 17 deletions(-) diff --git a/.github/workflows/publish_release.yml b/.github/workflows/publish_release.yml index 4901b63..e4b3fa3 100644 --- a/.github/workflows/publish_release.yml +++ b/.github/workflows/publish_release.yml @@ -1,19 +1,12 @@ -name: Publish GitHub Release +# This workflow will generate a release distribution and upload it to PyPI + +name: Publish Build and GitHub Release on: push: branches: - master jobs: - tag_release: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - name: Get Version - run: | - VERSION=$(python setup.py --version) - echo "VERSION=${VERSION}" >> $GITHUB_ENV - - uses: ncipollo/release-action@v1 - with: - token: ${{secrets.GITHUB_TOKEN}} - tag: ${{env.VERSION}} + build_and_publish_pypi_and_release: + uses: neongeckocom/.github/.github/workflows/publish_stable_release.yml@master + secrets: inherit diff --git a/.github/workflows/publish_test_build.yml b/.github/workflows/publish_test_build.yml index 3df6b1b..a5bdbe9 100644 --- a/.github/workflows/publish_test_build.yml +++ b/.github/workflows/publish_test_build.yml @@ -1,5 +1,3 @@ -# This workflow will generate a distribution and upload it to PyPI - name: Publish Alpha Build on: push: @@ -13,6 +11,6 @@ jobs: uses: neongeckocom/.github/.github/workflows/publish_alpha_release.yml@master secrets: inherit with: - publish_pypi: false + publish_pypi: true version_file: "brainforge_inference_utils/version.py" publish_prerelease: true diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index e65b252..ef8f860 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -11,7 +11,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12"] + python-version: ["3.11", "3.12", "3.13"] runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 From fda6a21aec68a9679eda9558e2cd367ede5d7364 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 15:54:04 -0700 Subject: [PATCH 17/23] Refactor so responses are always mutated in-place for consistency --- brainforge_inference_utils/pipeline.py | 7 ++----- brainforge_inference_utils/steps/base.py | 22 +++++++++------------- tests/helpers.py | 12 +++++------- tests/test_client.py | 12 ++---------- tests/test_pipeline.py | 4 ++-- 5 files changed, 20 insertions(+), 37 deletions(-) diff --git a/brainforge_inference_utils/pipeline.py b/brainforge_inference_utils/pipeline.py index 7c2282c..d9488f6 100644 --- a/brainforge_inference_utils/pipeline.py +++ b/brainforge_inference_utils/pipeline.py @@ -101,8 +101,7 @@ def __init__(self, steps: Sequence[ResponseStep]): async def run(self, transaction: UserTransaction) -> UserTransaction: """ - Run every applicable response step, threading the (possibly replaced) - response back onto the transaction. + Run every applicable response step, mutating the transaction in place. :param transaction: The transaction whose response to shape. :returns: The same `transaction`. @@ -112,9 +111,7 @@ async def run(self, transaction: UserTransaction) -> UserTransaction: if reason is not None: LOG.debug("Skipping response step %r: %s", step.id, reason) continue - revised = await step.process(transaction) - if revised is not None: - transaction.response = revised + await step.process(transaction) return transaction diff --git a/brainforge_inference_utils/steps/base.py b/brainforge_inference_utils/steps/base.py index 23f2661..da0f059 100644 --- a/brainforge_inference_utils/steps/base.py +++ b/brainforge_inference_utils/steps/base.py @@ -35,9 +35,9 @@ - `RequestStep` runs before generation. It shapes the request on the `UserTransaction` (RAG/MCP context, ambient context, request history validation, ...), mutating it in place. -- `ResponseStep` runs after generation. It has the complete response on the - `UserTransaction` and may return a revised response (or mutate it in place) to - annotate or rewrite it (citations, sanitizing, guardrails, ...). +- `ResponseStep` runs after generation. It shapes the response on the + `UserTransaction` (citations, sanitizing, guardrails, ...), mutating it + in place. Both share two conventions: @@ -59,8 +59,6 @@ from dataclasses import dataclass from typing import Any, ClassVar, List, Optional, Sequence -from neon_data_models.models.api.llm import LLMResponse - from brainforge_inference_utils.transaction import UserTransaction @@ -167,9 +165,8 @@ class ResponseStep(Step, ABC): """ A step that shapes the generated response. - Implementations read the response from `transaction.response` and either - mutate it in place or return a new `LLMResponse` to replace it. Returning - `None` leaves the current response unchanged. + Implementations mutate `transaction.response` in place (annotating or + rewriting it). They return nothing. """ def applicable(self, transaction: UserTransaction) -> Optional[str]: @@ -184,15 +181,14 @@ def applicable(self, transaction: UserTransaction) -> Optional[str]: return None @abstractmethod - async def process(self, - transaction: UserTransaction) -> Optional[LLMResponse]: + async def process(self, transaction: UserTransaction) -> None: """ - Annotate or rewrite the response. + Apply this step's effect to the response. Only called when `applicable(transaction)` returned `None`. + Implementations mutate `transaction.response` in place and return + nothing. :param transaction: The transaction whose `response` to shape. - :returns: A revised `LLMResponse` to replace the current one, or `None` - to keep it (including when the step mutated the response in place). """ raise NotImplementedError diff --git a/tests/helpers.py b/tests/helpers.py index 9721c5e..ab28a29 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -101,7 +101,7 @@ async def process(self, transaction: UserTransaction) -> None: class PrefixResponseStep(ResponseStep): - """Response step that returns a copy with `config["prefix"]` prepended.""" + """Response step that prepends `config["prefix"]` to the response.""" id = "prefix_response" @@ -110,20 +110,18 @@ def applicable(self, transaction: UserTransaction) -> Optional[str]: return "no response" return None - async def process(self, transaction: UserTransaction) -> Optional[LLMResponse]: - current = transaction.response + async def process(self, transaction: UserTransaction) -> None: prefix = self.config.get("prefix", "") - return current.model_copy(update={"response": prefix + current.response}) + transaction.response.response = prefix + transaction.response.response class MutateResponseInPlaceStep(ResponseStep): - """Response step that mutates `response.response` and returns None.""" + """Response step that appends `config["suffix"]` to the response.""" id = "mutate_in_place" - async def process(self, transaction: UserTransaction) -> Optional[LLMResponse]: + async def process(self, transaction: UserTransaction) -> None: transaction.response.response += self.config.get("suffix", "!") - return None TEST_STEPS = { diff --git a/tests/test_client.py b/tests/test_client.py index 0c2771d..ec5e541 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -29,9 +29,6 @@ # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import unittest -from typing import Optional - -from neon_data_models.models.api.llm import LLMResponse from brainforge_inference_utils.client import BrainForgeClient from brainforge_inference_utils.config import PipelineConfig @@ -78,20 +75,15 @@ class FlagStep(ResponseStep): RequestContextOption("location"), ) - async def process( - self, transaction: UserTransaction - ) -> Optional[LLMResponse]: + async def process(self, transaction: UserTransaction) -> None: return None class RecordResponseStep(ResponseStep): id = "record_response" - async def process( - self, transaction: UserTransaction - ) -> Optional[LLMResponse]: + async def process(self, transaction: UserTransaction) -> None: transaction.response_step_context[self.id] = self.config["value"] - return None class TestBrainForgeClient(unittest.IsolatedAsyncioTestCase): diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 0f2973e..b74cfc2 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -158,7 +158,7 @@ async def test_inapplicable_request_step_is_skipped(self): assert transaction.request.max_tokens == 512 assert transaction.request.max_history == 7 - async def test_response_steps_thread_returned_revisions(self): + async def test_response_steps_run_in_declared_order(self): pipeline = ResponsePipeline([ PrefixResponseStep({"prefix": "["}), PrefixResponseStep({"prefix": "<"}), @@ -168,7 +168,7 @@ async def test_response_steps_thread_returned_revisions(self): await pipeline.run(transaction) assert transaction.response.response == "<[ok" - async def test_response_step_none_keeps_in_place_mutation(self): + async def test_response_steps_mutate_the_same_response(self): pipeline = ResponsePipeline([ MutateResponseInPlaceStep({"suffix": "!"}), PrefixResponseStep({"prefix": ">"}), From 33d0b9804393ea9a1adb3de0bc54c9e29ab2e0aa Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 15:55:59 -0700 Subject: [PATCH 18/23] Update dates in license headers --- brainforge_inference_utils/__init__.py | 2 +- brainforge_inference_utils/client.py | 2 +- brainforge_inference_utils/config.py | 2 +- brainforge_inference_utils/mcp.py | 2 +- brainforge_inference_utils/pipeline.py | 2 +- brainforge_inference_utils/registry.py | 2 +- brainforge_inference_utils/steps/__init__.py | 2 +- brainforge_inference_utils/steps/base.py | 2 +- brainforge_inference_utils/transaction.py | 2 +- brainforge_inference_utils/version.py | 2 +- setup.py | 2 +- tests/helpers.py | 2 +- tests/test_client.py | 2 +- tests/test_config.py | 2 +- tests/test_mcp.py | 2 +- tests/test_pipeline.py | 2 +- tests/test_registry.py | 2 +- tests/test_steps.py | 2 +- tests/test_transaction.py | 2 +- 19 files changed, 19 insertions(+), 19 deletions(-) diff --git a/brainforge_inference_utils/__init__.py b/brainforge_inference_utils/__init__.py index 7418a0b..7e10002 100644 --- a/brainforge_inference_utils/__init__.py +++ b/brainforge_inference_utils/__init__.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/brainforge_inference_utils/client.py b/brainforge_inference_utils/client.py index af77b8d..d047e5a 100644 --- a/brainforge_inference_utils/client.py +++ b/brainforge_inference_utils/client.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py index 9830743..c8c6994 100644 --- a/brainforge_inference_utils/config.py +++ b/brainforge_inference_utils/config.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/brainforge_inference_utils/mcp.py b/brainforge_inference_utils/mcp.py index d0b6272..3da2f41 100644 --- a/brainforge_inference_utils/mcp.py +++ b/brainforge_inference_utils/mcp.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/brainforge_inference_utils/pipeline.py b/brainforge_inference_utils/pipeline.py index d9488f6..2ba6a6d 100644 --- a/brainforge_inference_utils/pipeline.py +++ b/brainforge_inference_utils/pipeline.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/brainforge_inference_utils/registry.py b/brainforge_inference_utils/registry.py index 41ab51b..fdff5d2 100644 --- a/brainforge_inference_utils/registry.py +++ b/brainforge_inference_utils/registry.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/brainforge_inference_utils/steps/__init__.py b/brainforge_inference_utils/steps/__init__.py index 19d1fc3..15c03f6 100644 --- a/brainforge_inference_utils/steps/__init__.py +++ b/brainforge_inference_utils/steps/__init__.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/brainforge_inference_utils/steps/base.py b/brainforge_inference_utils/steps/base.py index da0f059..f642610 100644 --- a/brainforge_inference_utils/steps/base.py +++ b/brainforge_inference_utils/steps/base.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/brainforge_inference_utils/transaction.py b/brainforge_inference_utils/transaction.py index fdab320..959a742 100644 --- a/brainforge_inference_utils/transaction.py +++ b/brainforge_inference_utils/transaction.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/brainforge_inference_utils/version.py b/brainforge_inference_utils/version.py index fcc63ed..0707f34 100644 --- a/brainforge_inference_utils/version.py +++ b/brainforge_inference_utils/version.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/setup.py b/setup.py index 5ca10d0..4402b66 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/tests/helpers.py b/tests/helpers.py index ab28a29..a6a7aeb 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/tests/test_client.py b/tests/test_client.py index ec5e541..1382f6f 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/tests/test_config.py b/tests/test_config.py index 38b066f..f0b4b62 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 1db1b59..c585bc1 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index b74cfc2..e235d0e 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/tests/test_registry.py b/tests/test_registry.py index 596fb7a..dd972d8 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/tests/test_steps.py b/tests/test_steps.py index 70887f8..3920ff0 100644 --- a/tests/test_steps.py +++ b/tests/test_steps.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, diff --git a/tests/test_transaction.py b/tests/test_transaction.py index 6011d4f..b9b004a 100644 --- a/tests/test_transaction.py +++ b/tests/test_transaction.py @@ -1,6 +1,6 @@ # NEON AI (TM) SOFTWARE, Software Development Kit & Application Framework # All trademark and other rights reserved by their respective owners -# Copyright 2008-2025 Neongecko.com Inc. +# Copyright 2008-2026 Neongecko.com Inc. # Contributors: Daniel McKnight, Richard Leeds, Bohdan Mykhailenko, # Dmitry Honcharenko, Mike Gray, Clary Gasper, Charlie Bailey, Ryan Venturi, # Roger Wiebe, Pierce Seigne, Sathvika Nadipalli, Guy Daniels, Elon Gasper, From 40c7cd6ce9327e27bd19fa576240aa7d7e1ce4e4 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 16:07:46 -0700 Subject: [PATCH 19/23] Refactor to deprecate references to bundled steps --- README.md | 2 +- brainforge_inference_utils/__init__.py | 2 - brainforge_inference_utils/client.py | 13 +- brainforge_inference_utils/config.py | 8 +- brainforge_inference_utils/pipeline.py | 43 +++---- brainforge_inference_utils/registry.py | 22 +++- brainforge_inference_utils/steps/__init__.py | 8 +- setup.py | 2 +- tests/helpers.py | 69 ++++++++--- tests/test_client.py | 66 ++++++----- tests/test_pipeline.py | 118 +++++++++++-------- tests/test_registry.py | 100 ++++++---------- tests/test_steps.py | 2 +- 13 files changed, 236 insertions(+), 219 deletions(-) diff --git a/README.md b/README.md index cee2823..4971572 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ response, and two per-side context dicts (`request_step_context` and | `RequestContextOption` | One `request_context` key a step reads, with optional allowed values for hosts to present. | | `RequestPipeline` / `ResponsePipeline` | Run an ordered list of steps over a transaction. | | `PipelineConfig` / `StepConfig` | Declarative, ordered step lists (see below). | -| `build_pipelines(config)` | Build both pipelines from config. Bundled ids come from `steps_by_id`; plugin ids come from entry points. | +| `build_pipelines(config)` | Build both pipelines from config. Step classes come from the entry point group for that pipeline. | ## Usage sketch diff --git a/brainforge_inference_utils/__init__.py b/brainforge_inference_utils/__init__.py index 7e10002..80bfd1d 100644 --- a/brainforge_inference_utils/__init__.py +++ b/brainforge_inference_utils/__init__.py @@ -48,7 +48,6 @@ REQUEST_STEP_ENTRY_POINT, RESPONSE_STEP_ENTRY_POINT, ) -from brainforge_inference_utils.steps import steps_by_id from brainforge_inference_utils.steps.base import ( RequestStep, RequestContextOption, @@ -71,7 +70,6 @@ "RequestStep", "ResponseStep", "RequestContextOption", - "steps_by_id", # Pipelines "RequestPipeline", "ResponsePipeline", diff --git a/brainforge_inference_utils/client.py b/brainforge_inference_utils/client.py index d047e5a..8f3b737 100644 --- a/brainforge_inference_utils/client.py +++ b/brainforge_inference_utils/client.py @@ -36,27 +36,22 @@ from __future__ import annotations -from typing import Dict, Mapping, Optional, Type +from typing import Dict from neon_data_models.models.api.llm import LLMRequest, LLMResponse from brainforge_inference_utils.config import PipelineConfig from brainforge_inference_utils.pipeline import build_pipelines -from brainforge_inference_utils.steps.base import RequestContextOption, Step +from brainforge_inference_utils.steps.base import RequestContextOption from brainforge_inference_utils.transaction import UserTransaction class BrainForgeClient: """Build pipelines from config and run them around host-owned inference.""" - def __init__( - self, - config: PipelineConfig, - steps: Optional[Mapping[str, Type[Step]]] = None, - ): + def __init__(self, config: PipelineConfig): self.config = config - self.request_pipeline, self.response_pipeline = build_pipelines( - config, steps) + self.request_pipeline, self.response_pipeline = build_pipelines(config) async def supported_request_context( self, timeout: float = 10.0 diff --git a/brainforge_inference_utils/config.py b/brainforge_inference_utils/config.py index c8c6994..1aaf391 100644 --- a/brainforge_inference_utils/config.py +++ b/brainforge_inference_utils/config.py @@ -32,8 +32,8 @@ Configuration is two explicitly ordered lists of step entries: `request` and `response`. Order is authoritative; the pipeline runs steps in exactly the -order given and does not reorder. Each entry names a step `id` (a bundled id -from `steps_by_id`, or a plugin id `author.step_name`) and optionally carries +order given and does not reorder. Each entry names a step `id` +(`author.step_name`, the entry point name) and optionally carries step-specific `config`. Example (YAML):: @@ -69,8 +69,8 @@ class StepConfig(BaseModel): """ A single step entry within a pipeline. - :param id: Step id to resolve and instantiate. Bundled steps use their - current ids. Plugin steps use `author.step_name`. + :param id: Step id to resolve and instantiate. This is the entry point + name, `author.step_name`. :param enabled: When false, the entry is skipped at build time (a convenient way to turn a step off without deleting its config). :param config: Step-specific configuration passed to the step constructor. diff --git a/brainforge_inference_utils/pipeline.py b/brainforge_inference_utils/pipeline.py index 2ba6a6d..3927f6f 100644 --- a/brainforge_inference_utils/pipeline.py +++ b/brainforge_inference_utils/pipeline.py @@ -44,15 +44,14 @@ from __future__ import annotations import logging -from typing import List, Mapping, Optional, Sequence, Tuple, Type +from typing import List, Mapping, Sequence, Tuple, Type from brainforge_inference_utils.config import PipelineConfig from brainforge_inference_utils.registry import ( REQUEST_STEP_ENTRY_POINT, RESPONSE_STEP_ENTRY_POINT, - merge_builtin_and_entry_points, + load_entry_point_steps, ) -from brainforge_inference_utils.steps import steps_by_id from brainforge_inference_utils.steps.base import ( RequestStep, ResponseStep, @@ -127,33 +126,27 @@ def _resolve_step(stage_id: str, def build_pipelines( config: PipelineConfig, - steps: Optional[Mapping[str, Type[Step]]] = None, ) -> Tuple[RequestPipeline, ResponsePipeline]: """ Instantiate both pipelines from configuration. - Step classes are looked up by id in `steps` when that mapping is passed. - Otherwise each side uses `steps_by_id` plus the entry point group for that - pipeline (`brainforge.inference.request_steps` or + Step classes come from the entry point group for that pipeline + (`brainforge.inference.request_steps` or `brainforge.inference.response_steps`). A response step placed in the - request list (or vice versa) fails fast at build time. + request list (or vice versa) is absent from that group and fails at + build time. :param config: Parsed pipeline configuration. - :param steps: Optional id -> `Step` class mapping. When omitted, bundled - steps are combined with installed entry points. :returns: `(request_pipeline, response_pipeline)`. - :raises KeyError: If a configured step id is not present in the mapping - used for that side. - :raises TypeError: If a configured step resolves to the wrong step type. + :raises KeyError: If a configured step id is not registered in that + pipeline's entry point group. + :raises TypeError: If an installed entry point is not the step type for + its group. """ - if steps is None: - request_step_classes = merge_builtin_and_entry_points( - steps_by_id, REQUEST_STEP_ENTRY_POINT, RequestStep) - response_step_classes = merge_builtin_and_entry_points( - steps_by_id, RESPONSE_STEP_ENTRY_POINT, ResponseStep) - else: - request_step_classes = steps - response_step_classes = steps + request_step_classes = load_entry_point_steps( + REQUEST_STEP_ENTRY_POINT, RequestStep) + response_step_classes = load_entry_point_steps( + RESPONSE_STEP_ENTRY_POINT, ResponseStep) request_steps: List[RequestStep] = [] for entry in config.request: @@ -161,10 +154,6 @@ def build_pipelines( LOG.info("Request step %r disabled by config", entry.id) continue cls = _resolve_step(entry.id, request_step_classes) - if not (isinstance(cls, type) and issubclass(cls, RequestStep)): - raise TypeError( - f"Step {entry.id!r} is configured as a request step but " - f"{getattr(cls, '__name__', cls)!r} is not a RequestStep") request_steps.append(cls(entry.config)) response_steps: List[ResponseStep] = [] @@ -173,10 +162,6 @@ def build_pipelines( LOG.info("Response step %r disabled by config", entry.id) continue cls = _resolve_step(entry.id, response_step_classes) - if not (isinstance(cls, type) and issubclass(cls, ResponseStep)): - raise TypeError( - f"Step {entry.id!r} is configured as a response step but " - f"{getattr(cls, '__name__', cls)!r} is not a ResponseStep") response_steps.append(cls(entry.config)) LOG.info("Built request pipeline: %s", diff --git a/brainforge_inference_utils/registry.py b/brainforge_inference_utils/registry.py index fdff5d2..0545bce 100644 --- a/brainforge_inference_utils/registry.py +++ b/brainforge_inference_utils/registry.py @@ -31,14 +31,15 @@ Discovery of plugin steps via setuptools entry points. Request and response steps are registered in separate groups, so the same -`author.step_name` id can exist on both sides. Bundled steps shipped in this -package stay in `steps_by_id` and are merged with whichever group is loaded. +`author.step_name` id can exist on both sides. Every step class comes from +the entry point group for its pipeline. """ from __future__ import annotations import logging import re +import warnings from importlib.metadata import entry_points from typing import Dict, Mapping, Type @@ -120,15 +121,24 @@ def merge_builtin_and_entry_points( expected_type: Type[Step], ) -> Dict[str, Type[Step]]: """ - Combine bundled steps with classes discovered from `group`. + Combine a caller-supplied step map with classes discovered from `group`. - :param builtin: In-package id -> step class mapping (`steps_by_id`). + Deprecated. This package no longer ships built-in steps. Call + `load_entry_point_steps` instead. The `builtin` map is still copied and + then extended with the entry point group so existing callers keep working. + + :param builtin: Caller-supplied id -> step class mapping. :param group: Entry point group for this pipeline side. :param expected_type: Step base class that group must provide. - :returns: A new mapping. `builtin` is copied unchanged, then plugin steps - are added. + :returns: A new mapping. `builtin` is copied, then plugin steps are added. :raises ValueError: If a plugin id is already present in `builtin`. """ + warnings.warn( + "merge_builtin_and_entry_points is deprecated; this package no longer " + "ships built-in steps. Use load_entry_point_steps instead.", + DeprecationWarning, + stacklevel=2, + ) resolved = dict(builtin) for step_id, step_cls in load_entry_point_steps(group, expected_type).items(): if step_id in resolved: diff --git a/brainforge_inference_utils/steps/__init__.py b/brainforge_inference_utils/steps/__init__.py index 15c03f6..5156507 100644 --- a/brainforge_inference_utils/steps/__init__.py +++ b/brainforge_inference_utils/steps/__init__.py @@ -31,14 +31,11 @@ Step base classes. Concrete steps ship in separate plugins and are discovered from entry points. -`steps_by_id` remains for any steps still bundled with this package. Plugin -step ids use `author.step_name`. +Step ids use `author.step_name`. """ from __future__ import annotations -from typing import Dict, Type - from brainforge_inference_utils.steps.base import ( RequestStep, RequestContextOption, @@ -46,12 +43,9 @@ Step, ) -steps_by_id: Dict[str, Type[Step]] = {} - __all__ = [ "Step", "RequestStep", "ResponseStep", "RequestContextOption", - "steps_by_id", ] diff --git a/setup.py b/setup.py index 4402b66..09182dc 100644 --- a/setup.py +++ b/setup.py @@ -102,7 +102,7 @@ def get_requirements(requirements_filename: str): ], # Plugin steps register under brainforge.inference.request_steps and # brainforge.inference.response_steps. Entry point names are - # author.step_name. Bundled steps stay in steps_by_id for now. + # author.step_name. This package registers no steps of its own. # entry_points={ # "brainforge.inference.request_steps": [ # "acme.guardrails=acme_guardrails:GuardrailRequestStep", diff --git a/tests/helpers.py b/tests/helpers.py index a6a7aeb..6e9df56 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -27,7 +27,8 @@ # LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING # NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -from typing import Optional +from types import SimpleNamespace +from typing import Optional, Type from neon_data_models.models.api.llm import ( LLMPersona, @@ -35,7 +36,11 @@ LLMResponse, ) -from brainforge_inference_utils.steps.base import RequestStep, ResponseStep +from brainforge_inference_utils.registry import ( + REQUEST_STEP_ENTRY_POINT, + RESPONSE_STEP_ENTRY_POINT, +) +from brainforge_inference_utils.steps.base import RequestStep, ResponseStep, Step from brainforge_inference_utils.transaction import UserTransaction @@ -70,7 +75,7 @@ def make_transaction(response: Optional[LLMResponse] = None, class SetMaxHistoryStep(RequestStep): """Request step that writes `config["max_history"]` onto the request.""" - id = "set_max_history" + id = "test.set_max_history" async def process(self, transaction: UserTransaction) -> None: transaction.request.max_history = self.config["max_history"] @@ -79,7 +84,7 @@ async def process(self, transaction: UserTransaction) -> None: class SkipUnlessQueryStep(RequestStep): """Request step skipped unless `request.query` matches `config["query"]`.""" - id = "skip_unless_query" + id = "test.skip_unless_query" def applicable(self, transaction: UserTransaction) -> Optional[str]: expected = self.config.get("query") @@ -94,7 +99,7 @@ async def process(self, transaction: UserTransaction) -> None: class AppendQueryStep(RequestStep): """Request step that appends `config["suffix"]` to `request.query`.""" - id = "append_query" + id = "test.append_query" async def process(self, transaction: UserTransaction) -> None: transaction.request.query += self.config.get("suffix", "") @@ -103,7 +108,7 @@ async def process(self, transaction: UserTransaction) -> None: class PrefixResponseStep(ResponseStep): """Response step that prepends `config["prefix"]` to the response.""" - id = "prefix_response" + id = "test.prefix_response" def applicable(self, transaction: UserTransaction) -> Optional[str]: if transaction.response is None: @@ -118,16 +123,52 @@ async def process(self, transaction: UserTransaction) -> None: class MutateResponseInPlaceStep(ResponseStep): """Response step that appends `config["suffix"]` to the response.""" - id = "mutate_in_place" + id = "test.mutate_in_place" async def process(self, transaction: UserTransaction) -> None: transaction.response.response += self.config.get("suffix", "!") -TEST_STEPS = { - SetMaxHistoryStep.id: SetMaxHistoryStep, - SkipUnlessQueryStep.id: SkipUnlessQueryStep, - AppendQueryStep.id: AppendQueryStep, - PrefixResponseStep.id: PrefixResponseStep, - MutateResponseInPlaceStep.id: MutateResponseInPlaceStep, -} +class FakeEntryPoint: + """Stand-in for an importlib entry point that loads a step class.""" + + def __init__(self, name, cls, dist="test-plugin"): + self.name = name + self._cls = cls + self.dist = SimpleNamespace(name=dist) + + def load(self): + return self._cls + + +def entry_points_for(request_points=(), response_points=()): + """Return an `entry_points(group=...)` replacement for the two step groups.""" + + def fake_entry_points(group): + if group == REQUEST_STEP_ENTRY_POINT: + return list(request_points) + if group == RESPONSE_STEP_ENTRY_POINT: + return list(response_points) + return [] + + return fake_entry_points + + +def entry_points_for_steps( + *step_classes: Type[Step], + dist: str = "test-plugin", +): + """ + Advertise `step_classes` on the entry point group matching each base class. + """ + request_points = [] + response_points = [] + for cls in step_classes: + point = FakeEntryPoint(cls.id, cls, dist=dist) + if issubclass(cls, RequestStep): + request_points.append(point) + elif issubclass(cls, ResponseStep): + response_points.append(point) + else: + raise TypeError(f"{cls} is not a request or response step") + return entry_points_for(request_points, response_points) diff --git a/tests/test_client.py b/tests/test_client.py index 1382f6f..05181ba 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -29,6 +29,7 @@ # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import unittest +from unittest.mock import patch from brainforge_inference_utils.client import BrainForgeClient from brainforge_inference_utils.config import PipelineConfig @@ -40,20 +41,21 @@ from brainforge_inference_utils.transaction import UserTransaction from helpers import ( PrefixResponseStep, + entry_points_for_steps, make_request, make_response, ) class RecordRequestStep(RequestStep): - id = "record_request" + id = "test.record_request" async def process(self, transaction: UserTransaction) -> None: transaction.request_step_context[self.id] = self.config["value"] class CatalogStep(RequestStep): - id = "catalog" + id = "test.catalog" async def process(self, transaction: UserTransaction) -> None: return None @@ -69,7 +71,7 @@ async def supported_request_context(self, timeout: float = 10.0): class FlagStep(ResponseStep): - id = "flag" + id = "test.flag" request_context_options = ( RequestContextOption("citations_enabled", (True, False)), RequestContextOption("location"), @@ -80,38 +82,47 @@ async def process(self, transaction: UserTransaction) -> None: class RecordResponseStep(ResponseStep): - id = "record_response" + id = "test.record_response" async def process(self, transaction: UserTransaction) -> None: transaction.response_step_context[self.id] = self.config["value"] class TestBrainForgeClient(unittest.IsolatedAsyncioTestCase): + def setUp(self): + self._entry_points = patch( + "brainforge_inference_utils.registry.entry_points", + entry_points_for_steps( + RecordRequestStep, + CatalogStep, + PrefixResponseStep, + RecordResponseStep, + FlagStep, + ), + ) + self._entry_points.start() + + def tearDown(self): + self._entry_points.stop() + def make_client(self): config = PipelineConfig.model_validate({ "request": [{ - "id": "record_request", + "id": "test.record_request", "config": {"value": {"documents": ["one"]}}, }], "response": [ { - "id": "prefix_response", + "id": "test.prefix_response", "config": {"prefix": "processed: "}, }, { - "id": "record_response", + "id": "test.record_response", "config": {"value": {"citations": [1]}}, }, ], }) - return BrainForgeClient( - config, - { - RecordRequestStep.id: RecordRequestStep, - RecordResponseStep.id: RecordResponseStep, - PrefixResponseStep.id: PrefixResponseStep, - }, - ) + return BrainForgeClient(config) async def test_run_request_pipeline_returns_transaction_with_context(self): client = self.make_client() @@ -122,7 +133,7 @@ async def test_run_request_pipeline_returns_transaction_with_context(self): self.assertIsNone(transaction.response) self.assertEqual( transaction.request_step_context, - {"record_request": {"documents": ["one"]}}, + {"test.record_request": {"documents": ["one"]}}, ) self.assertEqual(transaction.response_step_context, {}) @@ -137,30 +148,23 @@ async def test_run_response_pipeline_keeps_separate_step_namespaces(self): self.assertEqual(transaction.response.response, "processed: answer") self.assertEqual( transaction.request_step_context, - {"record_request": {"documents": ["one"]}}, + {"test.record_request": {"documents": ["one"]}}, ) self.assertEqual( transaction.response_step_context, - {"record_response": {"citations": [1]}}, + {"test.record_response": {"citations": [1]}}, ) async def test_supported_request_context_merges_pipeline_steps(self): config = PipelineConfig.model_validate({ "request": [ - {"id": "catalog", "config": {"providers": ["news", "docs"]}}, - {"id": "record_request", "config": {"value": {}}}, - {"id": "catalog", "config": {"providers": ["docs", "web"]}}, + {"id": "test.catalog", "config": {"providers": ["news", "docs"]}}, + {"id": "test.record_request", "config": {"value": {}}}, + {"id": "test.catalog", "config": {"providers": ["docs", "web"]}}, ], - "response": [{"id": "flag"}], + "response": [{"id": "test.flag"}], }) - client = BrainForgeClient( - config, - { - CatalogStep.id: CatalogStep, - RecordRequestStep.id: RecordRequestStep, - FlagStep.id: FlagStep, - }, - ) + client = BrainForgeClient(config) options = await client.supported_request_context() self.assertEqual( @@ -168,4 +172,4 @@ async def test_supported_request_context_merges_pipeline_steps(self): self.assertEqual( list(options["citations_enabled"].values), [True, False]) self.assertIsNone(options["location"].values) - self.assertNotIn("record_request", options) + self.assertNotIn("test.record_request", options) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index e235d0e..1073e13 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -29,6 +29,7 @@ # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import unittest +from unittest.mock import patch import pytest @@ -38,88 +39,101 @@ ResponsePipeline, build_pipelines, ) -from brainforge_inference_utils.steps import steps_by_id from helpers import ( AppendQueryStep, MutateResponseInPlaceStep, PrefixResponseStep, SetMaxHistoryStep, SkipUnlessQueryStep, - TEST_STEPS, + entry_points_for_steps, make_response, make_transaction, ) +def _patch_test_steps(monkeypatch): + monkeypatch.setattr( + "brainforge_inference_utils.registry.entry_points", + entry_points_for_steps( + AppendQueryStep, + SkipUnlessQueryStep, + PrefixResponseStep, + MutateResponseInPlaceStep, + ), + ) + + class TestBuildPipelines: - def test_empty_config_builds_empty_pipelines(self): - request_pipeline, response_pipeline = build_pipelines( - PipelineConfig(), TEST_STEPS) + def test_empty_config_builds_empty_pipelines(self, monkeypatch): + _patch_test_steps(monkeypatch) + request_pipeline, response_pipeline = build_pipelines(PipelineConfig()) assert request_pipeline.steps == [] assert response_pipeline.steps == [] - def test_default_registry_has_no_bundled_steps(self): - assert steps_by_id == {} - with pytest.raises(KeyError, match="append_query"): + def test_unregistered_id_raises(self, monkeypatch): + _patch_test_steps(monkeypatch) + with pytest.raises(KeyError, match="test.missing"): build_pipelines(PipelineConfig.model_validate({ - "request": ["append_query"], + "request": ["test.missing"], })) - def test_unknown_id_lists_known_steps(self): - with pytest.raises(KeyError, match="append_query"): + def test_unknown_id_lists_known_steps(self, monkeypatch): + _patch_test_steps(monkeypatch) + with pytest.raises(KeyError, match="test.append_query"): build_pipelines( - PipelineConfig.model_validate({"request": ["nope"]}), - TEST_STEPS, + PipelineConfig.model_validate({"request": ["test.missing"]}), ) - def test_disabled_steps_are_omitted(self): + def test_disabled_steps_are_omitted(self, monkeypatch): + _patch_test_steps(monkeypatch) request_pipeline, response_pipeline = build_pipelines( PipelineConfig.model_validate({ "request": [ - {"id": "append_query", "enabled": False, + {"id": "test.append_query", "enabled": False, "config": {"suffix": " skipped"}}, - {"id": "append_query", "config": {"suffix": " kept"}}, + {"id": "test.append_query", "config": {"suffix": " kept"}}, ], "response": [ - {"id": "prefix_response", "enabled": False, + {"id": "test.prefix_response", "enabled": False, "config": {"prefix": "X"}}, ], }), - TEST_STEPS, ) assert len(request_pipeline.steps) == 1 assert request_pipeline.steps[0].config == {"suffix": " kept"} assert response_pipeline.steps == [] - def test_response_step_in_request_list_raises(self): - with pytest.raises(TypeError, match="not a RequestStep"): + def test_response_step_id_is_unknown_on_the_request_pipeline( + self, monkeypatch): + _patch_test_steps(monkeypatch) + with pytest.raises(KeyError, match="test.prefix_response"): build_pipelines( PipelineConfig.model_validate({ - "request": ["prefix_response"]}), - TEST_STEPS, + "request": ["test.prefix_response"]}), ) - def test_request_step_in_response_list_raises(self): - with pytest.raises(TypeError, match="not a ResponseStep"): + def test_request_step_id_is_unknown_on_the_response_pipeline( + self, monkeypatch): + _patch_test_steps(monkeypatch) + with pytest.raises(KeyError, match="test.append_query"): build_pipelines( PipelineConfig.model_validate( - {"response": ["append_query"]}), - TEST_STEPS, + {"response": ["test.append_query"]}), ) - def test_duplicate_ids_instantiate_independent_instances(self): + def test_duplicate_ids_instantiate_independent_instances(self, monkeypatch): + _patch_test_steps(monkeypatch) request_pipeline, response_pipeline = build_pipelines( PipelineConfig.model_validate({ "request": [ - {"id": "append_query", "config": {"suffix": " A"}}, - {"id": "append_query", "config": {"suffix": " B"}}, + {"id": "test.append_query", "config": {"suffix": " A"}}, + {"id": "test.append_query", "config": {"suffix": " B"}}, ], "response": [ - {"id": "prefix_response", "config": {"prefix": "["}}, - {"id": "prefix_response", "config": {"prefix": "<"}}, + {"id": "test.prefix_response", "config": {"prefix": "["}}, + {"id": "test.prefix_response", "config": {"prefix": "<"}}, ], }), - TEST_STEPS, ) assert len(request_pipeline.steps) == 2 assert request_pipeline.steps[0] is not request_pipeline.steps[1] @@ -188,22 +202,30 @@ async def test_inapplicable_response_step_is_skipped(self): assert transaction.response is None async def test_build_then_run_duplicate_configured_steps(self): - request_pipeline, response_pipeline = build_pipelines( - PipelineConfig.model_validate({ - "request": [ - {"id": "append_query", "config": {"suffix": " A"}}, - {"id": "skip_unless_query", - "config": {"query": "nope", "max_tokens": 64}}, - {"id": "append_query", "config": {"suffix": " B"}}, - ], - "response": [ - {"id": "prefix_response", "config": {"prefix": "["}}, - {"id": "mutate_in_place", "config": {"suffix": "!"}}, - {"id": "prefix_response", "config": {"prefix": "<"}}, - ], - }), - TEST_STEPS, - ) + with patch( + "brainforge_inference_utils.registry.entry_points", + entry_points_for_steps( + AppendQueryStep, + SkipUnlessQueryStep, + PrefixResponseStep, + MutateResponseInPlaceStep, + ), + ): + request_pipeline, response_pipeline = build_pipelines( + PipelineConfig.model_validate({ + "request": [ + {"id": "test.append_query", "config": {"suffix": " A"}}, + {"id": "test.skip_unless_query", + "config": {"query": "nope", "max_tokens": 64}}, + {"id": "test.append_query", "config": {"suffix": " B"}}, + ], + "response": [ + {"id": "test.prefix_response", "config": {"prefix": "["}}, + {"id": "test.mutate_in_place", "config": {"suffix": "!"}}, + {"id": "test.prefix_response", "config": {"prefix": "<"}}, + ], + }), + ) transaction = make_transaction(query="hello") await request_pipeline.run(transaction) assert transaction.request.query == "hello A B" diff --git a/tests/test_registry.py b/tests/test_registry.py index dd972d8..af63f03 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -28,8 +28,6 @@ # NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -from types import SimpleNamespace - import pytest from brainforge_inference_utils.config import PipelineConfig @@ -38,19 +36,10 @@ REQUEST_STEP_ENTRY_POINT, RESPONSE_STEP_ENTRY_POINT, load_entry_point_steps, + merge_builtin_and_entry_points, ) from brainforge_inference_utils.steps.base import RequestStep, ResponseStep -from helpers import TEST_STEPS - - -class _EntryPoint: - def __init__(self, name, cls, dist="acme-plugin"): - self.name = name - self._cls = cls - self.dist = SimpleNamespace(name=dist) - - def load(self): - return self._cls +from helpers import FakeEntryPoint, entry_points_for class AcmeRequestGuard(RequestStep): @@ -74,23 +63,12 @@ async def process(self, transaction): return None -def _entry_points_for(request_points=(), response_points=()): - def fake_entry_points(group): - if group == REQUEST_STEP_ENTRY_POINT: - return list(request_points) - if group == RESPONSE_STEP_ENTRY_POINT: - return list(response_points) - return [] - - return fake_entry_points - - def test_same_author_step_name_loads_from_both_groups(monkeypatch): monkeypatch.setattr( "brainforge_inference_utils.registry.entry_points", - _entry_points_for( - [_EntryPoint("acme.guard", AcmeRequestGuard)], - [_EntryPoint("acme.guard", AcmeResponseGuard, dist="acme-plugin")], + entry_points_for( + [FakeEntryPoint("acme.guard", AcmeRequestGuard)], + [FakeEntryPoint("acme.guard", AcmeResponseGuard, dist="acme-plugin")], ), ) @@ -106,9 +84,9 @@ def test_same_author_step_name_loads_from_both_groups(monkeypatch): def test_build_pipelines_resolves_entry_point_steps(monkeypatch): monkeypatch.setattr( "brainforge_inference_utils.registry.entry_points", - _entry_points_for( - [_EntryPoint("acme.guard", AcmeRequestGuard)], - [_EntryPoint("acme.guard", AcmeResponseGuard)], + entry_points_for( + [FakeEntryPoint("acme.guard", AcmeRequestGuard)], + [FakeEntryPoint("acme.guard", AcmeResponseGuard)], ), ) @@ -124,29 +102,10 @@ def test_build_pipelines_resolves_entry_point_steps(monkeypatch): assert type(response_pipeline.steps[0]) is AcmeResponseGuard -def test_explicit_step_map_skips_entry_points(monkeypatch): - def fail_entry_points(group): - raise AssertionError(f"entry points consulted for {group}") - - monkeypatch.setattr( - "brainforge_inference_utils.registry.entry_points", - fail_entry_points, - ) - request_pipeline, response_pipeline = build_pipelines( - PipelineConfig.model_validate({ - "request": ["append_query"], - "response": ["prefix_response"], - }), - TEST_STEPS, - ) - assert [step.id for step in request_pipeline.steps] == ["append_query"] - assert [step.id for step in response_pipeline.steps] == ["prefix_response"] - - def test_entry_point_name_must_be_author_step_name(monkeypatch): monkeypatch.setattr( "brainforge_inference_utils.registry.entry_points", - _entry_points_for([_EntryPoint("guardrails", AcmeRequestGuard)]), + entry_points_for([FakeEntryPoint("guardrails", AcmeRequestGuard)]), ) with pytest.raises(ValueError, match="author.step_name"): load_entry_point_steps(REQUEST_STEP_ENTRY_POINT, RequestStep) @@ -155,9 +114,9 @@ def test_entry_point_name_must_be_author_step_name(monkeypatch): def test_duplicate_name_in_one_group_names_both_packages(monkeypatch): monkeypatch.setattr( "brainforge_inference_utils.registry.entry_points", - _entry_points_for([ - _EntryPoint("acme.guard", AcmeRequestGuard, dist="acme-a"), - _EntryPoint("acme.guard", AcmeRequestGuard, dist="acme-b"), + entry_points_for([ + FakeEntryPoint("acme.guard", AcmeRequestGuard, dist="acme-a"), + FakeEntryPoint("acme.guard", AcmeRequestGuard, dist="acme-b"), ]), ) with pytest.raises(ValueError, match="acme-a and acme-b"): @@ -167,7 +126,7 @@ def test_duplicate_name_in_one_group_names_both_packages(monkeypatch): def test_class_id_must_match_entry_point_name(monkeypatch): monkeypatch.setattr( "brainforge_inference_utils.registry.entry_points", - _entry_points_for([_EntryPoint("acme.guard", MismatchedId)]), + entry_points_for([FakeEntryPoint("acme.guard", MismatchedId)]), ) with pytest.raises(ValueError, match="does not match entry point"): load_entry_point_steps(REQUEST_STEP_ENTRY_POINT, RequestStep) @@ -176,24 +135,33 @@ def test_class_id_must_match_entry_point_name(monkeypatch): def test_response_class_in_request_group_raises(monkeypatch): monkeypatch.setattr( "brainforge_inference_utils.registry.entry_points", - _entry_points_for([_EntryPoint("acme.guard", AcmeResponseGuard)]), + entry_points_for([FakeEntryPoint("acme.guard", AcmeResponseGuard)]), ) with pytest.raises(TypeError, match="not a RequestStep"): load_entry_point_steps(REQUEST_STEP_ENTRY_POINT, RequestStep) -def test_plugin_id_that_collides_with_builtin_raises(monkeypatch): - monkeypatch.setattr( - "brainforge_inference_utils.pipeline.steps_by_id", - {"acme.guard": AcmeRequestGuard}, - ) +def test_merge_builtin_and_entry_points_is_deprecated(monkeypatch): monkeypatch.setattr( "brainforge_inference_utils.registry.entry_points", - _entry_points_for([ - _EntryPoint("acme.guard", AcmeRequestGuard, dist="acme-plugin"), + entry_points_for([ + FakeEntryPoint("acme.guard", AcmeRequestGuard, dist="acme-plugin"), ]), ) - with pytest.raises(ValueError, match="conflicts with a built-in step"): - build_pipelines(PipelineConfig.model_validate({ - "request": ["acme.guard"], - })) + + class CallerStep(RequestStep): + id = "test.caller" + + async def process(self, transaction): + return None + + with pytest.warns(DeprecationWarning, match="load_entry_point_steps"): + resolved = merge_builtin_and_entry_points( + {"test.caller": CallerStep}, + REQUEST_STEP_ENTRY_POINT, + RequestStep, + ) + + assert resolved["test.caller"] is CallerStep + assert resolved["acme.guard"] is AcmeRequestGuard + diff --git a/tests/test_steps.py b/tests/test_steps.py index 3920ff0..37c57be 100644 --- a/tests/test_steps.py +++ b/tests/test_steps.py @@ -56,7 +56,7 @@ def test_constructor_copies_config(self): def test_repr_includes_class_and_id(self): step = AppendQueryStep() - assert repr(step) == "" + assert repr(step) == "" def test_default_request_context_options_are_empty(self): assert list(AppendQueryStep().request_context_options) == [] From 778fec6a01d22bc25ece0f5436afe31738c6336a Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 16:40:07 -0700 Subject: [PATCH 20/23] Remove deprecated method for builtin entrypoints --- brainforge_inference_utils/registry.py | 37 +------------------------- tests/test_registry.py | 26 ------------------ 2 files changed, 1 insertion(+), 62 deletions(-) diff --git a/brainforge_inference_utils/registry.py b/brainforge_inference_utils/registry.py index 0545bce..d95acb4 100644 --- a/brainforge_inference_utils/registry.py +++ b/brainforge_inference_utils/registry.py @@ -39,9 +39,8 @@ import logging import re -import warnings from importlib.metadata import entry_points -from typing import Dict, Mapping, Type +from typing import Dict, Type from brainforge_inference_utils.steps.base import Step @@ -113,37 +112,3 @@ def load_entry_point_steps( LOG.debug("Loaded %s step %r from %s", expected_type.__name__, entry_point.name, dist_name) return loaded - - -def merge_builtin_and_entry_points( - builtin: Mapping[str, Type[Step]], - group: str, - expected_type: Type[Step], -) -> Dict[str, Type[Step]]: - """ - Combine a caller-supplied step map with classes discovered from `group`. - - Deprecated. This package no longer ships built-in steps. Call - `load_entry_point_steps` instead. The `builtin` map is still copied and - then extended with the entry point group so existing callers keep working. - - :param builtin: Caller-supplied id -> step class mapping. - :param group: Entry point group for this pipeline side. - :param expected_type: Step base class that group must provide. - :returns: A new mapping. `builtin` is copied, then plugin steps are added. - :raises ValueError: If a plugin id is already present in `builtin`. - """ - warnings.warn( - "merge_builtin_and_entry_points is deprecated; this package no longer " - "ships built-in steps. Use load_entry_point_steps instead.", - DeprecationWarning, - stacklevel=2, - ) - resolved = dict(builtin) - for step_id, step_cls in load_entry_point_steps(group, expected_type).items(): - if step_id in resolved: - raise ValueError( - f"Step id {step_id!r} from entry point group {group} conflicts " - f"with a built-in step") - resolved[step_id] = step_cls - return resolved diff --git a/tests/test_registry.py b/tests/test_registry.py index af63f03..b6f550e 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -36,7 +36,6 @@ REQUEST_STEP_ENTRY_POINT, RESPONSE_STEP_ENTRY_POINT, load_entry_point_steps, - merge_builtin_and_entry_points, ) from brainforge_inference_utils.steps.base import RequestStep, ResponseStep from helpers import FakeEntryPoint, entry_points_for @@ -140,28 +139,3 @@ def test_response_class_in_request_group_raises(monkeypatch): with pytest.raises(TypeError, match="not a RequestStep"): load_entry_point_steps(REQUEST_STEP_ENTRY_POINT, RequestStep) - -def test_merge_builtin_and_entry_points_is_deprecated(monkeypatch): - monkeypatch.setattr( - "brainforge_inference_utils.registry.entry_points", - entry_points_for([ - FakeEntryPoint("acme.guard", AcmeRequestGuard, dist="acme-plugin"), - ]), - ) - - class CallerStep(RequestStep): - id = "test.caller" - - async def process(self, transaction): - return None - - with pytest.warns(DeprecationWarning, match="load_entry_point_steps"): - resolved = merge_builtin_and_entry_points( - {"test.caller": CallerStep}, - REQUEST_STEP_ENTRY_POINT, - RequestStep, - ) - - assert resolved["test.caller"] is CallerStep - assert resolved["acme.guard"] is AcmeRequestGuard - From 310403a5a714f75bd85bef1102c3bf3c266644ee Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 16:49:40 -0700 Subject: [PATCH 21/23] Manually clean up README.md --- README.md | 33 +++++++++------------------------ 1 file changed, 9 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index 4971572..dac7c1f 100644 --- a/README.md +++ b/README.md @@ -9,15 +9,15 @@ safety guardrails, ...) can be implemented as **steps** that plug into one of two pipelines: - **Request pipeline** — runs before generation and shapes the request: - injecting ambient time/weather, retrieving RAG for host-selected - providers, then validating and truncating history last. + injecting informational context, RAG retrieval for host-selected + providers, validating and truncating history, etc. - **Response pipeline** — runs after generation and shapes the response: - citation rewrite via `ResponseParser`, sanitizing, guardrails, etc. + citation rewrite, sanitization, guardrails, etc. Primary generation stays with the host, between the two pipelines. The host keeps the `UserTransaction` from the request pipeline until the response is -ready. A response step may call a configured OpenAI-compatible endpoint -(for example citation rewrite) without the host owning that second request. +ready. A pipeline step may call an inference endpoint to modify the transaction +(for example citation rewrite) without the host owning those requests. Both pipelines operate on a `UserTransaction`, which carries the request, response, and two per-side context dicts (`request_step_context` and @@ -35,7 +35,6 @@ response, and two per-side context dicts (`request_step_context` and | `RequestContextOption` | One `request_context` key a step reads, with optional allowed values for hosts to present. | | `RequestPipeline` / `ResponsePipeline` | Run an ordered list of steps over a transaction. | | `PipelineConfig` / `StepConfig` | Declarative, ordered step lists (see below). | -| `build_pipelines(config)` | Build both pipelines from config. Step classes come from the entry point group for that pipeline. | ## Usage sketch @@ -61,7 +60,7 @@ between calls. - **Steps are independently defensive.** There is no dependency graph between steps. A step that needs something it does not find (e.g. a citation step with - no retrieved context) reports itself inapplicable via `applicable()` and is + no retrieved sources) reports itself inapplicable via `applicable()` and is skipped with a logged reason, rather than failing the pipeline. - **Ordering is explicit.** Step order is exactly the order declared in configuration; the pipeline does not reorder. @@ -73,9 +72,10 @@ between calls. A repeated step on the *same* side shares that key; a step id on the request side does not clash with the same id on the response side. Each step is responsible for how it stores and reads its own data across one or more - calls (for example appending to a list instead of overwriting a scalar). + calls (for example appending to a list instead of overwriting it). Response steps that consume request-step data should read - `transaction.request_step_context` and skip when nothing usable is present. + `transaction.request_step_context` and report inapplicable when required + context is absent. - **Format contracts live together.** Steps that share a parse format (e.g. a citation formatter and the sanitizer that strips it) belong in the same module and should derive from a shared constant so they cannot drift. @@ -123,18 +123,3 @@ pipeline: - id: acme.guardrails ``` -First-party steps ship in their own packages: - -| Package | Steps | -| --- | --- | -| `brainforge-inference-plugin-locale` | `brainforge.locale` | -| `brainforge-inference-plugin-rag` | `brainforge.rag`, `brainforge.citations` | -| `brainforge-inference-plugin-history` | `brainforge.history` | - -This package does not register those steps. Install the plugin to use the id. - -## Deferred - -- A **services** adapter for streaming, tool calls, and token counting. -- Richer persistent state or execution tracing beyond the transient - `request_step_context` / `response_step_context` dicts on `UserTransaction`. From 198653f7c9fdfbf0941aa4938b64e60c158b6e86 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 16:53:49 -0700 Subject: [PATCH 22/23] Loosen entrypoint naming requirements --- README.md | 5 ++--- brainforge_inference_utils/registry.py | 9 +++++---- tests/test_registry.py | 15 +++++++++++++++ 3 files changed, 22 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index dac7c1f..5bc57ee 100644 --- a/README.md +++ b/README.md @@ -95,9 +95,8 @@ Both groups use the step id `author.step_name`: - `author` is the plugin developer's slug (`acme`, `brainforge`). - `step_name` names that one step (`guardrails`, `rag`). -Each part is lowercase letters, digits, and underscores, and starts with a -letter. The class `id`, the config `id`, and the context-dict key are that -same string. +Each part may contain letters, digits, hyphens, and underscores. The class +`id`, the config `id`, and the context-dict key are that same string. The same `author.step_name` may be registered once in each group. Each pipeline reads only its own group, so the request class and the response class diff --git a/brainforge_inference_utils/registry.py b/brainforge_inference_utils/registry.py index d95acb4..7cb35b9 100644 --- a/brainforge_inference_utils/registry.py +++ b/brainforge_inference_utils/registry.py @@ -49,8 +49,8 @@ REQUEST_STEP_ENTRY_POINT = "brainforge.inference.request_steps" RESPONSE_STEP_ENTRY_POINT = "brainforge.inference.response_steps" -# author.step_name — one developer slug, one step name. -STEP_ID_PATTERN = re.compile(r"^[a-z][a-z0-9_]*\.[a-z][a-z0-9_]*$") +# author.step_name — letters, digits, hyphens, and underscores in each part. +STEP_ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+$") def _distribution_name(entry_point) -> str: @@ -66,8 +66,9 @@ def load_entry_point_steps( Load step classes advertised under `group`. Each entry point name must be `author.step_name` and must equal the - loaded class's `id`. Names must be unique within the group. The same name - may still be registered under the other pipeline's group. + loaded class's `id`. Each part may contain letters, digits, hyphens, and + underscores. Names must be unique within the group. The same name may + still be registered under the other pipeline's group. :param group: Entry point group to read. :param expected_type: `RequestStep` or `ResponseStep`. Loaded classes must diff --git a/tests/test_registry.py b/tests/test_registry.py index b6f550e..ab5d1a0 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -101,6 +101,21 @@ def test_build_pipelines_resolves_entry_point_steps(monkeypatch): assert type(response_pipeline.steps[0]) is AcmeResponseGuard +def test_entry_point_name_allows_alphanumerics_hyphen_and_underscore(monkeypatch): + class MixedName(RequestStep): + id = "Acme-corp.step_name-2" + + async def process(self, transaction): + return None + + monkeypatch.setattr( + "brainforge_inference_utils.registry.entry_points", + entry_points_for([FakeEntryPoint(MixedName.id, MixedName)]), + ) + loaded = load_entry_point_steps(REQUEST_STEP_ENTRY_POINT, RequestStep) + assert loaded[MixedName.id] is MixedName + + def test_entry_point_name_must_be_author_step_name(monkeypatch): monkeypatch.setattr( "brainforge_inference_utils.registry.entry_points", From b35d674915df63f7cc2da02f3d36f63df9264dd6 Mon Sep 17 00:00:00 2001 From: Daniel McKnight Date: Wed, 14 Oct 2026 16:54:42 -0700 Subject: [PATCH 23/23] Remove commented plugin examples from setup.py --- setup.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/setup.py b/setup.py index 09182dc..82ae5c2 100644 --- a/setup.py +++ b/setup.py @@ -100,15 +100,4 @@ def get_requirements(requirements_filename: str): "Intended Audience :: Developers", "Programming Language :: Python :: 3", ], - # Plugin steps register under brainforge.inference.request_steps and - # brainforge.inference.response_steps. Entry point names are - # author.step_name. This package registers no steps of its own. - # entry_points={ - # "brainforge.inference.request_steps": [ - # "acme.guardrails=acme_guardrails:GuardrailRequestStep", - # ], - # "brainforge.inference.response_steps": [ - # "acme.guardrails=acme_guardrails:GuardrailResponseStep", - # ], - # }, )