wehub/jina-ai--reader
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
/ test (push) Failing after 1s
Details
/ build-and-push-to-ghcr (push) Has been skipped
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:23:39 +08:00
commit f1ba9c6c36
177 changed files with 54678 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.c8rc.json
+15
View File
@@ -0,0 +1,15 @@
{
"all": true,
"src": ["src"],
"include": ["src/**/*.ts"],
"exclude": [
"src/**/*.d.ts",
"tests/**",
"tests-build/**",
"src/scripts/**"
],
"reporter": ["text", "html", "lcov"],
"reports-dir": "coverage",
"excludeAfterRemap": true,
"skip-full": false
}
.editorconfig
+17
View File
@@ -0,0 +1,17 @@
# EditorConfig is awesome: https://EditorConfig.org
# top-most EditorConfig file
root = true
# Unix-style newlines with a newline ending every file
[*]
end_of_line = lf
insert_final_newline = true
charset = utf-8
indent_style = space
indent_size = 4
[*.html]
[*.htm]
indent_style = space
indent_size = 2
.github/workflows/.keep
View File
.github/workflows/_test.yml
+41
View File
@@ -0,0 +1,41 @@
run-name: CI test (reusable)
on:
workflow_call:
inputs:
node-version:
type: string
default: '24.14.1'
description: 'Node.js version to use.'
jobs:
test:
runs-on: ubuntu-latest
permissions:
contents: read
steps:
- uses: actions/checkout@v4
with:
lfs: true
- name: Set up Node.js
uses: actions/setup-node@v4
with:
node-version: ${{ inputs.node-version }}
cache: npm
- name: npm install
run: npm ci
- name: get non-redistributable assets
run: npm run assets:download
- name: build application
run: npm run build
- name: tests with coverage
run: npm run test:coverage
- name: upload coverage to Codecov
uses: codecov/codecov-action@v6
with:
files: coverage/lcov.info
token: ${{ secrets.CODECOV_TOKEN }}
fail_ci_if_error: true
.github/workflows/ci.yml
+24
View File
@@ -0,0 +1,24 @@
run-name: CI (tests)
on:
pull_request:
paths-ignore:
- '**.md'
push:
branches-ignore:
- main
- ops
- ci-debug
- dev
paths-ignore:
- '**.md'
concurrency:
group: ci-${{ github.ref }}
cancel-in-progress: true
jobs:
test:
uses: ./.github/workflows/_test.yml
permissions:
contents: read
secrets: inherit
.github/workflows/oss-image.yml
+87
View File
@@ -0,0 +1,87 @@
run-name: Build push container image to GH
on:
push:
branches:
- main
tags:
- '*'
jobs:
test:
uses: ./.github/workflows/_test.yml
permissions:
contents: read
secrets: inherit
build-and-push-to-ghcr:
needs: test
runs-on: ubuntu-latest
concurrency:
group: ${{ github.ref_type == 'branch' && github.ref }}
cancel-in-progress: true
permissions:
contents: read
packages: write
steps:
- uses: actions/checkout@v5
with:
lfs: true
submodules: true
- name: Login to GitHub Container Registry
uses: docker/login-action@v3
with:
registry: ghcr.io
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Set controller release version
run: echo "RELEASE_VERSION=${GITHUB_REF#refs/*/}" >> $GITHUB_ENV
- name: Set up Node.js
uses: actions/setup-node@v5
with:
node-version: 24.14.1
cache: npm
- name: npm install
run: npm ci
- name: get maxmind mmdb
run: mkdir -p licensed && curl -o licensed/GeoLite2-City.mmdb https://raw.githubusercontent.com/P3TERX/GeoLite.mmdb/download/GeoLite2-City.mmdb
- name: get asn mmdb
run: curl -o licensed/geolite2-asn.mmdb https://cdn.jsdelivr.net/npm/@ip-location-db/geolite2-asn-mmdb/geolite2-asn.mmdb
- name: get source han sans font
run: curl -o licensed/SourceHanSansSC-Regular.otf https://raw.githubusercontent.com/adobe-fonts/source-han-sans/refs/heads/release/OTF/SimplifiedChinese/SourceHanSansSC-Regular.otf
- name: get gsa user agent list
run: curl -o licensed/gsa_useragents.txt https://raw.githubusercontent.com/searxng/searxng/refs/heads/master/searx/data/gsa_useragents.txt
- name: build application
run: npm run build
- name: Set package version
run: npm version --no-git-tag-version ${{ env.RELEASE_VERSION }}
if: github.ref_type == 'tag'
- name: Docker meta
id: meta
uses: docker/metadata-action@v5
with:
images: |
ghcr.io/jina-ai/reader
labels: |
org.opencontainers.image.created={{timestamp}}
org.opencontainers.image.source="https://github.com/jina-ai/reader"
tags: |
type=schedule
type=ref,event=branch
type=ref,event=tag
type=ref,event=pr
type=raw,value=oss,enable=${{ github.ref == format('refs/heads/{0}', 'main') }}
flavor: |
latest=true
- name: Set up QEMU
uses: docker/setup-qemu-action@v4
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v4
- name: Build and push
uses: docker/build-push-action@v7
with:
context: .
push: true
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
platforms: linux/amd64,linux/arm64
.gitignore
+81
View File
@@ -0,0 +1,81 @@
# Logs
logs
*.log
npm-debug.log*
yarn-debug.log*
yarn-error.log*
firebase-debug.log*
firebase-debug.*.log*
# Firebase cache
.firebase/
# Firebase config
# Uncomment this if you'd like others to create their own Firebase project.
# For a team working on the same Firebase project(s), it is recommended to leave
# it commented so all members can deploy to the same project(s) in .firebaserc.
# .firebaserc
# Runtime data
pids
*.pid
*.seed
*.pid.lock
# Directory for instrumented libs generated by jscoverage/JSCover
lib-cov
# Coverage directory used by tools like istanbul
coverage
# nyc test coverage
.nyc_output
# Grunt intermediate storage (http://gruntjs.com/creating-plugins#storing-task-files)
.grunt
# Bower dependency directory (https://bower.io/)
bower_components
# node-waf configuration
.lock-wscript
# Compiled binary addons (http://nodejs.org/api/addons.html)
build/Release
# Dependency directories
node_modules/
# Optional npm cache directory
.npm
# Optional eslint cache
.eslintcache
# Optional REPL history
.node_repl_history
# Output of 'npm pack'
*.tgz
# Yarn Integrity file
.yarn-integrity
# dotenv environment variables file
.env
.secret.local
toy*.ts
.DS_Store
build/
tests-build/
.firebase-emu/
*.log
.DS_Store
*.local
.secret.*
licensed/
.claude/
.gitmodules
View File
.nvmrc
+1
View File
@@ -0,0 +1 @@
24
.vscode/exensions.json
Vendored
+10
View File
@@ -0,0 +1,10 @@
{
"recommendations": [
"editorconfig.editorconfig",
"octref.vetur",
"redhat.vscode-yaml",
"dbaeumer.vscode-eslint",
"esbenp.prettier-vscode",
"streetsidesoftware.code-spell-checker"
]
}
.vscode/launch.json
Vendored
+144
View File
@@ -0,0 +1,144 @@
{
"version": "0.2.0",
"configurations": [
{
"name": "Debug Stand Alone Crawl",
"request": "launch",
"runtimeVersion": "24",
"runtimeArgs": [
],
"env": {
},
"cwd": "${workspaceFolder}",
"program": "build/stand-alone/crawl.js",
"skipFiles": [
"<node_internals>/**"
],
"type": "node",
"outputCapture": "std",
"preLaunchTask": "Backend:prepare",
"killBehavior": "forceful"
},
{
"name": "Debug Stand Alone Crawl + Browser",
"request": "launch",
"runtimeVersion": "24",
"runtimeArgs": [
],
"env": {
"DEBUG_BROWSER": "true"
},
"cwd": "${workspaceFolder}",
"program": "build/stand-alone/crawl.js",
"skipFiles": [
"<node_internals>/**"
],
"type": "node",
"outputCapture": "std",
"preLaunchTask": "Backend:prepare",
"killBehavior": "forceful"
},
{
"name": "Debug Stand Alone Crawl - EU",
"request": "launch",
"runtimeVersion": "24",
"runtimeArgs": [
],
"env": {
"GCLOUD_PROJECT": "reader-6b7dc",
"FIRESTORE_DATABASE": "reader-eu",
"GCP_STORAGE_BUCKET": "reader-eu"
},
"cwd": "${workspaceFolder}",
"program": "build/stand-alone/crawl.js",
"skipFiles": [
"<node_internals>/**"
],
"type": "node",
"outputCapture": "std",
"preLaunchTask": "Backend:prepare",
"killBehavior": "forceful"
},
{
"name": "Debug Stand Alone Search",
"request": "launch",
"runtimeVersion": "24",
"runtimeArgs": [
],
"env": {
// "GCLOUD_PROJECT": "reader-6b7dc",
// "PREFERRED_PROXY_COUNTRY": "us",
// "JINA_CRAWLER_OFFLOAD_ORIGIN": "https://r.jina.ai"
},
"cwd": "${workspaceFolder}",
"program": "build/stand-alone/search.js",
"skipFiles": [
"<node_internals>/**"
],
"type": "node",
"outputCapture": "std",
"preLaunchTask": "Backend:prepare",
"killBehavior": "forceful"
},
{
"name": "Debug Stand Alone Search + Offload",
"request": "launch",
"runtimeVersion": "24",
"runtimeArgs": [
],
"env": {
// "GCLOUD_PROJECT": "reader-6b7dc",
// "PREFERRED_PROXY_COUNTRY": "us",
// "JINA_CRAWLER_OFFLOAD_ORIGIN": "https://r.jina.ai"
},
"cwd": "${workspaceFolder}",
"program": "build/stand-alone/search.js",
"skipFiles": [
"<node_internals>/**"
],
"type": "node",
"outputCapture": "std",
"preLaunchTask": "Backend:prepare",
"killBehavior": "forceful"
},
{
"name": "Debug Stand Alone SERP",
"request": "launch",
"runtimeVersion": "24",
"runtimeArgs": [
],
"env": {
"GCLOUD_PROJECT": "reader-6b7dc",
// "OVERRIDE_GOOGLE_DOMAIN": "www.google.com.hk",
// "PREFERRED_PROXY_COUNTRY": "us"
},
"cwd": "${workspaceFolder}",
"program": "build/stand-alone/serp.js",
"skipFiles": [
"<node_internals>/**"
],
"type": "node",
"outputCapture": "std",
"preLaunchTask": "Backend:prepare",
"killBehavior": "forceful"
},
{
"name": "Attach",
"port": 9229,
"request": "attach",
"skipFiles": [
"<node_internals>/**"
],
"type": "node"
},
{
"name": "Attach by Process ID",
"processId": "${command:PickProcess}",
"request": "attach",
"skipFiles": [
"<node_internals>/**"
],
"type": "node"
},
]
}
.vscode/settings.json
Vendored
+32
View File
@@ -0,0 +1,32 @@
{
"editor.wordWrap": "on",
"editor.wordWrapColumn": 120,
"files.trimTrailingWhitespace": true,
"files.trimFinalNewlines": true,
"[javascript]": {
"editor.defaultFormatter": "vscode.typescript-language-features"
},
"[jsonc]": {
"editor.defaultFormatter": "vscode.json-language-features"
},
"[typescript]": {
"editor.defaultFormatter": "vscode.typescript-language-features"
},
"[json]": {
"editor.defaultFormatter": "vscode.json-language-features"
},
"[yaml]": {
"editor.defaultFormatter": "redhat.vscode-yaml"
},
"[markdown]": {
"files.trimTrailingWhitespace": false
},
"typescript.tsdk": "node_modules/typescript/lib",
"typescript.preferences.quoteStyle": "single",
"typescript.format.semicolons": "insert",
"typescript.preferences.importModuleSpecifier": "project-relative",
"typescript.locale": "en",
"cSpell.enabled": true,
"cSpell.words": [
],
}
.vscode/tasks.json
Vendored
+64
View File
@@ -0,0 +1,64 @@
{
"version": "2.0.0",
"tasks": [
{
"type": "npm",
"script": "build",
"group": "build",
"options": {
"cwd": "${workspaceFolder}"
},
"problemMatcher": [],
"label": "Backend:rebuild",
"detail": "Backend:rebuild"
},
{
"type": "typescript",
"options": {
"cwd": "${workspaceFolder}"
},
"tsconfig": "tsconfig.json",
"option": "watch",
"isBackground": true,
"problemMatcher": [
"$tsc-watch"
],
"group": "build",
"label": "Backend:build:watch"
},
{
"label": "Backend:docker-compose:up",
"type": "shell",
"command": "docker",
"args": [
"compose",
"up"
],
"group": "build",
"isBackground": true,
"presentation": {
"echo": true,
"reveal": "always",
"focus": false,
"panel": "shared"
},
"problemMatcher": {
"pattern": {
"regexp": ".*?"
},
"background": {
"activeOnStart": true,
"beginsPattern": "OKOKOK",
"endsPattern": ".*"
}
}
},
{
"label": "Backend:prepare",
"dependsOn": [
"Backend:docker-compose:up",
"Backend:build:watch"
],
}
]
}
CLAUDE.md
+126
View File
@@ -0,0 +1,126 @@
# CLAUDE.md
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
## Repository identity
This is the `oss` branch of `reader-saas`, the codebase behind `https://r.jina.ai` (URL → markdown) and `https://s.jina.ai` (search → markdown). It is published to https://github.com/jina-ai/reader. The MongoDB-backed SaaS storage layer is **not** in this branch — only the stateless and bucket-cached storage modes ship here.
Companion docs in this repo:
- `architecture.md` — engines, formatting profiles, abuse mitigation, deployment topology.
- `CONTRIBUTING.md` — full env var table, dev workflow, test policy.
- `cookbooks.md` — header recipes for common downstream pipelines (RAG, embedding, deep research, uploads, etc.).
## Common commands
Build / run:
- `npm run build` — runs `integrity-check.cjs` (requires `licensed/GeoLite2-City.mmdb` to exist) then `tsc -p .`. The integrity check is **not** optional — bare `tsc` will skip it.
- `npm run build:watch` — incremental build for the F5 debug flow.
- `npm start` — runs the compiled crawler entrypoint (`build/stand-alone/crawl.js`).
- `npm run dry-run` — `NODE_ENV=dry-run` boots `search.js`, resolves the DI graph, then exits via `finalizer.terminate()`. Used in the Dockerfile to warm `NODE_COMPILE_CACHE`.
Local services (`docker compose up -d`):
- `minio` on `:9000` (API) / `:9001` (console) — also aliased as `minio.dev.jina.ai`. Only needed when running with `BucketStorageLayer`.
VSCode F5 launches `Debug Stand Alone Crawl` (or `Search` / `SERP`), which runs `Backend:prepare` (docker compose + tsc watch) first. `--env-file=.secret.local` is loaded by Node directly — that file is gitignored and must be created locally.
Three independent stand-alone servers:
- `build/stand-alone/crawl.js` — mounts `CrawlerHost` (the `r.jina.ai` surface).
- `build/stand-alone/search.js` — mounts `SearcherHost`. At startup it deletes any registry entries tagged `'crawl'`, so search and crawl are mutually exclusive on a single process.
- `build/stand-alone/serp.js` — SERP-only.
Linting: `npm run lint` (eslint over `.js`/`.ts`).
## Tests
The repo uses the **Node.js built-in test runner**. Do not introduce Jest, Vitest, or similar.
- `npm run test:unit` — pure TS unit tests under `tests/unit/`. Compiled into `tests-build/` before running. No Docker required.
- `npm run test:e2e` — boots the real `CrawlStandAloneServer` via `serviceReady()` and hits `tests/e2e/*.test.ts` against it. Needs Docker services up and `.secret.local` configured.
- `npm test` — runs unit then e2e.
- `npm run test:coverage` — c8 across both suites; merges coverage from unit and e2e runs (`--no-clean` between them).
Single test:
```bash
tsc -p tests/tsconfig.json
node --test tests-build/unit/<file>.test.js
node --test --test-name-pattern '<regex>' tests-build/unit/<file>.test.js
```
For e2e, the test runner expects the crawl server already initialized — easier to run the full file via `node tests-build/run.js` after editing `run.ts` to filter, or just run the whole suite.
The e2e harness shuts down via `finalizer.teardown()` once the enqueued count matches the completed count; if you add async tests that don't go through `node:test`'s lifecycle, the process won't exit cleanly.
## Architecture
### DI: tsyringe + civkit
Every service is a `@singleton()` registered against the global `tsyringe` container. The graph is wired by side-effect: importing a module is what registers it. The conventional bootstrap is:
1. `import 'reflect-metadata';`
2. `import '../config';` — sets `AUTH_DTO_CLS` and `STORAGE_CLS` based on env vars (see below).
3. `container.resolve(...)` to get a fully-injected instance.
`AsyncService` (from `civkit`) is the base. Services emit `'ready'` after `dependencyReady()` resolves. `serviceReady()` waits until the entire graph is initialized. Always wait on `serviceReady()` before listening / running tests.
Don't construct services with `new` — go through `container.resolve` (or constructor injection). Adding a new service means: `@singleton()` + add it to a constructor that's already in the graph (or call `container.resolve` once).
### `src/config.ts` — runtime mode selection
The same code runs in two storage modes. `config.ts` swaps the implementation classes at import time:
- `AUTH_DTO_CLS`: `BaseAuthDTO` (the oss branch always uses the base DTO — no SaaS auth).
- `STORAGE_CLS`: defaults to `StorageLayer` (`db/noop-storage.ts` — every method returns nothing). With `GCP_STORAGE_*` → `BucketStorageLayer` (Stage 1: bucket-only cache).
This mirrors the "Progressive Clustering" stages in `architecture.md`. The noop layer means the same code paths exist in stateless mode — handlers always call `storage.findPageCache` etc., and the noop returns `undefined` so the request just falls through to a live fetch. Don't add `if (storage)` guards; rely on the noop.
When you change a method on the storage layer, you must update **both** implementations: `db/noop-storage.ts` and `db/bucket-storage.ts`.
### RPC routing: civkit + Koa
`src/services/registry.ts` is the routing core. It re-exports `Method`, `RPCMethod`, `Param`, `Ctx`, `RPCReflect` decorators and the `Context` type. Don't import these from `koa` or `civkit/civ-rpc` directly — go through `services/registry.ts` so you get the configured `RPCRegistry` and `ReaderEnvelope`.
`ReaderEnvelope` content-negotiates errors: JSON, text/markdown, and SSE branches all live there. If you add a new content type, add an error path here too — otherwise errors will fall back to the JSON envelope.
Endpoints are `RPCHost` subclasses in `src/api/`:
- `crawler.ts` — `CrawlerHost`, the `r.jina.ai` surface.
- `searcher.ts` — `SearcherHost`, the `s.jina.ai` surface (calls back into crawler).
- `serp.ts` — pure SERP.
### Worker threads
`ThreadedServiceRegistry` (`src/services/threaded.ts`) extends `civkit`'s `AbstractThreadedServiceRegistry`. CPU-heavy work (DOM manipulation, PDF parsing, markify) runs in worker threads. The registry sniffs hyperthreading and sets `maxWorkers` to `cpus.length` or `cpus.length / 2` accordingly. Workers share state via `PseudoTransfer`; do not use `postMessage` directly.
To make a method run in a worker, decorate it with `@Threaded()` — exported by `services/threaded.ts` alongside the RPC decorators. See `services/jsdom.ts` for the canonical pattern (its narrowing/parsing methods are all threaded).
### Engines and formatting profiles
The crawler is multi-engine and multi-profile:
- Engines (URL→HTML): `puppeteer.ts` (browser), `curl.ts` (curl-impersonate), `cf-browser-rendering.ts`. `auto` is the default and is implemented in `crawler.ts`.
- Formatting profiles (HTML→Markdown): `@mozilla/readability`, the `markify` rule-based engine (`services/markify.ts`), and `lm.ts` (ReaderLM v2).
- Header `x-respond-with` selects the profile; `x-engine` selects the engine.
The full surface — including `x-respond-timing`, `x-retain-images`, `x-retain-links`, `x-markdown-chunking`, and the body-field equivalents — lives in `src/dto/crawler-options.ts`. When users ask "what does header X do," that file is authoritative. `cookbooks.md` shows curated combinations.
When adding a new engine or profile, wire it into `CrawlerHost.crawl()` (the dispatch lives there) and register a singleton service for the implementation.
### Auth and rate limiting
Every public method takes an `auth: AUTH_DTO_CLS`. In this branch the base DTO accepts any request — there is no upstream auth check. Rate limiting goes through `storage.rateLimit(ctx, rpcReflect, auth)`, which the noop layer answers with an empty policy. The SaaS-mode enforcement lives outside this branch.
## Environment variables
The full table lives in `CONTRIBUTING.md`. Two shapes you'll see in practice:
- `.secret.local` — gitignored, loaded via Node's `--env-file` flag from `.vscode/launch.json`. Compiled / extracted by `bin/compile-secret.js` and `bin/extract-secret.js`.
- `SECRETS_COMBINED` — base64 JSON blob, the production shape. Decoded and merged by `services/envconfig.ts`.
Mode-selecting vars worth remembering when reading code: `GCP_STORAGE_ENDPOINT` + `GCP_STORAGE_BUCKET` (Stage 1 bucket cache), `NODE_ENV=dry-run` (DI-graph warmup, used in the Dockerfile), `DEBUG_BROWSER=true` (non-headless puppeteer), `JINA_CRAWLER_OFFLOAD_ORIGIN` (search → peer crawler cluster).
## Conventions specific to this repo
- Decorators (`@singleton`, `@Method`, `@Param`, `@Threaded`, etc.) require `experimentalDecorators` and `emitDecoratorMetadata` — both already on. Do not migrate to TC39 decorators.
- The `licensed/` folder holds external (non-redistributable) artifacts (GeoLite mmdbs, Source Han Sans font, gsa user-agent list). It's gitignored. CI fetches them inline; locally, run `npm run assets:download` (or `bash ./download-external-assets.sh` directly) — idempotent, skips files already present. The build's integrity check requires `licensed/GeoLite2-City.mmdb` to exist; if `npm run build` errors on it, you forgot to run the download. The folder is still named `licensed/` for source-code/Dockerfile compatibility — only the script and env vars use the "external" naming.
- Build output is `build/`, test build output is `tests-build/`. Both are gitignored. The Dockerfile copies pre-built `build/` rather than building inside the image.
- HTTP/2 cleartext (h2c) is used in production; `crawl.ts` / `search.ts` install an `http2.createServer` on top of the Koa callback. Each request gets its own `traceId` (random UUID) because h2c connections multiplex — don't rely on `x-cloud-trace-context` for per-request identity.
- Public assets in `public/` are served directly by the stand-alone server's `walkForAssets()` — drop a file in there, rebuild, it's served at `/<filename>`.
- The Reader API surface is entirely header- and body-driven; there are no path-based routes for options. When changing behavior, the right place is almost always `src/dto/crawler-options.ts` (parsing) plus `src/api/crawler.ts` (dispatch).
CONTRIBUTING.md
+132
View File
@@ -0,0 +1,132 @@
# Contributing to Reader
Thanks for your interest in contributing. This is the open source branch of the codebase that runs at `https://r.jina.ai` and `https://s.jina.ai`. The MongoDB-backed SaaS storage layer is not part of this branch — local development uses the stateless / bucket-cached modes only.
If you're not sure where to start, take a look at [architecture.md](./architecture.md) first.
## Local development
### Requirements
- **Node.js 22+** — earlier versions will not build.
- **Docker** *(optional)* — only needed if you want to run the bucket-cached storage mode against a local MinIO. Pure stateless mode needs nothing extra.
- **LibreOffice** *(optional)* — only needed if you want to test MS Office document handling locally.
### First-time setup
```bash
git clone git@github.com:jina-ai/reader.git
cd reader
npm install
# Optional: only if you want the local bucket cache
docker compose up -d
```
`docker compose up -d` starts:
| Service | Port(s) | Purpose |
| ------- | -------------- | ------------------------------------------------------------------ |
| `minio` | `9000`, `9001` | S3-compatible object storage for cached pages. Console on `:9001`. |
### Running the server
In VSCode, press `F5` to launch the debugger.
Or, after exporting the environment variables (see below):
```bash
docker compose up -d
npm run dev
```
### Useful scripts
- `npm run build` — TypeScript compile (also runs an integrity check).
- `npm run build:watch` — incremental build.
- `npm run start` — run the compiled `crawl` entrypoint.
- `npm run dry-run` — run `search.js` with `NODE_ENV=dry-run` to resolve the DI graph and exit. Used to warm `NODE_COMPILE_CACHE` in the Dockerfile.
- `npm run lint` — ESLint over `.js` / `.ts`.
## Environment variables
Reader picks up configuration from environment variables. The most relevant ones for local development are:
### Storage & data
| Variable | Notes |
| -------------------------- | ------------------------------------------------------------------------------------- |
| `GCP_STORAGE_ENDPOINT` | Object storage endpoint (use the local MinIO endpoint for dev). Enables Stage 1 bucket-cached mode. |
| `GCP_STORAGE_BUCKET` | Bucket name for cached objects. |
| `GCP_STORAGE_ACCESS_KEY` | MinIO root user locally. |
| `GCP_STORAGE_SECRET_KEY` | MinIO root password locally. |
| `GCP_STORAGE_REGION` | Optional; for parity with GCS. |
| `GCLOUD_PROJECT` | Alternative trigger for the bucket layer when combined with `GCP_STORAGE_ENDPOINT`. |
| `CACHE_LOCAL_STORAGE_ROOT` | Filesystem root for local cache (alternative to object storage in stateless modes). |
### Vendors & integrations
| Variable | Purpose |
| ------------------------------------------------- | -------------------------------------------------------------------- |
| `JINA_SERP_API_KEY` / `JINA_SERP_API_ORIGIN` | Jina SERP backend. |
| `JINA_SERP_API_POLICY` | SERP routing policy. |
| `SERPER_SEARCH_API_KEY` | serper.dev search backend. |
| `THORDATA_PROXY_URL` / `THORDATA_PROXY_URL_ALT` | Thordata residential proxy. |
| `THORDATA_SERP_API_KEY` | Thordata SERP API. |
| `BRIGHTDATA_PROXY_URL` / `BRIGHTDATA_ISP_PROXY_URL` / `BRIGHTDATA_SERP_API_KEY` | BrightData proxy + SERP. |
| `CLOUD_FLARE_API_KEY` | Required for the `cf-browser-rendering` engine. |
| `STRIPE_SECRET_KEY` / `STRIPE_WEBHOOK_KEY` | Billing integration. |
| `OPENAI_API_KEY` / `ANTHROPIC_API_KEY` / `OPENROUTER_API_KEY` / `GOOGLE_AI_STUDIO_API_KEY` / `REPLICATE_API_KEY` | LLM/VLM access. |
### Overrides & toggles
| Variable | Purpose |
| --------------------------------- | ------------------------------------------------------------------ |
| `PORT` | HTTP port. |
| `NODE_ENV` | `dry-run` is recognized for offline `search` runs. |
| `DEBUG_BROWSER` | Run headless Chrome with non-headless / debug behavior. |
| `OVERRIDE_CHROME_EXECUTABLE_PATH` | Use a specific Chrome binary instead of the bundled Puppeteer one. |
| `OVERRIDE_JINA_VLM_URL` | Point at a different VLM endpoint. |
| `OVERRIDE_READERLM_V` | Switch between ReaderLM versions. |
| `OVERRIDE_GOOGLE_DOMAIN` / `OVERRIDE_BING_DOMAIN` | Use a regional search domain. |
| `OVERRIDE_MANAGE_SERVER_URL` | Redirect calls to the management server. |
| `JINA_BOGO_SITES_RESORT_ORIGIN` | Origin for the bogo-sites resort list. |
| `JINA_CRAWLER_OFFLOAD_ORIGIN` | Offload crawler traffic to a peer cluster. |
| `PREFERRED_PROXY_COUNTRY` | Hint for proxy country selection. |
| `SLACK_REPORT_WEBHOOK_URL` | Slack channel for runtime reports. |
### `SECRETS_COMBINED`
You can pass a base64-encoded JSON object via `SECRETS_COMBINED` to bundle multiple variables into one. See `src/services/envconfig.ts`.
## Tests
The repo uses the Node.js built-in test runner (no Jest, no Vitest).
```bash
npm run test:unit # unit tests
npm run test:e2e # end-to-end tests (slower, hits docker services)
npm test # both
npm run test:unit:coverage
npm run test:e2e:coverage
npm run test:coverage # combined coverage report (c8)
```
Tests are written in TypeScript under `tests/` and compiled into `tests-build/` before running. The test entrypoints are `tests-build/run-unit.js` and `tests-build/run.js`.
## Submitting changes
1. Open an issue first if the change is non-trivial — it saves churn for both sides.
2. Keep PRs focused. A bug fix and a refactor in the same PR are harder to review and revert.
3. Run `npm run lint` and `npm test` before pushing.
4. Reference the issue in the PR description if one exists.
## Reporting issues
Bug reports are most useful when they include:
- The exact URL (for `r.jina.ai`) or query (for `s.jina.ai`).
- The request headers in use, especially any `x-*` overrides.
- The expected vs actual output.
Open an issue on GitHub and we'll take a look.
Dockerfile
+39
View File
@@ -0,0 +1,39 @@
# syntax=docker/dockerfile:1
FROM node:24 AS base
FROM base AS build-amd64
RUN apt-get update \
&& apt-get install -y wget gnupg \
&& wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | apt-key add - \
&& sh -c 'echo "deb [arch=amd64] http://dl.google.com/linux/chrome/deb/ stable main" >> /etc/apt/sources.list.d/google.list' \
&& apt-get update \
&& apt-get install -y libreoffice google-chrome-stable fonts-noto-cjk-extra fonts-noto-color-emoji fonts-liberation fonts-ipafont-gothic fonts-wqy-zenhei fonts-thai-tlwg fonts-kacst fonts-freefont-ttf libxss1 zstd libc++-dev \
--no-install-recommends \
&& rm -rf /var/lib/apt/lists/*
ENV OVERRIDE_CHROME_EXECUTABLE_PATH=/usr/bin/google-chrome-stable
FROM base AS build-arm64
RUN apt-get update \
&& apt-get install -y libreoffice chromium fonts-noto-cjk-extra fonts-noto-color-emoji fonts-liberation fonts-ipafont-gothic fonts-wqy-zenhei fonts-thai-tlwg fonts-kacst fonts-freefont-ttf libxss1 zstd libc++-dev \
--no-install-recommends \
&& rm -rf /var/lib/apt/lists/*
ENV OVERRIDE_CHROME_EXECUTABLE_PATH=/usr/bin/chromium
FROM build-${TARGETARCH} AS final
RUN groupadd -r jina
RUN useradd -g jina -G audio,video -m jina
USER jina
WORKDIR /app
COPY package.json package-lock.json ./
RUN npm ci
COPY build ./build
COPY public ./public
COPY licensed ./licensed
RUN rm -rf ~/.config/chromium && mkdir -p ~/.config/chromium
RUN NODE_COMPILE_CACHE=node_modules npm run dry-run
ENV NODE_COMPILE_CACHE=node_modules
ENV PORT=8080
EXPOSE 8080 8081
ENTRYPOINT ["node"]
CMD [ "build/stand-alone/crawl.js" ]
LICENSE
+193
View File
@@ -0,0 +1,193 @@
Copyright 2020-2024 Jina AI Limited. All rights reserved.
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
Copyright 2020-2021 Jina AI Limited
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
README.md
+288
View File
@@ -0,0 +1,288 @@
# Reader
[![codecov](https://codecov.io/gh/jina-ai/reader/branch/main/graph/badge.svg)](https://codecov.io/gh/jina-ai/reader)
[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/jina-ai/reader)
Your LLMs deserve better input.
Reader does two things:
- **Read**: It converts any URL to an **LLM-friendly** input with `https://r.jina.ai/https://your.url`. Get improved output for your agent and RAG systems at no cost.
- **Search**: It searches the web for a given query with `https://s.jina.ai/your+query`. This allows your LLMs to access the latest world knowledge from the web.
Check out [the live demo](https://jina.ai/reader#demo)
Or just visit these URLs (**Read**) https://r.jina.ai/https://github.com/jina-ai/reader, (**Search**) https://s.jina.ai/Who%20will%20win%202024%20US%20presidential%20election%3F and see yourself.
> Feel free to use Reader API in production. It is free, stable and scalable. We are maintaining it actively as one of the core products of Jina AI. [Check out rate limit](https://jina.ai/reader#pricing)
<img width="973" alt="image" src="https://github.com/jina-ai/reader/assets/2041322/2067c7a2-c12e-4465-b107-9a16ca178d41">
<img width="973" alt="image" src="https://github.com/jina-ai/reader/assets/2041322/675ac203-f246-41c2-b094-76318240159f">
> This repository is the open source branch of the codebase behind `https://r.jina.ai` and `https://s.jina.ai`. It runs in stateless or bucket-cached mode; the MongoDB-backed SaaS storage layer is not included here.
## Updates
- **2026-04** — Re-synchronized the open source branch with the SaaS code. The MongoDB-backed storage layer is stripped; the oss branch runs in stateless mode out of the box, with optional MinIO/S3-compatible bucket caching via `docker compose`. See [Local development](#local-development).
- **2025-12** — Storage layer decoupled and binary file uploads landed. PDFs and MS Office documents (Word, Excel, PowerPoint) can now be POSTed directly via the `file` body field — no need to host them first. See [cookbooks.md](./cookbooks.md#pdf-ms-office-and-raw-html-uploads).
- **2025-03** — Major refactor: Reader is no longer a Firebase application. The SaaS migrated off Firestore + Cloud Functions to a Cloud Run image with MongoDB Atlas, removing the platform-coupled bits and unblocking the local-Docker path above.
- **2024-05** — `s.jina.ai` launched, extending Reader from URL→markdown to search→markdown. PDFs added the same month — any URL ending in `.pdf` is parsed with PDF.js and returned as markdown.
- **2024-04** — Reader released and `r.jina.ai` went live as Jina AI's first SaaS API for converting URLs to LLM-friendly input.
## What Reader can read
- **Web pages** — rendered with headless Chrome, or fetched lightweight via `curl-impersonate`. Reader picks intelligently between the two.
- **PDFs** — any URL, parsed with PDF.js. [See this NASA PDF result](https://r.jina.ai/https://www.nasa.gov/wp-content/uploads/2023/01/55583main_vision_space_exploration2.pdf) vs [the original](https://www.nasa.gov/wp-content/uploads/2023/01/55583main_vision_space_exploration2.pdf).
- **MS Office documents** — Word, Excel, PowerPoint, converted via LibreOffice and then processed as HTML/PDF.
- **Images** — captioned by a vision-language model, so your downstream text-only LLM gets *just enough* hints to reason about them.
## Usage
### Using `r.jina.ai` for single URL fetching
Simply prepend `https://r.jina.ai/` to any URL. For example, to convert the URL `https://en.wikipedia.org/wiki/Artificial_intelligence` to an LLM-friendly input, use the following URL:
[https://r.jina.ai/https://en.wikipedia.org/wiki/Artificial_intelligence](https://r.jina.ai/https://en.wikipedia.org/wiki/Artificial_intelligence)
### [Using `r.jina.ai` for a full website fetching (Google Colab)](https://colab.research.google.com/drive/1uoBy6_7BhxqpFQ45vuhgDDDGwstaCt4P#scrollTo=5LQjzJiT9ewT)
### Using `s.jina.ai` for web search
Simply prepend `https://s.jina.ai/` to your search query. Note that if you are using this in the code, make sure to encode your search query first, e.g. if your query is `Who will win 2024 US presidential election?` then your url should look like:
[https://s.jina.ai/Who%20will%20win%202024%20US%20presidential%20election%3F](https://s.jina.ai/Who%20will%20win%202024%20US%20presidential%20election%3F)
Behind the scenes, Reader searches the web, fetches the top 5 results, visits each URL, and applies `r.jina.ai` to it. This is different from many `web search function-calling` in agent/RAG frameworks, which often return only the title, URL, and description provided by the search engine API. If you want to read one result more deeply, you have to fetch the content yourself from that URL. With Reader, `http://s.jina.ai` automatically fetches the content from the top 5 search result URLs for you (reusing the tech stack behind `http://r.jina.ai`). This means you don't have to handle browser rendering, blocking, or any issues related to JavaScript and CSS yourself.
### Using `s.jina.ai` for in-site search
Simply specify `site` in the query parameters such as:
```bash
curl 'https://s.jina.ai/When%20was%20Jina%20AI%20founded%3F?site=jina.ai&site=github.com'
```
### [Interactive Code Snippet Builder](https://jina.ai/reader#apiform)
We highly recommend using the code builder to explore different parameter combinations of the Reader API.
<a href="https://jina.ai/reader#apiform"><img width="973" alt="image" src="https://github.com/jina-ai/reader/assets/2041322/a490fd3a-1c4c-4a3f-a95a-c481c2a8cc8f"></a>
### Using request headers
You can control the behavior of the Reader API using request headers. The list below covers the most useful ones — for the full surface with up-to-date defaults and validation rules, see the live API docs at [https://r.jina.ai/docs](https://r.jina.ai/docs), or the source of truth in [`src/dto/crawler-options.ts`](./src/dto/crawler-options.ts).
- `x-respond-with` — select the output format.
- `markdown` returns markdown *without* going through `readability`
- `html` returns `documentElement.outerHTML`
- `text` returns `document.body.innerText`
- `screenshot` returns the URL of the webpage's screenshot
- `pageshot` similar to `screenshot` but tries to capture the whole page instead of just the viewport
- `frontmatter` returns **Markdown with a YAML frontmatter block**. The default plain-text response uses a custom `Title: …` / `URL Source: …` header format; `frontmatter` replaces that with a front matter block. Example:
```bash
curl -H 'X-Respond-With: frontmatter' 'https://r.jina.ai/https://example.com'
```
```markdown
---
title: "Example Domain"
description: "This domain is for use in illustrative examples."
url: "https://example.com/"
---
## Example Domain
This domain is for use in illustrative examples in documents. ...
```
- `markdown+frontmatter` — like `frontmatter` but covers the full page without readability filtering.
- `x-engine` — enforces a fetching engine: `browser` (headless Chrome), `curl` (lightweight, no JS), or `auto` (the default — Combined use of both browser and curl).
- `x-proxy-url` — route the traffic through your designated proxy.
- `x-cache-tolerance` — integer seconds; how stale a cached page is acceptable.
- `x-no-cache: true` — bypass the cached page (lifetime 3600s). Equivalent to `x-cache-tolerance: 0`.
- `x-target-selector` — a CSS selector. Reader returns content within the matched element instead of the full page. Useful when automatic content extraction misses what you want.
- `x-wait-for-selector` — a CSS selector. Reader waits until the matched element is rendered before returning. If `x-target-selector` is set, this can be omitted to wait for the same element.
- `x-timeout` — integer seconds (max 180). When set, Reader will not return early; it waits for network idle or until the timeout is reached.
- `x-max-tokens` — integer (≥500). Trim the response so it never exceeds this many tokens. Useful as a per-request guardrail when feeding a fixed-size context window — Reader truncates rather than rejects.
- `x-token-budget` — integer. Reject the request if the resulting content would exceed this many tokens. Use this when *over*-budget output is worse than no output (e.g. cost control). Ignored on the search endpoint.
- `x-respond-timing` — explicit control over *when* Reader is willing to return. Trade off latency against completeness:
- `html` — return as soon as the raw HTML lands. No JS execution, no waiting.
- `visible-content` — return the moment readable content is parseable. Lowest latency that still produces text.
- `mutation-idle` — wait for DOM mutations to settle for ≥0.2s. Good default for SPAs that lazy-render above the fold.
- `resource-idle` — wait for content-affecting resources to finish loading (≥0.5s quiet). The default heuristic for content-shaped requests.
- `media-idle` — wait for media (images, video, fonts) to also finish. Use with `screenshot` / `pageshot` / `vlm`.
- `network-idle` — full `networkidle0`. Slowest, most complete. Implied when `x-timeout` ≥ 20.
When omitted, Reader picks one based on `x-respond-with`, `x-timeout`, and `x-with-iframe`. See `presumedRespondTiming` in [src/dto/crawler-options.ts](./src/dto/crawler-options.ts) for the exact rules.
- `x-with-generated-alt: true` — caption images on the page with a VLM.
- `x-retain-images` — control how images survive into the output:
- `all` (default) — keep `![alt](url)` markdown for every image.
- `none` — drop images entirely.
- `alt` — keep alt text only, no URLs. Cheap on tokens; useful when the downstream LLM has no use for the image link.
- `x-retain-links` — control how links survive into the output:
- `all` (default) — keep `[text](url)` markdown.
- `none` — drop links entirely.
- `text` — keep link anchor text only, drop URLs. Best for embedding / semantic-index pipelines where URLs are noise.
- `gpt-oss` — emit citations in gpt-oss's `【{id}†...】` format and append a numbered URL footer (also auto-enables `x-with-links-summary`).
- `x-retain-media` — control how `<video>`, `<audio>`, and embedded video iframes (`<iframe>` from YouTube, Vimeo, Bilibili, etc.) appear in the output:
- `link` (default) — markdown link, e.g. `[Video 1](url)`. Embedded iframes are rewritten to their canonical watch URL. Respects `x-md-link-style`.
- `none` — drop media entirely; non-video iframes fall back to their inner text content.
- `text` — bare label only, e.g. `Video 1` or `Audio 1`. No URL.
- `image` — markdown image syntax, e.g. `![Video 1](url)`.
- `html` — the original HTML element with cosmetic attributes (`class`, `id`, `style`, `data-*`, `aria-*`) stripped. Embedded video iframes keep their original embed `src` rather than the canonical watch URL.
- `x-with-links-summary` / `x-with-images-summary` — append a deduplicated footer of all links / images to the output. Combine with `x-retain-links: text` or `x-retain-images: alt` to get inline anchor/alt text plus *one* canonical URL list at the end — convenient when you want the model to see URLs without paying for them inline. `x-with-links-summary: all` keeps every link instead of only the unique ones.
- `x-markdown-chunking` — opt-in semantic chunking of the markdown response. Returns a JSON array (or ``-delimited text) of chunks instead of one blob:
- `true` / `h1` … `h5` — heading-based split at the given heading level (e.g. `h3` chunks at `#`, `##`, and `###`).
- `structured` / `s1` … `s5` — block-level structured split. `s1` is coarsest, `s5` finest.
- `x-preset` — apply a pre-packaged option bundle for common scenarios. Preset values only take effect for options the caller does *not* set explicitly (via body or another header). See [cookbooks.md](./cookbooks.md#using-presets) for examples.
- `reader` — for displaying content to human users.
- `index` — for semantic indexing / embedding pipelines.
- `research` — for AI research agents needing structured, citable output.
- `agent` — for AI agents doing everyday browsing tasks.
- `spider` — for recursive site crawling with a full link inventory.
- `x-detach-invisibles` — detach elements with eventual `display:none` before snapshotting. Implies browser engine; disables caching.
- `x-set-cookie` — forward cookie settings. Requests with cookies are not cached.
- `x-md-*` — fine-tune markdown output (heading style, bullet markers, link style, etc.). See [src/dto/turndown-tweakable-options.ts](./src/dto/turndown-tweakable-options.ts).
### Using `r.jina.ai` for single page application (SPA) fetching
Many websites nowadays rely on JavaScript frameworks and client-side rendering, usually known as Single Page Applications (SPA). Thanks to [Puppeteer](https://github.com/puppeteer/puppeteer) and headless Chrome, Reader natively supports fetching these websites. However, due to specific approaches some SPAs are developed with, there may be some extra precautions to take.
#### SPAs with hash-based routing
By definition of the web standards, content after `#` in a URL is not sent to the server. To mitigate this, use `POST` with the `url` parameter in the body:
```bash
curl -X POST 'https://r.jina.ai/' -d 'url=https://example.com/#/route'
```
#### SPAs with preloading contents
Some SPAs (and even some non-SPAs) show preload content before later loading the main content dynamically. In this case, Reader may capture the preload content instead. Two ways to mitigate:
```bash
# wait for network idle or until timeout
curl 'https://r.jina.ai/https://example.com/' -H 'x-timeout: 10'
# wait for a specific element
curl 'https://r.jina.ai/https://example.com/' -H 'x-wait-for-selector: #content'
# combined use of both to wait for non-existent element (which means waiting for the full timeout duration)
curl 'https://r.jina.ai/https://example.com/' -H 'x-timeout: 30' -H 'x-wait-for-selector: non-existent-element'
```
### JSON mode
Use the accept-header to control the output format:
```bash
curl -H "Accept: application/json" https://r.jina.ai/https://en.m.wikipedia.org/wiki/Main_Page
```
### Generated alt
All images on a page that lack an `alt` tag can be auto-captioned by a VLM (vision-language model) and formatted as `![Image [idx]: [VLM_caption]](img_URL)`. This should give your downstream text-only LLM *just enough* hints to include those images in reasoning, selection, and summarization:
```bash
curl -H "X-With-Generated-Alt: true" https://r.jina.ai/https://en.m.wikipedia.org/wiki/Main_Page
```
## Cookbooks
For pipeline-specific recipes — RAG, semantic indexing, deep research, agentic browsing, visual snapshots, PDF/Office/HTML uploads, and more — see [cookbooks.md](./cookbooks.md). Each entry is a short curl example with the header combination that fits the use case and a paragraph explaining the trade-offs.
## Self-host with Docker
A prebuilt image of the open-source branch is published to GitHub Container Registry. It bundles headless Chrome, LibreOffice, and CJK fonts, so you can run Reader without building it yourself.
```bash
docker pull ghcr.io/jina-ai/reader:oss
```
### Run
The image exposes two ports:
- `8080` — **h2c** (HTTP/2 cleartext). Production-grade, multiplexed; this is what Cloud Run talks to. Plain `curl` won't speak it without `--http2-prior-knowledge`.
- `8081` — **HTTP/1.1** fallback. Same handler, same routes; use this from anything that doesn't speak h2c.
For a quick try-out from `curl` or a browser, map the HTTP/1.1 port:
```bash
docker run --rm -p 3000:8081 ghcr.io/jina-ai/reader:oss
# then: curl http://localhost:3000/https://example.com
```
For load-testing or production-shape traffic, map the h2c port instead (or both):
```bash
docker run --rm -p 3000:8080 -p 3001:8081 ghcr.io/jina-ai/reader:oss
```
With no extra config the container is fully stateless — every request hits the live URL, no cache, no rate limiting. That's the right default for a quick try-out, CI, or throwaway environments.
### Run with caching
Point Reader at an S3-compatible bucket to cache fetched pages and reuse them across requests:
```bash
docker run --rm -p 3000:8081 \
-e GCP_STORAGE_ENDPOINT=https://s3.example.com \
-e GCP_STORAGE_BUCKET=reader-cache \
-e GCP_STORAGE_ACCESS_KEY=... \
-e GCP_STORAGE_SECRET_KEY=... \
ghcr.io/jina-ai/reader:oss
```
See [CONTRIBUTING.md](./CONTRIBUTING.md#environment-variables) for the full env-var table.
## Local development
Requirements:
- nvm use
- Docker *(optional — only if you want a local MinIO bucket cache)*
```bash
git clone git@github.com:jina-ai/reader.git
cd reader
npm install
# Optional, for bucket-cached mode:
docker compose up -d
```
Then either press `F5` in VSCode to launch the debugger, or after setting up the appropriate environment variables:
```bash
npm run dev
```
For a deeper tour of the codebase — engines, formatting profiles, abuse alleviation, deployment topology — see [architecture.md](./architecture.md). For dev workflow, env vars, and tests, see [CONTRIBUTING.md](./CONTRIBUTING.md).
### Licensed assets
A few non-redistributable artifacts live in `licensed/` and are needed at build/runtime:
- `GeoLite2-City.mmdb` and `geolite2-asn.mmdb` — MaxMind GeoLite databases (geolocation + ASN lookups).
- `SourceHanSansSC-Regular.otf` — Source Han Sans (CJK rendering for PDFs/screenshots).
- `gsa_useragents.txt` — user-agent list used by the curl engine.
Fetch them in one shot:
```bash
npm run assets:download
```
The script (`download-external-assets.sh`) is idempotent — it skips files already present and exits 0 even on partial network failure. Set `FORCE_DOWNLOAD_EXTERNAL=1` to overwrite, or `SKIP_DOWNLOAD_EXTERNAL=1` to bypass entirely if you supply your own copies. The repo's CI fetches the same URLs inline; this script exists for local convenience.
## How it works
[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/jina-ai/reader)
## Having trouble on some websites?
Some sites push back against scrapers — bot challenges, geo blocks, stale CDN edges. A few knobs to try, in roughly increasing order of "this is bothering me":
- **Use an API key.** Anonymous traffic is the most aggressively rate-limited and lands in the lowest-trust pool. Authenticated requests get a higher quota and access to features like the internal proxy. Get one at [jina.ai/reader](https://jina.ai/reader#pricing).
- **Bypass the cache** with `-H 'x-no-cache: true'`. If a stale or already-blocked response got cached, this forces a fresh fetch.
- **Force the browser engine** with `-H 'x-engine: browser'`. The default `auto` engine prefers the lightweight curl path when it can; some sites only serve real content to a JS-capable browser.
- **Route through the SaaS proxy** with `-H 'x-proxy: auto'` (key required). Reader's hosted proxy pool rotates residential / datacenter IPs and handles common anti-bot challenges automatically. You can also pin a country, e.g. `x-proxy: us` (see [Geo- and locale-sensitive scraping](./cookbooks.md#geo--and-locale-sensitive-scraping)).
- **Bring your own proxy** with `-H 'x-proxy-url: <url>'`. As a last resort — when even the hosted proxy can't get through — buy a residential or ISP-grade proxy from a third-party provider (BrightData, Thordata, Oxylabs, etc.) and pass the URL directly. Supports `http`, `https`, `socks4`, `socks5`; for auth use `https://user:pass@host:port`.
If none of those help, please open an issue with the URL and the headers you tried — we'll take a look.
## License
Reader is backed by [Jina AI](https://jina.ai) and licensed under [Apache-2.0](./LICENSE).
README.wehub.md
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目：`jina-ai/reader`
- 原始仓库：https://github.com/jina-ai/reader
- 导入方式：上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
architecture.md
+83
View File
@@ -0,0 +1,83 @@
# Architecture
## Introduction
Jina Reader is an API-first SaaS application that turns URLs of web pages, PDFs, and other documents into markdown or images. It's built to help developers prepare data context for LLMs — now widely known as context engineering.
## Application Architecture
Jina Reader is a multi-threaded Node.js application.
- Web pages are rendered using a headless Chrome browser, with text content extracted via a stack of techniques (see [HTML to Markdown profiles](#multiple-html-to-markdown-profiles)).
- PDF parsing and rendering are done using PDF.js.
- MS Office documents are processed using LibreOffice.
## Stateless Core Features
### URL to Markdown / Image
Given a URL, Jina Reader fetches the content and renders it using headless Chrome if it's a web page (HTML/xHTML). If the content is a PDF, it uses PDF.js to parse and render. For MS Office documents, LibreOffice converts them to PDF + HTML first, after which the PDF/HTML path takes over.
Advanced options let you filter or manipulate the page content — CSS-selector-based filtering, custom JavaScript execution, custom proxy routing, and more.
### HTML to Markdown
Reader can also take raw HTML and convert it to markdown, using the same conversion pipeline as the URL-to-Markdown feature.
### PDF to Markdown / Image
Reader can take a PDF file, extract text content as markdown, and render each page as an image.
### MS Office to Markdown / Image
Reader can take MS Office documents (Word, Excel, PowerPoint) and convert them to markdown or images by first converting them to PDF/HTML using LibreOffice.
### Image to Text
Reader can take an image and produce a text description (captioning). This is built on the `jina-vlm` small vision-language model and can be extended to VQA tasks. Note: this is not exactly OCR.
## Multiple URL-to-HTML Engines
Reader supports several engines for fetching/rendering web pages to HTML.
### Browser
The most-used engine. The current implementation runs latest headless Chrome via the `puppeteer` library. It provides the most accurate rendering and can execute JavaScript on the page, which is essential for modern web pages.
### CURL
A lightweight engine that uses `curl-impersonate` to fetch the raw HTML of a web page. It does not execute JavaScript. Reader's implementation includes a simulated cookie layer to handle basic cookie-based redirection.
### CF-Browser-Rendering
Uses Cloudflare's Browser Rendering REST API for URL-to-HTML. Strict rate limits apply; this engine is meant for testing and as a fallback.
### Auto
The default. Reader intelligently uses the CURL and Browser engines in combination, based on content characteristics and request requirements.
## Multiple HTML-to-Markdown Profiles
Reader supports several profiles for converting HTML to markdown.
### `@mozilla/readability`
Readability is automatically used to clean HTML before converting to markdown. It produces a clean, readable version of the HTML content for many pages.
### Rule-based engine
A custom implementation inspired by the `turndown` library, with custom rules and plugins to convert HTML into markdown.
### ReaderLM v2
An experimental engine that uses a specifically trained small language model to convert HTML to markdown.
### ReaderLM v3 / JinaOCR / VLM
WIP / future engine that uses a vision-language model to convert webpage screenshots directly to markdown.
## Abuse Mitigation (SaaS)
- **Request filtering**: block requests targeting suspicious addresses.
- **Request throttling**: cap concurrent requests per page.
- **Anonymous-user pressure relief**: when one URL receives excessive anonymous traffic, temporarily block that website for anonymous users.
- **Excessive HTML nodes/depth**: fall back to HTML-to-text instead of markdown.
## Progressive Clustering
- **Stage 0**: fully stateless — no caching, no rate limit, no persistence.
- **Stage 1**: S3-like object storage for caching, no rate limit.
- **Stage 2**: MongoDB + S3-like object storage. MongoDB indexes the cached objects; rate limiting is available. This is the SaaS configuration and is not part of the open source branch.
## Vendor-Provided Features
- **Proxy**: Reader supports a built-in proxy provider for fetching content via a different IP.
- **SERP**: Reader primarily relies on external SERP providers for web search results.
- **VLM**: Reader relies on a vision-language model for image captioning. The current model is `gemini-2.5-flash-lite`, but it can be switched to any model with similar capabilities.
## Deployment Architecture
The SaaS version of Jina Reader is deployed as a Docker image on GCP Cloud Run. MongoDB Atlas is used for metadata indexing and rate limiting; Google Cloud Storage is used for cache data. Internal services and dependencies — such as billing, `jina-vlm`, and `readerlm-v2` — are reached over a private VPC peering link.
We run two independent clusters: **US** and **EU**. The US cluster spans 3 regions (`us-central1`, `us-east1`, `us-west1`); the EU cluster runs in 1 region (`europe-west1`).
Due to the high resource requirements of headless Chrome and LibreOffice, Reader is best deployed on serverless platforms that handle auto-scaling and resource management.
bin/compile-secret.js
Executable
+7
View File
@@ -0,0 +1,7 @@
#!/usr/bin/env node
const fs = require('fs');
const f = fs.readFileSync('.secret.local.json', 'utf8');
fs.writeFileSync('.secret.local', `SECRETS_COMBINED=${Buffer.from(JSON.stringify(JSON.parse(f)), 'utf-8').toString('base64')}`, 'utf8');
bin/extract-secret.js
Executable
+7
View File
@@ -0,0 +1,7 @@
#!/usr/bin/env node
const fs = require('fs');
const f = fs.readFileSync('.secret.local', 'utf8');
fs.writeFileSync('.secret.local.json', `${JSON.stringify(JSON.parse(Buffer.from(f.slice('SECRETS_COMBINED='.length), 'base64').toString('utf-8')), undefined, 2)}`, 'utf8');
cookbooks.md
+215
View File
@@ -0,0 +1,215 @@
# Reader Cookbooks
Recipes for shaping Reader's output to fit a specific downstream pipeline. The default output ("drop into an LLM and read") is fine for ad-hoc use; the recipes below trade defaults for token efficiency, latency, or compatibility with a specific consumer.
For the full list of headers and body fields these recipes pull from, see the [Using request headers](./README.md#using-request-headers) section of the README.
## Contents
- [Reader Cookbooks](#reader-cookbooks)
- [Contents](#contents)
- [Using presets](#using-presets)
- [RAG inference (the user will see what the LLM sees)](#rag-inference-the-user-will-see-what-the-llm-sees)
- [Semantic indexing (build embeddings; URLs are noise)](#semantic-indexing-build-embeddings-urls-are-noise)
- [Deep research (long-context model needs URLs, but only once)](#deep-research-long-context-model-needs-urls-but-only-once)
- [Visual snapshot / pageshot for multimodal reasoning](#visual-snapshot--pageshot-for-multimodal-reasoning)
- [Scrape a known template (article body only)](#scrape-a-known-template-article-body-only)
- [Inject a page script (click-to-reveal content)](#inject-a-page-script-click-to-reveal-content)
- [Iframes and shadow DOM](#iframes-and-shadow-dom)
- [Geo- and locale-sensitive scraping](#geo--and-locale-sensitive-scraping)
- [PDF, MS Office, and raw HTML uploads](#pdf-ms-office-and-raw-html-uploads)
- [PDF and Office files](#pdf-and-office-files)
- [Raw HTML](#raw-html)
## Using presets
`x-preset` is a one-header shortcut that bundles the options from the manual recipes below. Preset values only kick in for options the caller hasn't set explicitly — you can always override a single field.
| Preset | Best for | Key settings |
|---|---|---|
| `reader` | Displaying to humans | `respondWith: frontmatter`, `retainMedia: html`, `detachInvisibles`, `removeOverlay` |
| `index` | Embedding / vector stores | `retainLinks: text`, `retainImages: alt`, `retainMedia: none`, `markdownChunking: s3` |
| `research` | AI research agents | `respondWith: markdown+frontmatter`, `markdownChunking: h3`, all links/images/media |
| `agent` | Day-to-day AI agents | `respondWith: frontmatter`, `markdownChunking: h3`, `retainImages: alt` |
| `spider` | Recursive site crawling | `respondWith: markdown+frontmatter`, `markdownChunking: h3`, `withLinksSummary: all` |
```bash
# Semantic indexing with one header — equivalent to the manual recipe below
curl https://r.jina.ai/https://example.com/article \
-H 'x-preset: index'
# Override a single field (keep URLs even though the preset drops them)
curl https://r.jina.ai/https://example.com/article \
-H 'x-preset: index' \
-H 'x-retain-links: all'
```
## RAG inference (the user will see what the LLM sees)
The chat model needs to *cite* the source: image URLs render as inline images in the answer, link URLs become clickable references. Just use the default — no tuning needed.
```bash
curl https://r.jina.ai/https://example.com/article
```
Reader's defaults (`x-retain-images: all`, `x-retain-links: all`, `readability` filtering on) already produce the shape a chat model wants. Add `x-with-generated-alt: true` only if you need VLM captions for images that lack alt text.
## Semantic indexing (build embeddings; URLs are noise)
For chunk → embed → vector store, URLs are pure overhead — they bloat tokens and can pull semantically similar chunks apart by their slugs. Keep the readable text, drop the addresses, and chunk in one shot.
```bash
curl https://r.jina.ai/https://example.com/article \
-H 'Accept: application/json' \
-H 'x-retain-links: text' \
-H 'x-retain-images: alt' \
-H 'x-markdown-chunking: h3'
```
Anchor text and alt text are kept (they carry meaning); URLs are dropped. JSON + `x-markdown-chunking: h3` splits at `#`, `##`, and `###` — typically one chunk per subsection, which lines up well with embedding window sizes. Use `h2` for coarser chunks if the source has dense headings.
If headings are sparse or unreliable, switch to the structured (`s1` … `s5`) family — block-level splits that don't depend on heading discipline. `s3` is a reasonable starting point; `s2` is coarser, `s4` / `s5` are finer.
**Preset shortcut:** `x-preset: index` applies the same combination (with `s3` chunking).
## Deep research (long-context model needs URLs, but only once)
Deep-research / agentic models cite URLs but choke if every link reference appears inline a dozen times. They also rarely need image URLs — captions are enough. Move the URLs into a single canonical footer and keep only anchor text inline.
```bash
curl https://r.jina.ai/https://example.com/article \
-H 'x-retain-links: text' \
-H 'x-with-links-summary: true' \
-H 'x-retain-images: alt'
```
The model reads the article with clean anchor text inline, then consults the deduplicated link footer when it needs to attribute or fetch a source. For gpt-oss-style citation tokens (`【1†...】`) instead of plain anchor text, swap `x-retain-links: text` for `x-retain-links: gpt-oss` — it auto-enables the summary footer for you.
**Preset shortcut:** `x-preset: research` applies a similar setup (keeps all links inline with URLs and chunks at `h3`). Use `x-preset: spider` if you also want the full link inventory collected for further crawling.
## Visual snapshot / pageshot for multimodal reasoning
Capture the full rendered page — overlays gone, lazy-loaded media settled — for visual QA, archival, or feeding a multimodal model that prefers pixels over markdown.
```bash
curl https://r.jina.ai/https://example.com/article \
-H 'x-respond-with: pageshot' \
-H 'x-remove-overlay: true' \
-H 'x-timeout: 30'
```
`pageshot` captures the entire scrollable page (use `screenshot` for the viewport only). Reader automatically picks `media-idle` timing for shot/vlm requests, so images and fonts will be painted before the snapshot is taken. `x-remove-overlay: true` strips cookie banners and modal dialogs that would otherwise dominate the snapshot. The 30-second timeout bounds a misbehaving page without giving up too early.
## Scrape a known template (article body only)
When you control the source — a specific blog, docs site, or product catalog — you already know the DOM shape. Skip Reader's heuristics and target the element directly.
```bash
curl https://r.jina.ai/https://example.com/blog/post-slug \
-H 'x-target-selector: article.post-body' \
-H 'x-remove-selector: nav, .related-posts, .comments, footer'
```
`x-target-selector` doubles as a `wait-for-selector` — Reader won't return until the element appears, so this is also a guard against partial renders. `x-remove-selector` strips the navigation chrome and templated noise that repeats on every page of the same site.
## Inject a page script (click-to-reveal content)
Some pages hide the content you want behind a click — "Show transcript", "Read more", "Load comments". The `injectPageScript` body field runs JavaScript in the page on every navigation. Reader exposes a `window.waitForSelector(selector)` helper that resolves to the element once it appears, so the script reduces to one line.
YouTube's transcript is the canonical example. Wait for the transcript toggle, click it:
```bash
curl -F 'url=https://www.youtube.com/watch?v=dQw4w9WgXcQ' \
-F "injectPageScript=waitForSelector('ytd-video-description-transcript-section-renderer button').then((el) => el.click())" \
-H 'Accept: application/json' \
https://r.jina.ai/
```
After the click, Reader's own page-settle logic takes over — the headless browser dispatches a `mutationIdle` event on `document` whenever the `MutationObserver` has seen no DOM changes for 200ms, and Reader waits for that signal before extracting. So you don't need to chain anything onto the click; the transcript content will be in the DOM by the time Reader serializes the page.
A few notes:
- `injectPageScript` accepts an array — repeat the `-F` flag for multiple steps; each runs as a separate `frame.evaluate(...)` in order.
- Injecting a script disables Reader's early-return optimization, so set a sensible `x-timeout` if the click triggers slow async work.
- For content inside iframes (Twitter embeds, etc.), use `injectFrameScript` — same shape, runs in every frame rather than just the main one.
## Iframes and shadow DOM
By default Reader serializes only the main document. Iframe contents live in their own document, and shadow roots are separate node trees — neither is technically part of the page being serialized, so they're skipped. That's the right default for most articles, but it loses content on:
- Pages that embed real content via iframe (CodeSandbox / JSFiddle examples in docs, Notion / Airtable embeds, Twitter / YouTube cards inside articles).
- Modern component-heavy sites built on web components (Stencil, Lit, many design-system docs) where the actual text lives inside shadow roots.
Pull both into the main document for extraction:
```bash
curl 'https://r.jina.ai/https://example.com/docs-page' \
-H 'x-with-iframe: true' \
-H 'x-with-shadow-dom: true'
```
Use `x-with-iframe: quoted` instead of `true` to wrap iframe contents in a markdown blockquote — useful when you want the model to know "this came from an embed" rather than treat it as inline body copy.
A few notes:
- Both options violate the standard same-origin / encapsulation boundaries, so use them deliberately. They make the output longer and the request slower.
- Turning on either option forces `network-idle` timing — Reader waits for every iframe and resource to fully load. Combine with `x-timeout` (max 180s) to bound the wait.
- These flags are about *extraction*. If you only need to *interact* with a frame (click a button inside an iframe), reach for `injectFrameScript` instead — see the previous cookbook.
## Geo- and locale-sensitive scraping
Pricing pages, restricted content, and regional search results all change based on where the request appears to come from. Pin geography and language explicitly.
```bash
curl https://r.jina.ai/https://shop.example.com/product/123 \
-H 'x-proxy: de' \
-H 'x-locale: de-DE' \
-H 'x-set-cookie: country=DE; Path=/'
```
`x-proxy: de` exits through a German residential IP. `x-locale: de-DE` sets `navigator.language` and `Accept-Language` for the headless browser. `x-set-cookie` handles the third common gating mechanism — sites that override geo via a cookie. Note: any request with `x-set-cookie` skips the cache, so don't pair this with cache-warming pipelines.
Note: this requires a premium key.
## PDF, MS Office, and raw HTML uploads
Reader can ingest content you already have on hand — PDF, Word, Excel, PowerPoint, or raw HTML — no need to host it first. The same options that apply to web pages (`x-retain-images`, `x-markdown-chunking`, `Accept: application/json`) work on uploads too.
### PDF and Office files
Use the `file` body field — Reader sniffs the MIME type from the bytes, so the same field accepts `.pdf`, `.docx`, `.xlsx`, and `.pptx`. **Multipart** is the right default; it streams and avoids base64 overhead:
```bash
curl -X POST 'https://r.jina.ai/' \
-F 'file=@./report.pdf' \
-H 'Accept: application/json' \
-H 'x-markdown-chunking: s3'
```
**Selecting a single PDF page.** PDFs are returned as one document by default. To get one page, append `#N` (1-indexed) to the optional `url` field:
```bash
curl -X POST 'https://r.jina.ai/' \
-F 'file=@./long-report.pdf' \
-F 'page=7' \
-H 'Accept: application/json'
```
Subsequent page requests reuse the cached parse, so paging through a long document is cheap after the first call.
Notes:
- File requests are cached by sha256 of the bytes — re-uploading the same file is a hit.
- Office files round-trip through LibreOffice. If fidelity matters (complex Excel layouts, custom slide masters), export to PDF on your side first and upload that.
### Raw HTML
Send HTML you already have via the `html` body field — Reader skips the fetcher and runs the same conversion pipeline:
```bash
curl -X POST 'https://r.jina.ai/' \
-H 'Content-Type: application/json' \
-d '{"html": "<html>...</html>", "url": "https://example.com/source"}'
```
`url` is optional but recommended whenever the HTML contains relative links or images — Reader uses it as the base for resolving them.
docker-compose.yml
+31
View File
@@ -0,0 +1,31 @@
services:
minio:
image: minio/minio
hostname: minio
ports:
- 9001:9001
- 9000:9000
environment:
MINIO_ROOT_USER: minio
MINIO_ROOT_PASSWORD: minio123
volumes:
- minio-data:/data/minio
command: server /data/minio --console-address ":9001"
healthcheck:
test:
[
'CMD',
'curl',
'-f',
'http://localhost:9001/minio/health/live'
]
interval: 30s
timeout: 20s
retries: 3
networks:
default:
aliases:
- minio.dev.jina.ai
volumes:
minio-data:
download-external-assets.sh
Executable
+50
View File
@@ -0,0 +1,50 @@
#!/usr/bin/env bash
# Downloads external (non-redistributable) artifacts into ./licensed/. Run via
# `npm run assets:download` (or `bash ./download-external-assets.sh`).
# Idempotent: existing files are skipped unless FORCE_DOWNLOAD_EXTERNAL=1 is set.
# Set SKIP_DOWNLOAD_EXTERNAL=1 to bypass.
set -u
if [ "${SKIP_DOWNLOAD_EXTERNAL:-}" ]; then
echo "[download-external] SKIP_DOWNLOAD_EXTERNAL set, skipping."
exit 0
fi
cd "$(dirname "$0")"
mkdir -p licensed
ARTIFACTS=(
"GeoLite2-City.mmdb|https://raw.githubusercontent.com/P3TERX/GeoLite.mmdb/download/GeoLite2-City.mmdb"
"geolite2-asn.mmdb|https://cdn.jsdelivr.net/npm/@ip-location-db/geolite2-asn-mmdb/geolite2-asn.mmdb"
"SourceHanSansSC-Regular.otf|https://raw.githubusercontent.com/adobe-fonts/source-han-sans/refs/heads/release/OTF/SimplifiedChinese/SourceHanSansSC-Regular.otf"
"gsa_useragents.txt|https://raw.githubusercontent.com/searxng/searxng/refs/heads/master/searx/data/gsa_useragents.txt"
)
failed=0
for entry in "${ARTIFACTS[@]}"; do
name="${entry%%|*}"
url="${entry#*|}"
dest="licensed/$name"
if [ -z "${FORCE_DOWNLOAD_EXTERNAL:-}" ] && [ -s "$dest" ]; then
echo "[download-external] skip $name (already present)"
continue
fi
echo "[download-external] get $name"
if curl -fsSL --retry 3 --retry-delay 2 -o "$dest.partial" "$url"; then
mv "$dest.partial" "$dest"
else
rm -f "$dest.partial"
echo "[download-external] failed $name" >&2
failed=$((failed + 1))
fi
done
if [ "$failed" -gt 0 ]; then
echo "[download-external] $failed artifact(s) failed; \`npm run build\` will fail until they are present. Re-run with FORCE_DOWNLOAD_EXTERNAL=1 npm run assets:download." >&2
fi
# Don't fail npm install on transient network issues.
exit 0
fixtures/sample.html
+96
View File
@@ -0,0 +1,96 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>Test Article: Web Crawling Guide</title>
<meta name="description" content="A comprehensive guide to web crawling techniques">
<meta property="article:published_time" content="2024-01-15T10:00:00Z">
</head>
<body>
<header>
<nav id="site-nav">
<a href="/home">Home</a>
<a href="/about">About</a>
</nav>
</header>
<div class="ads banner">
<p>Buy our product!</p>
</div>
<article id="main-content">
<h1>Web Crawling Guide</h1>
<p>This is the <strong>main content</strong> of the article about web crawling.</p>
<h2>Section One: Basics</h2>
<p>Web crawling involves <em>fetching pages</em> from the internet automatically.</p>
<p>Here is a <a href="https://example.com/crawling">link to crawling docs</a> and
another <a href="https://example.org/robots">link about robots.txt</a>.</p>
<h2>Section Two: Images</h2>
<p>Below are some images used in crawling:</p>
<img src="https://example.com/spider.png" alt="A spider crawling the web">
<img src="https://example.com/network.jpg" alt="Network diagram">
<img src="https://example.com/noalt.png">
<h2>Section Three: Lists</h2>
<ul>
<li>Fetch HTML pages</li>
<li>Parse links</li>
<li>Follow links recursively</li>
</ul>
<blockquote>
<p>The web is a graph, not a tree.</p>
</blockquote>
<h2>Section Four: Code</h2>
<pre><code>const crawler = new Crawler();
crawler.start('https://example.com');</code></pre>
<h2>Section Five: Tables</h2>
<table>
<thead><tr><th>Method</th><th>Speed</th><th>Accuracy</th></tr></thead>
<tbody>
<tr><td>Browser</td><td>Slow</td><td>High</td></tr>
<tr><td>Curl</td><td>Fast</td><td>Medium</td></tr>
</tbody>
</table>
<h2>Section Six: Dividers</h2>
<p>Content before the rule.</p>
<hr>
<p>Content after the rule.</p>
<h2>Section Seven: More Resources</h2>
<img src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" alt="Tiny pixel">
<p>See also the <a href="/guide/advanced">advanced guide</a> and <a href="sibling-page">sibling page</a>.</p>
<h2>Section Eight: Deferred Content</h2>
<div id="deferred-content" class="loading" data-loaded="false">Loading deferred section, please wait…</div>
<script>
// Simulates a lazily-rendered widget: real browsers see the loaded
// state after ~100ms; HTML-only consumers see the static fallback
// above. Useful for exercising mutation-idle / wait-for-selector.
setTimeout(function () {
var el = document.getElementById('deferred-content');
if (!el) { return; }
el.classList.remove('loading');
el.classList.add('loaded');
el.setAttribute('data-loaded', 'true');
el.textContent = 'Deferred content fully loaded after 100ms.';
}, 100);
</script>
</article>
<aside class="sidebar">
<p>Related articles sidebar content</p>
<a href="https://example.com/related">Related Article</a>
</aside>
<footer>
<p>Copyright 2024</p>
<a href="/privacy">Privacy Policy</a>
</footer>
</body>
</html>
integrity-check.cjs
Executable
+11
View File
@@ -0,0 +1,11 @@
#!/usr/bin/env node
const fs = require('fs');
const path = require('path');
const file = path.resolve(__dirname, 'licensed/GeoLite2-City.mmdb');
if (!fs.existsSync(file)) {
console.error(`Integrity check failed: ${file} does not exist.`);
process.exit(1);
}
package-lock.json
Generated
+7996
View File
File diff suppressed because it is too large Load Diff
package.json
+81
View File
@@ -0,0 +1,81 @@
{
"name": "reader",
"version": "0.5.0",
"scripts": {
"lint": "eslint --ext .js,.ts .",
"assets:download": "bash ./download-external-assets.sh",
"build": "node ./integrity-check.cjs && tsc -p .",
"build:watch": "tsc --watch",
"build:clean": "rm -rf ./build",
"serve": "npm run build && npm run start",
"debug": "npm run build && npm run dev",
"start": "node ./build/stand-alone/crawl.js",
"dry-run": "NODE_ENV=dry-run node ./build/stand-alone/search.js",
"test:unit": "tsc -p tests/tsconfig.json && node tests-build/run-unit.js",
"test:e2e": "tsc -p tests/tsconfig.json && node tests-build/run.js",
"test": "npm run test:unit && npm run test:e2e",
"test:unit:coverage": "tsc -p tests/tsconfig.json && c8 node tests-build/run-unit.js",
"test:e2e:coverage": "tsc -p tests/tsconfig.json && c8 node tests-build/run.js",
"test:coverage": "tsc -p tests/tsconfig.json && c8 --reporter=none node tests-build/run-unit.js && c8 --no-clean --reporter=none node tests-build/run.js && c8 report"
},
"engines": {
"node": ">=22.15"
},
"main": "build/index.js",
"types": "./build/stand-alone/crawl.d.ts",
"exports": {
".": "./build/stand-alone/crawl.js",
"./crawl": "./build/stand-alone/crawl.js",
"./search": "./build/stand-alone/search.js",
"./serp": "./build/stand-alone/serp.js"
},
"dependencies": {
"@esm2cjs/normalize-url": "^8.0.0",
"@koa/bodyparser": "^6.1.0",
"@mozilla/readability": "^0.6.0",
"@napi-rs/canvas": "^0.1.68",
"@nomagick/mathml-to-latex": "^1.5.3",
"@nomagick/node-libcurl-impersonate": "^5.0.4",
"civkit": "^0.10.0-ca17c8a",
"dayjs": "^1.11.9",
"franc": "^6.2.0",
"generic-pool": "^3.9.0",
"google-auth-library": "^10.6.2",
"jose": "^5.1.0",
"koa": "^3.1.2",
"linkedom": "^0.18.4",
"lodash": "^4.18.1",
"lru-cache": "^11.0.2",
"mathml-to-latex": "^1.5.0",
"maxmind": "^4.3.18",
"minio": "^8.0.7",
"pdfjs-dist": "^4.10.38",
"puppeteer": "^24.42.0",
"robots-parser": "^3.0.1",
"set-cookie-parser": "^2.6.0",
"stripe": "^11.11.0",
"svg2png-wasm": "^1.4.1",
"tiktoken": "^1.0.16",
"tld-extract": "^2.1.0",
"tslib": "^2.8.1",
"tsyringe": "^4.10.0",
"undici": "^7.24.7"
},
"devDependencies": {
"@types/koa": "^2.15.0",
"@types/lodash": "^4.17.24",
"@types/node": "^24.12.0",
"@types/set-cookie-parser": "^2.4.7",
"@types/supertest": "^7.2.0",
"@typescript-eslint/eslint-plugin": "^5.12.0",
"@typescript-eslint/parser": "^5.12.0",
"c8": "^10.1.3",
"eslint": "^8.9.0",
"openai": "^4.20.0",
"pino-pretty": "^13.0.0",
"replicate": "^0.16.1",
"supertest": "^7.2.2",
"typescript": "^5.5.4"
},
"private": true
}
public/favicon.ico
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 14 KiB

public/robots.txt
+14
View File
@@ -0,0 +1,14 @@
# ───── Live user-initiated fetches (allow) ─────
User-agent: Claude-User
User-agent: ChatGPT-User
User-agent: Perplexity-User
User-agent: Meta-ExternalFetcher
User-agent: DuckAssistBot
User-agent: MistralAI-User
User-agent: kagi-fetcher
User-agent: Manus-User
Allow: /
# ───── Everything else (block) ─────
User-agent: *
Disallow: /
src/3rd-party/anthropic.ts
Vendored
+177
View File
@@ -0,0 +1,177 @@
import { HTTPService, HTTPServiceRequestOptions } from 'civkit/http';
import { InputServerEventStream } from '../lib/transform-server-event-stream';
import _ from 'lodash';
export interface ClaudeSamplingParameters {
prompt: string;
model: string;
stop_sequences: string[];
max_tokens_to_sample: number;
stream?: boolean;
temperature?: number;
top_k?: number;
top_p?: number;
tags?: {
[key: string]: string;
};
};
export interface ClaudeMessagingParameters {
model: string;
messages: {
role: string;
content: string | Array<
{ type: 'text', text: string; } |
{
type: 'image', source: {
type: 'base64';
media_type: string;
data: string;
};
} |
{ type: 'tool_use'; id: string; name: string; input: any; } |
{
type: 'tool_result';
tool_use_id: string;
is_error?: boolean;
content?: Array<{ type: 'text', text: string; } |
{
type: 'image', source: {
type: 'base64';
media_type: string;
data: string;
};
}>;
}
>;
}[];
stop_sequences: string[];
max_tokens: number;
stream?: boolean;
temperature?: number;
top_k?: number;
top_p?: number;
metadata?: {
[key: string]: string;
};
};
export interface ClaudeCompletionResponse {
completion: string;
stop: string | null;
stop_reason: "stop_sequence" | "max_tokens";
log_id: string;
};
export interface ClaudeMessagingResponse {
id: string;
type: 'message';
role: 'assistant';
content: Array<{
type: 'text';
text: string;
} | { type: 'tool_use', id: string, name: string, input: any; }>;
model: string;
stop_reason: string | 'end_turn' | 'max_tokens' | 'stop_sequence' | null;
stop_sequence: string | null;
usage: {
input_tokens: number;
output_tokens: number;
};
};
export const CLAUDE_MODELS = [
'claude-instant-1',
'claude-2',
'claude-v1',
'claude-v1-100k',
'claude-instant-v1',
'claude-instant-v1-100k',
'claude-v1.3',
'claude-v1.3-100k',
'claude-v1.2',
'claude-v1.0',
'claude-instant-v1.1',
'claude-instant-v1.1-100k',
'claude-instant-v1.0',
] as const;
export type CLAUDE_MODELS = typeof CLAUDE_MODELS[number];
export class AnthropicHTTP extends HTTPService {
name = 'Anthropic';
supportedCompletionModels = [
...CLAUDE_MODELS
] as string[];
constructor(public apiKey: string) {
super('https://api.anthropic.com');
this.baseHeaders['X-API-Key'] = `${apiKey}`;
this.baseHeaders['Anthropic-Version'] = `2023-06-01`;
this.baseOptions.timeout = 1000 * 60 * 30 * 0.5;
}
complete<T extends ClaudeSamplingParameters>(payload: T, opts?: typeof this['baseOptions']) {
return this.postJson<typeof payload['stream'] extends true ?
InputServerEventStream : ClaudeCompletionResponse>('/v1/complete', payload,
{ ...opts, responseType: payload.stream ? 'stream' : 'json' }
);
}
chatCompletions<T extends Partial<ClaudeSamplingParameters> & {
prompt: string;
system?: string;
}>(payload: T, opts?: typeof this['baseOptions']) {
const finalPayload: T = {
stop_sequences: ['\n\nHuman:'],
max_tokens_to_sample: 1000,
...(payload as any),
};
if (!finalPayload.prompt.startsWith('\n\nHuman:')) {
finalPayload.prompt = `\n\nHuman: ${finalPayload.prompt}\n\nAssistant:`;
}
if (finalPayload.system) {
finalPayload.prompt = `\n\nHuman: ${finalPayload.system}${finalPayload.prompt}`;
delete finalPayload.system;
}
return this.complete(finalPayload as ClaudeSamplingParameters, opts);
}
messageComplete<T extends Partial<ClaudeMessagingParameters>>(payload: T, opts?: typeof this['baseOptions']) {
return this.postJson<typeof payload['stream'] extends true ?
InputServerEventStream : ClaudeMessagingResponse>('/v1/messages', payload,
{ ...opts, responseType: payload.stream ? 'stream' : 'json' }
);
}
override async __processResponse(options: HTTPServiceRequestOptions, r: Response): Promise<any> {
if (r.status !== 200) {
throw await r.json();
}
const s = await super.__processResponse(options, r);
if (options.responseType === 'stream') {
const parseStream = new InputServerEventStream();
s.pipe(parseStream);
parseStream.once('end', () => {
if (!s.readableEnded) {
r.body?.cancel();
}
});
return parseStream;
}
return s;
}
}
src/3rd-party/cloud-flare.ts
Vendored
+22
View File
@@ -0,0 +1,22 @@
import { HTTPService } from 'civkit/http';
import _ from 'lodash';
export class CloudFlareHTTP extends HTTPService {
name = 'CloudFlare';
constructor(public accountId: string, public apiToken: string) {
super(`https://api.cloudflare.com/client/v4/accounts/${accountId}`);
this.baseHeaders['Authorization'] = `Bearer ${apiToken}`;
this.baseOptions.timeout = 1000 * 60 * 30 * 0.5;
}
fetchBrowserRenderedHTML(input: {
url: string;
rejectResourceTypes?: string[];
rejectRequestPattern?: string[];
}, opts?: typeof this['baseOptions']) {
return this.postJson<{ success: true; result: string; }>('/browser-rendering/content', input, { responseType: 'json', ...opts });
}
}
src/3rd-party/common-serp.ts
Vendored
+179
View File
@@ -0,0 +1,179 @@
import { HTTPService, HTTPServiceRequestOptions } from 'civkit/http';
import _ from 'lodash';
import { Agent, RetryAgent } from 'undici';
import secretExposer from '../services/envconfig';
import { parseJSONText } from 'civkit/vectorize';
export interface CommonSerpWebResponse {
organic: {
title: string;
link: string;
description: string;
display_link?: string;
siteLinks?: { title: string; link: string; }[];
position: number,
}[];
}
export interface CommonSerpImageResponse {
images: {
title: string;
image: string;
link: string;
source?: string;
image_alt?: string;
}[];
}
export interface CommonSerpNewsResponse {
news: {
title: string;
link: string;
date?: string;
description?: string;
image?: string;
source?: string;
}[];
}
export type CommonSerpSearchResponse = CommonSerpWebResponse | CommonSerpImageResponse | CommonSerpNewsResponse;
export abstract class CommonSerpBaseHTTP extends HTTPService {
get name() {
return this.constructor.name;
}
abstract queryJSON(url: string, opts?: typeof this['baseOptions']): Promise<CommonSerpSearchResponse>;
abstract queryHTML(url: string, opts?: typeof this['baseOptions']): Promise<string>;
}
// For retry of `Retry-After` header
const overrideAgent = new RetryAgent(
new Agent({
allowH2: true
}),
{
statusCodes: [429, 503, 500],
maxRetries: 3,
retryAfter: true,
minTimeout: 200,
}
);
export class ThorDataSerp extends CommonSerpBaseHTTP {
constructor(public apiKey: string) {
super('https://scraperapi.thordata.com');
Reflect.set(this.baseOptions, 'dispatcher', overrideAgent);
}
async queryJSON(url: string, opts?: this['baseOptions'] | undefined) {
const parsed = new URL(url);
parsed.searchParams.set('json', '1');
const r = await this.postJson('/request', {
url: parsed.href,
}, {
...opts,
responseType: 'json',
headers: {
...opts?.headers,
'Authorization': `Bearer ${this.apiKey}`,
// This is wrong but the server expects it
'Content-Type': 'application/json',
}
});
return r.parsed.data.result;
}
async queryHTML(url: string, opts?: this['baseOptions'] | undefined) {
const parsed = new URL(url);
parsed.searchParams.set('json', '0');
const r = await this.postJson('/request', {
url: parsed.href,
}, {
...opts,
responseType: 'json',
headers: {
...opts?.headers,
'Authorization': `Bearer ${this.apiKey}`,
// This is wrong but the server expects it
'Content-Type': 'text/html',
}
});
return r.parsed.html;
}
override async __processResponse(options: HTTPServiceRequestOptions, r: Response): Promise<any> {
if (r.status !== 200) {
throw await r.json();
}
return super.__processResponse(options, r);
}
}
export class BrightDataSerp extends CommonSerpBaseHTTP {
public apiKey: string;
public zone = 'serp_api1';
constructor(apiKey: string) {
super('https://api.brightdata.com');
this.apiKey = apiKey.split(':')[0];
this.zone = apiKey.split(':')[1] || this.zone;
Reflect.set(this.baseOptions, 'dispatcher', overrideAgent);
}
async queryJSON(url: string, opts?: this['baseOptions'] | undefined) {
const parsed = new URL(url);
parsed.searchParams.set('brd_json', '1');
const r = await this.postJson('/request', {
url: parsed.href,
zone: this.zone,
format: 'json'
}, {
...opts,
responseType: 'json',
headers: {
...opts?.headers,
'Authorization': `Bearer ${this.apiKey}`,
}
});
return parseJSONText(r.parsed.body);
}
async queryHTML(url: string, opts?: this['baseOptions'] | undefined) {
const parsed = new URL(url);
parsed.searchParams.set('json', '0');
const r = await this.postJson('/request', {
url: parsed.href,
zone: this.zone,
format: 'raw'
}, {
...opts,
responseType: 'json',
headers: {
...opts?.headers,
'Authorization': `Bearer ${this.apiKey}`,
// This is wrong but the server expects it
'Content-Type': 'text/html',
}
});
return r.parsed.body;
}
}
export const commonSerpClients: CommonSerpBaseHTTP[] = [];
if (secretExposer.THORDATA_SERP_API_KEY) {
commonSerpClients.push(new ThorDataSerp(secretExposer.THORDATA_SERP_API_KEY));
}
if (secretExposer.BRIGHTDATA_SERP_API_KEY) {
commonSerpClients.push(new BrightDataSerp(secretExposer.BRIGHTDATA_SERP_API_KEY));
}
export default commonSerpClients;
src/3rd-party/google-gemini.ts
Vendored
+774
View File
@@ -0,0 +1,774 @@
import { HTTPService, HTTPServiceRequestOptions } from 'civkit/http';
import _ from 'lodash';
import { ProxyAgent } from 'undici';
import { InputServerEventStream } from '../lib/transform-server-event-stream';
import { Readable } from 'stream';
/**
* @license
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
/**
* Possible roles.
* @public
*/
export const POSSIBLE_ROLES = ["user", "model", "function", "system"] as const;
/**
* Harm categories that would cause prompts or candidates to be blocked.
* @public
*/
export enum HarmCategory {
HARM_CATEGORY_UNSPECIFIED = "HARM_CATEGORY_UNSPECIFIED",
HARM_CATEGORY_HATE_SPEECH = "HARM_CATEGORY_HATE_SPEECH",
HARM_CATEGORY_SEXUALLY_EXPLICIT = "HARM_CATEGORY_SEXUALLY_EXPLICIT",
HARM_CATEGORY_HARASSMENT = "HARM_CATEGORY_HARASSMENT",
HARM_CATEGORY_DANGEROUS_CONTENT = "HARM_CATEGORY_DANGEROUS_CONTENT",
}
/**
* Threshold above which a prompt or candidate will be blocked.
* @public
*/
export enum HarmBlockThreshold {
// Threshold is unspecified.
HARM_BLOCK_THRESHOLD_UNSPECIFIED = "HARM_BLOCK_THRESHOLD_UNSPECIFIED",
// Content with NEGLIGIBLE will be allowed.
BLOCK_LOW_AND_ABOVE = "BLOCK_LOW_AND_ABOVE",
// Content with NEGLIGIBLE and LOW will be allowed.
BLOCK_MEDIUM_AND_ABOVE = "BLOCK_MEDIUM_AND_ABOVE",
// Content with NEGLIGIBLE, LOW, and MEDIUM will be allowed.
BLOCK_ONLY_HIGH = "BLOCK_ONLY_HIGH",
// All content will be allowed.
BLOCK_NONE = "BLOCK_NONE",
}
/**
* Probability that a prompt or candidate matches a harm category.
* @public
*/
export enum HarmProbability {
// Probability is unspecified.
HARM_PROBABILITY_UNSPECIFIED = "HARM_PROBABILITY_UNSPECIFIED",
// Content has a negligible chance of being unsafe.
NEGLIGIBLE = "NEGLIGIBLE",
// Content has a low chance of being unsafe.
LOW = "LOW",
// Content has a medium chance of being unsafe.
MEDIUM = "MEDIUM",
// Content has a high chance of being unsafe.
HIGH = "HIGH",
}
/**
* Reason that a prompt was blocked.
* @public
*/
export enum BlockReason {
// A blocked reason was not specified.
BLOCKED_REASON_UNSPECIFIED = "BLOCKED_REASON_UNSPECIFIED",
// Content was blocked by safety settings.
SAFETY = "SAFETY",
// Content was blocked, but the reason is uncategorized.
OTHER = "OTHER",
}
/**
* Reason that a candidate finished.
* @public
*/
export enum FinishReason {
// Default value. This value is unused.
FINISH_REASON_UNSPECIFIED = "FINISH_REASON_UNSPECIFIED",
// Natural stop point of the model or provided stop sequence.
STOP = "STOP",
// The maximum number of tokens as specified in the request was reached.
MAX_TOKENS = "MAX_TOKENS",
// The candidate content was flagged for safety reasons.
SAFETY = "SAFETY",
// The candidate content was flagged for recitation reasons.
RECITATION = "RECITATION",
// Unknown reason.
OTHER = "OTHER",
}
/**
* Task type for embedding content.
* @public
*/
export enum TaskType {
TASK_TYPE_UNSPECIFIED = "TASK_TYPE_UNSPECIFIED",
RETRIEVAL_QUERY = "RETRIEVAL_QUERY",
RETRIEVAL_DOCUMENT = "RETRIEVAL_DOCUMENT",
SEMANTIC_SIMILARITY = "SEMANTIC_SIMILARITY",
CLASSIFICATION = "CLASSIFICATION",
CLUSTERING = "CLUSTERING",
}
/**
* @public
*/
export enum FunctionCallingMode {
// Unspecified function calling mode. This value should not be used.
MODE_UNSPECIFIED = "MODE_UNSPECIFIED",
// Default model behavior, model decides to predict either a function call
// or a natural language repspose.
AUTO = "AUTO",
// Model is constrained to always predicting a function call only.
// If "allowed_function_names" are set, the predicted function call will be
// limited to any one of "allowed_function_names", else the predicted
// function call will be any one of the provided "function_declarations".
ANY = "ANY",
// Model will not predict any function call. Model behavior is same as when
// not passing any function declarations.
NONE = "NONE",
}
/**
* Content type for both prompts and response candidates.
* @public
*/
export interface Content {
role: string;
parts: Part[];
}
/**
* Content part - includes text or image part types.
* @public
*/
export type Part =
| TextPart
| InlineDataPart
| FunctionCallPart
| FunctionResponsePart
| FileDataPart;
/**
* Content part interface if the part represents a text string.
* @public
*/
export interface TextPart {
text: string;
inlineData?: never;
functionCall?: never;
functionResponse?: never;
fileData?: never;
}
/**
* Content part interface if the part represents an image.
* @public
*/
export interface InlineDataPart {
text?: never;
inlineData: GenerativeContentBlob;
functionCall?: never;
functionResponse?: never;
fileData?: never;
}
/**
* Content part interface if the part represents FunctionResponse.
* @public
*/
export interface FunctionCallPart {
text?: never;
inlineData?: never;
functionCall: FunctionCall;
functionResponse?: never;
fileData?: never;
}
/**
* Content part interface if the part represents FunctionResponse.
* @public
*/
export interface FunctionResponsePart {
text?: never;
inlineData?: never;
functionCall?: never;
functionResponse: FunctionResponse;
fileData?: never;
}
/**
* A predicted [FunctionCall] returned from the model
* that contains a string representing the [FunctionDeclaration.name]
* and a structured JSON object containing the parameters and their values.
* @public
*/
export interface FunctionCall {
name: string;
args: object;
}
/**
* The result output from a [FunctionCall] that contains a string
* representing the [FunctionDeclaration.name]
* and a structured JSON object containing any output
* from the function is used as context to the model.
* This should contain the result of a [FunctionCall]
* made based on model prediction.
* @public
*/
export interface FunctionResponse {
name: string;
response: object;
}
/**
* Interface for sending an image.
* @public
*/
export interface GenerativeContentBlob {
mimeType: string;
/**
* Image as a base64 string.
*/
data: string;
}
/**
* Content part interface if the part represents FunctionResponse.
* @public
*/
export interface FileDataPart {
text?: never;
inlineData?: never;
functionCall?: never;
functionResponse?: never;
fileData: FileData;
}
/**
* Data pointing to a file uploaded with the Files API.
* @public
*/
export interface FileData {
mimeType: string;
fileUri: string;
}
/**
* Base parameters for a number of methods.
* @public
*/
export interface BaseParams {
safetySettings?: SafetySetting[];
generationConfig?: GenerationConfig;
}
/**
* Params passed to {@link GoogleGenerativeAI.getGenerativeModel}.
* @public
*/
export interface ModelParams extends BaseParams {
model: string;
tools?: Tool[];
toolConfig?: ToolConfig;
systemInstruction?: Content;
}
/**
* Request sent to `generateContent` endpoint.
* @public
*/
export interface GenerateContentRequest extends BaseParams {
contents: Content[];
tools?: Tool[];
toolConfig?: ToolConfig;
systemInstruction?: Content;
}
/**
* Safety setting that can be sent as part of request parameters.
* @public
*/
export interface SafetySetting {
category: HarmCategory;
threshold: HarmBlockThreshold;
}
/**
* Config options for content-related requests
* @public
*/
export interface GenerationConfig {
candidateCount?: number;
stopSequences?: string[];
maxOutputTokens?: number;
temperature?: number;
topP?: number;
topK?: number;
}
/**
* Params for {@link GenerativeModel.startChat}.
* @public
*/
export interface StartChatParams extends BaseParams {
history?: Content[];
tools?: Tool[];
toolConfig?: ToolConfig;
systemInstruction?: Content;
}
/**
* Params for calling {@link GenerativeModel.countTokens}
* @public
*/
export interface CountTokensRequest {
contents: Content[];
}
/**
* Params for calling {@link GenerativeModel.embedContent}
* @public
*/
export interface EmbedContentRequest {
content: Content;
taskType?: TaskType;
title?: string;
}
/**
* Params for calling {@link GenerativeModel.batchEmbedContents}
* @public
*/
export interface BatchEmbedContentsRequest {
requests: EmbedContentRequest[];
}
/**
* Params passed to getGenerativeModel() or GoogleAIFileManager().
* @public
*/
export interface RequestOptions {
/**
* Request timeout in milliseconds.
*/
timeout?: number;
/**
* Version of API endpoint to call (e.g. "v1" or "v1beta"). If not specified,
* defaults to latest stable version.
*/
apiVersion?: string;
/**
* Additional attribution information to include in the x-goog-api-client header.
* Used by wrapper SDKs.
*/
apiClient?: string;
/**
* Base endpoint url. Defaults to "https://generativelanguage.googleapis.com"
*/
baseUrl?: string;
}
/**
* Defines a tool that model can call to access external knowledge.
* @public
*/
export declare type Tool = FunctionDeclarationsTool;
/**
* Structured representation of a function declaration as defined by the
* [OpenAPI 3.0 specification](https://spec.openapis.org/oas/v3.0.3). Included
* in this declaration are the function name and parameters. This
* FunctionDeclaration is a representation of a block of code that can be used
* as a Tool by the model and executed by the client.
* @public
*/
export declare interface FunctionDeclaration {
/**
* The name of the function to call. Must start with a letter or an
* underscore. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with
* a max length of 64.
*/
name: string;
/**
* Optional. Description and purpose of the function. Model uses it to decide
* how and whether to call the function.
*/
description?: string;
/**
* Optional. Describes the parameters to this function in JSON Schema Object
* format. Reflects the Open API 3.03 Parameter Object. string Key: the name
* of the parameter. Parameter names are case sensitive. Schema Value: the
* Schema defining the type used for the parameter. For function with no
* parameters, this can be left unset.
*
* @example with 1 required and 1 optional parameter: type: OBJECT properties:
* ```
* param1:
*
* type: STRING
* param2:
*
* type: INTEGER
* required:
*
* - param1
* ```
*/
parameters?: FunctionDeclarationSchema;
}
/**
* A FunctionDeclarationsTool is a piece of code that enables the system to
* interact with external systems to perform an action, or set of actions,
* outside of knowledge and scope of the model.
* @public
*/
export declare interface FunctionDeclarationsTool {
/**
* Optional. One or more function declarations
* to be passed to the model along with the current user query. Model may
* decide to call a subset of these functions by populating
* [FunctionCall][content.part.functionCall] in the response. User should
* provide a [FunctionResponse][content.part.functionResponse] for each
* function call in the next turn. Based on the function responses, Model will
* generate the final response back to the user. Maximum 64 function
* declarations can be provided.
*/
functionDeclarations?: FunctionDeclaration[];
}
/**
* Contains the list of OpenAPI data types
* as defined by https://swagger.io/docs/specification/data-models/data-types/
* @public
*/
export enum FunctionDeclarationSchemaType {
/** String type. */
STRING = "STRING",
/** Number type. */
NUMBER = "NUMBER",
/** Integer type. */
INTEGER = "INTEGER",
/** Boolean type. */
BOOLEAN = "BOOLEAN",
/** Array type. */
ARRAY = "ARRAY",
/** Object type. */
OBJECT = "OBJECT",
}
/**
* Schema for parameters passed to {@link FunctionDeclaration.parameters}.
* @public
*/
export interface FunctionDeclarationSchema {
/** The type of the parameter. */
type: FunctionDeclarationSchemaType;
/** The format of the parameter. */
properties: { [k: string]: FunctionDeclarationSchemaProperty; };
/** Optional. Description of the parameter. */
description?: string;
/** Optional. Array of required parameters. */
required?: string[];
}
/**
* Schema is used to define the format of input/output data.
* Represents a select subset of an OpenAPI 3.0 schema object.
* More fields may be added in the future as needed.
* @public
*/
export interface FunctionDeclarationSchemaProperty {
/**
* Optional. The type of the property. {@link
* FunctionDeclarationSchemaType}.
*/
type?: FunctionDeclarationSchemaType;
/** Optional. The format of the property. */
format?: string;
/** Optional. The description of the property. */
description?: string;
/** Optional. Whether the property is nullable. */
nullable?: boolean;
/** Optional. The items of the property. {@link FunctionDeclarationSchema} */
items?: FunctionDeclarationSchema;
/** Optional. The enum of the property. */
enum?: string[];
/** Optional. Map of {@link FunctionDeclarationSchema}. */
properties?: { [k: string]: FunctionDeclarationSchema; };
/** Optional. Array of required property. */
required?: string[];
/** Optional. The example of the property. */
example?: unknown;
}
/**
* Tool config. This config is shared for all tools provided in the request.
* @public
*/
export interface ToolConfig {
functionCallingConfig: FunctionCallingConfig;
[k: string]: unknown;
}
/**
* @public
*/
export interface FunctionCallingConfig {
mode?: FunctionCallingMode;
allowedFunctionNames?: string[];
}
/**
* Result object returned from generateContent() call.
*
* @public
*/
export interface GenerateContentResult {
response: EnhancedGenerateContentResponse;
}
/**
* Result object returned from generateContentStream() call.
* Iterate over `stream` to get chunks as they come in and/or
* use the `response` promise to get the aggregated response when
* the stream is done.
*
* @public
*/
export interface GenerateContentStreamResult {
stream: AsyncGenerator<EnhancedGenerateContentResponse>;
response: Promise<EnhancedGenerateContentResponse>;
}
/**
* Response object wrapped with helper methods.
*
* @public
*/
export interface EnhancedGenerateContentResponse
extends GenerateContentResponse {
/**
* Returns the text string assembled from all `Part`s of the first candidate
* of the response, if available.
* Throws if the prompt or candidate was blocked.
*/
text: () => string;
/**
* Deprecated: use `functionCalls()` instead.
* @deprecated - use `functionCalls()` instead
*/
functionCall: () => FunctionCall | undefined;
/**
* Returns function calls found in any `Part`s of the first candidate
* of the response, if available.
* Throws if the prompt or candidate was blocked.
*/
functionCalls: () => FunctionCall[] | undefined;
}
/**
* Individual response from {@link GenerativeModel.generateContent} and
* {@link GenerativeModel.generateContentStream}.
* `generateContentStream()` will return one in each chunk until
* the stream is done.
* @public
*/
export interface GenerateContentResponse {
candidates?: GenerateContentCandidate[];
promptFeedback?: PromptFeedback;
}
/**
* If the prompt was blocked, this will be populated with `blockReason` and
* the relevant `safetyRatings`.
* @public
*/
export interface PromptFeedback {
blockReason: BlockReason;
safetyRatings: SafetyRating[];
blockReasonMessage?: string;
}
/**
* A candidate returned as part of a {@link GenerateContentResponse}.
* @public
*/
export interface GenerateContentCandidate {
index: number;
content: Content;
finishReason?: FinishReason;
finishMessage?: string;
safetyRatings?: SafetyRating[];
citationMetadata?: CitationMetadata;
}
/**
* Citation metadata that may be found on a {@link GenerateContentCandidate}.
* @public
*/
export interface CitationMetadata {
citationSources: CitationSource[];
}
/**
* A single citation source.
* @public
*/
export interface CitationSource {
startIndex?: number;
endIndex?: number;
uri?: string;
license?: string;
}
/**
* A safety rating associated with a {@link GenerateContentCandidate}
* @public
*/
export interface SafetyRating {
category: HarmCategory;
probability: HarmProbability;
blocked?: boolean;
}
/**
* Response from calling {@link GenerativeModel.countTokens}.
* @public
*/
export interface CountTokensResponse {
totalTokens: number;
}
/**
* Response from calling {@link GenerativeModel.embedContent}.
* @public
*/
export interface EmbedContentResponse {
embedding: ContentEmbedding;
}
/**
* Response from calling {@link GenerativeModel.batchEmbedContents}.
* @public
*/
export interface BatchEmbedContentsResponse {
embeddings: ContentEmbedding[];
}
/**
* A single content embedding.
* @public
*/
export interface ContentEmbedding {
values: number[];
}
export class GoogleGeminiHTTP extends HTTPService {
name = 'GoogleGemini';
constructor(apiKey: string, public baseUri: string = 'https://generativelanguage.googleapis.com/v1beta') {
super(baseUri);
this.baseHeaders['x-goog-api-client'] = `Developers know-better/1.0`;
// this.baseHeaders['x-goog-api-key'] = `${apiKey}`;
this.apiKey = apiKey;
this.baseOptions.timeout = 1000 * 60 * 30 * 0.5;
let agent!: ProxyAgent;
const proxyUri = process.env.http_proxy;
if (proxyUri) {
agent = new ProxyAgent({
uri: proxyUri
});
}
if (agent) {
Reflect.set(this.baseOptions, 'dispatcher', agent);
}
}
set apiKey(apiKey: string) {
this.baseHeaders['x-goog-api-key'] = apiKey;
}
get apiKey() {
return this.baseHeaders['x-goog-api-key'] as string;
}
complete<T extends GenerateContentRequest>(model: string, payload: T, opts?: typeof this['baseOptions']) {
return this.postJson<GenerateContentResponse>(`/models/${model}:generateContent`, payload,
{ ...opts, responseType: 'json' }
);
}
streamComplete<T extends GenerateContentRequest>(model: string, payload: T, opts?: typeof this['baseOptions']) {
return this.postJsonWithSearchParams<InputServerEventStream>(`/models/${model}:generateContent`, {
alt: 'sse'
}, payload,
{ ...opts, responseType: 'stream' }
);
}
countTokens<T extends CountTokensRequest>(model: string, payload: T, opts?: typeof this['baseOptions']) {
return this.postJson<CountTokensResponse>(`/models/${model}:countTokens`, payload,
{ ...opts, responseType: 'json' }
);
}
override async __processResponse(options: HTTPServiceRequestOptions, r: Response): Promise<any> {
if (r.status === 401) {
this.emit('unauthorized', r);
}
if (r.status !== 200) {
throw await r.json();
}
const s = await super.__processResponse(options, r) as any as Readable;
if (options.responseType === 'stream') {
const parseStream = new InputServerEventStream();
s.pipe(parseStream);
parseStream.once('end', () => {
if (!s.readableEnded) {
r.body?.cancel();
}
});
return parseStream;
}
return s;
}
}
export class VertexGeminiHTTP extends GoogleGeminiHTTP {
override name = 'VertexGemini';
constructor(
apiKey: string,
public project: string,
public region: string = 'us-central1',
) {
super(apiKey, `https://${region}-aiplatform.googleapis.com/v1/projects/${project}/locations/${region}/publishers/google`);
}
override get apiKey() {
return this.baseHeaders['x-goog-api-key'] as string;
}
override set apiKey(input: string) {
this.baseHeaders['authorization'] = `Bearer ${input}`;
this.baseHeaders['x-goog-api-key'] = input;
}
override streamComplete<T extends GenerateContentRequest>(model: string, payload: T, opts?: typeof this['baseOptions']) {
return this.postJsonWithSearchParams<InputServerEventStream>(`/models/${model}:streamGenerateContent`, { alt: 'sse' }, payload,
{ ...opts, responseType: 'stream' }
);
}
}
src/3rd-party/internal-cloudrun.ts
Vendored
+35
View File
@@ -0,0 +1,35 @@
// import { OpenAICompatHTTP } from './openai-compat';
// export class InternalCloudRunReaderLM2 extends OpenAICompatHTTP {
// name = 'InternalReaderLMCloudRun';
// supportedCompletionModels = [
// ];
// supportedChatCompletionModels = [
// 'readerlm-v2',
// ];
// supportedImageGenerationModels = [
// ];
// constructor(apiKey: string, organization?: string) {
// super(apiKey, process.env.OVERRIDE_READERLM_V2_URL', organization);
// }
// }
// export class InternalCloudRunJinaVLM extends OpenAICompatHTTP {
// name = 'InternalJinaVLMCloudRun';
// supportedCompletionModels = [
// ];
// supportedChatCompletionModels = [
// 'jina-vlm',
// ];
// supportedImageGenerationModels = [
// ];
// constructor(apiKey: string, organization?: string) {
// super(apiKey, process.env.OVERRIDE_JINA_VLM_URL, organization);
// }
// }
src/3rd-party/jina-embeddings.ts
Vendored
+102
View File
@@ -0,0 +1,102 @@
import { HTTPService } from 'civkit/http';
import _ from 'lodash';
export interface JinaWallet {
trial_balance: number;
trial_start: Date;
trial_end: Date;
regular_balance: number;
total_balance: number;
}
export interface JinaUserBrief {
user_id: string;
email: string | null;
full_name: string | null;
customer_id: string | null;
avatar_url?: string;
billing_address: Partial<{
address: string;
city: string;
state: string;
country: string;
postal_code: string;
}>;
payment_method: Partial<{
brand: string;
last4: string;
exp_month: number;
exp_year: number;
}>;
wallet: JinaWallet;
metadata: {
[k: string]: any;
};
}
export interface JinaUsageReport {
model_name: string;
api_endpoint: string;
consumer: {
user_id: string;
customer_plan?: string;
[k: string]: any;
};
usage: {
total_tokens: number;
};
labels: {
user_type?: string;
model_name?: string;
[k: string]: any;
};
}
export class JinaEmbeddingsDashboardHTTP extends HTTPService {
name = 'JinaEmbeddingsDashboardHTTP';
constructor(
public apiKey: string,
public baseUri: string = process.env.OVERRIDE_MANAGE_SERVER_URL || 'https://dash.jina.ai/api'
) {
super(baseUri);
this.baseOptions.timeout = 30_000; // 30 sec
}
async authorization(token: string) {
const r = await this.get<JinaUserBrief>('/v1/authorization', {
headers: {
Authorization: `Bearer ${token}`
},
responseType: 'json',
});
return r;
}
async validateToken(token: string) {
const r = await this.getWithSearchParams<JinaUserBrief>('/v1/api_key/user', {
api_key: token,
}, {
responseType: 'json',
});
return r;
}
async reportUsage(token: string, query: JinaUsageReport) {
const r = await this.postJson('/v1/usage', query, {
headers: {
Authorization: `Bearer ${token}`,
'x-api-key': this.apiKey,
},
responseType: 'text',
});
return r;
}
}
src/3rd-party/open-router.ts
Vendored
+122
View File
@@ -0,0 +1,122 @@
import { HTTPService, HTTPServiceRequestOptions } from 'civkit/http';
import type OpenAI from 'openai';
import { InputServerEventStream } from '../lib/transform-server-event-stream';
import _ from 'lodash';
import { ProxyAgent, Agent } from 'undici';
import { Readable } from 'stream';
export class OpenRouterHTTP extends HTTPService {
name: string = 'OpenRouter';
constructor(
public apiKey: string,
public userTitle?: string,
public userUrl?: string,
) {
super('https://openrouter.ai/api/v1');
this.baseHeaders['Authorization'] = `Bearer ${apiKey}`;
if (userTitle) {
this.baseHeaders['X-Title'] = userTitle;
}
if (userUrl) {
this.baseHeaders['HTTP-Referer'] = userUrl;
}
let agent: ProxyAgent | Agent;
const proxyUri = process.env.http_proxy;
const proxyOptions: Omit<ProxyAgent.Options, 'uri'> = {
connectTimeout: 1000 * 30,
headersTimeout: 1000 * 60 * 15,
};
if (proxyUri) {
agent = new ProxyAgent({
...proxyOptions,
uri: proxyUri
});
} else {
agent = new Agent(proxyOptions);
}
this.baseOptions.timeout = 1000 * 60 * 15;
Reflect.set(this.baseOptions, 'dispatcher', agent);
}
listModels() {
return this.get<{
data: {
id: string;
canonical_slug: string;
hugging_face_id: string;
name: string;
created: number;
description: string;
context_length: number;
architecture: {
modality: string;
input_modalities: string[];
output_modalities: string[];
tokenizer: string;
instruct_type: null;
};
pricing: {
prompt: string;
completion: string;
request: string;
image: string;
web_search: string;
internal_reasoning: string;
};
top_provider: {
context_length: number;
max_completion_tokens: number;
is_moderated: boolean;
};
per_request_limits: null;
supported_parameters: string[];
}[];
}>('/models');
}
completions<T extends OpenAI.CompletionCreateParams>(payload: T, opts?: typeof this['baseOptions']) {
return this.postJson<typeof payload['stream'] extends true ?
InputServerEventStream : OpenAI.Completion>('/completions', payload,
{ ...opts, responseType: payload.stream ? 'stream' : 'json' }
);
}
chatCompletions<T extends OpenAI.ChatCompletionCreateParams>(payload: T, opts?: typeof this['baseOptions']) {
return this.postJson<typeof payload['stream'] extends true ?
InputServerEventStream : OpenAI.Completion
>('/chat/completions', payload,
{ ...opts, responseType: payload.stream ? 'stream' : 'json' }
);
}
override async __processResponse(options: HTTPServiceRequestOptions, r: Response): Promise<any> {
if (r.status !== 200) {
if (r.headers.get('content-type')?.includes('application/json')) {
throw await r.json();
}
throw await r.text();
}
const s = await super.__processResponse(options, r) as any as Readable;
if (options.responseType === 'stream') {
const parseStream = new InputServerEventStream();
s.pipe(parseStream);
parseStream.once('end', () => {
if (!s.readableEnded) {
r.body?.cancel();
}
});
return parseStream;
}
return s;
}
}
src/3rd-party/openai-compat.ts
Vendored
+119
View File
@@ -0,0 +1,119 @@
import { HTTPService, HTTPServiceRequestOptions } from 'civkit/http';
import type OpenAI from 'openai';
import { InputServerEventStream } from '../lib/transform-server-event-stream';
import _ from 'lodash';
import { ProxyAgent, Agent } from 'undici';
import { Readable } from 'stream';
export abstract class OpenAICompatHTTP extends HTTPService {
abstract name: string;
abstract supportedCompletionModels: string[];
abstract supportedChatCompletionModels: string[];
abstract supportedImageGenerationModels: string[];
functionCallingSupported = true;
constructor(public apiKey: string, baseUri: string, public organization?: string) {
super(baseUri);
this.baseHeaders['Authorization'] = `Bearer ${apiKey}`;
if (organization) {
this.baseHeaders['OpenAI-Organization'] = organization;
}
let agent: ProxyAgent | Agent;
const proxyUri = process.env.http_proxy;
const proxyOptions: Omit<ProxyAgent.Options, 'uri'> = {
connectTimeout: 1000 * 30,
headersTimeout: 1000 * 60 * 15,
};
if (proxyUri) {
agent = new ProxyAgent({
...proxyOptions,
uri: proxyUri
});
} else {
agent = new Agent(proxyOptions);
}
this.baseOptions.timeout = 1000 * 60 * 15;
Reflect.set(this.baseOptions, 'dispatcher', agent);
}
listModels() {
return this.get<OpenAI.Models>('models');
}
getModelDetail(model: string, opts?: typeof this['baseOptions']) {
return this.get<OpenAI.Model>(`/models/${model}`, opts);
}
completions<T extends OpenAI.CompletionCreateParams>(payload: T, opts?: typeof this['baseOptions']) {
return this.postJson<typeof payload['stream'] extends true ?
InputServerEventStream : OpenAI.Completion>('/completions', payload,
{ ...opts, responseType: payload.stream ? 'stream' : 'json' }
);
}
chatCompletions<T extends OpenAI.ChatCompletionCreateParams>(payload: T, opts?: typeof this['baseOptions']) {
return this.postJson<typeof payload['stream'] extends true ?
InputServerEventStream : OpenAI.Completion
>('/chat/completions', payload,
{ ...opts, responseType: payload.stream ? 'stream' : 'json' }
);
}
imagesGenerations<T extends OpenAI.ImageGenerateParams>(payload: T, opts?: typeof this['baseOptions']) {
return this.postJson<OpenAI.ImagesResponse>('/images/generations', payload,
{ ...opts, responseType: 'json' }
);
}
imagesEdits<T extends OpenAI.ImageGenerateParams & { image: File; mask?: File; prompt?: string; }>(payload: T, opts?: typeof this['baseOptions']) {
const req: any[] = [];
for (const [k, v] of Object.entries(payload)) {
req.push([k, v]);
}
return this.postMultipart<OpenAI.ImagesResponse>('/images/edits', req,
{ ...opts, responseType: 'json' }
);
}
imagesVariants<T extends OpenAI.ImageGenerateParams>(payload: T & { image: File; prompt?: string; }, opts?: typeof this['baseOptions']) {
const req: any[] = [];
for (const [k, v] of Object.entries(payload)) {
req.push([k, v]);
}
return this.postMultipart<OpenAI.ImagesResponse>('/images/variations', req,
{ ...opts, responseType: 'json' }
);
}
override async __processResponse(options: HTTPServiceRequestOptions, r: Response): Promise<any> {
if (r.status !== 200) {
if (r.headers.get('content-type')?.includes('application/json')) {
throw await r.json();
}
throw await r.text();
}
const s = await super.__processResponse(options, r) as any as Readable;
if (options.responseType === 'stream') {
const parseStream = new InputServerEventStream();
s.pipe(parseStream);
parseStream.once('end', () => {
if (!s.readableEnded) {
r.body?.cancel();
}
});
return parseStream;
}
return s;
}
}
src/3rd-party/openai.ts
Vendored
+54
View File
@@ -0,0 +1,54 @@
import { OpenAICompatHTTP } from './openai-compat';
export class OpenAIHTTP extends OpenAICompatHTTP {
name = 'OfficialOpenAI';
supportedCompletionModels = [
'gpt-3.5-turbo-instruct',
'babbage-002',
'davinci-002',
];
supportedChatCompletionModels = [
'gpt-3.5-turbo',
'gpt-3.5-turbo-0125',
'gpt-3.5-turbo-1106',
'gpt-3.5-turbo-0613',
'gpt-3.5-turbo-0301',
'gpt-3.5-turbo-16k',
'gpt-3.5-turbo-16k-1106',
'gpt-3.5-turbo-16k-0613',
'gpt-4',
'gpt-4-0613',
'gpt-4-vision-preview',
'gpt-4-turbo-preview',
'gpt-4-0125-preview',
'gpt-4-1106-preview',
'gpt-4-32k',
'gpt-4-32k-0613',
'gpt-4-turbo',
'gpt-4-turbo-2024-04-09',
'gpt-4o',
'gpt-4o-2024-05-13',
'gpt-4o-mini',
'gpt-4o-mini-2024-07-18',
];
supportedImageGenerationModels = [
'dall-e-2',
'dall-e-3'
];
constructor(apiKey: string, organization?: string) {
super(apiKey, 'https://api.openai.com/v1', organization);
}
}
src/3rd-party/replicate.ts
Vendored
+226
View File
@@ -0,0 +1,226 @@
import { HTTPService, HTTPServiceRequestOptions, } from 'civkit/http';
import { DownstreamServiceFailureError, JSONPrimitive } from 'civkit/civ-rpc';
import { mimeOf } from 'civkit/mime';
import { TimeoutError } from 'civkit/defer';
import { delay } from 'civkit/timeout';
import _ from 'lodash';
import type { Model, ModelVersion, Page, Prediction, WebhookEventType } from 'replicate';
import { InputServerEventStream } from '../lib/transform-server-event-stream';
export class ReplicateHTTP extends HTTPService {
name = 'Replicate';
versionMap = new Map<string, string>();
constructor(public token: string, public baseUri: string = 'https://api.replicate.com/v1') {
super(baseUri);
this.baseHeaders['Authorization'] = `Token ${token}`;
this.baseOptions.timeout = 1000 * 30;
}
getModel(modelPath: string) {
return this.get<Model>(`/models/${modelPath}`);
}
getLatestVersionId(modelPath: string) {
return this.getModel(modelPath).then((r) => r.data.latest_version?.id);
}
listModelVersions(modelPath: string) {
return this.get<Page<ModelVersion>>(`/models/${modelPath}/versions`);
}
getVersionDetail(modelPath: string, version: string) {
return this.get<ModelVersion>(`/models/${modelPath}/versions/${version}`);
}
async encodeInput(input: object) {
const mapped: { [k: string]: JSONPrimitive; } = {};
for (const [k, v] of Object.entries(input)) {
if (Buffer.isBuffer(v)) {
const mimeVec = await mimeOf(v);
mapped[k] = `data:${mimeVec.mediaType}/${mimeVec.subType};base64,${v.toString('base64')}`;
} else if (v instanceof Blob) {
mapped[k] = `data:${v.type};base64,${Buffer.from(await v.arrayBuffer()).toString('base64')}`;
} else {
mapped[k] = v;
}
}
return mapped;
}
async resolveVersion(text: string) {
if (!text.includes('/')) {
return text;
}
const r = this.versionMap.get(text);
if (r) {
return r;
}
const latestVersion = await this.getLatestVersionId(text);
if (!latestVersion) {
throw new Error(`Replicate model ${text} cannot be resolved`);
}
this.versionMap.set(text, latestVersion);
setTimeout(() => {
this.versionMap.delete(text);
}, 3600 * 1000);
return latestVersion;
}
async createPrediction(version: string,
input: object, webhook?: string, webhook_events_filter?: WebhookEventType[], opts?: typeof this['baseOptions']
) {
const resolvedVersion = await this.resolveVersion(version);
return this.postJson<Prediction>(`/predictions`, {
version: resolvedVersion,
input: await this.encodeInput(input),
webhook,
webhook_events_filter
}, opts);
}
async createPredictionPreferringWait(version: string,
input: object, webhook?: string, webhook_events_filter?: WebhookEventType[], opts?: typeof this['baseOptions']
) {
const resolvedVersion = await this.resolveVersion(version);
return this.postJson<Prediction>(`/predictions`, {
version: resolvedVersion,
input: await this.encodeInput(input),
webhook,
webhook_events_filter
}, {
headers: {
'Prefer': 'wait'
},
timeout: 1000 * 70,
...opts
});
}
async createStreamPrediction(version: string,
input: object, webhook?: string, webhook_events_filter?: WebhookEventType[], opts?: typeof this['baseOptions']
) {
const resolvedVersion = await this.resolveVersion(version);
return this.postJson<Prediction>(`/predictions`, {
version: resolvedVersion,
input: await this.encodeInput(input),
stream: true,
webhook,
webhook_events_filter
}, opts);
}
checkPrediction(predictionId: string, opts?: typeof this['baseOptions']) {
return this.get<Prediction>(`/predictions/${predictionId}`, opts);
}
async *iterPollPrediction(predictionId: string, interval: number = 1000, maxAttempts: number = 210, opts?: typeof this['baseOptions']) {
let lastCheckpoint = Date.now();
const deadline = lastCheckpoint + maxAttempts * interval;
while (lastCheckpoint <= deadline) {
const now = Date.now();
const dt = now - lastCheckpoint;
if (dt < interval) {
await delay(interval - dt);
}
lastCheckpoint = now;
const prediction = (await this.checkPrediction(predictionId, opts)).data;
switch (prediction.status as typeof prediction['status'] | 'queued') {
case 'starting':
case 'queued': {
continue;
}
case 'processing': {
yield prediction;
continue;
}
case 'succeeded': {
yield prediction;
return prediction;
}
case 'failed': {
throw new DownstreamServiceFailureError({ message: 'Prediction failed', prediction });
}
case 'canceled': {
throw new DownstreamServiceFailureError({ message: 'Prediction canceled', prediction });
}
default: {
throw new Error('Unknown status');
}
}
}
throw new TimeoutError('Prediction timeout');
}
async pollPrediction(predictionId: string, interval: number = 2000, maxAttempts: number = 105, opts?: typeof this['baseOptions']) {
let r: Prediction;
for await (const prediction of this.iterPollPrediction(predictionId, interval, maxAttempts, opts)) {
r = prediction;
}
return r!;
}
async expectStream(url: string, opts?: typeof this['baseOptions']) {
const r = await this.get(url, {
headers: {
...this.baseHeaders,
'Accept': 'text/event-stream'
},
responseType: 'stream',
timeout: 1000 * 3600,
...opts,
});
const sseStream = new InputServerEventStream();
r.parsed.pipe(sseStream, { end: true });
r.parsed.once('error', (err: any) => {
sseStream.destroy(err);
});
sseStream.on('data', (sse) => {
if (sse.event === 'error') {
sseStream.destroy(new this.Error(sse.data?.detail));
return;
}
if (sse.event === 'done') {
sseStream.end();
return;
}
});
return sseStream;
}
override async __processResponse(options: HTTPServiceRequestOptions, r: Response): Promise<any> {
if (!r.ok) {
throw await r.json();
}
const s = await super.__processResponse(options, r);
return s;
}
}
src/3rd-party/serper-search.ts
Vendored
+194
View File
File diff suppressed because one or more lines are too long
src/api/crawler.ts
+1808
View File
File diff suppressed because it is too large Load Diff
src/api/searcher.ts
+950
View File
@@ -0,0 +1,950 @@
import { singleton } from 'tsyringe';
import {
assignTransferProtocolMeta, RPCHost, RPCReflection, AssertionFailureError, assignMeta, RawString,
DownstreamServiceFailureError,
AuthenticationRequiredError,
ArrayOf,
} from 'civkit/civ-rpc';
import { marshalErrorLike } from 'civkit/lang';
import { objHashMd5B64Of } from 'civkit/hash';
import _ from 'lodash';
import { CrawlerHost, ExtraScrappingOptions } from './crawler';
import { CrawlerOptions, RESPOND_TIMING } from '../dto/crawler-options';
import { SnapshotFormatter, FormattedPage as RealFormattedPage, FormattedPageDto } from '../services/snapshot-formatter';
import { GoogleSearchExplicitOperatorsDto } from '../services/serper-search';
import { GlobalLogger } from '../services/logger';
import { AsyncLocalContext } from '../services/async-context';
import { Context, Ctx, Method, Param, RPCReflect } from '../services/registry';
import { InputServerEventStream, OutputServerEventStream } from '../lib/transform-server-event-stream';
import { SerperBingSearchService, SerperGoogleSearchService } from '../services/serp/serper';
import { consumeAsyncGenerator, delayGenerator, finalYield, raceAsyncGenerators, timeoutGenerator, toAsyncGenerator } from '../utils/misc';
import { SerperSearchQueryParams, WORLD_COUNTRIES, WORLD_LANGUAGES } from '../3rd-party/serper-search';
import { WebSearchEntry } from '../services/serp/compat';
import { CommonGoogleSERP } from '../services/serp/common-serp';
import { GoogleSERP, GoogleSERPOldFashion } from '../services/serp/google';
import { EnvConfig } from '../services/envconfig';
import { Readable } from 'stream';
import { once } from 'events';
import { Defer } from 'civkit/defer';
import { BingSERP } from '../services/serp/bing';
import { bcp47ToIso639_3 } from '../utils/languages';
import { parseSearchQuery } from '../utils/search-query';
import { JSDomControl } from '../services/jsdom';
import { delay } from 'civkit/timeout';
import { BaseAuthDTO } from '../dto/base-auth';
import { SERPResult } from '../db/models';
import { StorageLayer } from '../db/noop-storage';
import { AUTH_DTO_CLS } from '../config';
const WORLD_COUNTRY_CODES = Object.keys(WORLD_COUNTRIES).map((x) => x.toLowerCase());
interface FormattedPage extends RealFormattedPage {
favicon?: string;
date?: string;
}
@singleton()
export class SearcherHost extends RPCHost {
logger = this.globalLogger.child({ service: this.constructor.name });
cacheRetentionMs = 1000 * 3600 * 24 * 7;
cacheValidMs = 1000 * 3600;
pageCacheToleranceMs = 1000 * 3600 * 24;
reasonableDelayMs = 25_000;
targetResultCount = 6;
constructor(
protected globalLogger: GlobalLogger,
protected envConfig: EnvConfig,
protected threadLocal: AsyncLocalContext,
protected crawler: CrawlerHost,
protected snapshotFormatter: SnapshotFormatter,
protected jsdomControl: JSDomControl,
protected serperGoogle: SerperGoogleSearchService,
protected serperBing: SerperBingSearchService,
protected commonGoogleSerp: CommonGoogleSERP,
protected googleSERP: GoogleSERP,
protected googleSERPOld: GoogleSERPOldFashion,
protected bingSERP: BingSERP,
protected storageLayer: StorageLayer,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
@Method({
name: 'searchIndex',
ext: {
http: {
action: ['get', 'post'],
path: '/search'
}
},
tags: ['search'],
returnType: [RawString, ArrayOf(FormattedPageDto), OutputServerEventStream],
})
@Method({
ext: {
http: {
action: ['get', 'post'],
path: '::q'
}
},
tags: ['search'],
returnType: [RawString, ArrayOf(FormattedPageDto), OutputServerEventStream],
})
async search(
@RPCReflect() rpcReflect: RPCReflection,
@Ctx() ctx: Context,
@Param({ type: AUTH_DTO_CLS }) auth: BaseAuthDTO,
crawlerOptions: CrawlerOptions,
searchExplicitOperators: GoogleSearchExplicitOperatorsDto,
@Param('type', { type: new Set(['web', 'images', 'news']), default: 'web' })
variant: 'web' | 'images' | 'news',
@Param('count', { validate: (v: number) => v >= 0 && v <= 20 })
@Param('num', { validate: (v: number) => v >= 0 && v <= 20 })
count: number = 10,
@Param('provider', { type: new Set(['google', 'bing', 'reader']) })
@Param('engine', { type: new Set(['google', 'bing', 'reader']) })
searchEngine?: 'google' | 'bing' | 'reader',
@Param('gl', { validate: (v: string) => WORLD_COUNTRY_CODES.includes(v?.toLowerCase()) }) gl?: string,
@Param('hl', { validate: (v: string) => WORLD_LANGUAGES.some(l => l.code === v) }) hl?: string,
@Param('location') location?: string,
@Param('page') page?: number,
@Param('fallback', { type: Boolean, default: false }) fallback?: boolean,
@Param('q') q?: string,
@Param('nfpr') nfpr?: boolean,
) {
// Return content by default
const crawlWithoutContent = crawlerOptions.respondWith.includes('no-content');
const withFavicon = Boolean(ctx.get('X-With-Favicons'));
this.threadLocal.set('collect-favicon', withFavicon);
crawlerOptions.respondTiming ??= RESPOND_TIMING.VISIBLE_CONTENT;
let chargeAmount = 0;
const noSlashPath = decodeURIComponent(ctx.path).slice(1);
if (!noSlashPath && !q) {
const index = await this.crawler.getIndex(auth);
if (!auth.isInternal && !auth.bearerToken) {
index.note = 'Authentication is required to use this endpoint. Please provide a valid API key via Authorization header.';
}
if (!ctx.accepts('text/plain') && (ctx.accepts('text/json') || ctx.accepts('application/json'))) {
return index;
}
return assignTransferProtocolMeta(`${index}`,
{ contentType: 'text/plain', envelope: null }
);
}
const {
uid,
reportOptions,
reportUsage,
isAnonymous,
} = await this.storageLayer.rateLimit(ctx, rpcReflect, auth as any);
if (isAnonymous && !auth.isInternal) {
throw new AuthenticationRequiredError('Authentication is required to use this endpoint. Please provide a valid API key via Authorization header.');
}
rpcReflect.finally(() => {
reportOptions?.(crawlerOptions.customizedProps());
reportUsage?.(chargeAmount, 'reader-search');
});
delete crawlerOptions.html;
delete crawlerOptions.file;
delete crawlerOptions.pdf;
const crawlOpts = await this.crawler.configure(crawlerOptions);
crawlOpts.eligibleForPageIndex = true;
const searchQuery = searchExplicitOperators.addTo(q || noSlashPath);
this.logger.info(`Accepting request from ${uid || ctx.ip}`, { opts: crawlerOptions, searchQuery });
let chargeAmountScaler = 1;
if (searchEngine === 'bing') {
this.threadLocal.set('bing-preferred', true);
chargeAmountScaler = 3;
}
if (variant !== 'web') {
chargeAmountScaler = 5;
}
const searchParams = {
variant,
provider: searchEngine,
q: searchQuery,
num: count,
gl,
hl,
location,
page,
nfpr,
};
const timeoutMs = crawlOpts.timeoutMs || this.reasonableDelayMs;
const t0 = performance.now();
let localSearchIterator: AsyncGenerator<FormattedPage[], void, undefined> | undefined;
let localResultsPromise: Promise<FormattedPage[] | undefined> | undefined;
let isDelayDemanding = false;
let it;
let results: FormattedPage[];
if (searchEngine === 'reader') {
this.logger.debug(`Preparing local cache search`, { timeoutMs });
localSearchIterator = this.readerLocalSearch(searchParams, crawlOpts, crawlerOptions);
localResultsPromise = localSearchIterator.next().then((r) => r.value!);
results = await localResultsPromise || [];
it = localSearchIterator;
} else if (variant === 'web' && !crawlerOptions.noCache) {
const cachedSearchIterator = this.cachedSearch(searchParams, crawlOpts, crawlerOptions);
localSearchIterator = this.readerLocalSearch(searchParams, crawlOpts, crawlerOptions);
let raceIsOver = false;
if (timeoutMs <= 5000) {
localResultsPromise = localSearchIterator.next().then((r) => r.value!);
isDelayDemanding = true;
this.logger.debug(`Preparing local cache search`, { timeoutMs });
} else {
localResultsPromise = delay(timeoutMs - 5000).then(() => {
if (raceIsOver) {
localSearchIterator?.return();
return;
}
this.logger.debug(`Preparing local cache search`, { timeoutMs });
return localSearchIterator!.next();
}).then((r) => r?.value || undefined);
}
const liveResults = cachedSearchIterator.next().then((r) => r.value!);
const t0 = performance.now();
const r = await Promise.race([
delay(timeoutMs)
.then(() => localResultsPromise)
.then((r) => ({ results: r, it: localSearchIterator! }))
.catch((err) => {
this.logger.warn(`Error happened during consumption of local search generator`, { err });
return liveResults.then((r) => ({ results: r, it: cachedSearchIterator }));
}),
liveResults.then((r) => ({ results: r, it: cachedSearchIterator }))
]).finally(() => raceIsOver = true);
it = r.it;
const dt = performance.now() - t0;
if (it === localSearchIterator) {
this.logger.debug(`Local search won the race after ${dt.toFixed(1)}ms`, { timeoutMs });
delete crawlerOptions.timeout;
delete crawlOpts.timeoutMs;
consumeAsyncGenerator(cachedSearchIterator).catch((err) => {
this.logger.warn(`Error happened during consumption of live search generator`, { err });
});
} else {
this.logger.debug(`Cached search won the race after ${dt.toFixed(1)}ms`, { timeoutMs });
}
results = r.results || [];
} else {
it = this.cachedSearch(searchParams, crawlOpts, crawlerOptions);
results = (await it.next()).value;
}
if (!results?.length && fallback) {
this.logger.debug(`No results found, falling back to local search`, { searchQuery, timeoutMs });
const localResults = await localResultsPromise;
if (localResults?.length) {
results = localResults;
it = localSearchIterator!;
this.logger.debug(`Fallback to local search results`, { resultsNum: results.length, searchQuery });
}
}
if (!results?.length) {
throw new AssertionFailureError(`No search results available for query ${searchQuery}`);
}
let lastScrapped: any[] | undefined;
if (!crawlerOptions.respondWith.includes('no-content') &&
['html', 'text', 'shot', 'markdown', 'content'].some((x) => crawlerOptions.respondWith.includes(x))
) {
for (const x of results) {
x.content ??= '';
}
}
if (crawlWithoutContent || count === 0) {
if (localSearchIterator) {
localSearchIterator.return();
}
delete crawlerOptions.timeout;
delete crawlOpts.timeoutMs;
if (isDelayDemanding) {
consumeAsyncGenerator(it);
} else {
it.return();
}
it = toAsyncGenerator(results);
} else if (localSearchIterator && it !== localSearchIterator && (!page || page <= 1)) {
const dt = performance.now() - t0;
it = raceAsyncGenerators(it, delayGenerator(timeoutMs - dt, localSearchIterator));
} else {
const dt = performance.now() - t0;
it = raceAsyncGenerators(it, timeoutGenerator(Math.max(timeoutMs - dt, 2000)));
}
if (!ctx.accepts('text/plain') && ctx.accepts('text/event-stream')) {
const sseStream = new OutputServerEventStream();
rpcReflect.return(sseStream);
try {
for await (const scrapped of it) {
if (!scrapped) {
continue;
}
if (rpcReflect.signal.aborted) {
break;
}
chargeAmount = this.assignChargeAmount(scrapped, count, chargeAmountScaler);
lastScrapped = scrapped.slice(0, count);
sseStream.write({
event: 'data',
data: scrapped,
});
}
} catch (err: any) {
this.logger.error(`Failed to collect search result for query ${searchQuery}`,
{ err: marshalErrorLike(err) }
);
sseStream.write({
event: 'error',
data: marshalErrorLike(err),
});
}
sseStream.end();
return sseStream;
}
if (!ctx.accepts('text/plain') && (ctx.accepts('text/json') || ctx.accepts('application/json'))) {
for await (const scrapped of it) {
lastScrapped = scrapped;
if (rpcReflect.signal.aborted) {
break;
}
if (!lastScrapped || !this.searchResultsQualified(lastScrapped, count)) {
continue;
}
await this.assignGeneralMixin(lastScrapped, count);
chargeAmount = this.assignChargeAmount(lastScrapped, count, chargeAmountScaler);
return lastScrapped;
}
if (!lastScrapped) {
throw new AssertionFailureError(`No content available for query ${searchQuery}`);
}
await this.assignGeneralMixin(lastScrapped, count);
chargeAmount = this.assignChargeAmount(lastScrapped, count, chargeAmountScaler);
return lastScrapped;
}
for await (const scrapped of it) {
lastScrapped = scrapped;
if (rpcReflect.signal.aborted) {
break;
}
if (!lastScrapped || !this.searchResultsQualified(lastScrapped, count)) {
continue;
}
await this.assignGeneralMixin(lastScrapped, count);
chargeAmount = this.assignChargeAmount(lastScrapped, count, chargeAmountScaler);
return assignTransferProtocolMeta(`${lastScrapped}`, { contentType: 'text/plain', envelope: null });
}
if (!lastScrapped) {
throw new AssertionFailureError(`No content available for query ${searchQuery}`);
}
await this.assignGeneralMixin(lastScrapped, count);
chargeAmount = this.assignChargeAmount(lastScrapped, count, chargeAmountScaler);
return assignTransferProtocolMeta(`${lastScrapped}`, { contentType: 'text/plain', envelope: null });
}
async *fetchSearchResults(
mode: string | 'markdown' | 'html' | 'text' | 'screenshot' | 'favicon' | 'content',
searchResults?: FormattedPage[],
options?: ExtraScrappingOptions,
crawlerOptions?: CrawlerOptions,
count?: number,
) {
if (!searchResults) {
return;
}
const urls = searchResults.map((x) => new URL(x.url!));
if (crawlerOptions) {
let offloaded = false;
for await (const x of this.offloadScrapMany(urls, crawlerOptions, options)) {
offloaded = true;
for (const [i, v] of x.entries()) {
if (v) {
Object.assign(searchResults[i], v);
}
}
yield this.reOrganizeSearchResults(searchResults, count);
}
if (offloaded) {
return;
}
}
const snapshotMap = new WeakMap();
for await (const scrapped of this.crawler.scrapMany(urls, options, crawlerOptions)) {
const mapped = scrapped.map((x, i) => {
if (!x) {
return {};
}
if (snapshotMap.has(x)) {
return snapshotMap.get(x);
}
return this.snapshotFormatter.formatSnapshot(mode, x, urls[i]).then((r) => {
snapshotMap.set(x, r);
return r;
}).catch((err) => {
this.logger.error(`Failed to format snapshot for ${urls[i].href}`, { err: marshalErrorLike(err) });
return {};
});
});
const resultArray = await Promise.all(mapped) as FormattedPage[];
for (const [i, v] of resultArray.entries()) {
if (v) {
Object.assign(searchResults[i], v);
}
}
yield this.reOrganizeSearchResults(searchResults, count);
}
}
async *offloadScrapMany(urls: URL[], crawlerOptions: CrawlerOptions, scrappingOptions?: ExtraScrappingOptions) {
if (!process.env.JINA_CRAWLER_OFFLOAD_ORIGIN) {
return undefined;
}
this.logger.debug(`Offloading ${urls.length} URLs to internal crawler API`, { urls });
const results = new Array(urls.length);
let nextDeferred = Defer<any>();
const abortCtrl = new AbortController();
const allJobs = urls.map(async (x, i) => {
try {
const cacheIt = this.crawler.queryCache(x, this.pageCacheToleranceMs);
const cacheMeta = await finalYield(cacheIt);
if (cacheMeta?.isFresh) {
const snapshot = cacheMeta.snapshot;
const patchedSnapshot = await this.jsdomControl.narrowSnapshot(snapshot, scrappingOptions);
if (!patchedSnapshot) {
return undefined;
}
const finalSnapshot = await this.snapshotFormatter.formatSnapshot(
crawlerOptions.respondWith,
patchedSnapshot,
new URL(patchedSnapshot.href)
);
results[i] = finalSnapshot;
return;
}
const r = await fetch(`${process.env.JINA_CRAWLER_OFFLOAD_ORIGIN}/`, {
method: 'POST',
headers: {
'Accept': 'text/event-stream',
'Content-Type': 'application/json',
},
body: JSON.stringify({
...crawlerOptions,
url: x.href,
proxy: 'none',
}),
signal: abortCtrl.signal,
});
if (r.status !== 200) {
const err = await r.json();
this.logger.warn(`Failed to offload scrap for ${x.href}, status: ${r.status}`, { status: r.status });
throw new DownstreamServiceFailureError({
message: `Failed to offload scrap for ${x.href}, status: ${r.status}`,
cause: err,
});
}
const streamNormalized = Readable.fromWeb(r.body as any);
const parsed = streamNormalized.pipe(new InputServerEventStream());
streamNormalized.on('error', (err) => parsed.destroy(err));
parsed.on('error', () => 'dont crash anything');
parsed.on('data', (event) => {
if (event.error) {
const err = new DownstreamServiceFailureError({
message: `Failed to offload scrap for ${x.href}`,
cause: event.error,
});
parsed.destroy(err);
return;
}
if (event.data && typeof event.data === 'object') {
results[i] = event.data;
nextDeferred.resolve();
nextDeferred = Defer<any>();
}
});
await once(parsed, 'end');
} catch (err: any) {
if (err?.name === 'AbortError') {
return;
}
this.logger.warn(`Failed to offload scrap for ${x.href}`, { err });
}
});
let done;
Promise.allSettled(allJobs).then(() => {
done = true;
nextDeferred.resolve();
});
yield results;
try {
do {
const r = await nextDeferred.promise.catch(() => 'dont crash anything');
if (typeof r === 'string') {
continue;
}
yield results;
} while (!done);
yield results;
} finally {
abortCtrl.abort();
this.logger.debug(`Offloaded scrap for ${urls.length} URLs done`, { urls });
}
}
reOrganizeSearchResults(searchResults: FormattedPage[], count?: number) {
const targetResultCount = count || this.targetResultCount;
const [qualifiedPages, unqualifiedPages] = _.partition(searchResults, (x) => this.pageQualified(x));
const acceptSet = new Set(qualifiedPages);
const n = targetResultCount - qualifiedPages.length;
for (const x of unqualifiedPages.slice(0, n >= 0 ? n : 0)) {
acceptSet.add(x);
}
const filtered = searchResults.filter((x) => acceptSet.has(x)).slice(0, targetResultCount);
const resultArray = filtered;
resultArray.toString = searchResults.toString;
return resultArray;
}
assignChargeAmount(formatted: FormattedPage[], num: number, scaler: number) {
let contentCharge = 0;
for (const x of formatted) {
const itemAmount = this.crawler.assignChargeAmount(x) || 0;
if (!itemAmount) {
continue;
}
contentCharge += itemAmount;
}
const numCharge = Math.ceil(formatted.length / 10) * 10000 * scaler;
let final = Math.max(contentCharge, numCharge);
if (final > 2_000_000) {
// We decided it's not quite fair to charge above 2M tokens even though the pages contains more.
final = 2_000_000;
}
if (final === numCharge) {
for (const x of formatted) {
x.usage = { tokens: Math.ceil(numCharge / formatted.length) };
}
}
const metadata: Record<string, any> = { usage: { tokens: final } };
assignMeta(formatted, metadata);
assignTransferProtocolMeta(formatted, { headers: { 'X-Usage-Tokens': final.toString() } });
return final;
}
pageQualified(formattedPage: FormattedPage) {
return formattedPage.title &&
formattedPage.content ||
formattedPage.screenshotUrl ||
formattedPage.pageshotUrl ||
formattedPage.text ||
formattedPage.html;
}
searchResultsQualified(results: FormattedPage[], targetResultCount = this.targetResultCount) {
return _.every(results, (x) => this.pageQualified(x)) && results.length >= targetResultCount;
}
async getFavicon(domain: string) {
const url = `https://www.google.com/s2/favicons?sz=32&domain_url=${domain}`;
try {
const response = await fetch(url);
if (!response.ok) {
return '';
}
const ab = await response.arrayBuffer();
const buffer = Buffer.from(ab);
const base64 = buffer.toString('base64');
return `data:image/png;base64,${base64}`;
} catch (error: any) {
this.logger.warn(`Failed to get favicon base64 string`, { err: marshalErrorLike(error) });
return '';
}
}
* iterProviders(preference?: string) {
if (preference === 'bing') {
if (this.envConfig.SERPER_SEARCH_API_KEY) {
yield this.serperBing;
}
yield this.bingSERP;
return;
}
if (preference === 'google') {
yield this.googleSERP;
if (this.envConfig.SERPER_SEARCH_API_KEY) {
yield this.serperGoogle;
}
yield this.commonGoogleSerp;
return;
}
if (this.envConfig.SERPER_SEARCH_API_KEY) {
yield this.serperGoogle;
}
yield this.googleSERP;
yield this.bingSERP;
yield this.commonGoogleSerp;
}
async *cachedSearch(query: Record<string, any> & { variant: 'web' | 'news' | 'images'; }, scrappingOptions: ExtraScrappingOptions, crawlerOptions: CrawlerOptions) {
const queryDigest = objHashMd5B64Of(query);
const variant = query.variant || 'web';
const provider = query.provider;
Reflect.deleteProperty(query, 'provider');
let cache;
let results;
let cacheUsed = false;
if (!crawlerOptions.noCache) {
cache = await this.storageLayer.findSERPResult({ queryDigest });
if (cache) {
const age = Date.now() - cache.createdAt.valueOf();
const stale = cache.createdAt.valueOf() < (Date.now() - this.cacheValidMs);
this.logger.info(`${stale ? 'Stale cache exists' : 'Cache hit'} for search query "${query.q}", normalized digest: ${queryDigest}, ${age}ms old`, {
query, digest: queryDigest, age, stale
});
if (!stale && Array.isArray(cache.response)) {
results = cache.response.filter((x) => x.link).map((x) => this.mapSearchEntryToPartialFormattedPage(x));
cacheUsed = true;
yield results;
}
}
}
if (!results) {
try {
let r: WebSearchEntry[] | undefined;
let lastError;
outerLoop:
for (const client of this.iterProviders(provider)) {
const t0 = Date.now();
let func;
switch (variant) {
case 'images': {
func = Reflect.get(client, 'imageSearch');
break;
}
case 'news': {
func = Reflect.get(client, 'newsSearch');
break;
}
case 'web':
default: {
func = Reflect.get(client, 'webSearch');
break;
}
}
if (!func) {
continue;
}
try {
r = await Reflect.apply(func, client, [query]);
const dt = Date.now() - t0;
this.logger.info(`Search took ${dt}ms, ${client.constructor.name}(${variant})`, { searchDt: dt, variant, client: client.constructor.name });
break outerLoop;
} catch (err) {
lastError = err;
const dt = Date.now() - t0;
this.logger.warn(`Failed to do ${variant} search using ${client.constructor.name}`, { err, variant, searchDt: dt, });
}
}
if (r?.length) {
const nowDate = new Date();
const record = SERPResult.from({
query,
queryDigest,
response: r,
createdAt: nowDate,
expireAt: new Date(nowDate.valueOf() + this.cacheRetentionMs)
});
Reflect.deleteProperty(record, '_id');
this.storageLayer.storeSERPResult(record).catch((err) => {
this.logger.warn(`Failed to save serp search result`, { err });
});
if (variant === 'web') {
r.map((x) => {
this.storageLayer.indexWebSearchEntry(x, {
geolocation: query.gl?.toLowerCase(),
language: bcp47ToIso639_3(query.hl),
}).catch((err) => {
this.logger.warn(`Failed to index SERP result`, { err });
});
});
}
} else if (lastError) {
throw lastError;
} else if (!r) {
throw new AssertionFailureError(`No provider can do ${variant} search atm.`);
}
results = r.map((x) => this.mapSearchEntryToPartialFormattedPage(x));
yield results;
} catch (err: any) {
if (cache) {
this.logger.warn(`Failed to fetch search result, but a stale cache is available. falling back to stale cache`, { err: marshalErrorLike(err) });
cacheUsed = true;
yield cache.response as any;
} else {
throw err;
}
}
}
if (cacheUsed && crawlerOptions.respondWith?.includes('no-content')) {
yield results;
return;
}
yield* this.fetchSearchResults(crawlerOptions.respondWith, results, scrappingOptions,
CrawlerOptions.from({ ...crawlerOptions, cacheTolerance: crawlerOptions.cacheTolerance ?? this.pageCacheToleranceMs }),
query.num,
);
}
mapSearchEntryToPartialFormattedPage(input: WebSearchEntry): FormattedPage {
const whitelistedProps = [
'imageUrl', 'imageWidth', 'imageHeight', 'source', 'date', 'siteLinks'
];
const result = {
title: input.title,
url: input.link,
description: Reflect.get(input, 'snippet'),
..._.pick(input, whitelistedProps),
};
return result;
}
async assignGeneralMixin(results: FormattedPage[], num?: number) {
const collectFavicon = this.threadLocal.get('collect-favicon');
await Promise.allSettled(results.map(async (x) => {
if (collectFavicon && x.url) {
const url = new URL(x.url);
Reflect.set(x, 'favicon', await this.getFavicon(url.origin));
}
x.content ??= '';
Object.setPrototypeOf(x, searchResultProto);
return x;
}));
if (num) {
results.length = Math.min(results.length, num);
}
results.toString = function (this: FormattedPage[]) {
let r = this.map((x, i) => x ? Reflect.apply(x.toString, x, [i]) : '').join('\n\n').trimEnd() + '\n';
return r;
};
return results;
}
async *readerLocalSearch(inputQuery: SerperSearchQueryParams, scrappingOptions: ExtraScrappingOptions, crawlerOptions: CrawlerOptions) {
const query = parseSearchQuery(inputQuery);
this.logger.debug(`Running local search`, { inputQuery, query });
const results = await this.storageLayer.searchLocalIndex(query.query, {
queryMixins: query.queryMixins,
num: inputQuery.num,
lang: bcp47ToIso639_3(inputQuery.hl),
});
const mapped = results.map((x) => ({
title: x.title,
url: x.url,
description: x.highlights?.map((x) => x.texts.map((y) => y.value)).flat().join(' ') || x.description,
date: x.publishedAt ? `${x.publishedAt.toISOString()}` : undefined,
})) as FormattedPage[];
this.logger.debug(`Local search returned ${mapped.length} results`, { resultsNum: mapped.length, inputQuery });
yield mapped;
this.logger.debug(`Extracting local cache for full-text`, { resultsNum: mapped.length, inputQuery });
const r = await Promise.allSettled(results.map((x, i) => {
return this.storageLayer.readFile(`snapshots/${x._id}`).then(async (r) => {
const snapshot = JSON.parse(r.toString());
const patchedSnapshot = await this.jsdomControl.narrowSnapshot(snapshot, scrappingOptions);
if (!patchedSnapshot) {
return undefined;
}
const finalSnapshot = await this.snapshotFormatter.formatSnapshot(
crawlerOptions.respondWith,
patchedSnapshot,
new URL(patchedSnapshot.href)
);
const lite = mapped[i] as any;
const pageDescription = lite.description?.replaceAll(x.description || '', '') || '';
finalSnapshot.description = `${x.description ? `${x.description}\n` : ''}${pageDescription}`;
finalSnapshot.title ??= lite.title;
// @ts-ignore
finalSnapshot.source ??= lite.source;
// @ts-ignore
finalSnapshot.date ??= lite.date;
return finalSnapshot;
});
}));
this.logger.debug(`Done extracting reader local cache for full-text`, { resultsNum: mapped.length, inputQuery });
yield r.map((x) => {
if (x.status === 'fulfilled') {
if (!x.value) {
return null;
}
return x.value;
}
this.logger.warn(`Failed to read snapshot`, { err: x.reason });
return null;
}).filter(Boolean) as FormattedPage[];
}
}
const dataItems = [
{ key: 'title', label: 'Title' },
{ key: 'source', label: 'Source' },
{ key: 'url', label: 'URL Source' },
{ key: 'imageUrl', label: 'Image URL' },
{ key: 'description', label: 'Description' },
{ key: 'publishedTime', label: 'Published Time' },
{ key: 'imageWidth', label: 'Image Width' },
{ key: 'imageHeight', label: 'Image Height' },
{ key: 'date', label: 'Date' },
{ key: 'favicon', label: 'Favicon' },
];
const searchResultProto = {
toString(this: FormattedPage, i?: number) {
const chunks = [];
for (const item of dataItems) {
const v = Reflect.get(this, item.key);
if (typeof v !== 'undefined') {
if (i === undefined) {
chunks.push(`[${item.label}]: ${v}`);
} else {
chunks.push(`[${i + 1}] ${item.label}: ${v}`);
}
}
}
const content = this.content || this.text;
if (content) {
chunks.push(`\n${content}`);
}
if (this.images) {
const imageSummaryChunks = [`${i === undefined ? '' : `[${i + 1}] `}Images:`];
for (const [k, v] of Object.entries(this.images)) {
imageSummaryChunks.push(`- ![${k}](${v})`);
}
if (imageSummaryChunks.length === 1) {
imageSummaryChunks.push('This page does not seem to contain any images.');
}
chunks.push(imageSummaryChunks.join('\n'));
}
if (this.links) {
const linkSummaryChunks = [`${i === undefined ? '' : `[${i + 1}] `}Links/Buttons:`];
if (Array.isArray(this.links)) {
for (const [k, v] of this.links) {
linkSummaryChunks.push(`- [${k}](${v})`);
}
} else {
for (const [k, v] of Object.entries(this.links)) {
linkSummaryChunks.push(`- [${k}](${v})`);
}
}
if (linkSummaryChunks.length === 1) {
linkSummaryChunks.push('This page does not seem to contain any buttons/links.');
}
chunks.push(linkSummaryChunks.join('\n'));
}
return chunks.join('\n');
}
};
src/api/serp.ts
+451
View File
@@ -0,0 +1,451 @@
import { singleton } from 'tsyringe';
import {
RPCHost, RPCReflection, assignMeta, RawString,
ParamValidationError,
assignTransferProtocolMeta,
AssertionFailureError,
} from 'civkit/civ-rpc';
import { marshalErrorLike } from 'civkit/lang';
import _ from 'lodash';
import { GlobalLogger } from '../services/logger';
import { AsyncLocalContext } from '../services/async-context';
import { Context, Ctx, Method, Param, RPCReflect } from '../services/registry';
import { OutputServerEventStream } from '../lib/transform-server-event-stream';
// import { JinaEmbeddingsAuthDTO } from '../dto/jina-embeddings-auth';
import { WORLD_COUNTRIES, WORLD_LANGUAGES } from '../3rd-party/serper-search';
import { GoogleSERP, GoogleSERPOldFashion } from '../services/serp/google';
import { WebSearchEntry } from '../services/serp/compat';
import { CrawlerOptions } from '../dto/crawler-options';
import { ScrappingOptions } from '../services/serp/puppeteer';
import { objHashMd5B64Of } from 'civkit/hash';
import { SerperBingSearchService, SerperGoogleSearchService } from '../services/serp/serper';
import { CommonGoogleSERP } from '../services/serp/common-serp';
import { BingSERP } from '../services/serp/bing';
import { bcp47ToIso639_3 } from '../utils/languages';
import { StorageLayer } from '../db/noop-storage';
import { BaseAuthDTO } from '../dto/base-auth';
import { SERPResult } from '../db/models';
import { EnvConfig } from '../services/envconfig';
const WORLD_COUNTRY_CODES = Object.keys(WORLD_COUNTRIES).map((x) => x.toLowerCase());
const indexProto = {
toString: function (): string {
return _(this)
.toPairs()
.map(([k, v]) => k ? `[${_.upperFirst(_.lowerCase(k))}] ${v}` : '')
.value()
.join('\n') + '\n';
}
};
@singleton()
export class SerpHost extends RPCHost {
logger = this.globalLogger.child({ service: this.constructor.name });
cacheRetentionMs = 1000 * 3600 * 24 * 7;
cacheValidMs = 1000 * 3600;
pageCacheToleranceMs = 1000 * 3600 * 24;
reasonableDelayMs = 15_000;
targetResultCount = 5;
async getIndex(ctx: Context, auth?: BaseAuthDTO) {
const indexObject: Record<string, string | number | undefined> = Object.create(indexProto);
Object.assign(indexObject, {
usage1: 'https://r.jina.ai/YOUR_URL',
usage2: 'https://s.jina.ai/YOUR_SEARCH_QUERY',
usage3: `${ctx.origin}/?q=YOUR_SEARCH_QUERY`,
homepage: 'https://jina.ai/reader',
});
if (auth && auth.user) {
// if (auth instanceof JinaEmbeddingsAuthDTO) {
// indexObject[''] = undefined;
// indexObject.authenticatedAs = `${auth.user.user_id} (${auth.user.full_name})`;
// indexObject.balanceLeft = auth.user.wallet.total_balance;
// }
} else {
indexObject.note = 'Authentication is required to use this endpoint. Please provide a valid API key via Authorization header.';
}
return indexObject;
}
constructor(
protected globalLogger: GlobalLogger,
protected threadLocal: AsyncLocalContext,
protected envConfig: EnvConfig,
protected googleSerp: GoogleSERP,
protected googleSerpOld: GoogleSERPOldFashion,
protected serperGoogle: SerperGoogleSearchService,
protected serperBing: SerperBingSearchService,
protected commonGoogleSerp: CommonGoogleSERP,
protected bingSERP: BingSERP,
protected storageLayer: StorageLayer,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
@Method({
name: 'searchIndex',
ext: {
http: {
action: ['get', 'post'],
path: '/'
}
},
tags: ['search'],
returnType: [String, OutputServerEventStream, RawString],
})
@Method({
ext: {
http: {
action: ['get', 'post'],
}
},
tags: ['search'],
returnType: [String, OutputServerEventStream, RawString],
})
async search(
@RPCReflect() rpcReflect: RPCReflection,
@Ctx() ctx: Context,
crawlerOptions: CrawlerOptions,
auth: BaseAuthDTO,
@Param('type', { type: new Set(['web', 'images', 'news']), default: 'web' })
variant: 'web' | 'images' | 'news',
@Param('q') q?: string,
@Param('provider', { type: new Set(['google', 'bing']) })
@Param('engine', { type: new Set(['google', 'bing']) })
searchEngine?: 'google' | 'bing',
@Param('num', { validate: (v: number) => v >= 0 && v <= 20 })
num?: number,
@Param('gl', { validate: (v: string) => WORLD_COUNTRY_CODES.includes(v?.toLowerCase()) }) gl?: string,
@Param('hl', { validate: (v: string) => WORLD_LANGUAGES.some(l => l.code === v) }) hl?: string,
@Param('location') location?: string,
@Param('page') page?: number,
@Param('fallback') fallback?: boolean,
@Param('nfpr') nfpr?: boolean,
) {
if (!q) {
if (ctx.path === '/') {
const indexObject = await this.getIndex(ctx, auth);
if (!ctx.accepts('text/plain') && (ctx.accepts('text/json') || ctx.accepts('application/json'))) {
return indexObject;
}
return assignTransferProtocolMeta(`${indexObject}`,
{ contentType: 'text/plain; charset=utf-8', envelope: null }
);
}
throw new ParamValidationError({
path: 'q',
message: `Required but not provided`
});
}
let chargeAmount = 0;
const {
reportOptions,
reportUsage
} = await this.storageLayer.rateLimit(ctx, rpcReflect, auth as any);
rpcReflect.finally(() => {
reportOptions?.(crawlerOptions.customizedProps());
reportUsage?.(chargeAmount, 'reader-search');
});
let chargeAmountScaler = 1;
if (searchEngine === 'bing') {
chargeAmountScaler = 3;
}
if (variant !== 'web') {
chargeAmountScaler = 5;
}
let realQuery = q;
let queryTerms = q.split(/\s+/g).filter((x) => !!x);
let results = await this.cachedSearch(variant, {
provider: searchEngine,
q,
num,
gl,
hl,
location,
page,
nfpr,
}, crawlerOptions);
if (fallback && !results?.length && (!page || page === 1)) {
let tryTimes = 1;
const containsRTL = /[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDFF\uFE70-\uFEFF\u0590-\u05FF\uFB1D-\uFB4F\u0700-\u074F\u0780-\u07BF\u07C0-\u07FF]/.test(q);
const lastResort = (containsRTL ? queryTerms.slice(queryTerms.length - 2) : queryTerms.slice(0, 2)).join(' ');
const n = 4;
let terms: string[] = [];
while (tryTimes < n) {
const delta = Math.ceil(queryTerms.length / n) * tryTimes;
terms = containsRTL ? queryTerms.slice(delta) : queryTerms.slice(0, queryTerms.length - delta);
const query = terms.join(' ');
if (!query) {
break;
}
if (realQuery === query) {
continue;
}
tryTimes += 1;
realQuery = query;
this.logger.info(`Retrying search with fallback query: "${realQuery}"`);
results = await this.cachedSearch(variant, {
provider: searchEngine,
q: realQuery,
num,
gl,
// hl,
location,
}, crawlerOptions);
if (results?.length) {
break;
}
}
if (!results?.length && realQuery.length > lastResort.length) {
realQuery = lastResort;
this.logger.info(`Retrying search with fallback query: "${realQuery}"`);
tryTimes += 1;
results = await this.cachedSearch(variant, {
provider: searchEngine,
q: realQuery,
num,
gl,
// hl,
location,
}, crawlerOptions);
}
chargeAmountScaler *= tryTimes;
}
if (!results?.length) {
results = [];
}
const finalResults = results.map((x: any) => this.mapToFinalResults(x));
await Promise.all(finalResults.map((x: any) => this.assignGeneralMixin(x)));
chargeAmount = this.assignChargeAmount(finalResults, chargeAmountScaler);
assignMeta(finalResults, {
query: realQuery,
fallback: realQuery === q ? undefined : realQuery,
});
return finalResults;
}
assignChargeAmount(items: unknown[], scaler: number) {
const numCharge = Math.ceil(items.length / 10) * 10000 * scaler;
assignMeta(items, { usage: { tokens: numCharge } });
return numCharge;
}
async getFavicon(domain: string) {
const url = `https://www.google.com/s2/favicons?sz=32&domain_url=${domain}`;
try {
const response = await fetch(url);
if (!response.ok) {
return '';
}
const ab = await response.arrayBuffer();
const buffer = Buffer.from(ab);
const base64 = buffer.toString('base64');
return `data:image/png;base64,${base64}`;
} catch (error: any) {
this.logger.warn(`Failed to get favicon base64 string`, { err: marshalErrorLike(error) });
return '';
}
}
async configure(opts: CrawlerOptions) {
const crawlOpts: ScrappingOptions = {
proxyUrl: opts.proxyUrl,
cookies: opts.setCookies,
overrideUserAgent: opts.userAgent,
timeoutMs: opts.timeout ? opts.timeout * 1000 : undefined,
locale: opts.locale,
referer: opts.referer,
viewport: opts.viewport,
proxyResources: (opts.proxyUrl || opts.proxy?.endsWith('+')) ? true : false,
allocProxy: opts.proxy?.endsWith('+') ? opts.proxy.slice(0, -1) : opts.proxy,
};
if (opts.locale) {
crawlOpts.extraHeaders ??= {};
crawlOpts.extraHeaders['Accept-Language'] = opts.locale;
}
return crawlOpts;
}
mapToFinalResults(input: WebSearchEntry) {
const whitelistedProps = [
'imageUrl', 'imageWidth', 'imageHeight', 'source', 'date', 'siteLinks'
];
const result = {
title: input.title,
url: input.link,
description: Reflect.get(input, 'snippet'),
..._.pick(input, whitelistedProps),
};
return result;
}
*iterProviders(preference?: string) {
if (preference === 'bing') {
if (this.envConfig.SERPER_SEARCH_API_KEY) {
yield this.serperBing;
}
yield this.bingSERP;
return;
}
if (preference === 'google') {
yield this.googleSerp;
if (this.envConfig.SERPER_SEARCH_API_KEY) {
yield this.serperGoogle;
}
yield this.commonGoogleSerp;
return;
}
if (this.envConfig.SERPER_SEARCH_API_KEY) {
yield this.serperGoogle;
}
yield this.googleSerp;
yield this.bingSERP;
yield this.commonGoogleSerp;
}
async cachedSearch(variant: 'web' | 'news' | 'images', query: Record<string, any>, opts: CrawlerOptions) {
const queryDigest = objHashMd5B64Of({ ...query, variant });
const provider = query.provider;
Reflect.deleteProperty(query, 'provider');
const noCache = opts.noCache;
let cache;
if (!noCache) {
cache = await this.storageLayer.findSERPResult({ queryDigest });
if (cache) {
const age = Date.now() - cache.createdAt.valueOf();
const stale = cache.createdAt.valueOf() < (Date.now() - this.cacheValidMs);
this.logger.info(`${stale ? 'Stale cache exists' : 'Cache hit'} for search query "${query.q}", normalized digest: ${queryDigest}, ${age}ms old`, {
query, digest: queryDigest, age, stale
});
if (!stale) {
return cache.response as any;
}
}
}
const scrappingOptions = await this.configure(opts);
try {
let r: any[] | undefined;
let lastError;
outerLoop:
for (const client of this.iterProviders(provider)) {
const t0 = Date.now();
let func;
switch (variant) {
case 'images': {
func = Reflect.get(client, 'imageSearch');
break;
}
case 'news': {
func = Reflect.get(client, 'newsSearch');
break;
}
case 'web':
default: {
func = Reflect.get(client, 'webSearch');
break;
}
}
if (!func) {
continue;
}
try {
r = await Reflect.apply(func, client, [query, scrappingOptions]);
const dt = Date.now() - t0;
this.logger.info(`Search took ${dt}ms, ${client.constructor.name}(${variant})`, { searchDt: dt, variant, client: client.constructor.name });
break outerLoop;
} catch (err) {
lastError = err;
const dt = Date.now() - t0;
this.logger.warn(`Failed to do ${variant} search using ${client.constructor.name}`, { err, variant, searchDt: dt, });
}
}
if (r?.length) {
const nowDate = new Date();
const record = SERPResult.from({
query,
queryDigest,
response: r,
createdAt: nowDate,
expireAt: new Date(nowDate.valueOf() + this.cacheRetentionMs)
});
Reflect.deleteProperty(record, '_id');
this.storageLayer.storeSERPResult(record).catch((err) => {
this.logger.warn(`Failed to save search result`, { err: marshalErrorLike(err) });
});
if (variant === 'web') {
r.map((x) => {
this.storageLayer.indexWebSearchEntry(x, {
geolocation: query.gl?.toLowerCase(),
language: bcp47ToIso639_3(query.hl),
}).catch((err) => {
this.logger.warn(`Failed to index SERP result`, { err });
});
});
}
} else if (lastError) {
throw lastError;
} else if (!r) {
throw new AssertionFailureError(`No provider can do ${variant} search atm.`);
}
return r;
} catch (err: any) {
if (cache) {
this.logger.warn(`Failed to fetch search result, but a stale cache is available. falling back to stale cache`, { err: marshalErrorLike(err) });
return cache.response as any;
}
throw err;
}
}
async assignGeneralMixin(result: Partial<WebSearchEntry>) {
const collectFavicon = this.threadLocal.get('collect-favicon');
if (collectFavicon && result.link) {
const url = new URL(result.link);
Reflect.set(result, 'favicon', await this.getFavicon(url.origin));
}
}
}
src/config.ts
+13
View File
@@ -0,0 +1,13 @@
import { container } from 'tsyringe';
import { StorageLayer } from './db/noop-storage';
import { BaseAuthDTO } from './dto/base-auth';
export let AUTH_DTO_CLS = BaseAuthDTO;
export let STORAGE_CLS = StorageLayer;
if (process.env.GCP_STORAGE_ENDPOINT && (process.env.GCP_STORAGE_BUCKET || process.env.GCLOUD_PROJECT)) {
STORAGE_CLS = require('./db/bucket-storage').BucketStorageLayer;
}
container.registerSingleton(StorageLayer, STORAGE_CLS);
src/db/bucket-storage.ts
+315
View File
@@ -0,0 +1,315 @@
import { singleton } from 'tsyringe';
import { ASNBlockade, ConsecutiveError, Crawled, DomainBlockade, ImgAlt, IndexedPage, SERPResult } from './models';
import { Context } from '../services/registry';
import { RPCReflection } from 'civkit/civ-rpc';
import { HashManager } from 'civkit/hash';
// import { JinaEmbeddingsAuthDTO } from '../dto/jina-embeddings-auth';
import _ from 'lodash';
import type { PageSnapshot } from '../services/puppeteer';
import { WebSearchEntry } from '../services/serp/compat';
import { StorageLayer } from './noop-storage';
import { DefaultBucket } from '../services/default-bucket';
import { BaseAuthDTO } from '../dto/base-auth';
const md5Hasher = new HashManager('md5', 'hex');
@singleton()
export class BucketStorageLayer extends StorageLayer {
constructor(protected defaultBucket: DefaultBucket) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
override async findImageAlt(query: Partial<ImgAlt>): Promise<ImgAlt | undefined> {
const id = query.urlDigest || query._id;
if (!id) {
throw new Error('ID or URL digest is required when finding image alt text');
}
try {
const buff = await this.defaultBucket.readSingleFile(`image-alts/${id}`);
return ImgAlt.from(JSON.parse(buff.toString()));
} catch (err) {
return undefined;
}
}
override async storeImageAlt(draft: ImgAlt): Promise<unknown> {
const id = draft.urlDigest || draft._id;
if (!id) {
throw new Error('ID or URL digest is required when storing image alt text');
}
await this.defaultBucket.putBuffer(`image-alts/${id}`, Buffer.from(JSON.stringify(draft)), {
contentType: 'application/json',
});
return undefined;
}
override async findDomainBlockade(query: Partial<DomainBlockade>): Promise<DomainBlockade | undefined> {
const id = query.domain;
if (!id) {
throw new Error('ID or URL digest is required when finding domain blockade');
}
try {
const buff = await this.defaultBucket.readSingleFile(`domain-blockade/${md5Hasher.hash(id.toString())}`);
return DomainBlockade.from(JSON.parse(buff.toString()));
} catch (err) {
return undefined;
}
}
override async storeDomainBlockade(draft: DomainBlockade): Promise<unknown> {
const id = draft.domain;
if (!id) {
throw new Error('ID or URL digest is required when storing domain blockade');
}
await this.defaultBucket.putBuffer(`domain-blockade/${md5Hasher.hash(id.toString())}`, Buffer.from(JSON.stringify(draft)), {
contentType: 'application/json',
});
return undefined;
}
override async findASNBlockade(asn: number): Promise<ASNBlockade | undefined> {
try {
const buff = await this.defaultBucket.readSingleFile(`asn-blockade/${asn}`);
const r = ASNBlockade.from(JSON.parse(buff.toString()));
if (r.expireAt && new Date(r.expireAt).valueOf() <= Date.now()) {
return undefined;
}
return r;
} catch (err) {
return undefined;
}
}
override async storeASNBlockade(draft: ASNBlockade): Promise<unknown> {
if (draft._id === undefined || draft._id === null) {
throw new Error('ASN is required when storing ASN blockade');
}
await this.defaultBucket.putBuffer(`asn-blockade/${draft._id}`, Buffer.from(JSON.stringify(draft)), {
contentType: 'application/json',
});
return undefined;
}
override async rateLimit(ctx: Context, _rpcReflect: RPCReflection, auth: BaseAuthDTO): Promise<{
isAnonymous: boolean;
record: object;
uid?: string;
rateLimitPolicies?: object[];
reportOptions?: (opts: object) => void;
reportUsage?: (amount: number, desc?: string) => void;
}> {
const uid = await auth.solveUID();
let isAnonymous = uid ? false : true;
let reportOptions;
let reportUsage;
const apiRoll: Record<string, any> = {
uid,
ip: ctx.ip,
};
reportOptions = (opts: object) => apiRoll.crawlerOptions = opts;
reportUsage = (amount: number, _desc?: string) => {
apiRoll.chargeAmount = amount;
};
return {
isAnonymous,
record: apiRoll,
uid,
rateLimitPolicies: [],
reportOptions,
reportUsage,
};
}
override async findConsecutiveError(query: Partial<ConsecutiveError>): Promise<ConsecutiveError | undefined> {
if (!query.count) {
throw new Error('Count is required when finding consecutive error record');
}
let id = query._id;
if (!id && query.url) {
id = md5Hasher.hash<string>(query.url.toLowerCase());
}
if (!id) {
throw new Error('URL is required when finding consecutive error record');
}
try {
const buff = await this.defaultBucket.readSingleFile(`consecutive-error/${id}`);
const record = ConsecutiveError.from(JSON.parse(buff.toString()));
if (record.count > query.count) {
return record;
}
return undefined;
} catch (err) {
return undefined;
}
}
override async storeConsecutiveError(draft: ConsecutiveError): Promise<unknown> {
let id = draft._id;
if (!id && draft.url) {
id = md5Hasher.hash<string>(draft.url.toLowerCase());
}
if (!id) {
throw new Error('URL is required when finding consecutive error record');
}
const r = await this.findConsecutiveError(draft);
await this.defaultBucket.putBuffer(`consecutive-error/${id}`, Buffer.from(JSON.stringify({
...r, ...draft,
count: (r?.count || 0) + draft.count
})), {
contentType: 'application/json',
});
return undefined;
}
override async clearConsecutiveError(draft: Partial<ConsecutiveError>): Promise<unknown> {
let id = draft._id;
if (!id && draft.url) {
id = md5Hasher.hash<string>(draft.url.toLowerCase());
}
if (!id) {
throw new Error('URL is required when clearing consecutive error record');
}
await this.defaultBucket.removeSingleFile(`consecutive-error/${id}`);
return undefined;
}
override async findPageCache(query: Partial<Crawled>): Promise<Crawled | undefined> {
const id = query.urlPathDigest || query._id;
if (!id) {
throw new Error('ID or URL digest is required when finding page cache record');
}
try {
const buff = await this.defaultBucket.readSingleFile(`page-cache/${id}`);
return Crawled.from(JSON.parse(buff.toString()));
} catch (err) {
return undefined;
}
}
override async storePageCache(draft: Crawled): Promise<unknown> {
const id = draft.urlPathDigest || draft._id;
if (!id) {
throw new Error('URL digest or id is required when storing page cache record');
}
await this.defaultBucket.putBuffer(`page-cache/${id}`, Buffer.from(JSON.stringify(draft)), {
contentType: 'application/json',
});
return undefined;
}
override async clearPageCache(draft: Partial<Crawled>): Promise<unknown> {
const id = draft.urlDigest || draft._id;
if (!id) {
throw new Error('URL digest or id is required when clearing page cache record');
}
await this.defaultBucket.removeSingleFile(`page-cache/${id}`);
return undefined;
}
override async indexSnapshot(_digest: string, _snapshot: PageSnapshot): Promise<unknown> {
return undefined;
}
override async indexWebSearchEntry(webSearchEntry: WebSearchEntry, additional?: Partial<IndexedPage>): Promise<unknown> {
return undefined;
}
override async searchLocalIndex(_query: string,
_options: {
queryMixins?: { filter?: any[], should?: any[], must?: any[], mustNot?: any[]; };
num?: number;
lang?: string;
} = {}
): Promise<{
_id: string;
title: string;
url: string;
domain: string;
description?: string;
publishedAt?: Date;
highlights?: {
path: string;
texts: { value: string, type: string; }[];
score: number;
}[];
score: number;
seq: string;
}[]> {
return [];
}
override async storeSERPResult(draft: SERPResult): Promise<unknown> {
const id = draft.queryDigest || draft._id;
if (!id) {
throw new Error('URL digest or id is required when storing SERP result');
}
await this.defaultBucket.putBuffer(`serp-results/${id}`, Buffer.from(JSON.stringify(draft)), {
contentType: 'application/json',
});
return undefined;
}
override async findSERPResult(query: Partial<SERPResult>): Promise<SERPResult | undefined> {
const id = query.queryDigest || query._id;
if (!id) {
throw new Error('ID or URL digest is required when finding SERP result');
}
try {
const buff = await this.defaultBucket.readSingleFile(`serp-results/${id}`);
return SERPResult.from(JSON.parse(buff.toString()));
} catch (err) {
return undefined;
}
}
override async storeFile(path: string, buffer: Buffer, metadata?: Record<string, any>): Promise<unknown> {
await this.defaultBucket.putBuffer(path, buffer, metadata);
return undefined;
}
override async readFile(path: string): Promise<Buffer> {
try {
return await this.defaultBucket.readSingleFile(path);
} catch (err) {
throw new Error(`Failed to read file ${path}: ${err}`);
}
}
override async fileExists(path: string): Promise<boolean> {
try {
await this.defaultBucket.getSingleFileStat(path);
return true;
} catch (_err) {
return false;
}
}
override async signDownloadUrl(path: string, expirySec?: number): Promise<string> {
const url = await this.defaultBucket.signDownloadObject(path, expirySec);
return url;
}
}
src/db/models.ts
+293
View File
@@ -0,0 +1,293 @@
import { Also, Coercible, Prop } from 'civkit/coercible';
import type { PageSnapshot } from '../services/puppeteer';
import { UUID } from 'bson';
@Also({
dictOf: Object
})
export class ConsecutiveError extends Coercible {
@Prop({
required: true,
})
_id!: string;
@Prop({
required: true
})
url!: string;
@Prop({
required: true,
default: 0,
})
count!: number;
@Prop()
lastError?: string;
@Prop()
createdAt!: Date;
@Prop()
updatedAt!: Date;
@Prop()
expireAt?: Date;
[k: string]: any;
}
@Also({
dictOf: Object
})
export class Crawled extends Coercible {
@Prop({
required: true
})
_id!: string;
@Prop({
required: true
})
url!: string;
@Prop({
required: true
})
urlPathDigest!: string;
@Prop()
htmlSignificantlyModifiedByJs?: boolean;
@Prop()
coerced?: boolean;
@Prop({ arrayOf: String })
traits?: string[];
@Prop()
snapshot?: PageSnapshot & { screenshot: never; pageshot: never; };
@Prop()
screenshotAvailable?: boolean;
@Prop()
pageshotAvailable?: boolean;
@Prop()
snapshotAvailable?: boolean;
@Prop()
createdAt!: Date;
@Prop()
expireAt!: Date;
[k: string]: any;
}
@Also({
dictOf: Object
})
export class DomainBlockade extends Coercible {
@Prop({
defaultFactory: () => new UUID()
})
_id!: UUID;
@Prop({
required: true
})
domain!: string;
@Prop({ required: true })
triggerReason!: string;
@Prop()
triggerUrl?: string;
@Prop()
createdAt!: Date;
@Prop()
expireAt?: Date;
[k: string]: any;
}
@Also({
dictOf: Object
})
export class ASNBlockade extends Coercible {
@Prop({
required: true
})
_id!: number;
@Prop()
note?: string;
@Prop()
createdAt!: Date;
@Prop()
expireAt?: Date;
[k: string]: any;
}
@Also({
dictOf: Object
})
export class ImgAlt extends Coercible {
@Prop({
required: true,
})
_id!: string;
@Prop({
required: true
})
src!: string;
@Prop({
required: true
})
urlDigest!: string;
@Prop()
width?: number;
@Prop()
height?: number;
@Prop()
generatedAlt?: string;
@Prop()
originalAlt?: string;
@Prop()
createdAt!: Date;
@Prop()
expireAt?: Date;
[k: string]: any;
}
export class IndexedPage extends Coercible {
@Prop({ required: true })
_id!: string;
@Prop({ required: true })
url!: string;
@Prop({ required: true })
urlPathDigest!: string;
@Prop({ required: true })
domain!: string;
@Prop({ required: true })
tld!: string;
@Prop()
language?: string;
@Prop()
geolocation?: string;
@Prop()
title?: string;
@Prop()
description?: string;
@Prop()
text?: Record<string, string>;
@Prop()
semanticText?: string;
@Prop()
createdAt!: Date;
@Prop()
publishedAt?: Date;
@Prop()
scrappedAt?: Date;
@Prop()
expireAt!: Date;
[k: string]: any;
}
@Also({
dictOf: Object
})
export class PDFContent extends Coercible {
@Prop({
required: true
})
_id!: string;
@Prop({
required: true
})
src!: string;
@Prop({
required: true
})
urlDigest!: string;
@Prop()
meta?: { [k: string]: any; };
@Prop()
text?: string;
@Prop()
content?: string;
@Prop()
createdAt!: Date;
@Prop()
expireAt?: Date;
[k: string]: any;
}
@Also({
dictOf: Object
})
export class SERPResult extends Coercible {
@Prop({
defaultFactory: () => new UUID()
})
_id!: UUID;
@Prop({
required: true
})
query!: any;
@Prop({
required: true
})
queryDigest!: string;
@Prop()
response?: any;
@Prop()
createdAt!: Date;
@Prop()
expireAt?: Date;
[k: string]: any;
}
src/db/noop-storage.ts
+169
View File
@@ -0,0 +1,169 @@
import { AsyncService } from 'civkit/async-service';
import { singleton } from 'tsyringe';
import { ASNBlockade, ConsecutiveError, Crawled, DomainBlockade, ImgAlt, IndexedPage, SERPResult } from './models';
import { Context } from '../services/registry';
import { RPCReflection } from 'civkit/civ-rpc';
// import { JinaEmbeddingsAuthDTO } from '../dto/jina-embeddings-auth';
import _ from 'lodash';
import type { PageSnapshot } from '../services/puppeteer';
import { WebSearchEntry } from '../services/serp/compat';
import { BaseAuthDTO } from '../dto/base-auth';
@singleton()
export class StorageLayer extends AsyncService {
override async init() {
await this.dependencyReady();
this.emit('ready');
}
async findImageAlt(_query: Partial<ImgAlt>): Promise<ImgAlt | undefined> {
return undefined;
}
async storeImageAlt(_draft: ImgAlt): Promise<unknown> {
return undefined;
}
async findDomainBlockade(_query: Partial<DomainBlockade>): Promise<DomainBlockade | undefined> {
return undefined;
}
async storeDomainBlockade(_draft: DomainBlockade): Promise<unknown> {
return undefined;
}
async findASNBlockade(_asn: number): Promise<ASNBlockade | undefined> {
return undefined;
}
async storeASNBlockade(_draft: ASNBlockade): Promise<unknown> {
return undefined;
}
async rateLimit(ctx: Context, _rpcReflect: RPCReflection, auth: BaseAuthDTO): Promise<{
isAnonymous: boolean;
record: object;
uid?: string;
rateLimitPolicies?: object[];
reportOptions?: (opts: object) => void;
reportUsage?: (amount: number, desc?: string) => void;
}> {
const uid = await auth.solveUID();
let isAnonymous = uid ? false : true;
let reportOptions;
let reportUsage;
const apiRoll: Record<string, any> = {
uid,
ip: ctx.ip,
};
reportOptions = (opts: object) => apiRoll.crawlerOptions = opts;
reportUsage = (amount: number, _desc?: string) => {
apiRoll.chargeAmount = amount;
};
return {
isAnonymous,
record: apiRoll,
uid,
rateLimitPolicies: [],
reportOptions,
reportUsage,
};
}
async findConsecutiveError(query: Partial<ConsecutiveError>): Promise<ConsecutiveError | undefined> {
if (!query.count) {
throw new Error('Count is required when finding consecutive error record');
}
return undefined;
}
async storeConsecutiveError(_draft: ConsecutiveError): Promise<unknown> {
return undefined;
}
async clearConsecutiveError(draft: Partial<ConsecutiveError>): Promise<unknown> {
if (!draft._id) {
throw new Error('ID is required to clear consecutive error record');
}
return undefined;
}
async findPageCache(_query: Partial<Crawled>): Promise<Crawled | undefined> {
return undefined;
}
async storePageCache(_draft: Crawled): Promise<unknown> {
return undefined;
}
async clearPageCache(draft: Partial<Crawled>): Promise<unknown> {
if (!draft._id) {
throw new Error('ID is required to clear page cache record');
}
return undefined;
}
async indexSnapshot(_digest: string, _snapshot: PageSnapshot): Promise<unknown> {
return undefined;
}
async indexWebSearchEntry(webSearchEntry: WebSearchEntry, additional?: Partial<IndexedPage>): Promise<unknown> {
return undefined;
}
async searchLocalIndex(_query: string,
_options: {
queryMixins?: { filter?: any[], should?: any[], must?: any[], mustNot?: any[]; };
num?: number;
lang?: string;
} = {}
): Promise<{
_id: string;
title: string;
url: string;
domain: string;
description?: string;
publishedAt?: Date;
highlights?: {
path: string;
texts: { value: string, type: string; }[];
score: number;
}[];
score: number;
seq: string;
}[]> {
return [];
}
async storeSERPResult(_draft: SERPResult): Promise<unknown> {
return undefined;
}
async findSERPResult(_query: Partial<SERPResult>): Promise<SERPResult | undefined> {
return undefined;
}
async storeFile(_path: string, _buffer: Buffer, _metadata?: Record<string, any>): Promise<unknown> {
return undefined;
}
async readFile(_path: string): Promise<Buffer> {
throw new Error('Not implemented');
}
async fileExists(_path: string): Promise<boolean> {
return false;
}
async signDownloadUrl(_path: string, _expirySec?: number): Promise<string> {
return '';
}
}
src/dto/base-auth.ts
+138
View File
@@ -0,0 +1,138 @@
import _ from 'lodash';
import {
Also, AuthenticationRequiredError,
RPC_CALL_ENVIRONMENT,
Coercible,
} from 'civkit/civ-rpc';
import { htmlEscape } from 'civkit/escape';
import type { Context } from 'koa';
import { InjectProperty } from '../services/registry';
import { AsyncLocalContext } from '../services/async-context';
import { TierFeatureConstraintError } from '../services/errors';
import { isIPInNonPublicRange } from '../utils/ip';
@Also({
openapi: {
operation: {
parameters: {
'Authorization': {
description: htmlEscape`Token for authentication.\n\n` +
htmlEscape`- Member of <BaseAuthDTO>\n\n` +
`- Authorization: Bearer {YOUR_JINA_TOKEN}`
,
in: 'header',
schema: {
anyOf: [
{ type: 'string', format: 'token' }
]
}
}
}
}
}
})
export class BaseAuthDTO<U extends object = object> extends Coercible {
uid?: string;
bearerToken?: string;
user?: U;
ip?: string;
isInternal?: boolean;
@InjectProperty(AsyncLocalContext)
ctxMgr!: AsyncLocalContext;
static override from(input: any) {
const instance = super.from(input) as BaseAuthDTO;
instance.isInternal = false;
const ctx = input[RPC_CALL_ENVIRONMENT] as Context;
if (ctx) {
const authorization = ctx.get('authorization');
if (authorization) {
const authToken = authorization.split(' ')[1] || authorization;
instance.bearerToken = authToken;
}
if (ctx.ip) {
instance.ip = ctx.ip;
instance.isInternal = ctx.ips.length <= 2 && isIPInNonPublicRange(ctx.ip);
}
}
if (!instance.bearerToken && input._token) {
instance.bearerToken = input._token;
}
return instance;
}
async solveUID() {
if (this.uid) {
this.ctxMgr.set('uid', this.uid);
this.ctxMgr.set('bearerToken', this.bearerToken);
return this.uid;
}
if (this.isInternal) {
return undefined;
}
if (this.bearerToken) {
throw new Error('Not implemented');
}
return undefined;
}
async assertUID() {
const uid = await this.solveUID();
if (!uid) {
throw new AuthenticationRequiredError('Authentication failed');
}
return uid;
}
async assertUser() {
if (this.user) {
return this.user;
}
throw new Error('Not implemented');
}
async assertTier(n: number, feature?: string) {
if (this.isInternal) {
return true;
}
let user;
try {
user = await this.assertUser();
} catch (err) {
if (err instanceof AuthenticationRequiredError) {
throw new AuthenticationRequiredError({
message: `Authentication is required to use this feature${feature ? ` (${feature})` : ''}. Please provide a valid API key.`
});
}
throw err;
}
const tier = user ? 2 : 0;
if (isNaN(tier) || tier < n) {
throw new TierFeatureConstraintError({
message: `Your current plan does not support this feature${feature ? ` (${feature})` : ''}. Please upgrade your plan.`
});
}
return true;
}
}
src/dto/crawler-options.ts
+1027
View File
File diff suppressed because it is too large Load Diff
src/dto/turndown-tweakable-options.ts
+61
View File
@@ -0,0 +1,61 @@
import { Coercible, Prop } from 'civkit/coercible';
import {Context} from '../services/registry';
import _ from 'lodash';
export class TurnDownTweakableOptions extends Coercible {
@Prop({
desc: 'Turndown options > headingStyle',
type: new Set(['setext', 'atx']),
})
headingStyle?: 'setext' | 'atx';
@Prop({
desc: 'Turndown options > hr',
validate: (v: string) => v.length > 0 && v.length <= 128
})
hr?: string;
@Prop({
desc: 'Turndown options > bulletListMarker',
type: new Set(['-', '+', '*']),
})
bulletListMarker?: '-' | '+' | '*';
@Prop({
desc: 'Turndown options > emDelimiter',
type: new Set(['_', '*']),
})
emDelimiter?: '_' | '*';
@Prop({
desc: 'Turndown options > strongDelimiter',
type: new Set(['__', '**']),
})
strongDelimiter?: '__' | '**';
@Prop({
desc: 'Turndown options > linkStyle',
type: new Set(['inlined', 'referenced', 'discarded']),
})
linkStyle?: 'inlined' | 'referenced' | 'discarded';
@Prop({
desc: 'Turndown options > linkReferenceStyle',
type: new Set(['full', 'collapsed', 'shortcut', 'discarded']),
})
linkReferenceStyle?: 'full' | 'collapsed' | 'shortcut' | 'discarded';
static fromCtx(ctx: Context, prefix= 'x-md-') {
const draft: Record<string, string> = {};
for (const [k, v] of Object.entries(ctx.headers)) {
if (k.startsWith(prefix)) {
const prop = k.slice(prefix.length);
const sk = _.camelCase(prop);
draft[sk] = v as string;
}
}
return this.from(draft);
}
}
src/fetch.d.ts
Vendored
+13
View File
@@ -0,0 +1,13 @@
declare global {
export const {
fetch,
FormData,
Headers,
Request,
Response,
File,
}: typeof import('undici');
export type { FormData, Headers, Request, RequestInit, Response, RequestInit, File } from 'undici';
}
export { };
src/lib/filtered-stream.ts
+59
View File
@@ -0,0 +1,59 @@
import { Transform, TransformCallback } from "stream";
import _ from 'lodash';
class FilteredTextResponseStream extends Transform {
constructor(public path: string, public predicate?: (data: any) => boolean) {
super({ writableObjectMode: true, encoding: 'utf-8' });
}
override _transform(chunk: any, _encoding: string, callback: TransformCallback) {
if (this.predicate) {
if (!this.predicate(chunk)) {
callback();
return;
}
}
const interest = _.get(chunk, this.path);
if (typeof interest === 'string') {
if (interest) {
this.push(interest);
}
}
callback();
}
}
export function getFilteredTextStream(path: string = 'data.choices[0].delta.content', predicate?: (data: any) => boolean) {
return new FilteredTextResponseStream(path, predicate);
}
class FilteredResponseStream extends Transform {
constructor(public path: string, public predicate?: (data: any) => boolean) {
super({ objectMode: true });
}
override _transform(chunk: any, _encoding: string, callback: TransformCallback) {
if (this.predicate) {
if (!this.predicate(chunk)) {
callback();
return;
}
}
const interest = _.get(chunk, this.path);
if (interest) {
this.push(interest);
}
callback();
}
}
export function getFilteredStream(path: string = 'data.choices[0].delta.content', predicate?: (data: any) => boolean) {
return new FilteredResponseStream(path, predicate);
}
src/lib/json-parse-stream.ts
+1535
View File
File diff suppressed because it is too large Load Diff
src/lib/koa-compress.ts
+243
View File
@@ -0,0 +1,243 @@
import zlib from 'zlib';
import { Stream, Transform, Readable } from 'stream';
import type { Context, Next, Middleware } from 'koa';
type Encoding = 'zstd' | 'br' | 'gzip' | 'deflate';
type EncoderFactory = (opts?: any) => Transform;
export interface CompressOptions {
filter?: (mimeType: string) => boolean;
threshold?: number;
encodingPreference?: Encoding[];
br?: zlib.BrotliOptions | false | null;
gzip?: zlib.ZlibOptions | false | null;
deflate?: zlib.ZlibOptions | false | null;
zstd?: object | false | null;
}
const encodingMethods: Record<Encoding, EncoderFactory> = {
zstd: zlib.createZstdCompress,
br: zlib.createBrotliCompress,
gzip: zlib.createGzip,
deflate: zlib.createDeflate,
};
type OneShot = (buf: Buffer | string, opts: any) => Promise<Buffer>;
function makeOneShot(fn: any): OneShot {
return (buf, opts) => new Promise((resolve, reject) => {
fn(buf, opts, (err: Error | null, result: Buffer) => {
if (err) reject(err);
else resolve(result);
});
});
}
const oneShotEncoders: Record<Encoding, OneShot> = {
zstd: makeOneShot(zlib.zstdCompress),
br: makeOneShot(zlib.brotliCompress),
gzip: makeOneShot(zlib.gzip),
deflate: makeOneShot(zlib.deflate),
};
const defaultPreference: Encoding[] = ['zstd', 'br', 'gzip', 'deflate'];
const encoderDefaults: Record<Encoding, object> = {
gzip: {},
deflate: {},
br: { params: { [zlib.constants.BROTLI_PARAM_QUALITY]: 4 } },
zstd: {},
};
const NO_TRANSFORM_REGEX = /(?:^|,)\s*?no-transform\s*?(?:,|$)/;
const EMPTY_BODY_STATUSES = new Set([204, 205, 304]);
const COMPRESSIBLE_DEFAULT_RE = /^(?:text\/|application\/(?:json|xml|javascript|ld\+json|graphql|x-ndjson)\b|[^;]+\+(?:json|xml|text))/i;
function defaultFilter(type: string): boolean {
return !!type && COMPRESSIBLE_DEFAULT_RE.test(type);
}
interface ParsedAccept {
explicit: Record<string, number>;
wildcard?: number;
}
function parseAcceptEncoding(header: string): ParsedAccept {
const explicit: Record<string, number> = {};
let wildcard: number | undefined;
for (const part of header.split(',')) {
const seg = part.trim();
if (!seg) continue;
const [rawName, ...params] = seg.split(';').map((s) => s.trim());
const name = rawName.toLowerCase();
let q = 1;
for (const p of params) {
const m = /^q\s*=\s*(.+)$/i.exec(p);
if (m) {
const v = parseFloat(m[1]);
if (!isNaN(v)) q = Math.max(0, Math.min(1, v));
}
}
if (name === '*') wildcard = q;
else explicit[name] = q;
}
return { explicit, wildcard };
}
// RFC 9110 §8.4.1: "identity" is the no-transformation coding. It's the only place
// in this module where the literal matters — we read its q to decide whether the
// client strictly prefers no compression over the best compressible option.
const IDENTITY = 'identity';
function negotiateEncoding(
header: string | undefined,
available: Encoding[],
preference: Encoding[],
): Encoding | undefined {
if (!header) return undefined;
const { explicit, wildcard } = parseAcceptEncoding(header);
// RFC 9110: identity is implicitly acceptable unless excluded.
const identityQ = explicit[IDENTITY] ?? wildcard ?? 1;
let best: { enc: Encoding; q: number; prefIdx: number; } | undefined;
for (const enc of available) {
const q = explicit[enc] ?? wildcard;
if (q === undefined || q === 0) continue;
const prefIdx = preference.indexOf(enc);
if (
!best ||
q > best.q ||
(q === best.q && prefIdx !== -1 && (best.prefIdx === -1 || prefIdx < best.prefIdx))
) {
best = { enc, q, prefIdx };
}
}
// No compressible encoding acceptable, or client strictly prefers identity.
if (!best || identityQ > best.q) return undefined;
return best.enc;
}
export default function compress(options: CompressOptions = {}): Middleware {
const filter = options.filter || defaultFilter;
const preference = options.encodingPreference || defaultPreference;
const threshold = options.threshold ?? 1024;
const enabled = preference.filter((enc) => {
const opt = (options as any)[enc];
return opt !== false && opt !== null;
});
return async function compressMiddleware(ctx: Context, next: Next) {
ctx.vary('Accept-Encoding');
await next();
const body = ctx.body;
const type = ctx.response.type;
const size = ctx.response.length;
const forced = (ctx as any).compress;
if (
!body ||
ctx.res.headersSent ||
!ctx.writable ||
forced === false ||
ctx.request.method === 'HEAD' ||
EMPTY_BODY_STATUSES.has(+ctx.response.status) ||
ctx.response.get('Content-Encoding') ||
!(forced === true || filter(type || '')) ||
NO_TRANSFORM_REGEX.test(ctx.response.get('Cache-Control'))
) {
return;
}
if (threshold && typeof size === 'number' && size < threshold) return;
const encoding = negotiateEncoding(
ctx.get('accept-encoding') as string | undefined,
enabled,
preference,
);
if (!encoding) {
const cc = ctx.response.get('Cache-Control') || '';
ctx.set('Cache-Control', cc ? `${cc}, no-transform` : 'no-transform');
return;
};
const encoderOpts = { ...encoderDefaults[encoding], ...((options as any)[encoding] || {}) };
const source = await normalizeBody(body);
if (!source) return;
if (source.kind === 'stream') {
// Length unknown until the stream ends — chunked transfer.
ctx.set('Content-Encoding', encoding);
ctx.res.removeHeader('Content-Length');
const stream = encodingMethods[encoding](encoderOpts);
ctx.body = stream;
source.value.pipe(stream);
return;
}
if (!source.value.length) {
// Don't bother compressing an empty body, and some encoders error on zero-length input.
return;
}
// One-shot path: compress fully so we can emit an accurate Content-Length.
const compressed = await oneShotEncoders[encoding](source.value, encoderOpts);
ctx.set('Content-Encoding', encoding);
if (source.size !== undefined) {
// Only set Content-Length if we know the original size, to avoid misleading clients about compression ratio.
ctx.set('X-Decompressed-Content-Length', source.size.toString());
}
const cc = ctx.response.get('Cache-Control') || '';
ctx.set('Cache-Control', cc ? `${cc}, no-transform` : 'no-transform');
ctx.body = compressed; // Koa's body setter writes Content-Length = compressed.length
};
}
type NormalizedBody =
| { kind: 'buffer'; value: Buffer | string; size?: number; }
| { kind: 'stream'; value: Readable; };
const WebReadableStream: typeof globalThis.ReadableStream | undefined =
(globalThis as any).ReadableStream;
async function normalizeBody(body: unknown): Promise<NormalizedBody | undefined> {
if (body == null) return undefined;
if (typeof body === 'string') return { kind: 'buffer', value: body, size: Buffer.byteLength(body) };
if (Buffer.isBuffer(body)) return { kind: 'buffer', value: body, size: body.byteLength };
if (body instanceof ArrayBuffer) return { kind: 'buffer', value: Buffer.from(body), size: body.byteLength };
if (ArrayBuffer.isView(body)) {
// Any TypedArray (Uint8/16/32, Int8/16/32, Float32/64, BigInt64/Uint64,
// Uint8Clamped) or DataView. Buffer was matched above to avoid a copy.
// Slice via byteOffset/byteLength so views over a larger buffer stay correct.
return { kind: 'buffer', value: Buffer.from(body.buffer, body.byteOffset, body.byteLength), size: body.byteLength };
}
if (body instanceof Stream) return { kind: 'stream', value: body as Readable };
// Web ReadableStream — must come before the Symbol.asyncIterator branch (it is one).
if (WebReadableStream && body instanceof WebReadableStream) {
return { kind: 'stream', value: Readable.fromWeb(body as any) };
}
if (body instanceof Blob) {
const ab = await body.arrayBuffer();
return { kind: 'buffer', value: Buffer.from(ab), size: ab.byteLength };
}
// Async iterables / iterators — must come after Stream and ReadableStream, both of
// which also expose Symbol.asyncIterator.
if (typeof (body as any)[Symbol.asyncIterator] === 'function') {
return { kind: 'stream', value: Readable.from(body as AsyncIterable<unknown>) };
}
if (typeof body === 'object') {
const json = Buffer.from(JSON.stringify(body));
return { kind: 'buffer', value: json, size: json.byteLength };
}
return undefined;
}
src/lib/pseudo-boolean.ts
+49
View File
@@ -0,0 +1,49 @@
import { AutoConstructor } from 'civkit/civ-rpc';
export class PseudoBoolean {
@AutoConstructor
static from(input: any) {
if (input === undefined || input === null) {
return false;
}
if (typeof input === 'boolean') {
return input;
}
if (typeof input === 'string') {
if (['', 'false', 'none', 'null', 'nan', 'nil', '0', 'no', 'undefined', 'disabled', 'f', 'n/a'].includes(input.toLowerCase().trim())) {
return false;
}
if (['true', 'yes', '1', 'ok', 'enabled', 't'].includes(input.toLowerCase().trim())) {
return true;
}
}
throw new TypeError(`Cannot convert ${input} to boolean`);
}
}
export class PseudoBooleanLoose {
@AutoConstructor
static from(input: any) {
if (input === undefined || input === null) {
return false;
}
if (typeof input === 'boolean') {
return input;
}
if (typeof input === 'string') {
if (['', 'false', 'none', 'null', 'nan', 'nil', '0', 'no', 'undefined', 'disabled', 'f', 'n/a'].includes(input.toLowerCase().trim())) {
return false;
}
return true;
}
throw new TypeError(`Cannot convert ${input} to boolean`);
}
}
src/lib/transform-server-event-stream.ts
+170
View File
@@ -0,0 +1,170 @@
import { TPM } from 'civkit/civ-rpc';
import { parseJSONText } from 'civkit/vectorize';
import { Transform, TransformCallback, TransformOptions } from 'stream';
export class InputServerEventStream extends Transform {
cache: string[] = [];
constructor(options?: TransformOptions) {
super({
...options,
readableObjectMode: true
});
}
decodeRoutine() {
if (!this.cache.length) {
return;
}
const vecs = this.cache.join('').split(/\r?\n\r?\n/);
this.cache.length = 0;
const lastVec = vecs.pop();
if (lastVec) {
this.cache.push(lastVec);
}
for (const x of vecs) {
const lines: string[] = x.split(/\r?\n/);
const event: {
event?: string;
data?: string;
id?: string;
retry?: number;
} = {};
for (const l of lines) {
const columnPos = l.indexOf(':');
if (columnPos <= 0) {
continue;
}
const key = l.substring(0, columnPos);
const rawValue = l.substring(columnPos + 1);
const value = rawValue.startsWith(' ') ? rawValue.slice(1) : rawValue;
if (key === 'data') {
if (event.data) {
event.data += value || '\n';
} else if (event.data === '') {
event.data += '\n';
event.data += value || '\n';
} else {
event.data = value;
}
} else if (key === 'retry') {
event.retry = parseInt(value, 10);
} else {
Reflect.set(event, key, value);
}
}
if (event.data) {
const parsed = parseJSONText(event.data);
if (parsed && typeof parsed === 'object') {
event.data = parsed;
}
}
if (Object.keys(event).length) {
this.push(event);
}
}
}
override _transform(chunk: any, encoding: BufferEncoding, callback: TransformCallback): void {
if (chunk === null) {
this.push(null);
}
this.cache.push(chunk.toString());
this.decodeRoutine();
callback();
}
override _final(callback: (error?: Error | null | undefined) => void): void {
this.decodeRoutine();
callback();
}
}
@TPM({
contentType: 'text/event-stream',
})
export class OutputServerEventStream extends Transform {
n: number = 0;
constructor(options?: TransformOptions) {
super({
...options, writableObjectMode: true, encoding: 'utf-8'
});
}
encodeRoutine(chunk: {
event?: string;
data?: any;
id?: string;
retry?: number;
} | string) {
if (typeof chunk === 'object') {
const lines: string[] = [];
if (chunk.event) {
lines.push(`event: ${chunk.event}`);
}
if (chunk.data) {
if (typeof chunk.data === 'string') {
for (const x of chunk.data.split(/\r?\n/)) {
lines.push(`data: ${x}`);
}
} else {
lines.push(`data: ${JSON.stringify(chunk.data)}`);
}
}
if (chunk.id) {
lines.push(`id: ${chunk.id}`);
}
if (chunk.retry) {
lines.push(`retry: ${chunk.retry}`);
}
if (!lines.length) {
lines.push(`data: ${JSON.stringify(chunk)}`);
}
this.push(lines.join('\n'));
this.push('\n\n');
this.n++;
return;
} else if (typeof chunk === 'string') {
const lines: string[] = [];
for (const x of chunk.split(/\r?\n/)) {
lines.push(`data: ${x}`);
}
this.push(lines.join('\n'));
this.push('\n\n');
this.n++;
}
}
override _transform(chunk: any, encoding: BufferEncoding, callback: TransformCallback): void {
if (chunk === null) {
this.push(null);
}
this.encodeRoutine(chunk);
callback();
}
}
export interface OutputServerEventStream extends Transform {
write(chunk: string | {
event?: string;
data?: any;
id?: string;
retry?: number;
}, callback?: (error: Error | null | undefined) => void): boolean;
write(chunk: any, callback?: (error: Error | null | undefined) => void): boolean;
write(chunk: any, encoding: BufferEncoding, callback?: (error: Error | null | undefined) => void): boolean;
}
src/scripts/smoke-test-llm.ts
+55
View File
@@ -0,0 +1,55 @@
import 'reflect-metadata';
import { injectable, container } from 'tsyringe';
import { AsyncService } from 'civkit/async-service';
import { GlobalLogger } from '../services/logger';
import { LLMManager } from '../services/common-llm';
import { AltTextService } from '../services/alt-text';
@injectable()
export class Hello extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
constructor(
protected globalLogger: GlobalLogger,
protected llmManager: LLMManager,
protected altTextService: AltTextService,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
await this.llmManager.importOpenRouterModel('google/gemini-2.5-flash-lite');
this.emit('ready');
}
async main() {
await this.serviceReady();
const r = await this.altTextService.caption('https://jina-ai-gmbh.ghost.io/content/images/2025/09/banner_code_embeddings.png');
this.logger.info(`Alt text: ${r}`);
// const r = await this.llmManager.run('google/gemini-2.5-flash-lite', {
// prompt: interleavedPrompt`${new URL('https://img.freepik.com/free-photo/closeup-scarlet-macaw-from-side-view-scarlet-macaw-closeup-head_488145-3540.jpg?semt=ais_incoming&w=740&q=80')}\nWhat is this`,
// options: {
// stream: true,
// }
// });
// r.on('data', (chunk) => {
// this.logger.info(`LLM: ${chunk}`);
// });
// await once(r, 'end');
}
}
const instance = container.resolve(Hello);
export default instance;
if (require.main === module) {
instance.main().finally(() => process.exit(0));
}
src/services/alt-text.ts
+190
View File
@@ -0,0 +1,190 @@
import { AsyncService } from 'civkit/async-service';
import { HashManager } from 'civkit/hash';
import { AssertionFailureError } from 'civkit/civ-rpc';
import { singleton } from 'tsyringe';
import { GlobalLogger } from './logger';
import { CanvasService } from './canvas';
import { ImageInterrogationManager } from './common-iminterrogate';
import { ImgBrief } from './puppeteer';
import { AsyncLocalContext } from './async-context';
import { ImgAlt } from '../db/models';
import { StorageLayer } from '../db/noop-storage';
import { LLMManager } from './common-llm';
const md5Hasher = new HashManager('md5', 'hex');
@singleton()
export class AltTextService extends AsyncService {
altsToIgnore = 'image,img,photo,picture,pic,alt,figure,fig,图片'.split(',');
logger = this.globalLogger.child({ service: this.constructor.name });
model = 'google/gemini-2.5-flash-lite';
constructor(
protected globalLogger: GlobalLogger,
protected imageInterrogator: ImageInterrogationManager,
protected canvasService: CanvasService,
protected asyncLocalContext: AsyncLocalContext,
protected storageLayer: StorageLayer,
protected llmManager: LLMManager,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
if (process.env.GCLOUD_PROJECT) {
this.model = 'vertex-gemini-2.5-flash-lite';
} else if (!this.llmManager.hasModel(this.model)) {
await this.llmManager.importOpenRouterModel(this.model);
}
this.emit('ready');
}
async caption(input: string | Buffer, overrideModel?: string) {
try {
const img = await this.canvasService.loadImage(input);
const contentTypeHint = Reflect.get(img, 'contentType');
if (Math.min(img.naturalHeight, img.naturalWidth) <= 1) {
return `A ${img.naturalWidth}x${img.naturalHeight} image, likely be a tacker probe`;
}
if (Math.min(img.naturalHeight, img.naturalWidth) < 64) {
return `A ${img.naturalWidth}x${img.naturalHeight} small image, likely a logo, icon or avatar`;
}
const resized = this.canvasService.fitImageToSquareBox(img, 1024);
const exported = await this.canvasService.canvasToBuffer(resized, 'image/png');
const svgHint = contentTypeHint.includes('svg') ? `Beware this image is a SVG rendered on a gray background, the gray background is not part of the image.\n\n` : '';
const svgSystemHint = contentTypeHint.includes('svg') ? ` Sometimes the system renders SVG on a gray background. When this happens, you must not include the gray background in the description.` : '';
const r = await this.imageInterrogator.interrogate(overrideModel || this.model, {
image: exported,
prompt: `${svgHint}Give a concise image caption descriptive sentence in third person. Start directly with the subject.`,
system: `You are BLIP2, an image captioning model. You will generate Alt Text (in web pages) for any image for a11y purposes. You must not start with "This image is sth...", instead, start direly with "sth..."${svgSystemHint}`,
});
return r.replaceAll(/[\n\"]|(\.\s*$)/g, '').trim();
} catch (err) {
throw new AssertionFailureError({
message: `Could not generate alt text for ${typeof input === 'string' ? `url ${input}` : 'buffer'}`,
cause: err
});
}
}
async vqa(
input: string | Buffer,
instruction: string = 'Give a professional, accurate, concise description for this image.',
overrideModel?: string
) {
try {
const img = await this.canvasService.loadImage(input);
const contentTypeHint = Reflect.get(img, 'contentType');
if (Math.min(img.naturalHeight, img.naturalWidth) <= 1) {
return `A ${img.naturalWidth}x${img.naturalHeight} image, likely be a tacker probe`;
}
if (Math.min(img.naturalHeight, img.naturalWidth) < 64) {
return `A ${img.naturalWidth}x${img.naturalHeight} small image, likely a logo, icon or avatar`;
}
const resized = this.canvasService.fitImageToSquareBox(img, 1024);
const exported = await this.canvasService.canvasToBuffer(resized, 'image/png');
const svgHint = contentTypeHint.includes('svg') ? `Beware this image is a SVG rendered on a gray background, the gray background is not part of the image.\n\n` : '';
const r = await this.imageInterrogator.interrogate(overrideModel || this.model, {
image: exported,
prompt: `${svgHint}${instruction}`,
max_tokens: 768,
temperature: 0,
modelSpecific: {
stop: ['.', '\n']
}
});
return r.replaceAll(/[\n\"]|(\.\s*$)/g, '').trim();
} catch (err) {
throw new AssertionFailureError({
message: `Could not generate alt text for ${typeof input === 'string' ? `url ${input}` : 'buffer'}`,
cause: err
});
}
}
async getAltText(imgBrief: ImgBrief, overrideModel?: string) {
if (!imgBrief.src) {
return undefined;
}
if (imgBrief.alt && !this.altsToIgnore.includes(imgBrief.alt.trim().toLowerCase())) {
return imgBrief.alt;
}
const digest = md5Hasher.hash(imgBrief.src);
let dims: number[] = [];
do {
if (imgBrief.loaded) {
if (imgBrief.naturalWidth && imgBrief.naturalHeight) {
if (Math.min(imgBrief.naturalWidth, imgBrief.naturalHeight) < 64) {
dims = [imgBrief.naturalWidth, imgBrief.naturalHeight];
break;
}
}
}
if (imgBrief.width && imgBrief.height) {
if (Math.min(imgBrief.width, imgBrief.height) < 64) {
dims = [imgBrief.width, imgBrief.height];
break;
}
}
} while (false);
if (Math.min(...dims) <= 1) {
return `A ${dims[0]}x${dims[1]} image, likely be a tacker probe`;
}
if (Math.min(...dims) < 64) {
return `A ${dims[0]}x${dims[1]} small image, likely a logo, icon or avatar`;
}
const existing = await this.storageLayer.findImageAlt({ urlDigest: digest });
if (existing) {
return existing.generatedAlt || existing.originalAlt || '';
}
let generatedCaption = '';
try {
generatedCaption = await this.caption(imgBrief.buff || imgBrief.src, overrideModel);
} catch (err) {
this.logger.warn(`Unable to generate alt text for ${imgBrief.src}`, { err });
}
if (this.asyncLocalContext.ctx.DNT) {
// Don't cache alt text if DNT is set
return generatedCaption;
}
// Don't try again until the next day
const expireMixin = generatedCaption ? {} : { expireAt: new Date(Date.now() + 1000 * 3600 * 24) };
const cache = ImgAlt.from({
_id: digest,
src: imgBrief.src || '',
width: imgBrief.naturalWidth || 0,
height: imgBrief.naturalHeight || 0,
urlDigest: digest,
originalAlt: imgBrief.alt || '',
generatedAlt: generatedCaption || '',
createdAt: new Date(),
...expireMixin
});
await this.storageLayer.storeImageAlt(cache).catch((err) => {
this.logger.warn(`Unable to cache alt text for ${imgBrief.src}`, { err });
});
return generatedCaption;
}
};
src/services/async-context.ts
+10
View File
@@ -0,0 +1,10 @@
import { GlobalAsyncContext } from 'civkit/async-context';
import { container, singleton } from 'tsyringe';
@singleton()
export class AsyncLocalContext extends GlobalAsyncContext { }
const instance = container.resolve(AsyncLocalContext);
Reflect.set(process, 'asyncLocalContext', instance);
export default instance;
src/services/binary-extractor.ts
+375
View File
@@ -0,0 +1,375 @@
import { singleton } from 'tsyringe';
import { AsyncService } from 'civkit/async-service';
import type { Blob } from 'buffer';
import { GlobalLogger } from './logger';
import { readBlob, readFile as readFileText, WEB_SUPPORTED_ENCODINGS } from '../utils/encoding';
import type { PageSnapshot } from './puppeteer';
import { AssertionFailureError, DataStreamBrokenError } from 'civkit/civ-rpc';
import { PDFExtractor } from './pdf-extract';
import { TempFileManager } from './temp-file';
import { mkdir, readFile, writeFile } from 'fs/promises';
import { pathToFileURL } from 'url';
import { SOffice } from './soffice';
import { Threaded } from './threaded';
import { AsyncLocalContext } from './async-context';
import { CanvasService } from './canvas';
import { AltTextService } from './alt-text';
import { stat } from 'fs/promises';
import { JSDomControl } from './jsdom';
@singleton()
export class BinaryExtractorService extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
constructor(
protected globalLogger: GlobalLogger,
protected tempFileManager: TempFileManager,
protected asyncLocalContext: AsyncLocalContext,
protected pdfExtractor: PDFExtractor,
protected sOfficeControl: SOffice,
protected canvasService: CanvasService,
protected altTextService: AltTextService,
protected jsdomControl: JSDomControl,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
async createSnapshotFromBlob(url: URL, blob: Blob, overrideContentType?: string, overrideFileName?: string) {
if (overrideContentType === 'application/octet-stream') {
overrideContentType = undefined;
}
const contentType: string = (overrideContentType || blob.type).toLowerCase();
const fileName = overrideFileName || `${url.origin || ''}${url.pathname || '-'}`;
const urlCopy = new URL(url.href);
urlCopy.hash = '';
const snapshot: PageSnapshot = {
title: '',
href: urlCopy.href,
html: '',
text: '',
traits: ['blob'],
};
// Text-based files
try {
let encoding: string | undefined = contentType.includes('charset=') ? contentType.split('charset=')[1]?.trim().toLowerCase().replaceAll(/["';]/gi, '') : 'utf-8';
if (!WEB_SUPPORTED_ENCODINGS.has(encoding)) {
encoding = 'utf-8';
}
if (!contentType || contentType.startsWith('text/html') || contentType.startsWith('application/xhtml+xml')) {
if (blob.size > 1024 * 1024 * 32) {
throw new AssertionFailureError(`Failed to access ${url}: file too large`);
}
snapshot.html = await readBlob(blob, encoding);
let innerCharset;
const peek = snapshot.html.slice(0, 1024);
innerCharset ??= peek.match(/<meta[^>]+text\/html;\s*?charset=([^>"]+)/i)?.[1]?.toLowerCase();
innerCharset ??= peek.match(/<meta[^>]+charset="([^>"]+)\"/i)?.[1]?.toLowerCase();
if (innerCharset && innerCharset !== encoding && WEB_SUPPORTED_ENCODINGS.has(innerCharset)) {
snapshot.html = await readBlob(blob, innerCharset);
}
snapshot.traits = [];
return snapshot;
}
if ((contentType.startsWith('text/') && !contentType.startsWith('text/xml')) || contentType.startsWith('application/json')) {
if (blob.size > 1024 * 1024 * 32) {
throw new AssertionFailureError(`Failed to access ${url}: file too large`);
}
snapshot.text = await readBlob(blob, encoding);
snapshot.html = `<html><head><meta name="color-scheme" content="light dark"></head><body><pre style="word-wrap: break-word; white-space: pre-wrap;">${snapshot.text}</pre></body></html>`;
return snapshot;
}
} catch (err: any) {
this.logger.warn(`Failed to read from file: ${url}`, { err, url });
throw new DataStreamBrokenError(`Failed to access ${url}: ${err?.message}`);
}
if (blob.size === 0) {
throw new AssertionFailureError(`Unexpected empty file: ${url}`);
}
// Binary files
const filePath = this.tempFileManager.alloc();
this.tempFileManager.bindPathTo(this.asyncLocalContext.ctx, filePath);
await writeFile(filePath, blob.stream() as any, { flush: true });
const dirPath = this.tempFileManager.alloc();
this.tempFileManager.bindPathTo(this.asyncLocalContext.ctx, dirPath);
await mkdir(dirPath);
const snapshotFromBinary = await this.localFileToSnapshot({
filePath,
contentType,
fileName,
outPath: dirPath,
url,
});
Object.assign(snapshot, snapshotFromBinary);
return snapshot;
}
argumentContentType(fileName: string, contentType?: string) {
if (contentType && contentType !== 'application/octet-stream') {
return contentType;
}
const ext = fileName.split('.').pop()?.toLowerCase();
if (!ext) {
return 'application/octet-stream';
}
switch (ext) {
case 'html':
case 'htm':
return 'text/html';
case 'json':
return 'application/json';
case 'txt':
return 'text/plain';
case 'pdf':
return 'application/pdf';
case 'doc':
case 'docx':
return 'application/vnd.openxmlformats-officedocument.wordprocessingml.document';
case 'xls':
case 'xlsx':
return 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet';
case 'ppt':
case 'pptx':
return 'application/vnd.openxmlformats-officedocument.presentationml.presentation';
case 'jpg':
case 'jpeg':
return 'image/jpeg';
case 'png':
return 'image/png';
case 'gif':
return 'image/gif';
case 'bmp':
return 'image/bmp';
case 'svg':
return 'image/svg+xml';
case 'webp':
return 'image/webp';
default:
return 'application/octet-stream';
}
}
@Threaded()
async localFileToSnapshot(options: {
filePath: string;
contentType: string;
fileName: string;
outPath: string;
url: URL;
}) {
const { filePath, contentType: initialContentType, fileName, url, outPath } = options;
const contentType = this.argumentContentType(fileName, initialContentType);
const pageNumber = url.hash ? parseInt(url.hash.slice(1), 10) : undefined;
const urlCopy = new URL(url.href);
urlCopy.hash = '';
const snapshot: PageSnapshot = {
title: '',
href: url.href,
html: '',
text: '',
traits: ['blob'],
};
if (contentType.startsWith('application/pdf')) {
const pagesToRender = pageNumber ? [pageNumber, pageNumber + 1, pageNumber + 2] : [1, 2, 3];
const extracted = await this.pdfExtractor.extractRendered(filePath, outPath, pagesToRender);
snapshot.title = extracted.meta.title || fileName;
snapshot.text = extracted.text;
snapshot.parsed = {
content: extracted.content,
byline: extracted.meta?.Author,
lang: extracted.meta?.Language || undefined,
title: extracted.meta?.Title,
publishedTime: this.pdfExtractor.parsePdfDate(extracted.meta?.ModDate || extracted.meta?.CreationDate)?.toISOString(),
};
snapshot.metadata = extracted.meta;
snapshot.childFrames = extracted.pages.map((page) => {
const childUrl = new URL(urlCopy.href);
childUrl.hash = `#${page.page}`;
return {
title: `${snapshot.title}#${page.page}`,
href: childUrl.href,
html: '',
text: page.text,
parsed: {
content: page.content,
},
screenshotUrl: page.pngPath ? pathToFileURL(page.pngPath).href : undefined,
} as PageSnapshot;
});
snapshot.traits!.push('pdf');
return snapshot;
}
if (
contentType.startsWith('application/msword') ||
contentType.startsWith('application/vnd.ms-word') ||
contentType.startsWith('application/vnd.openxmlformats-officedocument.wordprocessingml.document')
) {
const pagesToRender = pageNumber ? [pageNumber, pageNumber + 1, pageNumber + 2] : [1, 2, 3];
const extracted = await this.sOfficeControl.extractFromWriter(filePath, outPath, { pagesToRender });
snapshot.title = extracted.meta.title || fileName;
snapshot.html = extracted.html;
snapshot.text = extracted.text;
snapshot.parsed = {
content: extracted.content,
};
snapshot.metadata = extracted.meta;
snapshot.childFrames = extracted.pages.map((page, idx) => {
const childUrl = new URL(urlCopy.href);
childUrl.hash = `#${idx + 1}`;
return {
title: `${snapshot.title}#${idx + 1}`,
href: childUrl.href,
html: page.html,
text: page.text,
parsed: {
content: page.content,
},
screenshotUrl: page.pngPath ? pathToFileURL(page.pngPath).href : undefined,
} as PageSnapshot;
});
snapshot.traits!.push('writer');
return snapshot;
}
if (
contentType.startsWith('application/vnd.ms-excel') ||
contentType.startsWith('application/ms-excel') ||
contentType.startsWith('application/vnd.openxmlformats-officedocument.spreadsheetml.sheet')
) {
const pagesToRender = pageNumber ? [pageNumber, pageNumber + 1] : [1];
const extracted = await this.sOfficeControl.extractFromCalc(filePath, outPath, { pagesToRender });
snapshot.title = extracted.meta.title || fileName;
snapshot.html = extracted.html;
snapshot.text = extracted.text;
snapshot.parsed = {
content: extracted.content,
};
snapshot.metadata = extracted.meta;
snapshot.childFrames = extracted.pages.map((page, idx) => {
const childUrl = new URL(urlCopy.href);
childUrl.hash = `#${idx + 1}`;
return {
title: `${snapshot.title}#${idx + 1}`,
href: childUrl.href,
html: page.html,
text: page.text,
parsed: {
content: page.content,
},
screenshotUrl: page.pngPath ? pathToFileURL(page.pngPath).href : undefined,
} as PageSnapshot;
});
snapshot.traits!.push('calc');
return snapshot;
}
if (
contentType.startsWith('application/ms-powerpoint') ||
contentType.startsWith('application/vnd.ms-powerpoint') ||
contentType.startsWith('application/vnd.openxmlformats-officedocument.presentationml.presentation')
) {
const pagesToRender = pageNumber ? [pageNumber, pageNumber + 1, pageNumber + 2] : [1, 2, 3];
const extracted = await this.sOfficeControl.extractFromImpress(filePath, outPath, { pagesToRender });
snapshot.title = extracted.meta.title || fileName;
snapshot.html = extracted.html;
snapshot.text = extracted.text;
snapshot.parsed = {
content: extracted.content,
};
snapshot.metadata = extracted.meta;
snapshot.childFrames = extracted.pages.map((page, idx) => {
const childUrl = new URL(urlCopy.href);
childUrl.hash = `#${idx + 1}`;
return {
title: `${snapshot.title}#${idx + 1}`,
href: childUrl.href,
html: page.html,
text: page.text,
parsed: {
content: page.content,
},
screenshotUrl: page.pngPath ? pathToFileURL(page.pngPath).href : undefined,
} as PageSnapshot;
});
snapshot.traits!.push('impress');
return snapshot;
}
if (contentType.startsWith('image/')) {
snapshot.title = fileName;
snapshot.imgs = [{ src: url.href }];
const viewPort = this.asyncLocalContext.get('viewport');
const img = await this.canvasService.loadImage(await readFile(filePath));
const resized = this.canvasService.fitImageToBox(img, viewPort?.width || 1024, viewPort?.height || 1024);
const imageBuff = await this.canvasService.canvasToBuffer(resized);
// const altText = await this.altTextService.getAltText({
// src: url.href,
// buff: imageBuff,
// }, 'jina-vlm');
const altText = await this.altTextService.vqa(imageBuff, this.asyncLocalContext.get('instruction'), 'jina-vlm');
snapshot.text = altText || '';
snapshot.parsed = {
content: altText || '',
textContent: altText || '',
};
snapshot.screenshot = imageBuff;
snapshot.html = `<html style="height: 100%;"><head><meta name="viewport" content="width=device-width, minimum-scale=0.1"><title>${fileName}</title></head><body style="margin: 0px; height: 100%; background-color: rgb(14, 14, 14);"><p>${altText}</p></body></html>`;
snapshot.traits!.push('image');
return snapshot;
}
if (contentType.startsWith('text/xml') || contentType.startsWith('application/xml') || contentType.includes('+xml')) {
const fileSize = await stat(filePath).then((stat) => stat.size);
if (fileSize > 1024 * 1024 * 32) {
throw new AssertionFailureError(`Failed to access ${url}: file too large`);
}
const encoding = contentType.includes('charset=') ? contentType.split('charset=')[1]?.trim().toLowerCase().replaceAll(/["';]/gi, '') : 'utf-8';
const xmlText = await readFileText(filePath, encoding);
const snapshot = await this.jsdomControl.xmlTextToSnapshot(xmlText, url);
return snapshot;
}
throw new AssertionFailureError(`Failed to interpret ${url}: unexpected content type ${contentType}`);
}
}
src/services/blackhole-detector.ts
+82
View File
@@ -0,0 +1,82 @@
import { singleton } from 'tsyringe';
import { AsyncService } from 'civkit/async-service';
import { GlobalLogger } from './logger';
import { delay } from 'civkit/timeout';
@singleton()
export class BlackHoleDetector extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
lastWorkedTs?: number;
lastDoneRequestTs?: number;
lastIncomingRequestTs?: number;
maxDelay = 1000 * 30;
concurrentRequests = 0;
strikes = 0;
constructor(protected globalLogger: GlobalLogger) {
super(...arguments);
if (process.env.NODE_ENV?.startsWith('prod')) {
setInterval(() => {
this.routine();
}, 1000 * 30).unref();
}
}
override async init() {
await this.dependencyReady();
if (!process.env.GCLOUD_PROJECT) {
this.emit('ready');
return;
}
this.logger.debug('BlackHoleDetector started');
this.emit('ready');
}
async routine() {
// We give routine a 3s grace period for potentially paused CPU to spin up and process some requests
await delay(3000);
const now = Date.now();
const lastWorked = this.lastWorkedTs;
if (!lastWorked) {
return;
}
const dt = (now - lastWorked);
if (this.concurrentRequests > 1 &&
this.lastIncomingRequestTs && lastWorked &&
this.lastIncomingRequestTs >= lastWorked &&
(dt > (this.maxDelay * (this.strikes + 1)))
) {
this.logger.warn(`BlackHole detected, last worked: ${Math.ceil(dt / 1000)}s ago, concurrentRequests: ${this.concurrentRequests}`);
this.strikes += 1;
}
if (this.strikes >= 3) {
this.logger.error(`BlackHole detected for ${this.strikes} strikes, last worked: ${Math.ceil(dt / 1000)}s ago, concurrentRequests: ${this.concurrentRequests}`);
process.nextTick(() => {
this.emit('error', new Error(`BlackHole detected for ${this.strikes} strikes, last worked: ${Math.ceil(dt / 1000)}s ago, concurrentRequests: ${this.concurrentRequests}`));
// process.exit(1);
});
}
}
incomingRequest() {
this.lastIncomingRequestTs = Date.now();
this.lastWorkedTs ??= Date.now();
this.concurrentRequests++;
}
doneWithRequest() {
this.concurrentRequests--;
this.lastDoneRequestTs = Date.now();
}
itWorked() {
this.lastWorkedTs = Date.now();
this.strikes = 0;
}
};
src/services/bogo-sites.ts
+105
View File
@@ -0,0 +1,105 @@
import { AsyncService } from 'civkit/async-service';
import { GlobalLogger } from './logger';
import { HTTPService } from 'civkit/http';
import { EnvConfig } from './envconfig';
import { singleton } from 'tsyringe';
type IntegrityEnvelopeWrapped<T> = {
code: number;
status: number;
data: T;
meta: { [k: string]: any; };
};
export class BogoSiteClient extends HTTPService {
async listHandlerRegExps() {
const r = await this.get<IntegrityEnvelopeWrapped<{
handlers: string[];
}>>('/listHandlers');
return r.data.data.handlers.map((h) => new RegExp(h, 'i'));
}
async access(url: string) {
const r = await this.postJson<Blob>('/access', { url }, { responseType: 'blob' });
return r;
}
}
@singleton()
export class BogoSitesControl extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
clients: BogoSiteClient[] = [];
regExpMap: Map<RegExp, BogoSiteClient> = new Map();
protected __manifestSyncInterval?: ReturnType<typeof setInterval>;
constructor(
protected globalLogger: GlobalLogger,
protected envConfig: EnvConfig,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
if (process.env.JINA_BOGO_SITES_RESORT_ORIGIN) {
this.clients.push(new BogoSiteClient(process.env.JINA_BOGO_SITES_RESORT_ORIGIN));
}
if (this.clients.length) {
await this.synchronizeHandlers().catch((err) => {
this.logger.warn(`Failed to synchronize bogus site handlers on startup`, { error: err });
});
}
if (this.__manifestSyncInterval) {
clearInterval(this.__manifestSyncInterval);
delete this.__manifestSyncInterval;
}
if (this.clients.length) {
this.__manifestSyncInterval = setInterval(
() => this.synchronizeHandlers().catch((err) => {
this.logger.warn(`Failed to synchronize bogus site handlers`, { error: err });
}),
5 * 60 * 1000 + Math.random() * 60 * 1000
).unref();
}
this.emit('ready');
}
async synchronizeHandlers() {
for (const client of this.clients) {
const handlers = await client.listHandlerRegExps();
for (const handler of handlers) {
this.regExpMap.set(handler, client);
}
}
}
async attempt(url: string) {
for (const [regExp, client] of this.regExpMap.entries()) {
if (regExp.test(url)) {
try {
this.logger.debug(`Accessing URL with bogo site client: ${url}`);
const r = await client.access(url);
return r;
} catch (e) {
this.logger.warn('Error accessing URL with bogo site client', { url, error: e });
}
break;
}
}
return;
}
}
src/services/canvas.ts
+225
View File
@@ -0,0 +1,225 @@
import { singleton, container } from 'tsyringe';
import { ParamValidationError, SubmittedDataMalformedError, /* downloadFile */ } from 'civkit/civ-rpc';
import { AsyncService } from 'civkit/async-service';
import { mimeOf } from 'civkit/mime';
import { readFile } from 'fs/promises';
import type canvas from '@napi-rs/canvas';
export type { Canvas, Image } from '@napi-rs/canvas';
import { GlobalLogger } from './logger';
import { TempFileManager } from './temp-file';
import { isMainThread } from 'worker_threads';
import type { svg2png } from 'svg2png-wasm' with { 'resolution-mode': 'import' };
import path from 'path';
import { Threaded } from './threaded';
const downloadFile = async (uri: string) => {
const resp = await fetch(uri);
if (!(resp.ok && resp.body)) {
throw new Error(`Unexpected response ${resp.statusText}`);
}
const contentLength = parseInt(resp.headers.get('content-length') || '0');
if (contentLength > 1024 * 1024 * 100) {
throw new Error('File too large');
}
const buff = await resp.arrayBuffer();
return { buff, contentType: resp.headers.get('content-type') };
};
@singleton()
export class CanvasService extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
svg2png!: typeof svg2png;
canvas!: typeof canvas;
constructor(
protected temp: TempFileManager,
protected globalLogger: GlobalLogger,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
if (!isMainThread) {
const { createSvg2png, initialize } = require('svg2png-wasm');
const wasmBuff = await readFile(path.resolve(path.dirname(require.resolve('svg2png-wasm')), '../svg2png_wasm_bg.wasm'));
const fontBuff = await readFile(path.resolve(__dirname, '../../licensed/SourceHanSansSC-Regular.otf'));
await initialize(wasmBuff);
this.svg2png = createSvg2png({
fonts: [Uint8Array.from(fontBuff)],
defaultFontFamily: {
serifFamily: 'Source Han Sans SC',
sansSerifFamily: 'Source Han Sans SC',
cursiveFamily: 'Source Han Sans SC',
fantasyFamily: 'Source Han Sans SC',
monospaceFamily: 'Source Han Sans SC',
}
});
}
this.canvas = require('@napi-rs/canvas');
this.emit('ready');
}
@Threaded()
async renderSvgToPng(svgContent: string,) {
return this.svg2png(svgContent, { backgroundColor: '#D3D3D3' });
}
protected async _loadImage(input: string | Buffer) {
let buff;
let contentType;
do {
if (typeof input === 'string') {
if (input.startsWith('data:')) {
const firstComma = input.indexOf(',');
const header = input.slice(0, firstComma);
const data = input.slice(firstComma + 1);
const encoding = header.split(';')[1];
contentType = header.split(';')[0].split(':')[1];
if (encoding?.startsWith('base64')) {
buff = Buffer.from(data, 'base64');
} else {
buff = Buffer.from(decodeURIComponent(data), 'utf-8');
}
break;
}
if (input.startsWith('http')) {
const r = await downloadFile(input);
buff = Buffer.from(r.buff);
contentType = r.contentType;
break;
}
}
if (Buffer.isBuffer(input)) {
buff = input;
const mime = await mimeOf(buff);
contentType = `${mime.mediaType}/${mime.subType}`;
break;
}
throw new ParamValidationError('Invalid input');
} while (false);
if (!buff) {
throw new ParamValidationError('Invalid input');
}
if (contentType?.includes('svg')) {
buff = await this.renderSvgToPng(buff.toString('utf-8'));
}
const img = await this.canvas.loadImage(buff);
Reflect.set(img, 'contentType', contentType);
return img;
}
async loadImage(uri: string | Buffer) {
const t0 = Date.now();
try {
const theImage = await this._loadImage(uri);
const t1 = Date.now();
this.logger.debug(`Image loaded in ${t1 - t0}ms`);
return theImage;
} catch (err: any) {
if (err?.message?.includes('Unsupported image type') || err?.message?.includes('unsupported')) {
this.logger.warn(`Failed to load image ${uri.slice(0, 128)}`, { err });
throw new SubmittedDataMalformedError(`Unknown image format for ${uri.slice(0, 128)}`);
}
throw err;
}
}
fitImageToSquareBox(image: canvas.Image | canvas.Canvas, size: number = 1024) {
// this.logger.debug(`Fitting image(${ image.width }x${ image.height }) to ${ size } box`);
// const t0 = Date.now();
if (image.width <= size && image.height <= size) {
if (image instanceof this.canvas.Canvas) {
return image;
}
const canvasInstance = this.canvas.createCanvas(image.width, image.height);
const ctx = canvasInstance.getContext('2d');
ctx.drawImage(image, 0, 0, image.width, image.height, 0, 0, canvasInstance.width, canvasInstance.height);
// this.logger.debug(`No need to resize, copied to canvas in ${ Date.now() - t0 } ms`);
return canvasInstance;
}
const aspectRatio = image.width / image.height;
const resizedWidth = Math.round(aspectRatio > 1 ? size : size * aspectRatio);
const resizedHeight = Math.round(aspectRatio > 1 ? size / aspectRatio : size);
const canvasInstance = this.canvas.createCanvas(resizedWidth, resizedHeight);
const ctx = canvasInstance.getContext('2d');
ctx.drawImage(image, 0, 0, image.width, image.height, 0, 0, resizedWidth, resizedHeight);
// this.logger.debug(`Resized to ${ resizedWidth }x${ resizedHeight } in ${ Date.now() - t0 } ms`);
return canvasInstance;
}
fitImageToBox(image: canvas.Image | canvas.Canvas, maxWidth: number, maxHeight: number) {
// this.logger.debug(`Fitting image(${ image.width }x${ image.height }) to ${ maxWidth }x${ maxHeight } box`);
// const t0 = Date.now();
if (image.width <= maxWidth && image.height <= maxHeight) {
if (image instanceof this.canvas.Canvas) {
return image;
}
const canvasInstance = this.canvas.createCanvas(image.width, image.height);
const ctx = canvasInstance.getContext('2d');
ctx.drawImage(image, 0, 0, image.width, image.height, 0, 0, canvasInstance.width, canvasInstance.height);
// this.logger.debug(`No need to resize, copied to canvas in ${ Date.now() - t0 } ms`);
return canvasInstance;
}
const aspectRatio = image.width / image.height;
let resizedWidth = maxWidth;
let resizedHeight = Math.round(maxWidth / aspectRatio);
if (resizedHeight > maxHeight) {
resizedHeight = maxHeight;
resizedWidth = Math.round(maxHeight * aspectRatio);
}
const canvasInstance = this.canvas.createCanvas(resizedWidth, resizedHeight);
const ctx = canvasInstance.getContext('2d');
ctx.drawImage(image, 0, 0, image.width, image.height, 0, 0, resizedWidth, resizedHeight);
// this.logger.debug(`Resized to ${ resizedWidth }x${ resizedHeight } in ${ Date.now() - t0 } ms`);
return canvasInstance;
}
corpImage(image: canvas.Image | canvas.Canvas, x: number, y: number, w: number, h: number) {
// this.logger.debug(`Cropping image(${ image.width }x${ image.height }) to ${ w }x${ h } at ${ x },${ y } `);
// const t0 = Date.now();
const canvasInstance = this.canvas.createCanvas(w, h);
const ctx = canvasInstance.getContext('2d');
ctx.drawImage(image, x, y, w, h, 0, 0, w, h);
// this.logger.debug(`Crop complete in ${ Date.now() - t0 } ms`);
return canvasInstance;
}
canvasToDataUrl(canvas: canvas.Canvas, mimeType?: 'image/png' | 'image/jpeg') {
// this.logger.debug(`Exporting canvas(${ canvas.width }x${ canvas.height })`);
// const t0 = Date.now();
return canvas.toDataURLAsync((mimeType || 'image/png') as 'image/png');
}
async canvasToBuffer(canvas: canvas.Canvas, mimeType?: 'image/png' | 'image/jpeg') {
// this.logger.debug(`Exporting canvas(${ canvas.width }x${ canvas.height })`);
// const t0 = Date.now();
return canvas.toBuffer((mimeType || 'image/png') as 'image/png');
}
}
const instance = container.resolve(CanvasService);
export default instance;
src/services/cf-browser-rendering.ts
+54
View File
@@ -0,0 +1,54 @@
import { container, singleton } from 'tsyringe';
import { AsyncService } from 'civkit/async-service';
import { GlobalLogger } from './logger';
import { CloudFlareHTTP } from '../3rd-party/cloud-flare';
import { HTTPServiceError } from 'civkit/http';
import { ServiceNodeResourceDrainError } from './errors';
import { EnvConfig } from './envconfig';
@singleton()
export class CFBrowserRendering extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
client!: CloudFlareHTTP;
constructor(
protected globalLogger: GlobalLogger,
protected envConfig: EnvConfig,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
const [account, key] = this.envConfig.CLOUD_FLARE_API_KEY?.split(':');
this.client = new CloudFlareHTTP(account, key);
this.emit('ready');
}
async fetchContent(url: string) {
try {
const r = await this.client.fetchBrowserRenderedHTML({ url });
return r.parsed.result;
} catch (err) {
if (err instanceof HTTPServiceError) {
if (err.status === 429) {
// Rate limit exceeded, return empty result
this.logger.warn('Cloudflare browser rendering rate limit exceeded', { url });
throw new ServiceNodeResourceDrainError(`Cloudflare browser rendering (our account) is at capacity, please try again later or switch to another engine.`,);
}
}
throw err;
}
}
}
const instance = container.resolve(CFBrowserRendering);
export default instance;
src/services/common-iminterrogate/base.ts
+241
View File
@@ -0,0 +1,241 @@
import { Blob } from 'buffer';
import { AsyncService } from 'civkit/async-service';
import { ApplicationError, Coercible } from 'civkit/civ-rpc';
import { HTTPService, HTTPServiceRequestOptions } from 'civkit/http';
import { LoggerInterface } from 'civkit/logger';
import { mimeOf } from 'civkit/mime';
import { delay } from 'civkit/timeout';
type ImageData = Blob | Buffer | string;
export interface ImageInterrogationOptions<T> {
image: ImageData;
prompt?: string;
system?: string;
max_tokens?: number;
steps?: number;
n?: number;
modelSpecific?: Partial<T>;
mask?: ImageData;
seed?: number;
}
export abstract class AbstractImageInterrogationModel<T = any, C = unknown, R = string> extends AsyncService {
static aliases?: string[];
static description?: string;
static modelOptionsType?: typeof Coercible;
abstract clients: C[];
abstract logger: LoggerInterface;
protected clientBlockade: Map<C, Date> = new Map();
get aliases() {
return (this.constructor as typeof AbstractImageInterrogationModel).aliases || [];
}
get name() {
return this.aliases[0];
}
get description() {
return (this.constructor as typeof AbstractImageInterrogationModel).description;
}
abstract _run<U extends ImageInterrogationOptions<T>>(
client: this['clients'][number],
modelOpts: U,
requestOptions?: HTTPServiceRequestOptions
): Promise<R>;
async withClient<U extends (client: this['clients'][number]) => any>(func: U, options: {
maxTry?: number;
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: (client: any) => boolean;
} = {}): Promise<ReturnType<U>> {
await this.serviceReady();
const filteredClients = options?.filterClients ? this.clients.filter(options.filterClients) : this.clients;
const gentleClients = this.filterClientsByAvailability(filteredClients);
const clients = gentleClients.length ? gentleClients : filteredClients;
if (!clients.length) {
throw new Error('Model client not ready');
}
const maxTry = options?.maxTry ?? (clients.length + 1);
let i = 0;
while (true) {
for (const client of options.reverseProviderPreference ? [...clients].reverse() : clients) {
if (typeof client !== 'object' || !client) {
throw new Error('Invalid client');
}
try {
this.logger.debug(`Calling ImageInterrogation model ${this.constructor.name} with ${(client as any).name || client.constructor.name}`);
const t0 = Date.now();
const r = await func(client);
this.logger.debug(`Call succeeded. ImageInterrogation model ${this.constructor.name} with ${(client as any).name || client.constructor.name} in ${Date.now() - t0}ms.`);
return r;
} catch (err: any) {
if (err?.cause instanceof ApplicationError) {
this.logger.error(`Failed to run ImageInterrogation model ${this.constructor.name} with ${(client as any).name || client.constructor.name}. Non recoverable error: Code ${err.code}.`, { error: err });
throw err.cause;
}
// const code = err.status || (err as any).code;
// if ((((client as any).Error && err instanceof (client as any).Error) || err instanceof HarmfulContentError) && (typeof code === 'number')) {
// if (code !== 20 && code !== 429 && code !== 401 && !(code >= 500 && code < 600)) {
// this.logger.error(`Failed to run ImageInterrogation model ${this.constructor.name} with ${(client as any).name || client.constructor.name}. Non recoverable error: Code ${code}.`, { error: marshalErrorLike(err) });
// throw new DownstreamServiceFailureError({
// message: `Failed to run ImageInterrogation model ${this.constructor.name} with ${(client as any).name || client.constructor.name}. Non recoverable error: Code ${code}.`,
// err
// });
// }
// }
this.clientAvailabilityRoutine(client, err);
if (++i >= maxTry) {
this.logger.error(`Failed to run ImageInterrogation model ${this.constructor.name} with ${(client as any).name || client.constructor.name}. No tries left.`, { error: err });
throw err;
}
const delayMs = Math.floor((options?.delayFactor ?? 1) * (1 + Math.random() * 0.4 - 0.2) * 100);
this.logger.warn(`Failed to run ImageInterrogation model ${this.constructor.name} with ${(client as any).name || client.constructor.name}, retrying after ${delayMs}ms...`, { error: err });
await delay(delayMs);
}
}
}
throw new Error('unreachable');
}
interrogate<U extends ImageInterrogationOptions<T>>(
modelOpts: U,
execOpts?: {
maxTry?: number;
requestOptions?: HTTPServiceRequestOptions,
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: <T extends HTTPService>(client: T) => boolean;
}
): Promise<string> {
const requestOptions = execOpts?.requestOptions;
this.logger.info(`Executing ${this.constructor.name} Image Interrogation...`);
return this.withClient((client) => this._run(client, modelOpts, requestOptions), execOpts) as any;
}
protected clientAvailabilityRoutine(client: C, err: Error) {
if (typeof err !== 'object' || !client) {
return;
}
const safeErr = err as any;
const retryAfter = safeErr?.retryAfter || safeErr?.retry_after || safeErr?.response?.headers?.get('retry-after');
if (!retryAfter) {
return;
}
const retryAfterSec = Number.parseInt(retryAfter);
if (retryAfterSec) {
this.logger.warn(`Client ${client.constructor.name}(${this.clients.indexOf(client)}) informed blockade for ${retryAfterSec} seconds.`);
this.clientBlockade.set(client, new Date(Date.now() + retryAfterSec * 1000));
return;
}
const retryAfterDate = new Date(retryAfter);
if (isNaN(retryAfterDate as any)) {
return;
}
this.logger.warn(`Client ${client.constructor.name}(${this.clients.indexOf(client)}) informed blockade until ${retryAfterDate}.`);
this.clientBlockade.set(client, retryAfterDate);
}
protected filterClientsByAvailability(clients: C[]) {
const now = new Date();
const filtered = clients.filter((client) => {
const blockade = this.clientBlockade.get(client);
if (!blockade) {
return true;
}
if (blockade < now) {
return true;
}
return false;
});
return filtered.length ? filtered : clients;
}
protected async toUrl(input: ImageData) {
if (Buffer.isBuffer(input)) {
const mimeInfo = await mimeOf(input);
return `data:${mimeInfo.mediaType}/${mimeInfo.subType};base64,${input.toString('base64')}`;
}
if (input instanceof Blob) {
return `data:${input.type};base64,${Buffer.from(await input.arrayBuffer()).toString('base64')}`;
}
return input;
}
protected async toBuffer(input: ImageData) {
if (Buffer.isBuffer(input)) {
return input;
}
if (input instanceof Blob) {
return Buffer.from(await input.arrayBuffer());
}
if (input.startsWith('data:')) {
const [_mime, base64] = input.split(';base64,');
if (!base64) {
throw new Error('Invalid data url');
}
return Buffer.from(base64, 'base64');
}
const r = await fetch(input);
return Buffer.from(await r.arrayBuffer());
}
protected async toBlob(input: ImageData) {
if (Buffer.isBuffer(input)) {
const mimeInfo = await mimeOf(input);
return new Blob([input as any], { type: `${mimeInfo.mediaType}/${mimeInfo.subType}` });
}
if (input instanceof Blob) {
return input;
}
if (input.startsWith('data:')) {
const [mime, base64] = input.split(';base64,');
if (!base64) {
throw new Error('Invalid data url');
}
return new Blob([Buffer.from(base64, 'base64')], { type: mime });
}
const r = await fetch(input);
return await r.blob();
}
}
src/services/common-iminterrogate/index.ts
+5
View File
@@ -0,0 +1,5 @@
export * from './base';
export * from './registry';
import instance from './registry';
export default instance;
src/services/common-iminterrogate/instruct-blip.ts
+129
View File
@@ -0,0 +1,129 @@
import { Coercible, Prop } from 'civkit/coercible';
import { ReplicateHTTP } from '../../3rd-party/replicate';
import _ from 'lodash';
import { injectable } from 'tsyringe';
import { AbstractImageInterrogationModel, ImageInterrogationOptions } from './base';
import { EnvConfig } from '../envconfig';
import { GlobalLogger } from '../logger';
export class ReplicateInstructBlipVicuna13bModelOptions extends Coercible {
@Prop({
desc: `Image prompt to send to the model.`,
required: true
})
img!: string | Blob | Buffer;
@Prop({
desc: `Text prompt to send to the model.`,
required: true,
default: `Describe the image in great detail.`
})
prompt!: string;
@Prop({
desc: `Maximum number of tokens to generate. A word is generally 2-3 tokens`,
})
max_tokens?: number;
@Prop({
desc: `Adjusts randomness of outputs, greater than 1 is random and 0 is deterministic, 0.75 is a good starting value.`,
})
temperature?: number;
@Prop({
desc: `When decoding text, samples from the top p percentage of most likely tokens; lower to ignore less likely tokens`,
})
top_p?: number;
@Prop({
desc: `When decoding text, samples from the top k most likely tokens; lower to ignore less likely tokens. Defaults to 0 (no top-k sampling).`,
})
top_k?: number;
@Prop({
desc: `When > 0 and top_k > 1, penalizes new tokens based on their similarity to previous tokens. Can help minimize repitition while maintaining semantic coherence. Set to 0 to disable.`,
})
penalty_alpha?: number;
@Prop({
desc: `Penalty for repeated words in generated text; 1 is no penalty, values greater than 1 discourage repetition, less than 1 encourage it.`,
})
repetition_penalty?: number;
@Prop({
desc: `Increasing the length_penalty parameter above 1.0 will cause the model to favor longer sequences, while decreasing it below 1.0 will cause the model to favor shorter sequences.`,
})
length_penalty?: number;
@Prop({
desc: `If set to int > 0, all ngrams of size no_repeat_ngram_size can only occur once.`,
})
no_repeat_ngram_size?: number;
@Prop({
desc: `Set seed for reproducible outputs. Set to -1 for random seed.`,
})
seed?: number;
@Prop({
desc: `provide debugging output in logs`,
})
debug?: boolean;
}
@injectable()
export class InstructBLIP extends AbstractImageInterrogationModel<ReplicateInstructBlipVicuna13bModelOptions> {
static override description = 'An instruction-tuned multi-modal model based on BLIP-2 and Vicuna-13B';
static override aliases = [
'instructblip'
];
static override modelOptionsType = ReplicateInstructBlipVicuna13bModelOptions;
// https://replicate.com/joehoover/instructblip-vicuna13b/versions
static replicateModel: string = 'c4c54e3c8c97cd50c2d2fec9be3b6065563ccf7d43787fb99f84151b867178fe';
override clients: Array<ReplicateHTTP> = [];
logger = this.globalLogger.child({ service: this.constructor.name });
constructor(
protected envConfig: EnvConfig,
protected globalLogger: GlobalLogger,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.clients = [
new ReplicateHTTP(this.envConfig.REPLICATE_API_KEY),
];
this.emit('ready');
}
override async _run<U extends ImageInterrogationOptions<ReplicateInstructBlipVicuna13bModelOptions>>(client: this['clients'][number], modelOpts: U): Promise<string> {
const payload: any = ReplicateInstructBlipVicuna13bModelOptions.from({
...modelOpts.modelSpecific,
img: await this.toUrl(modelOpts.image),
prompt: modelOpts.prompt,
max_length: modelOpts.max_tokens,
});
if (modelOpts.prompt) {
payload.question = modelOpts.prompt;
} else {
payload.caption = true;
}
const rp = await client.createPrediction((this.constructor as typeof InstructBLIP).replicateModel,
payload
);
const predictionId = rp.data.id;
const r = await client.pollPrediction(predictionId);
return r.output.join('');
}
}
src/services/common-iminterrogate/llms.ts
+85
View File
@@ -0,0 +1,85 @@
import _ from "lodash";
import { injectable } from "tsyringe";
import { AbstractImageInterrogationModel, ImageInterrogationOptions } from './base';
import { EnvConfig } from '../envconfig';
import { GlobalLogger } from '../logger';
import type { AbstractLLM, PromptChunk } from '../common-llm/base';
export function imageInterrogationWithVisualLLM(llm: AbstractLLM<unknown, unknown>): any {
const cls = llm.constructor as typeof AbstractLLM;
if (!llm.interleavedPromptSupported) {
throw new Error(`LLM ${llm.name} does not support image interrogation`);
}
@injectable()
class LLMDerivedImageInterrogator extends AbstractImageInterrogationModel<typeof AbstractLLM['modelOptionsType']> {
static override description = cls.description;
static override aliases = cls.aliases;
static override modelOptionsType = cls.modelOptionsType;
override clients: Array<InstanceType<typeof cls>> = [];
logger = this.globalLogger.child({ service: this.constructor.name });
constructor(
protected envConfig: EnvConfig,
protected globalLogger: GlobalLogger,
) {
super(...arguments);
this.clients.push(llm);
this.dependsOn(llm);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
override async withClient<U extends (client: InstanceType<typeof cls>) => any>(func: U, _options: {
maxTry?: number;
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: (client: any) => boolean;
} = {}): Promise<ReturnType<U>> {
await this.serviceReady();
const client = this.clients[0];
if (!this.clients.length) {
throw new Error('Model client not ready');
}
if (typeof client !== 'object' || !client) {
throw new Error('Invalid client');
}
const r = await func(client);
return r;
}
override async _run<U extends ImageInterrogationOptions<typeof cls['modelOptionsType']>>(client: this['clients'][number], modelOpts: U): Promise<string> {
this.logger.debug(`Calling ${cls.name} Image Caption...`);
const promptChunks: PromptChunk[] = [await this.toBlob(modelOpts.image)];
if (modelOpts.prompt) {
promptChunks.push(modelOpts.prompt);
}
const r = await client.exec(promptChunks, {
stream: false,
n: modelOpts.n,
seed: modelOpts.seed,
max_tokens: modelOpts.max_tokens,
system: modelOpts.system,
modelSpecific: modelOpts.modelSpecific
});
return r as string;
}
}
Object.defineProperty(LLMDerivedImageInterrogator, 'name', { value: `${cls.name}DerivedImageInterrogator`, writable: false });
return LLMDerivedImageInterrogator as any;
}
src/services/common-iminterrogate/registry.ts
+136
View File
@@ -0,0 +1,136 @@
import { HTTPService, HTTPServiceRequestOptions } from 'civkit/http';
import { container, singleton } from 'tsyringe';
import { readdirSync } from 'fs';
import _ from 'lodash';
import { GlobalLogger } from '../logger';
import { AbstractImageInterrogationModel, ImageInterrogationOptions } from './base';
import { AsyncService } from 'civkit/async-service';
import { AssertionFailureError } from 'civkit/civ-rpc';
import { LLMManager } from '../common-llm/registry';
import { imageInterrogationWithVisualLLM } from './llms';
function loadModulesDynamically(path: string) {
const moduleDir = readdirSync(path,
{ withFileTypes: true, encoding: 'utf-8' });
const modules = moduleDir.filter((x) => x.isFile() && x.name.endsWith('.js') || x.isDirectory()).map((x) => x.name);
const imClasses: (typeof AbstractImageInterrogationModel<unknown>)[] = [];
for (const m of modules) {
try {
if (m === 'index.js') {
continue;
}
// if (process.env.hasOwnProperty(`IMGEN_DISABLE_${m.toUpperCase()}`)) {
// continue;
// }
// FIXME: Does not work with esm
// eslint-disable-next-line @typescript-eslint/no-var-requires
const mod = require(`${path}/${m}`);
for (const [_k, v] of Object.entries<Function>(mod)) {
if (v.prototype instanceof AbstractImageInterrogationModel) {
imClasses.push(v as any);
}
}
} catch (err) {
// ignore
console.warn(`Failed to load module ${m}`, err);
}
}
return imClasses;
}
@singleton()
export class ImageInterrogationManager extends AsyncService {
modelClasses: (typeof AbstractImageInterrogationModel<unknown>)[];
modelMap: Record<string, AbstractImageInterrogationModel<unknown>> = {};
logger = this.globalLogger.child({ service: this.constructor.name });
constructor(
protected globalLogger: GlobalLogger,
protected llmManager: LLMManager,
) {
super(...arguments);
this.modelClasses = loadModulesDynamically(__dirname);
const instances: AsyncService[] = [];
for (const x of this.modelClasses) {
const instance: AbstractImageInterrogationModel<unknown> = container.resolve(x as any);
const names = _.uniq([x.name, x.name.toLowerCase(), ...(x.aliases || [])]);
for (const n of names) {
this.modelMap[n] = instance;
}
instances.push(instance);
}
this.dependsOn(...instances);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
getModel(name: string) {
let mdl = this.modelMap[name];
if (!mdl && this.llmManager.getModel(name)?.interleavedPromptSupported) {
mdl = container.resolve(imageInterrogationWithVisualLLM(this.llmManager.getModel(name)!));
this.modelMap[name] = mdl;
}
if (!mdl) {
return;
}
return mdl;
}
assertModel(name: string) {
const mdl = this.getModel(name);
if (!mdl) {
throw new AssertionFailureError(`Unknown model: ${name}`);
}
return mdl;
}
hasModel(name: string) {
return Boolean(this.getModel(name));
}
interrogate<U extends ImageInterrogationOptions<T>, T>(
model: string,
modelOpts: U,
execOpts?: {
maxTry?: number;
requestOptions?: HTTPServiceRequestOptions,
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: <T extends HTTPService>(client: T) => boolean;
}
): Promise<string> {
return this.assertModel(model).interrogate(modelOpts, execOpts);
}
brief() {
return Object.fromEntries(this.modelClasses.map((x) => [x.aliases![0], x.description!]));
}
}
const instance = container.resolve(ImageInterrogationManager);
export default instance;
src/services/common-llm/base.ts
+1098
View File
File diff suppressed because it is too large Load Diff
src/services/common-llm/google-gemini.ts
+1022
View File
File diff suppressed because it is too large Load Diff
src/services/common-llm/gpt-1106.ts
+239
View File
@@ -0,0 +1,239 @@
import { DownstreamServiceFailureError, isCoercibleClass } from 'civkit/civ-rpc';
import _ from "lodash";
import { injectable } from "tsyringe";
import {
CHAT_GPT_ROLE,
ChatGPT0613,
ChatGPTModelOptions,
} from './gpt-35';
import { OpenAIHTTP } from '../../3rd-party/openai';
import { DetectFunctions, LLMDto } from './base';
import { FunctionCallingAwareLLMMessage, FunctionCallingAwareLLMModelOptions, LLMFunctionCallRequest, LLMMessage, LLMModelOptions, LLMPeakStream, PromptChunk } from './misc';
import { isReadable, once, Readable } from 'stream';
import { HTTPServiceRequestOptions } from 'civkit/http';
import { isPrimitiveType } from 'civkit/lang';
export class ChatGPTModelOptions1106 extends ChatGPTModelOptions {
static override windowSize = 16384;
}
@injectable()
export class ChatGPT1106 extends ChatGPT0613 {
static override description = 'OpenAI GPT 3.5 Turbo Model with tools';
static override aliases = [
'gpt-3.5-turbo-1106', 'chatgpt-1106', 'gpt3.5-1106', 'gpt35-1106',
];
static override streamingSupported = true;
static override chatOptimized = true;
static override systemSupported = true;
static override modelOptionsType = ChatGPTModelOptions1106;
static override windowSize = ChatGPTModelOptions1106.windowSize;
static override nSupported = true;
static override functionCalling = 'native' as const;
static override modelName = 'gpt-3.5-turbo-1106';
static override visionEnabled: boolean = false;
override async structured<U extends (typeof LLMDto), MO extends FunctionCallingAwareLLMModelOptions<ChatGPTModelOptions>>(
expectOutputClass: U,
input: string | FunctionCallingAwareLLMMessage[],
modelOpts?: MO,
execOpts?: {
maxTry?: number;
requestOptions?: HTTPServiceRequestOptions,
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: (client: OpenAIHTTP) => boolean;
asyncValidate?: (parsed: InstanceType<U>, raw: string) => Promise<boolean>;
peekStream?: LLMPeakStream;
peekChannel?: string;
collectRawOutput?: (rawOutput: any) => void;
collectRawInput?: (
options: FunctionCallingAwareLLMModelOptions<unknown>,
prompt?: string | PromptChunk[],
) => void;
}
): Promise<DetectFunctions<MO, InstanceType<U>>> {
const messages: LLMModelOptions<ChatGPTModelOptions>['messages'] = [
...(
typeof input === 'string' ?
[{ role: 'user', content: input }] :
await Promise.all(
input.map(
(x) => {
if ((x as LLMMessage).role) {
return x as LLMMessage;
}
return this.getFunctionCallMessage(x);
}
)
)
)
];
if (modelOpts?.messages?.length) {
messages.push(...await Promise.all(
modelOpts.messages.map(
(x) => {
if ((x as LLMMessage).role) {
return x as LLMMessage;
}
return this.getFunctionCallMessage(x);
}
)
));
}
const expectedObjectLike = isCoercibleClass(expectOutputClass);
let lastError = '';
let triesLeft = execOpts?.maxTry || 3;
let trick: undefined | 'jsonMode' = expectedObjectLike ? 'jsonMode' : undefined;
while (triesLeft > 0) {
triesLeft -= 1;
if (lastError) {
execOpts?.peekStream?.write({
event: 'retry',
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
}
let ret: any;
if (trick === undefined) {
const opts = {
...modelOpts,
messages,
n: 1
};
ret = await this.exec('', opts, execOpts as any);
execOpts?.collectRawInput?.(opts, '');
} else if (trick === 'jsonMode') {
let system = modelOpts?.system || '';
if (
!system.includes('JSON') &&
!((modelOpts?.messages?.find((x) => (x as LLMMessage).role === 'system') as LLMMessage)?.content?.includes('JSON'))
) {
system = `${system ? `${system}\n` : ''}${expectedObjectLike ? expectOutputClass?.toJSONModeSystemPromptFragment() : `Provide your response in the following JSON Schema: ${JSON.stringify({
type: 'object',
properties: {
data: {
type: expectOutputClass.name.toLowerCase(),
description: `The expected output`
}
}
})}`}`;
}
const opts = {
...modelOpts,
modelSpecific: {
response_format: { type: 'json_object' }
},
messages,
n: 1,
system,
};
ret = await this.exec('', opts, execOpts as any) as string;
execOpts?.collectRawInput?.(opts, '');
}
execOpts?.collectRawOutput?.(ret);
const textChunks: string[] = [];
if (isReadable(ret as Readable)) {
let lastNonTextChunk;
(ret as Readable).on('data', (c: any) => {
if (typeof c === 'string') {
textChunks.push(c);
execOpts?.peekStream?.write({
event: 'chunk',
data: c,
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
} else {
lastNonTextChunk = c;
}
});
await once(ret, 'end');
ret = lastNonTextChunk || textChunks.join('');
}
execOpts?.peekStream?.write({
event: 'response',
data: ret,
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
if (ret instanceof LLMFunctionCallRequest) {
return ret as any;
} else if (typeof ret === 'string') {
messages.push({
role: 'assistant',
content: ret,
});
}
let structure: InstanceType<U>;
try {
structure = expectOutputClass.fromString ? expectOutputClass.fromString(ret as string) as InstanceType<U> :
isPrimitiveType(expectOutputClass) ?
expectOutputClass.call(undefined, ret) as InstanceType<U> :
Reflect.construct(expectOutputClass, [ret]) as InstanceType<U>;
} catch (err: any) {
lastError = err?.readableMessage || err?.message || 'Invalid response';
const errorMsg = `[SYSTEM] Error: ${lastError}\nTry again. Note: this is a system message. DO NOT REPLY.`;
execOpts?.peekStream?.write({
event: 'exception',
data: errorMsg,
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
messages.push({ role: CHAT_GPT_ROLE.USER, content: errorMsg });
continue;
}
if (execOpts?.asyncValidate) {
try {
execOpts?.peekStream?.write({
event: 'check',
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
const validated = await execOpts.asyncValidate(
structure as any,
ret || ''
);
if (!validated) {
throw new Error('Invalid response');
}
} catch (err: any) {
lastError = err?.readableMessage || err?.message || 'Invalid response';
const errorMsg = `[SYSTEM] Error: ${lastError}\nTry again. Note: this is a system message. DO NOT REPLY.`;
execOpts?.peekStream?.write({
event: 'exception',
data: errorMsg,
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
messages.push({ role: 'user', content: errorMsg });
continue;
}
}
execOpts?.peekStream?.write({
event: 'result',
data: structure,
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
return structure as any;
}
throw new DownstreamServiceFailureError(`Failed to parse ${expectOutputClass.name} from LLM: ${lastError}`);
}
}
@injectable()
export class ChatGPT0125 extends ChatGPT1106 {
static override description = 'GPT-3.5 Turbo model with higher accuracy at responding in requested formats and a fix for a bug which caused a text encoding issue for non-English language function calls';
static override aliases = [
'gpt-3.5-turbo-0125', 'chatgpt-0125', 'gpt3.5-0125', 'gpt35-0125',
];
static override modelName = 'gpt-3.5-turbo-0125';
}
src/services/common-llm/gpt-35.ts
+946
View File
@@ -0,0 +1,946 @@
import { Also, ArrayOf, Coercible, ParamValidationError, Pick, Prop } from 'civkit/civ-rpc';
import { AbstractLLM, DependsOnOptions } from './base';
import { LLMFunctionCallRequest, LLMFunctionCallResponse, LLMMessage, LLMModelOptions, PromptChunk, chatMLEncode, parseJSON, stringPromptChunks } from './misc';
import { OpenAIHTTP } from '../../3rd-party/openai';
import { EnvConfig } from '../envconfig';
import _ from 'lodash';
import { Readable, Transform, TransformCallback } from 'stream';
import { injectable } from 'tsyringe';
import { countGPTToken } from "../../utils/openai";
import { GlobalLogger } from "../logger";
import type { ChatCompletion, ChatCompletionChunk, ChatCompletionMessageToolCall } from 'openai/resources/index';
import { TempFileManager } from '../temp-file';
import { readFile } from 'fs/promises';
// import { OpenAICompatHTTP } from '../../3rd-party/openai-compat';
import { AsyncLocalContext } from '../async-context';
import { HTTPServiceRequestOptions } from 'civkit/http';
import { downloadFile } from 'civkit/download';
import { mimeOf } from 'civkit/mime';
import { HarmfulContentError } from '../errors';
export enum CHAT_GPT_ROLE {
USER = 'user',
SYSTEM = 'system',
ASSISTANT = 'assistant',
FUNCTION = 'function',
TOOL = 'tool',
}
export enum CHAT_GPT_FUNCTION_CALL_MODE {
NONE = 'none',
AUTO = 'auto',
}
export interface ChatGPTFunctionCall {
name: string;
// The arguments to call the function with, as generated by the model in JSON format. Note that the model does not always generate valid JSON, and may hallucinate parameters not defined by your function schema. Validate the arguments in your code before calling your function.
arguments: string;
}
export class ChatGPTFunctionCall extends Coercible {
@Prop({
desc: 'The name of the function to call.',
required: true,
})
name!: string;
@Prop({
desc: 'The arguments to call the function with, as generated by the model in JSON format. Note that the model does not always generate valid JSON, and may hallucinate parameters not defined by your function schema. Validate the arguments in your code before calling your function.',
required: true,
})
arguments!: string;
}
export class ChatGPTToolCall_Function extends Coercible {
@Prop({
desc: 'The ID of the tool call',
required: true,
})
id!: string;
@Prop({
desc: 'The type of the tool. Currently, only function is supported.',
required: true,
default: 'function',
})
type!: string;
@Prop({
desc: `The function that the model called.`,
required: true,
})
function!: ChatGPTFunctionCall;
@Prop()
index?: number;
}
export class ChatGPTMessageTextChunk extends Coercible {
@Prop({
desc: 'Type of the message chunk',
default: 'text',
required: true,
validate: (x) => x === 'text'
})
type!: 'text';
@Prop({
desc: 'The content of the message chunk',
required: true,
})
text!: string;
}
export enum CHAT_GPT_IMAGE_DETAIL {
AUTO = 'auto',
LOW = 'low',
HIGH = 'high',
}
export class ChatGPTMessageImageUrl extends Coercible {
@Prop({
desc: 'Either a URL of the image or the base64 encoded image data.',
})
url?: string;
@Prop({
desc: 'Either a URL of the image or the base64 encoded image data.',
type: CHAT_GPT_IMAGE_DETAIL,
default: CHAT_GPT_IMAGE_DETAIL.AUTO,
})
detail?: CHAT_GPT_IMAGE_DETAIL | 'string';
}
export class ChatGPTMessageImageUrlChunk extends Coercible {
@Prop({
desc: 'Type of the message chunk',
default: 'image_url',
required: true,
validate: (x) => x === 'image_url'
})
type!: 'image_url';
@Prop({
desc: 'The content of the image url chunk',
required: true,
})
image_url!: ChatGPTMessageImageUrl;
}
export class ChatGPTMessage extends Coercible {
@Prop({
desc: 'The role of the author of this message.',
required: true,
default: 'user',
type: CHAT_GPT_ROLE,
})
role!: CHAT_GPT_ROLE | string;
@Prop({
desc: 'The contents of the message',
default: '',
nullable: true,
type: [ArrayOf(ChatGPTMessageTextChunk, ChatGPTMessageImageUrlChunk), String],
})
content!: string | (ChatGPTMessageTextChunk | ChatGPTMessageImageUrlChunk)[];
@Prop({
desc: 'The name of the user in a multi-user chat'
})
name?: string;
@Prop({
desc: 'The tool calls generated by the model, such as function calls.',
arrayOf: [ChatGPTToolCall_Function],
})
tool_calls?: (ChatGPTToolCall_Function)[];
@Prop({
desc: 'The name and arguments of a function that should be called, as generated by the model.',
})
function_call?: ChatGPTFunctionCall;
@Prop({
desc: 'Tool call that this message is responding to.',
})
tool_call_id?: string;
}
@Also({
validate: (i: ChatGPTMessage) => (i.role === CHAT_GPT_ROLE.FUNCTION) ? Boolean(i.name) : true
})
export class ClassicChatGPTMessage extends ChatGPTMessage {
@Prop({
desc: 'The contents of the message',
default: '',
nullable: true,
})
override content!: string;
}
export class ChatGPTFunctionDescriber extends Coercible {
@Prop({
desc: 'The name of the function to be called. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 64.',
validate: (x: string) => /[a-zA-Z0-9_]{1,64}/.test(x),
required: true
})
name!: string;
@Prop({
desc: 'The description of what the function does.',
default: ''
})
description!: string;
@Prop({
desc: 'The parameters the functions accepts, described as a JSON Schema object. See the guide for examples, and the JSON Schema reference for documentation about the format.',
required: true,
})
parameters!: object;
}
export enum CHAT_GPT_RESPONSE_FORMAT {
TEXT = 'text',
JSON_OBJECT = 'json_object',
}
export class ChatGPTResponseFormat extends Coercible {
@Prop({
required: true,
default: CHAT_GPT_RESPONSE_FORMAT.TEXT,
type: CHAT_GPT_RESPONSE_FORMAT,
})
type!: CHAT_GPT_RESPONSE_FORMAT | string;
}
export enum CHAT_GPT_TOOL_TYPE {
FUNCTION = 'function',
}
export class ChatGPTToolDescriber_Function extends Coercible {
@Prop({
required: true,
default: CHAT_GPT_TOOL_TYPE.FUNCTION,
type: CHAT_GPT_TOOL_TYPE,
})
type!: CHAT_GPT_TOOL_TYPE.FUNCTION;
@Prop({
required: true,
})
function!: ChatGPTFunctionDescriber;
}
export class ChatGPTToolReference_Function extends Coercible {
@Prop({
required: true,
default: CHAT_GPT_TOOL_TYPE.FUNCTION,
type: CHAT_GPT_TOOL_TYPE,
})
type!: CHAT_GPT_TOOL_TYPE.FUNCTION;
@Prop({
required: true,
type: Pick(ChatGPTFunctionDescriber, 'name'),
})
function!: Pick<ChatGPTFunctionDescriber, 'name'>;
}
export class ChatGPTModelOptions extends Coercible {
static windowSize = 4096;
@Prop({
default: 'gpt-3.5-turbo',
})
model!: string;
@Prop({
required: true,
desc: 'The messages to generate chat completions for, in the [chat format](/docs/guides/chat/introduction).',
arrayOf: ChatGPTMessage,
validateCollection(this: typeof ChatGPTModelOptions, messages: ChatGPTMessage[]) {
const totalTokens = _.sum(messages.map((m) => countGPTToken(m.content, 'gpt-3.5-turbo')));
if (totalTokens > this.windowSize) {
throw new Error(`The total number of tokens in the messages (${totalTokens}) exceeds the maximum window size of ${this.windowSize}`);
}
return true;
},
})
messages!: ChatGPTMessage[];
@Prop({
desc: 'The maximum number of [tokens](/tokenizer) to generate in the completion. The token count of your prompt plus `max_tokens` cannot exceed the model\'s context length. Most models have a context length of 2048 tokens (except for the newest models, which support 4096).',
validate(this: typeof ChatGPTModelOptions, val: number) {
return val >= 1 && val <= this.windowSize;
},
})
max_tokens?: number;
@Prop({
desc: 'What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic. We generally recommend altering this or `top_p` but not both.',
validate: (val: number) => val >= 0.0 && val <= 2.0,
})
temperature?: number;
@Prop({
desc: 'An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered. We generally recommend altering this or `temperature` but not both.',
validate: (val: number) => val >= 0.0 && val <= 1.0,
})
top_p?: number;
@Prop({
desc: 'How many chat completion choices to generate for each input message. **Note:** Because this parameter generates many completions, it can quickly consume your token quota. Use carefully and ensure that you have reasonable settings for `max_tokens` and `stop`.',
validate: (val: number) => val >= 1 && val <= 64,
})
n?: number;
@Prop({
desc: 'If set, partial message deltas will be sent, like in ChatGPT. Tokens will be sent as data-only [server-sent events](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#Event_stream_format) as they become available, with the stream terminated by a `data: [DONE]` message.'
})
stream?: boolean;
@Prop({
type: [ArrayOf(String), String],
})
stop?: string | string[];
@Prop({
desc: 'Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far, increasing the model\'s likelihood to talk about new topics. [See more information about frequency and presence penalties.](/docs/api-reference/parameter-details)',
validate: (val: number) => val >= -2.0 && val <= 2.0,
})
presence_penalty?: number;
@Prop({
desc: 'Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model\'s likelihood to repeat the same line verbatim. [See more information about frequency and presence penalties.](/docs/api-reference/parameter-details)',
validate: (val: number) => val >= -2.0 && val <= 2.0,
})
frequency_penalty?: number;
@Prop({
desc: 'Modify the likelihood of specified tokens appearing in the completion. Accepts a json object that maps tokens (specified by their token ID in the GPT tokenizer) to an associated bias value from -100 to 100. You can use this [tokenizer tool](/tokenizer?view=bpe) (which works for both GPT-2 and GPT-3) to convert text to token IDs. Mathematically, the bias is added to the logits generated by the model prior to sampling. The exact effect will vary per model, but values between -1 and 1 should decrease or increase likelihood of selection; values like -100 or 100 should result in a ban or exclusive selection of the relevant token. As an example, you can pass `{\"50256\": -100}` to prevent the <|endoftext|> token from being generated.',
dictOf: Number,
})
logit_bias?: Record<string, number>;
@Prop({
desc: 'A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse. [Learn more](/docs/guides/safety-best-practices/end-user-ids).',
})
user?: string;
@Prop({
desc: 'An object specifying the format that the model must output. Used to enable JSON mode.',
})
response_format?: ChatGPTResponseFormat;
@Prop({
desc: 'This feature is in Beta. If specified, our system will make a best effort to sample deterministically, such that repeated requests with the same seed and parameters should return the same result. Determinism is not guaranteed, and you should refer to the system_fingerprint response parameter to monitor changes in the backend.',
nullable: true
})
seed?: number;
@Prop({
desc: 'A list of tools the model may call. Currently, only functions are supported as a tool. Use this to provide a list of functions the model may generate JSON inputs for.',
arrayOf: [ChatGPTToolDescriber_Function],
})
tools?: ChatGPTToolDescriber_Function[];
@Prop({
desc: `Controls which (if any) function is called by the model. none means the model will not call a function and instead generates a message. auto means the model can pick between generating a message or calling a function.`,
type: [ChatGPTToolReference_Function, CHAT_GPT_FUNCTION_CALL_MODE]
})
tool_choice?: ChatGPTToolReference_Function | CHAT_GPT_FUNCTION_CALL_MODE;
@Prop({
desc: 'A list of functions the model may generate JSON inputs for.',
arrayOf: ChatGPTFunctionDescriber,
})
functions?: ChatGPTFunctionDescriber[];
@Prop({
desc: `Controls how the model responds to function calls. "none" means the model does not call a function, and responds to the end-user. "auto" means the model can pick between an end-user or calling a function. Specifying a particular function via {"name":\ "my_function"} forces the model to call that function. "none" is the default when no functions are present. "auto" is the default if functions are present.`,
type: [Pick(ChatGPTFunctionDescriber, 'name'), CHAT_GPT_FUNCTION_CALL_MODE]
})
function_call?: { name: string; } | CHAT_GPT_FUNCTION_CALL_MODE;
}
export class ChatGPT16KModelOptions extends ChatGPTModelOptions {
static override windowSize = 16384;
@Prop({
default: 'gpt-3.5-turbo-16k',
})
override model!: string;
}
export class ChatGPTResponseStream extends Transform {
lastEvent?: ChatCompletionChunk;
get id() {
return this.lastEvent?.id;
}
get crated() {
return this.lastEvent ? new Date(this.lastEvent.created * 1000) : undefined;
}
get model() {
return this.lastEvent?.model;
}
get systemFingerprint() {
return this.lastEvent?.system_fingerprint;
}
get choice() {
return this.accumulatedObject?.choices?.find((x) => x.index === this.choiceIndex);
}
get role(): string | undefined {
return this.choice?.delta.role;
}
history?: LLMMessage;
finishReason?: string;
accumulatedObject: ChatCompletionChunk = {} as any;
constructor(protected choiceIndex: number = 0) {
super({
objectMode: true,
highWaterMark: 4 * 1024,
});
}
applyDelta(ptr: object = this.accumulatedObject, delta?: any, stack: string[] = []) {
if (typeof delta !== 'object' || delta === null) {
return;
}
for (const [k, v] of Object.entries(delta)) {
if (Array.isArray(v)) {
let cur = Reflect.get(ptr, k);
if (!Array.isArray(cur)) {
cur = [];
Reflect.set(ptr, k, cur);
}
for (const inner of v) {
if (typeof inner.index !== 'number') {
continue;
}
let nextHop = Reflect.get(cur, inner.index);
if (!nextHop) {
nextHop = {};
Reflect.set(cur, inner.index, nextHop);
}
this.applyDelta(nextHop, inner, [...stack, k, inner.index]);
}
continue;
}
if (typeof v === 'object' && v !== null) {
let nextHop = Reflect.get(ptr, k);
if (!nextHop || typeof nextHop !== 'object') {
nextHop = {};
Reflect.set(ptr, k, nextHop);
}
this.applyDelta(nextHop, v, [...stack, k]);
continue;
}
const curVal = Reflect.get(ptr, k);
if (curVal && typeof v !== typeof curVal) {
// Bug in OpenAI API or Azure screwed up the protocol
continue;
}
if (typeof v === 'string' && stack.includes('delta')) {
Reflect.set(ptr, k, `${curVal || ''}${v}`);
continue;
}
Reflect.set(ptr, k, v);
}
}
override _flush(callback: TransformCallback) {
const choice = this.choice;
if (choice?.delta.tool_calls) {
for (const x of choice?.delta.tool_calls) {
this.emit('tool', x, x.id);
if (x.type === 'function') {
this.emit('call', x.function, x.id);
}
}
this.emit('tools', choice.delta.tool_calls);
this.push(LLMFunctionCallRequest.from({
text: this.accumulatedObject.choices[this.choiceIndex].delta.content || '',
functionCalls: choice.delta.tool_calls.map((x) => {
return {
id: x.id,
name: x.function!.name,
arguments: parseJSON(x.function!.arguments as string)
};
}),
}));
}
this.emit('history', choice?.delta);
this.history = choice?.delta as any;
this.emit('final', this.accumulatedObject);
this.push(null);
callback(null);
}
override _transform(input: { data: ChatCompletionChunk | '[DONE]'; } | null, _encoding: string, callback: TransformCallback) {
if (input === null) {
return callback(null);
}
const chunk = input.data;
if (typeof chunk === 'string') {
// '[DONE]'
return callback(null);
}
if (!chunk) {
return callback(null);
}
this.lastEvent = chunk;
this.emit('event', chunk);
const choices = [...(chunk.choices || [])].filter((x) => x.index === this.choiceIndex);
this.applyDelta(this.accumulatedObject, chunk);
for (const x of choices) {
if (x.delta.content) {
this.emit('text', x.delta.content);
this.push(x.delta.content);
}
if (x.finish_reason) {
this.finishReason = x.finish_reason;
this.emit('reason', this.finishReason);
}
}
callback(null);
return;
}
}
export interface ChatGPTResponseStream extends Transform {
on(event: 'event', handler: (event: ChatCompletionChunk) => void): this;
on(event: 'role', handler: (role: NonNullable<ChatCompletionChunk['choices'][0]['delta']['role']>) => void): this;
on(event: 'text', handler: (textChunk: string) => void): this;
on(event: 'call', handler: (functionCall: ChatCompletionMessageToolCall['function'], id?: string) => void): this;
on(event: 'tools', handler: (toolCalls: ChatCompletionMessageToolCall[]) => void): this;
on(event: 'tool', handler: (toolCall: ChatCompletionMessageToolCall, id: string) => void): this;
on(event: 'reason', handler: (reason: string) => void): this;
on(event: 'history', handler: (history: ChatCompletionChunk['choices'][0]['delta']) => void): this;
once(event: 'event', handler: (event: ChatCompletionChunk) => void): this;
once(event: 'role', handler: (role: NonNullable<ChatCompletionChunk['choices'][0]['delta']['role']>) => void): this;
once(event: 'text', handler: (textChunk: string) => void): this;
once(event: 'call', handler: (functionCall: ChatCompletionMessageToolCall['function'], id?: string) => void): this;
once(event: 'tools', handler: (toolCalls: ChatCompletionMessageToolCall[]) => void): this;
once(event: 'tool', handler: (toolCall: ChatCompletionMessageToolCall, id: string) => void): this;
once(event: 'reason', handler: (reason: string) => void): this;
once(event: 'history', handler: (history: ChatCompletionChunk['choices'][0]['delta']) => void): this;
on(event: 'close', listener: () => void): this;
on(event: 'data', listener: (chunk: any) => void): this;
on(event: 'drain', listener: () => void): this;
on(event: 'end', listener: () => void): this;
on(event: 'error', listener: (err: Error) => void): this;
on(event: 'finish', listener: () => void): this;
on(event: 'pause', listener: () => void): this;
on(event: 'pipe', listener: (src: Readable) => void): this;
on(event: 'readable', listener: () => void): this;
on(event: 'resume', listener: () => void): this;
on(event: 'unpipe', listener: (src: Readable) => void): this;
on(event: string | symbol, listener: (...args: any[]) => void): this;
once(event: 'close', listener: () => void): this;
once(event: 'data', listener: (chunk: any) => void): this;
once(event: 'drain', listener: () => void): this;
once(event: 'end', listener: () => void): this;
once(event: 'error', listener: (err: Error) => void): this;
once(event: 'finish', listener: () => void): this;
once(event: 'pause', listener: () => void): this;
once(event: 'pipe', listener: (src: Readable) => void): this;
once(event: 'readable', listener: () => void): this;
once(event: 'resume', listener: () => void): this;
once(event: 'unpipe', listener: (src: Readable) => void): this;
once(event: string | symbol, listener: (...args: any[]) => void): this;
}
@injectable()
export class ChatGPT0613 extends AbstractLLM<ChatGPTModelOptions> {
static override description = 'OpenAI GPT-3.5 (ChatGPT) Model';
static override aliases = [
'gpt-3.5-turbo-0613', 'chatgpt-0613', 'gpt3.5-0613', 'gpt35-0613'
];
static override streamingSupported = true;
static override chatOptimized = true;
static override systemSupported = true;
static override modelOptionsType = ChatGPTModelOptions;
static override windowSize = ChatGPTModelOptions.windowSize;
static modelName = 'gpt-3.5-turbo-0613';
static override nSupported = true;
static override functionCalling = 'native';
static visionEnabled: boolean = false;
override clients: Array<OpenAIHTTP> = [];
logger = this.globalLogger.child({ service: this.constructor.name });
constructor(
protected envConfig: EnvConfig,
protected globalLogger: GlobalLogger,
protected tempFileManager: TempFileManager,
protected threadLocal: AsyncLocalContext,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.clients = [
new OpenAIHTTP(this.envConfig.OPENAI_API_KEY),
].filter((x) => x?.supportedChatCompletionModels.includes((this.constructor as typeof ChatGPT0613).modelName));
this.emit('ready');
}
override async getFunctionCallMessage(thing: unknown) {
if (thing instanceof LLMFunctionCallRequest) {
return {
role: 'assistant',
content: `${thing.text || ''}`,
tool_calls: thing.functionCalls.map((x) => ({
id: x.id,
type: 'function',
function: {
name: x.name,
arguments: x.arguments,
}
}))
} as LLMMessage;
}
if (thing instanceof LLMFunctionCallResponse) {
return {
role: 'tool',
content: `${typeof thing.content === 'string' ? thing.content : JSON.stringify(thing.content)}`,
name: thing.name,
tool_call_id: thing.functionCallId
} as LLMMessage;
}
return super.getFunctionCallMessage(thing);
}
override async _run<U extends LLMModelOptions<ChatGPTModelOptions>>(
client: this['clients'][number],
prompt: string | PromptChunk[],
modelOpts?: U,
requestOptions?: HTTPServiceRequestOptions
): Promise<DependsOnOptions<U>> {
const modelName = (this.constructor as typeof ChatGPT0613).modelName;
const draftOptions: Partial<ChatGPTModelOptions> = {
top_p: modelOpts?.top_p,
max_tokens: modelOpts?.max_tokens,
temperature: modelOpts?.temperature,
seed: modelOpts?.seed,
...modelOpts?.modelSpecific,
n: modelOpts?.n || 1,
stream: modelOpts?.stream,
model: modelName,
};
if (modelOpts?.functions) {
draftOptions.tools = modelOpts.functions.map((x) => ({
type: CHAT_GPT_TOOL_TYPE.FUNCTION,
function: x
}));
}
if (modelOpts?.function_call) {
draftOptions.tool_choice = typeof modelOpts.function_call === 'string' ?
modelOpts.function_call as CHAT_GPT_FUNCTION_CALL_MODE : {
type: CHAT_GPT_TOOL_TYPE.FUNCTION,
function: modelOpts.function_call
};
}
draftOptions.messages ??= [];
if (modelOpts?.system) {
draftOptions.messages!.unshift({ role: CHAT_GPT_ROLE.SYSTEM, content: modelOpts.system });
}
if (modelOpts?.messages) {
const ps = modelOpts.messages.map(async (x) => {
const mixin = {
content: await this.encodePromptChunks(x.content),
role: x.role,
name: x.name,
};
if (Array.isArray(mixin.content) && mixin.content?.find((x) => x.type !== 'text')) {
if (mixin.role === 'assistant') {
mixin.role = 'user';
mixin.name ??= 'special-assistant';
}
}
if (mixin.role !== 'user' && Array.isArray(mixin.content)) {
mixin.content = mixin.content.map((x) => x.text).join('\n');
}
return {
...x,
...mixin,
};
});
const rs = await Promise.all(ps);
draftOptions.messages?.push(...rs);
}
if (prompt) {
draftOptions.messages!.push({ role: CHAT_GPT_ROLE.USER, content: await this.encodePromptChunks(prompt) });
}
const opts = (this.constructor as typeof ChatGPT0613).modelOptionsType.from(draftOptions);
const numTokens = countGPTToken(chatMLEncode(stringPromptChunks(prompt), modelOpts));
this.logger.debug(`Calling ${modelName} by ${client.name}(${client.baseURL}) approx ${numTokens} tokens`);
if ((numTokens + (opts.max_tokens || 0)) >= this.windowSize) {
const newMaxTokens = Math.floor(this.windowSize - numTokens - 2);
if (newMaxTokens < 0) {
throw new ParamValidationError(`Prompt too long for ${modelName} ${numTokens}/${this.windowSize}`);
}
this.logger.warn(`Prompt has ${numTokens} tokens, max_tokens is ${opts.max_tokens}, which is greater than the window size of ${this.windowSize}. Shrinking max_tokens to ${newMaxTokens}.`);
opts.max_tokens = newMaxTokens;
}
const r = await client.chatCompletions(opts as any, requestOptions);
if (this.threadLocal.ctx?.tokenUsage === undefined) {
this.threadLocal.ctx.tokenUsage = {};
}
if (this.threadLocal.ctx?.tokenUsage[modelName] === undefined) {
this.threadLocal.ctx.tokenUsage[modelName] = {};
}
const usageTrack = this.threadLocal.ctx.tokenUsage[modelName];
if (opts.stream) {
this.logger.debug(`Streaming response from ${modelName} by ${client.name}`);
(r.data as Readable).once('end', () => {
this.logger.debug(`Streaming completed from ${modelName} by ${client.name}`);
});
const streams = _.range(0, opts.n || 1).map(
(i) =>
(r.data as any as Readable).pipe(
new ChatGPTResponseStream(i)
)
);
if (opts.n && opts.n > 1) {
streams[0].once('end', () => {
if (!(r.data as Readable)?.readableEnded) {
(r.data as Readable).destroy();
}
usageTrack.inputTokens ??= 0;
usageTrack.outputTokens ??= 0;
usageTrack.inputTokens += _.get(streams[0], 'accumulatedObject.usage.prompt_tokens', 0);
usageTrack.outputTokens += _.get(streams[0], 'accumulatedObject.usage.completion_tokens', 0);
});
return streams as any;
}
const onlyStream: ChatGPTResponseStream = streams[0];
onlyStream.once('end', () => {
if (!(r.data as Readable)?.readableEnded) {
(r.data as Readable).destroy();
}
usageTrack.inputTokens ??= 0;
usageTrack.outputTokens ??= 0;
usageTrack.inputTokens += _.get(onlyStream.accumulatedObject, 'usage.prompt_tokens', 0);
usageTrack.outputTokens += _.get(onlyStream.accumulatedObject, 'usage.completion_tokens', 0);
});
return onlyStream as any;
}
const rData = r.data as any as ChatCompletion;
usageTrack.inputTokens ??= 0;
usageTrack.outputTokens ??= 0;
usageTrack.inputTokens += _.get(rData, 'usage.prompt_tokens', 0);
usageTrack.outputTokens += _.get(rData, 'usage.completion_tokens', 0);
this.logger.debug(`Received response from ${modelName} by ${client.name}, ${rData.usage?.completion_tokens} tokens.`);
const texts = _(rData).get('choices', []).map((x) => {
if (x.finish_reason === 'content_filter') {
const filterDetails = (x as any).content_filter_result;
if (filterDetails) {
for (const [k, v] of Object.entries<any>(filterDetails)) {
if (v?.filtered) {
throw new HarmfulContentError(`Your input is filtered because it provokes ${k}. Please change your input and try again.`);
}
}
}
this.logger.warn(`Received content triggers 'content_filter' restriction`);
throw new HarmfulContentError('Output possibly includes harmful contents, please check your input.');
}
const tool_calls = _.get(x, 'message.tool_calls');
if (tool_calls && !_.isEmpty(tool_calls)) {
return LLMFunctionCallRequest.from({
functionCalls: tool_calls.map((x) => ({
id: x.id,
name: x.function.name,
arguments: parseJSON(x.function!.arguments as string)
})),
text: _.get(x, 'message.content') || ''
});
}
return _.get(x, 'message.content') || '';
});
if (opts.n && opts.n > 1) {
return texts as any;
}
return texts[0] || '' as any;
}
override async encodePromptChunks(promptChunks: string | PromptChunk[] | null): Promise<any> {
if (!this.interleavedPromptSupported) {
return super.encodePromptChunks(promptChunks);
}
if (typeof promptChunks === 'string') {
return promptChunks;
}
if (!promptChunks) {
return '';
}
const ps = promptChunks.map(async (x) => {
if (typeof x === 'string') {
return { type: 'text', text: x };
}
if (x instanceof URL) {
if (x.hostname === 'localhost' || x.hostname.startsWith('127.')) {
const tmpPath = this.tempFileManager.alloc();
await downloadFile(x.toString(), tmpPath);
const buff = await readFile(tmpPath);
const mimeVec = await mimeOf(buff);
const r = {
type: 'image_url',
image_url: {
url: `data:${mimeVec.mediaType}/${mimeVec.subType};base64,${buff.toString('base64')}`,
detail: 'auto',
}
};
this.tempFileManager.bindPathTo(r, tmpPath);
return r;
}
return {
type: 'image_url', image_url: {
url: `${x}`,
detail: 'auto',
}
};
}
if (Buffer.isBuffer(x)) {
const mimeInfo = await mimeOf(x);
return {
type: 'image_url',
image_url: {
url: `data:${mimeInfo.mediaType}/${mimeInfo.subType};base64,${x.toString('base64')}`,
detail: 'auto',
}
};
}
if (x instanceof Blob) {
return {
type: 'image_url',
image_url: {
url: `data:${x.type};base64,${Buffer.from(await x.arrayBuffer()).toString('base64')}`,
detail: 'auto',
}
};
}
return { type: 'text', text: JSON.stringify(x) };
});
return await Promise.all(ps);
}
}
@injectable()
export class ChatGPT16K0613 extends ChatGPT0613 implements AbstractLLM<ChatGPT16KModelOptions> {
static override description = 'OpenAI GPT-3.5 (ChatGPT) Model With 16k Window';
static override aliases = [
'gpt3.5-16k-0613', 'chatgpt-16k-0613', 'gpt35-16k-0613', 'gpt-3.5-turbo-16k-0613'
];
static override streamingSupported = true;
static override chatOptimized = true;
static override modelOptionsType = ChatGPT16KModelOptions;
static override windowSize = ChatGPT16KModelOptions.windowSize;
static override modelName = 'gpt-3.5-turbo-16k';
static override nSupported = true;
override clients: Array<OpenAIHTTP> = [];
override async init() {
await this.dependencyReady();
this.clients = [
new OpenAIHTTP(this.envConfig.OPENAI_API_KEY),
].filter((x) => x.supportedChatCompletionModels.includes((this.constructor as typeof ChatGPT16K0613).modelName));
this.emit('ready');
}
}
@injectable()
export class ChatGPT0301 extends ChatGPT0613 {
static override description = 'OpenAI GPT-3.5 (ChatGPT) Model With function call feature';
static override aliases = ['gpt-3.5-turbo-0301', 'chatgpt-0301', 'gpt3.5-0301'];
static override streamingSupported = true;
static override chatOptimized = true;
static override windowSize = 4096;
static override modelOptionsType = ChatGPTModelOptions;
static override modelName = 'gpt-3.5-turbo-0301';
static override nSupported = true;
static override functionCalling = 'none';
override clients: Array<OpenAIHTTP> = [];
override async init() {
await this.dependencyReady();
this.clients = [
new OpenAIHTTP(this.envConfig.OPENAI_API_KEY),
].filter((x) => x.supportedChatCompletionModels.includes((this.constructor as typeof ChatGPT0301).modelName));
this.emit('ready');
}
}
src/services/common-llm/gpt-4.ts
+61
View File
@@ -0,0 +1,61 @@
import { Prop } from 'civkit/civ-rpc';
import _ from "lodash";
import { injectable } from "tsyringe";
import { ChatGPTMessage, ChatGPTModelOptions } from './gpt-35';
import { ChatGPT1106, } from './gpt-1106';
import { countGPTToken } from '../../utils/openai';
export class GPT4oModelOptions extends ChatGPTModelOptions {
static override windowSize = 128_000;
@Prop({
default: 'gpt-4o',
})
override model!: string;
@Prop({
required: true,
desc: 'The messages to generate chat completions for, in the [chat format](/docs/guides/chat/introduction).',
arrayOf: ChatGPTMessage,
validateCollection(this: typeof ChatGPTModelOptions, messages: ChatGPTMessage[]) {
return _.sum(messages.map((m) => countGPTToken(m.content, 'gpt-4o'))) <= this.windowSize;
},
})
override messages!: ChatGPTMessage[];
}
@injectable()
export class GPT4o_20240513 extends ChatGPT1106 {
static override description = 'OpenAI GPT 4 Omni Model 2024-05-13';
static override aliases = [
'gpt-4o-2024-05-13',
];
static override modelOptionsType = GPT4oModelOptions;
static override windowSize = GPT4oModelOptions.windowSize;
static override visionEnabled = true;
static override interleavedPromptSupported = true;
static override modelName = 'gpt-4o-2024-05-13';
}
@injectable()
export class GPT4o extends GPT4o_20240513 {
static override description = 'OpenAI GPT 4 Omni Model 2024-08-06';
static override aliases = [
'gpt-4o-2024-08-06',
'gpt-4o',
];
static override modelName = 'gpt-4o-2024-08-06';
}
@injectable()
export class GPT4oMini extends GPT4o_20240513 {
static override description = 'OpenAI GPT 4 Omni Model mini 2024-07-18';
static override aliases = [
'gpt-4o-mini-2024-07-18',
'gpt-4o-mini',
];
static override modelName = 'gpt-4o-mini-2024-07-18';
}
src/services/common-llm/index.ts
+6
View File
@@ -0,0 +1,6 @@
export * from './misc';
export * from './base';
export * from './prompt-profile';
export * from './registry';
import instance from './registry';
export default instance;
src/services/common-llm/jina-vlm.ts
+40
View File
@@ -0,0 +1,40 @@
import { injectable } from 'tsyringe';
import { ChatGPT0613, ChatGPTModelOptions } from './gpt-35';
import { Prop } from 'civkit/civ-rpc';
export class ReaderLMOptions extends ChatGPTModelOptions {
static override windowSize = Math.floor(512000 * 0.8);
@Prop({
default: 1.05,
})
repetition_penalty?: number;
}
@injectable()
export class JinaVLM extends ChatGPT0613 {
static override description = 'JinaAI Jina-VLM';
static override aliases = [
'jina-vlm'
];
static override streamingSupported = true;
static override chatOptimized = true;
static override modelOptionsType = ReaderLMOptions;
static override windowSize = ReaderLMOptions.windowSize;
static override nSupported = false;
static override functionCalling = 'none' as const;
static override interleavedPromptSupported = true;
static override modelName = 'jina-vlm';
static override visionEnabled: boolean = true;
override async init() {
await this.dependencyReady();
this.clients = [
// new InternalCloudRunJinaVLM(this.envConfig.JINA_VLM_DEPLOYMENT_API_KEY)
];
this.emit('ready');
}
}
src/services/common-llm/misc.ts
+315
View File
@@ -0,0 +1,315 @@
import { randomBytes, randomUUID } from 'crypto';
import path from 'path';
import fsp from 'fs/promises';
import { countGPTToken } from '../../utils/openai';
import _ from 'lodash';
import { Writable } from 'stream';
import { JSONAccumulation, JSONParserStream, JSONParserStreamOptions } from '../../lib/json-parse-stream';
import { FancyFile } from 'civkit/fancy-file';
import { Coercible, Prop } from 'civkit/coercible';
export type PromptChunk = string | URL | Buffer | File | object;
export interface LLMMessage {
role: 'user' | 'system' | 'assistant' | string,
content: string | PromptChunk[] | null,
name?: string;
[k: string]: any;
}
export type FunctionCallingAwareLLMMessage = LLMMessage | LLMFunctionCallRequest | LLMFunctionCallResponse;
export interface FunctionCallingAwareLLMModelOptions<T> {
system?: string;
messages?: Array<FunctionCallingAwareLLMMessage>;
stop?: string[];
stream?: boolean;
max_tokens?: number;
temperature?: number;
top_p?: number;
top_k?: number;
n?: number;
seed?: number;
functions?: { name: string, description: string, parameters: any; }[];
function_call?: { name: string; } | string;
attachments?: (FancyFile | File | string)[];
modelSpecific?: Partial<T>;
}
export function chatMLEncode(prompt?: string, opts?: {
system?: string;
messages?: {
role: string;
content: string | PromptChunk[] | null;
[k: string]: any;
}[];
}) {
const chunks: string[] = [];
const imStart = '<|im_start|>';
const imEnd = '<|im_end|>';
if (opts?.system) {
chunks.push(`${imStart}system\n${opts.system}${imEnd}`);
}
if (opts?.messages?.length) {
for (const m of opts.messages) {
const params = _.omit(m, 'content', 'role');
const paramParts = _.map(params, (v, k) => `${k}=${typeof v === 'object' ? JSON.stringify(v) : v}`).join(', ');
const content = typeof m.content === 'string' ? m.content : m.content?.filter((x) => typeof x === 'string').join('');
chunks.push(`${imStart}${m.role || 'user'}${paramParts ? ` ${paramParts}` : ''}\n${content}${imEnd}`);
}
}
if (prompt) {
if (prompt.startsWith(imStart)) {
chunks.push(prompt);
} else {
chunks.push(`${imStart}user\n${prompt}${imEnd}\n${imStart}assistant\n`);
}
}
return chunks.join('\n');
}
const DATAURL_REGEXP = /^data:([a-z]+\/[a-z0-9-+.]+)?;base64,(.*)$/i;
function parseDataUrl(s: string) {
return DATAURL_REGEXP.exec(s);
}
export async function parseFileFromString(input: string, overrideName?: string) {
const dataurl = parseDataUrl(input);
if (dataurl) {
return new File([Buffer.from(dataurl[2], 'base64')], overrideName || randomUUID(), { type: dataurl[1] });
}
if (input.startsWith('http://') || input.startsWith('https://')) {
const parsedUrl = new URL(input);
const name = overrideName || path.basename(parsedUrl.pathname) || randomUUID();
return await fetch(input).then((r) => r.blob()).then((r) => new File([r], name, { type: r.type }));
}
if (Buffer.isBuffer(input)) {
return new File([input], overrideName || randomUUID());
}
return new File([await fsp.readFile(input) as any], overrideName || path.basename(input));
}
export function interleavedPrompt(strs: TemplateStringsArray | string[], ...args: (PromptChunk | PromptChunk[])[]) {
const chunks = _.zip(strs, args).flat() as PromptChunk[];
const expandedChunks: PromptChunk[] = [];
for (const x of chunks) {
if (Array.isArray(x)) {
expandedChunks.push(...x);
continue;
}
expandedChunks.push(x);
}
// Merge text chunks
const finalChunks: PromptChunk[] = [];
const strStack: PromptChunk[] = [];
for (const x of expandedChunks) {
if (typeof x === 'string') {
if (x) {
strStack.push(x);
}
continue;
} else if (typeof x === 'boolean' || (x as any) === 0) {
strStack.push(`${x}`);
continue;
} else if (!x) {
continue;
}
if (strStack.length) {
finalChunks.push(strStack.join(''));
strStack.length = 0;
}
finalChunks.push(x);
}
if (strStack.length) {
finalChunks.push(strStack.join(''));
strStack.length = 0;
}
return finalChunks;
}
export function stringPromptChunks(promptChunks: string | PromptChunk[] | null | undefined): string {
if (!Array.isArray(promptChunks)) {
return promptChunks || '';
}
const chunks: PromptChunk[] = [];
for (const x of promptChunks) {
if (typeof x === 'string') {
chunks.push(x);
} else if (x instanceof URL || x instanceof Blob || Buffer.isBuffer(x)) {
chunks.push(`[Image]`);
} else {
chunks.push(`${JSON.stringify(x)}`);
}
}
return chunks.join('\n');
}
export function trimMessages(prompt: string | PromptChunk[], opts: LLMModelOptions<any>, targetSize: number = 3096) {
const stringPrompt = stringPromptChunks(prompt);
const initialEncoded = chatMLEncode(stringPrompt, opts);
let tokenCount = countGPTToken(initialEncoded);
const tailMessages = [...(opts.messages || [])];
const headMessages: any[] = [];
while (tokenCount > targetSize) {
const firstMessage = tailMessages.shift();
if (!firstMessage) {
break;
}
if (firstMessage.role === 'system') {
headMessages.push(firstMessage);
continue;
}
tokenCount -= countGPTToken(stringPromptChunks(firstMessage?.content));
tokenCount -= 4;
const rest = _.omit(firstMessage, 'role', 'content');
if (!_.isEmpty(rest)) {
const restChunks: string[] = [];
for (const [k, v] of Object.entries(rest)) {
restChunks.push(`${k}=${typeof v === 'object' ? JSON.stringify(v) : v}`);
}
tokenCount -= countGPTToken(restChunks.join(', '));
}
}
return [...headMessages, ...tailMessages];
}
export interface LLMPeakEvent {
event: 'chunk' | 'response' | 'check' | 'exception' | 'retry' | 'result' | 'error';
data?: string | object | any[];
channel?: string;
}
export class LLMPeakStream extends Writable {
constructor() {
super({ objectMode: true, highWaterMark: 4 * 1024 });
}
}
export class LLMFunctionCall extends Coercible {
@Prop({
desc: 'Call id',
required: true,
defaultFactory() {
return randomBytes(6).toString('base64url');
}
})
id!: string;
@Prop({
desc: 'name of the function',
required: true
})
name!: string;
@Prop({
desc: 'The function call arguments.',
default: ''
})
arguments!: string | any;
}
export class LLMFunctionCallRequest extends Coercible {
@Prop({
default: 'LLMFunctionCallRequest'
})
$$typeof!: 'LLMFunctionCallRequest';
@Prop({
desc: 'Optional text preceding the alternative response.',
})
text?: string;
@Prop({
desc: 'The function call response model generated.',
required: true,
arrayOf: LLMFunctionCall,
})
functionCalls!: LLMFunctionCall[];
}
export class LLMFunctionCallResponse extends Coercible {
@Prop({
default: 'LLMFunctionCallResponse'
})
$$typeof!: 'LLMFunctionCallResponse';
@Prop({
desc: 'Function call id',
required: true,
defaultFactory() {
return randomBytes(6).toString('base64url');
}
})
functionCallId!: string;
@Prop({
desc: 'Function name',
})
name?: string;
@Prop({
desc: 'The function call response.',
required: true
})
content: any;
@Prop({
desc: 'The response was an error.',
})
isError?: boolean;
requestCall?: LLMFunctionCall;
}
export function parseJSON(i: string, opts: JSONParserStreamOptions & { withOffsetMetadata?: boolean; } = {
expectAbruptTerminationOfInput: true,
expectContaminated: 'object',
expectControlCharacterInString: true,
expectCasingInLiteral: true,
}) {
let parsedArg: any = i;
try {
const s = JSONParserStream.parse(i, opts);
parsedArg = JSONAccumulation.parse(s, undefined, opts.withOffsetMetadata);
} catch (_err) {
void 0;
}
return parsedArg;
}
export interface LLMModelOptions<T> extends FunctionCallingAwareLLMModelOptions<T> {
messages?: Array<LLMMessage>;
};
src/services/common-llm/open-router.ts
+919
View File
@@ -0,0 +1,919 @@
import { ArrayOf, Coercible, DownstreamServiceFailureError, isCoercibleClass, ParamValidationError, Pick, Prop } from 'civkit/civ-rpc';
import { injectable } from 'tsyringe';
import { isReadable, once, Readable } from 'stream';
import { chatMLEncode, FunctionCallingAwareLLMMessage, FunctionCallingAwareLLMModelOptions, LLMFunctionCallRequest, LLMFunctionCallResponse, LLMMessage, LLMModelOptions, LLMPeakStream, parseJSON, PromptChunk, stringPromptChunks } from './misc';
import { AbstractLLM, DependsOnOptions, DetectFunctions, LLMDto } from './base';
import { OpenRouterHTTP } from '../../3rd-party/open-router';
import { EnvConfig } from '../envconfig';
import { GlobalLogger } from '../logger';
import { TempFileManager } from '../temp-file';
import { AsyncLocalContext } from '../async-context';
import { HTTPServiceRequestOptions, downloadFile, isPrimitiveType, mimeOf } from 'civkit';
import { readFile } from 'fs/promises';
import _ from 'lodash';
import { ChatCompletion } from 'openai/resources/index.mjs';
import { countGPTToken } from '../../utils/openai';
import { HarmfulContentError } from '../errors';
import { ChatGPTResponseStream } from './gpt-35';
export enum CHAT_GPT_IMAGE_DETAIL {
AUTO = 'auto',
LOW = 'low',
HIGH = 'high',
}
export class ChatGPTMessageImageUrl extends Coercible {
@Prop({
desc: 'Either a URL of the image or the base64 encoded image data.',
})
url?: string;
@Prop({
desc: 'Either a URL of the image or the base64 encoded image data.',
type: CHAT_GPT_IMAGE_DETAIL,
default: CHAT_GPT_IMAGE_DETAIL.AUTO,
})
detail?: CHAT_GPT_IMAGE_DETAIL | 'string';
}
export class ChatGPTMessageImageUrlChunk extends Coercible {
@Prop({
desc: 'Type of the message chunk',
default: 'image_url',
required: true,
validate: (x) => x === 'image_url'
})
type!: 'image_url';
@Prop({
desc: 'The content of the image url chunk',
required: true,
})
image_url!: ChatGPTMessageImageUrl;
}
export enum CHAT_GPT_ROLE {
USER = 'user',
SYSTEM = 'system',
ASSISTANT = 'assistant',
FUNCTION = 'function',
TOOL = 'tool',
}
export class ChatGPTMessageTextChunk extends Coercible {
@Prop({
desc: 'Type of the message chunk',
default: 'text',
required: true,
validate: (x) => x === 'text'
})
type!: 'text';
@Prop({
desc: 'The content of the message chunk',
required: true,
})
text!: string;
}
export interface ChatGPTFunctionCall {
name: string;
// The arguments to call the function with, as generated by the model in JSON format. Note that the model does not always generate valid JSON, and may hallucinate parameters not defined by your function schema. Validate the arguments in your code before calling your function.
arguments: string;
}
export class ChatGPTToolCall_Function extends Coercible {
@Prop({
desc: 'The ID of the tool call',
required: true,
})
id!: string;
@Prop({
desc: 'The type of the tool. Currently, only function is supported.',
required: true,
default: 'function',
})
type!: string;
@Prop({
desc: `The function that the model called.`,
required: true,
})
function!: ChatGPTFunctionCall;
@Prop()
index?: number;
}
export class ChatGPTMessage extends Coercible {
@Prop({
desc: 'The role of the author of this message.',
required: true,
default: 'user',
type: CHAT_GPT_ROLE,
})
role!: CHAT_GPT_ROLE | string;
@Prop({
desc: 'The contents of the message',
default: '',
nullable: true,
type: [ArrayOf(ChatGPTMessageTextChunk, ChatGPTMessageImageUrlChunk), String],
})
content!: string | (ChatGPTMessageTextChunk | ChatGPTMessageImageUrlChunk)[];
@Prop({
desc: 'The name of the user in a multi-user chat'
})
name?: string;
@Prop({
desc: 'The tool calls generated by the model, such as function calls.',
arrayOf: [ChatGPTToolCall_Function],
})
tool_calls?: (ChatGPTToolCall_Function)[];
@Prop({
desc: 'The name and arguments of a function that should be called, as generated by the model.',
})
function_call?: ChatGPTFunctionCall;
@Prop({
desc: 'Tool call that this message is responding to.',
})
tool_call_id?: string;
}
export enum CHAT_GPT_TOOL_TYPE {
FUNCTION = 'function',
}
export class ChatGPTFunctionDescriber extends Coercible {
@Prop({
desc: 'The name of the function to be called. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 64.',
validate: (x: string) => /[a-zA-Z0-9_]{1,64}/.test(x),
required: true
})
name!: string;
@Prop({
desc: 'The description of what the function does.',
default: ''
})
description!: string;
@Prop({
desc: 'The parameters the functions accepts, described as a JSON Schema object. See the guide for examples, and the JSON Schema reference for documentation about the format.',
required: true,
})
parameters!: object;
}
export class ChatGPTToolDescriber_Function extends Coercible {
@Prop({
required: true,
default: CHAT_GPT_TOOL_TYPE.FUNCTION,
type: CHAT_GPT_TOOL_TYPE,
})
type!: CHAT_GPT_TOOL_TYPE.FUNCTION;
@Prop({
required: true,
})
function!: ChatGPTFunctionDescriber;
}
export class ChatGPTToolReference_Function extends Coercible {
@Prop({
required: true,
default: CHAT_GPT_TOOL_TYPE.FUNCTION,
type: CHAT_GPT_TOOL_TYPE,
})
type!: CHAT_GPT_TOOL_TYPE.FUNCTION;
@Prop({
required: true,
type: Pick(ChatGPTFunctionDescriber, 'name'),
})
function!: Pick<ChatGPTFunctionDescriber, 'name'>;
}
export enum CHAT_GPT_FUNCTION_CALL_MODE {
NONE = 'none',
AUTO = 'auto',
}
export enum OPENROUTER_REASONING_EFFORT {
HIGH = 'high',
MEDIUM = 'medium',
LOW = 'low',
}
export class OpenRouterReasoningOptions extends Coercible {
@Prop({
type: OPENROUTER_REASONING_EFFORT,
})
effort?: OPENROUTER_REASONING_EFFORT;
@Prop()
max_tokens?: number;
@Prop()
exclude?: boolean;
@Prop()
enabled?: boolean;
}
// Definitions of subtypes are below
export class OpenRouterRequestOptions extends Coercible {
static contextLength = 8192;
// Either "messages" or "prompt" is required
@Prop({
arrayOf: ChatGPTMessage
})
messages?: ChatGPTMessage[];
@Prop()
prompt?: string;
@Prop()
reasoning?: OpenRouterReasoningOptions;
// If "model" is unspecified, uses the user's default
@Prop()
model?: string; // See "Supported Models" section
// Allows to force the model to produce specific output format.
// See models page and note on this docs page for which models support it.
@Prop()
response_format?: { type: 'json_object' | string; };
@Prop({
type: [ArrayOf(String), String],
})
stop?: string | string[];
@Prop()
stream?: boolean; // Enable streaming
// See LLM Parameters (openrouter.ai/docs/api-reference/parameters)
@Prop({
validate(this: typeof OpenRouterRequestOptions, v: number) { return v >= 1 && v <= this.contextLength; }
})
max_tokens?: number; // Range: [1, context_length)
@Prop({
validate: (v: number) => v >= 0 && v <= 2,
})
temperature?: number; // Range: [0, 2]
// Tool calling
// Will be passed down as-is for providers implementing OpenAI's interface.
// For providers with custom interfaces, we transform and map the properties.
// Otherwise, we transform the tools into a YAML template. The model responds with an assistant message.
// See models supporting tool calling: openrouter.ai/models?supported_parameters=tools
@Prop({
arrayOf: ChatGPTToolDescriber_Function,
})
tools?: ChatGPTToolDescriber_Function[];
@Prop({
type: [ChatGPTToolReference_Function, CHAT_GPT_FUNCTION_CALL_MODE]
})
tool_choice?: ChatGPTToolReference_Function | CHAT_GPT_FUNCTION_CALL_MODE;
// Advanced optional parameters
@Prop({
validate: (v: number) => Number.isInteger(v)
})
seed?: number; // Integer only
@Prop({
validate: (v: number) => v > 0 && v < 1
})
top_p?: number; // Range: (0, 1]
@Prop({
validate: (v: number) => v >= 1
})
top_k?: number; // Range: [1, Infinity) Not available for OpenAI models
@Prop({
validate: (v: number) => v >= -2 && v <= 2
})
frequency_penalty?: number; // Range: [-2, 2]
@Prop({
validate: (v: number) => v >= -2 && v <= 2
})
presence_penalty?: number; // Range: [-2, 2]
@Prop({
validate: (v: number) => v > 0 && v < 2
})
repetition_penalty?: number; // Range: (0, 2]
@Prop({
dictOf: Number
})
logit_bias?: { [key: string]: number; };
@Prop({
validate: (v: number) => Number.isInteger(v) && v > 0 && v < 20
})
top_logprobs?: number; // Integer only
@Prop({
validate: (v: number) => v >= 0 && v <= 1
})
min_p?: number; // Range: [0, 1]
@Prop({
validate: (v: number) => v >= 0 && v <= 1
})
top_a?: number; // Range: [0, 1]
// Reduce latency by providing the model with a predicted output
// https://platform.openai.com/docs/guides/latency-optimization#use-predicted-outputs
@Prop()
prediction?: { type: 'content'; content: string; };
// OpenRouter-only parameters
// See "Prompt Transforms" section: openrouter.ai/docs/transforms
@Prop({ arrayOf: String })
transforms?: string[];
// See "Model Routing" section: openrouter.ai/docs/model-routing
@Prop({
arrayOf: String
})
models?: string[];
// See "Provider Routing" section: openrouter.ai/docs/provider-routing
@Prop()
provider?: ProviderPreferences;
@Prop()
user?: string; // A stable identifier for your end-users. Used to help detect and prevent abuse.
};
type ProviderPreferences = {
order?: string[];
allow_fallbacks?: boolean;
require_parameters?: boolean;
data_collection?: 'allow' | 'deny';
only?: string[];
ignore?: string[];
quantizations?: string[];
sort?: string;
max_price?: {
prompt?: number;
completion?: number;
image?: number;
};
};
// Definitions of subtypes are below
export type OpenRouterChatCompletionResponse = {
id: string;
// Depending on whether you set "stream" to "true" and
// whether you passed in "messages" or a "prompt", you
// will get a different output shape
choices: (NonStreamingChoice | StreamingChoice | NonChatChoice)[];
created: number; // Unix timestamp
model: string;
object: 'chat.completion' | 'chat.completion.chunk';
system_fingerprint?: string; // Only present if the provider supports it
// Usage data is always returned for non-streaming.
// When streaming, you will get one usage object at
// the end accompanied by an empty choices array.
usage?: ResponseUsage;
};
// If the provider returns usage, we pass it down
// as-is. Otherwise, we count using the GPT-4 tokenizer.
type ResponseUsage = {
/** Including images and tools if any */
prompt_tokens: number;
/** The tokens generated */
completion_tokens: number;
/** Sum of the above two fields */
total_tokens: number;
};
// Subtypes:
type NonChatChoice = {
index: number;
finish_reason: string | null;
text: string;
error?: ErrorResponse;
};
type NonStreamingChoice = {
index: number;
finish_reason: string | null;
native_finish_reason: string | null;
message: {
content: string | null;
role: string;
tool_calls?: ToolCall[];
};
error?: ErrorResponse;
};
type StreamingChoice = {
index: number;
finish_reason: string | null;
native_finish_reason: string | null;
delta: {
content: string | null;
role?: string;
tool_calls?: ToolCall[];
};
error?: ErrorResponse;
};
type ErrorResponse = {
code: number; // See "Error Handling" section
message: string;
metadata?: Record<string, unknown>; // Contains additional error information such as provider details, the raw error message, etc.
};
type ToolCall = {
id: string;
type: 'function';
function: FunctionCall;
};
type FunctionCall = {
name: string;
// The arguments to call the function with, as generated by the model in JSON format. Note that the model does not always generate valid JSON, and may hallucinate parameters not defined by your function schema. Validate the arguments in your code before calling your function.
arguments: string;
};
@injectable()
export class OpenRouterLLM extends AbstractLLM<OpenRouterRequestOptions> {
static override description = 'OpenRouter LLM';
static override aliases = [
'open-router'
];
static override streamingSupported = true;
static override chatOptimized = true;
static override systemSupported = true;
static override modelOptionsType = OpenRouterRequestOptions;
static override windowSize = OpenRouterRequestOptions.contextLength;
static modelName = 'open-router';
static override functionCalling = 'native';
override clients: Array<OpenRouterHTTP> = [];
logger = this.globalLogger.child({ service: this.constructor.name });
constructor(
protected envConfig: EnvConfig,
protected globalLogger: GlobalLogger,
protected tempFileManager: TempFileManager,
protected threadLocal: AsyncLocalContext,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.clients = [
new OpenRouterHTTP(this.envConfig.OPENROUTER_API_KEY),
];
this.emit('ready');
}
override async getFunctionCallMessage(thing: unknown) {
if (thing instanceof LLMFunctionCallRequest) {
return {
role: 'assistant',
content: `${thing.text || ''}`,
tool_calls: thing.functionCalls.map((x) => ({
id: x.id,
type: 'function',
function: {
name: x.name,
arguments: x.arguments,
}
}))
} as LLMMessage;
}
if (thing instanceof LLMFunctionCallResponse) {
return {
role: 'tool',
content: `${typeof thing.content === 'string' ? thing.content : JSON.stringify(thing.content)}`,
name: thing.name,
tool_call_id: thing.functionCallId
} as LLMMessage;
}
return super.getFunctionCallMessage(thing);
}
override async _run<U extends LLMModelOptions<OpenRouterRequestOptions>>(
client: this['clients'][number],
prompt: string | PromptChunk[],
modelOpts?: U,
requestOptions?: HTTPServiceRequestOptions
): Promise<DependsOnOptions<U>> {
const modelName = (this.constructor as typeof OpenRouterLLM).modelName;
const draftOptions: Partial<OpenRouterRequestOptions> = {
top_p: modelOpts?.top_p,
max_tokens: modelOpts?.max_tokens,
temperature: modelOpts?.temperature,
seed: modelOpts?.seed,
...modelOpts?.modelSpecific,
stream: modelOpts?.stream,
model: modelName,
};
if (modelOpts?.functions) {
draftOptions.tools = modelOpts.functions.map((x) => ({
type: CHAT_GPT_TOOL_TYPE.FUNCTION,
function: x
}));
}
if (modelOpts?.function_call) {
draftOptions.tool_choice = typeof modelOpts.function_call === 'string' ?
modelOpts.function_call as CHAT_GPT_FUNCTION_CALL_MODE : {
type: CHAT_GPT_TOOL_TYPE.FUNCTION,
function: modelOpts.function_call
};
}
draftOptions.messages ??= [];
if (modelOpts?.system) {
draftOptions.messages!.unshift({ role: CHAT_GPT_ROLE.SYSTEM, content: modelOpts.system });
}
if (modelOpts?.messages) {
const ps = modelOpts.messages.map(async (x) => {
const mixin = {
content: await this.encodePromptChunks(x.content),
role: x.role,
name: x.name,
};
if (Array.isArray(mixin.content) && mixin.content?.find((x) => x.type !== 'text')) {
if (mixin.role === 'assistant') {
mixin.role = 'user';
mixin.name ??= 'special-assistant';
}
}
if (mixin.role !== 'user' && Array.isArray(mixin.content)) {
mixin.content = mixin.content.map((x) => x.text).join('\n');
}
return {
...x,
...mixin,
};
});
const rs = await Promise.all(ps);
draftOptions.messages?.push(...rs);
}
if (prompt) {
draftOptions.messages!.push({ role: CHAT_GPT_ROLE.USER, content: await this.encodePromptChunks(prompt) });
}
const opts = (this.constructor as typeof OpenRouterLLM).modelOptionsType.from(draftOptions);
const numTokens = countGPTToken(chatMLEncode(stringPromptChunks(prompt), modelOpts));
this.logger.debug(`Calling ${modelName} by ${client.name}(${client.baseURL}) approx ${numTokens} tokens`);
if ((numTokens + (opts.max_tokens || 0)) >= this.windowSize) {
const newMaxTokens = Math.floor(this.windowSize - numTokens - 2);
if (newMaxTokens < 0) {
throw new ParamValidationError(`Prompt too long for ${modelName} ${numTokens}/${this.windowSize}`);
}
this.logger.warn(`Prompt has ${numTokens} tokens, max_tokens is ${opts.max_tokens}, which is greater than the window size of ${this.windowSize}. Shrinking max_tokens to ${newMaxTokens}.`);
opts.max_tokens = newMaxTokens;
}
const r = await client.chatCompletions(opts as any, requestOptions);
let usageTrack: Record<string, any> = {};
if (this.threadLocal.available()) {
if (this.threadLocal.ctx?.tokenUsage === undefined) {
this.threadLocal.ctx.tokenUsage = {};
}
if (this.threadLocal.ctx?.tokenUsage[modelName] === undefined) {
this.threadLocal.ctx.tokenUsage[modelName] = {};
}
usageTrack = this.threadLocal.ctx.tokenUsage[modelName];
}
if (opts.stream) {
this.logger.debug(`Streaming response from ${modelName} by ${client.name}`);
(r.data as Readable).once('end', () => {
this.logger.debug(`Streaming completed from ${modelName} by ${client.name}`);
});
const onlyStream: ChatGPTResponseStream = (r.data as any as Readable).pipe(
new ChatGPTResponseStream(0)
);
onlyStream.once('end', () => {
if (!(r.data as Readable)?.readableEnded) {
(r.data as Readable).destroy();
}
usageTrack.inputTokens ??= 0;
usageTrack.outputTokens ??= 0;
usageTrack.inputTokens += _.get(onlyStream.accumulatedObject, 'usage.prompt_tokens', 0);
usageTrack.outputTokens += _.get(onlyStream.accumulatedObject, 'usage.completion_tokens', 0);
});
return onlyStream as any;
}
const rData = r.data as any as ChatCompletion;
usageTrack.inputTokens ??= 0;
usageTrack.outputTokens ??= 0;
usageTrack.inputTokens += _.get(rData, 'usage.prompt_tokens', 0);
usageTrack.outputTokens += _.get(rData, 'usage.completion_tokens', 0);
this.logger.debug(`Received response from ${modelName} by ${client.name}, ${rData.usage?.completion_tokens} tokens.`);
const texts = _(rData).get('choices', []).map((x) => {
if (x.finish_reason === 'content_filter') {
const filterDetails = (x as any).content_filter_result;
if (filterDetails) {
for (const [k, v] of Object.entries<any>(filterDetails)) {
if (v?.filtered) {
throw new HarmfulContentError(`Your input is filtered because it provokes ${k}. Please change your input and try again.`);
}
}
}
this.logger.warn(`Received content triggers 'content_filter' restriction`);
throw new HarmfulContentError('Output possibly includes harmful contents, please check your input.');
}
const tool_calls = _.get(x, 'message.tool_calls');
if (tool_calls) {
return LLMFunctionCallRequest.from({
functionCalls: tool_calls.map((x) => ({
id: x.id,
name: x.function.name,
arguments: parseJSON(x.function!.arguments as string)
})),
text: _.get(x, 'message.content') || ''
});
}
return _.get(x, 'message.content') || '';
});
return texts[0] || '' as any;
}
override async encodePromptChunks(promptChunks: string | PromptChunk[] | null): Promise<any> {
if (!this.interleavedPromptSupported) {
return super.encodePromptChunks(promptChunks);
}
if (typeof promptChunks === 'string') {
return promptChunks;
}
if (!promptChunks) {
return '';
}
const ps = promptChunks.map(async (x) => {
if (typeof x === 'string') {
return { type: 'text', text: x };
}
if (x instanceof URL) {
if (x.hostname === 'localhost' || x.hostname.startsWith('127.')) {
const tmpPath = this.tempFileManager.alloc();
await downloadFile(x.toString(), tmpPath);
const buff = await readFile(tmpPath);
const mimeVec = await mimeOf(buff);
const r = {
type: 'image_url',
image_url: {
url: `data:${mimeVec.mediaType}/${mimeVec.subType};base64,${buff.toString('base64')}`,
detail: 'auto',
}
};
this.tempFileManager.bindPathTo(r, tmpPath);
return r;
}
return {
type: 'image_url', image_url: {
url: `${x}`,
detail: 'auto',
}
};
}
if (Buffer.isBuffer(x)) {
const mimeInfo = await mimeOf(x);
return {
type: 'image_url',
image_url: {
url: `data:${mimeInfo.mediaType}/${mimeInfo.subType};base64,${x.toString('base64')}`,
detail: 'auto',
}
};
}
if (x instanceof Blob) {
return {
type: 'image_url',
image_url: {
url: `data:${x.type};base64,${Buffer.from(await x.arrayBuffer()).toString('base64')}`,
detail: 'auto',
}
};
}
return { type: 'text', text: JSON.stringify(x) };
});
return await Promise.all(ps);
}
override async structured<U extends (typeof LLMDto), MO extends FunctionCallingAwareLLMModelOptions<OpenRouterRequestOptions>>(
expectOutputClass: U,
input: string | FunctionCallingAwareLLMMessage[],
modelOpts?: MO,
execOpts?: {
maxTry?: number;
requestOptions?: HTTPServiceRequestOptions,
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: (client: OpenRouterHTTP) => boolean;
asyncValidate?: (parsed: InstanceType<U>, raw: string) => Promise<boolean>;
peekStream?: LLMPeakStream;
peekChannel?: string;
collectRawOutput?: (rawOutput: any) => void;
collectRawInput?: (
options: FunctionCallingAwareLLMModelOptions<unknown>,
prompt?: string | PromptChunk[],
) => void;
}
): Promise<DetectFunctions<MO, InstanceType<U>>> {
const messages: LLMModelOptions<OpenRouterRequestOptions>['messages'] = [
...(
typeof input === 'string' ?
[{ role: 'user', content: input }] :
await Promise.all(
input.map(
(x) => {
if ((x as LLMMessage).role) {
return x as LLMMessage;
}
return this.getFunctionCallMessage(x);
}
)
)
)
];
if (modelOpts?.messages?.length) {
messages.push(...await Promise.all(
modelOpts.messages.map(
(x) => {
if ((x as LLMMessage).role) {
return x as LLMMessage;
}
return this.getFunctionCallMessage(x);
}
)
));
}
const expectedObjectLike = isCoercibleClass(expectOutputClass);
let lastError = '';
let triesLeft = execOpts?.maxTry || 3;
let trick: undefined | 'jsonMode' = expectedObjectLike ? 'jsonMode' : undefined;
while (triesLeft > 0) {
triesLeft -= 1;
if (lastError) {
execOpts?.peekStream?.write({
event: 'retry',
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
}
let ret: any;
if (trick === undefined) {
const opts = {
...modelOpts,
messages,
n: 1
};
ret = await this.exec('', opts, execOpts as any);
execOpts?.collectRawInput?.(opts, '');
} else if (trick === 'jsonMode') {
let system = modelOpts?.system || '';
if (
!system.includes('JSON') &&
!((modelOpts?.messages?.find((x) => (x as LLMMessage).role === 'system') as LLMMessage)?.content?.includes('JSON'))
) {
system = `${system ? `${system}\n` : ''}${expectedObjectLike ? expectOutputClass?.toJSONModeSystemPromptFragment() : `Provide your response in the following JSON Schema: ${JSON.stringify({
type: 'object',
properties: {
data: {
type: expectOutputClass.name.toLowerCase(),
description: `The expected output`
}
}
})}`}`;
}
const opts = {
...modelOpts,
modelSpecific: {
response_format: { type: 'json_object' }
},
messages,
n: 1,
system,
};
ret = await this.exec('', opts, execOpts as any) as string;
execOpts?.collectRawInput?.(opts, '');
}
execOpts?.collectRawOutput?.(ret);
const textChunks: string[] = [];
if (isReadable(ret as Readable)) {
let lastNonTextChunk;
(ret as Readable).on('data', (c: any) => {
if (typeof c === 'string') {
textChunks.push(c);
execOpts?.peekStream?.write({
event: 'chunk',
data: c,
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
} else {
lastNonTextChunk = c;
}
});
await once(ret, 'end');
ret = lastNonTextChunk || textChunks.join('');
}
execOpts?.peekStream?.write({
event: 'response',
data: ret,
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
if (ret instanceof LLMFunctionCallRequest) {
return ret as any;
} else if (typeof ret === 'string') {
messages.push({
role: 'assistant',
content: ret,
});
}
let structure: InstanceType<U>;
try {
structure = expectOutputClass.fromString ? expectOutputClass.fromString(ret as string) as InstanceType<U> :
isPrimitiveType(expectOutputClass) ?
expectOutputClass.call(undefined, ret) as InstanceType<U> :
Reflect.construct(expectOutputClass, [ret]) as InstanceType<U>;
} catch (err: any) {
lastError = err?.readableMessage || err?.message || 'Invalid response';
const errorMsg = `[SYSTEM] Error: ${lastError}\nTry again. Note: this is a system message. DO NOT REPLY.`;
execOpts?.peekStream?.write({
event: 'exception',
data: errorMsg,
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
messages.push({ role: CHAT_GPT_ROLE.USER, content: errorMsg });
continue;
}
if (execOpts?.asyncValidate) {
try {
execOpts?.peekStream?.write({
event: 'check',
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
const validated = await execOpts.asyncValidate(
structure as any,
ret || ''
);
if (!validated) {
throw new Error('Invalid response');
}
} catch (err: any) {
lastError = err?.readableMessage || err?.message || 'Invalid response';
const errorMsg = `[SYSTEM] Error: ${lastError}\nTry again. Note: this is a system message. DO NOT REPLY.`;
execOpts?.peekStream?.write({
event: 'exception',
data: errorMsg,
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
messages.push({ role: 'user', content: errorMsg });
continue;
}
}
execOpts?.peekStream?.write({
event: 'result',
data: structure,
id: execOpts.peekChannel ? `${execOpts.peekChannel}-${Reflect.get(execOpts.peekStream, 'n') || 0 + 1}` : undefined
});
return structure as any;
}
throw new DownstreamServiceFailureError(`Failed to parse ${expectOutputClass.name} from LLM: ${lastError}`);
}
}
src/services/common-llm/prompt-profile.ts
+77
View File
@@ -0,0 +1,77 @@
import { COERCIBLE_OPTIONS_SYMBOL, Coercible, inputSingle } from 'civkit/coercible';
import { AbstractLLM, LLMDto, LLMModelOptions } from './base';
import { Readable } from 'stream';
import { FunctionCallingAwareLLMMessage, FunctionCallingAwareLLMModelOptions, LLMFunctionCallRequest } from './misc';
export interface PromptProfileRuntimeMetadata {
modelName: string;
model: AbstractLLM<unknown>;
prompt: string | FunctionCallingAwareLLMMessage[];
modelOptions: FunctionCallingAwareLLMModelOptions<unknown>;
iterations: { input: FunctionCallingAwareLLMModelOptions<unknown>, output?: Readable | string | LLMDto | LLMFunctionCallRequest; }[];
}
export abstract class PromptProfile<T = unknown> extends Coercible {
abstract modelOutput?: LLMDto | string | string[] | number | boolean | Readable | Readable[];
runtime?: PromptProfileRuntimeMetadata;
selectModel(): string | Promise<string> {
return 'default';
}
renderSystemPrompt(): string | undefined | Promise<string> | Promise<undefined> {
return undefined;
}
renderModelOptions(): LLMModelOptions<T> | Promise<LLMModelOptions<T>> | Promise<undefined> | undefined {
return undefined;
}
abstract renderPrompt(): string | FunctionCallingAwareLLMMessage[] | Promise<string | FunctionCallingAwareLLMMessage[]>;
get modelOutputDto() {
const theConstructor = this.constructor as typeof Coercible;
const opts = theConstructor[COERCIBLE_OPTIONS_SYMBOL];
const expectedType = opts?.['modelOutput']?.type;
if (!expectedType) {
throw new Error('Invalid modelOutput type');
}
if (Array.isArray(expectedType)) {
throw new Error('Invalid modelOutput type');
}
return expectedType;
}
async acceptModelOutput(parsed: typeof this.modelOutput, raw: string): Promise<boolean> {
const theConstructor = this.constructor as typeof Coercible;
const opts = theConstructor[COERCIBLE_OPTIONS_SYMBOL];
const propOpts = opts?.['modelOutput'];
if (!propOpts) {
throw new Error('Invalid modelOutput type');
}
const final = inputSingle(this.constructor.name, parsed, undefined, {
type: propOpts.type, desc: propOpts.desc
});
this.modelOutput = final;
return true;
}
get modelOutputJSONSchema() {
const s = this.modelOutputDto?.JSONSchema;
if (s) {
return s;
}
return { type: s.name.toLowerCase() };
}
}
src/services/common-llm/reader-lm.ts
+41
View File
@@ -0,0 +1,41 @@
import { injectable } from 'tsyringe';
// import { InternalCloudRunReaderLM2 } from '../../3rd-party/internal-cloudrun';
import { ChatGPT0613, ChatGPTModelOptions } from './gpt-35';
import { Prop } from 'civkit/civ-rpc';
export class ReaderLMOptions extends ChatGPTModelOptions {
static override windowSize = Math.floor(512000 * 0.8);
@Prop({
default: 1.05,
})
repetition_penalty?: number;
}
@injectable()
export class ReaderLM2 extends ChatGPT0613 {
static override description = 'JinaAI ReaderLM-v2';
static override aliases = [
'readerlm-v2'
];
static override streamingSupported = true;
static override chatOptimized = true;
static override systemSupported = true;
static override modelOptionsType = ReaderLMOptions;
static override windowSize = ReaderLMOptions.windowSize;
static override nSupported = false;
static override functionCalling = 'none' as const;
static override modelName = 'readerlm-v2';
static override visionEnabled: boolean = false;
override async init() {
await this.dependencyReady();
this.clients = [
// new InternalCloudRunReaderLM2(this.envConfig.READER_LM2_DEPLOYMENT_API_KEY)
];
this.emit('ready');
}
}
src/services/common-llm/registry.ts
+547
View File
@@ -0,0 +1,547 @@
import { ApplicationError, AssertionFailureError, isCoercibleClass } from 'civkit/civ-rpc';
import { container, injectable, singleton } from 'tsyringe';
import { readdirSync } from 'fs';
import _ from 'lodash';
import { GlobalLogger } from '../logger';
import { AbstractLLM, DependsOnOptions, DetectFunctions, LLMDto, PromptChunk } from './base';
import { Readable, isReadable } from 'stream';
import { PromptProfile, PromptProfileRuntimeMetadata } from './prompt-profile';
import { FunctionCallingAwareLLMMessage, FunctionCallingAwareLLMModelOptions, LLMFunctionCallRequest, LLMPeakStream } from './misc';
import { AsyncService } from 'civkit/async-service';
import { HarmfulContentError } from '../errors';
import { delay } from 'civkit/timeout';
import { HTTPServiceRequestOptions } from 'civkit/http';
import { isScalarLike } from '../../utils/misc';
import { EnvConfig } from '../envconfig';
import { OpenRouterHTTP } from '../../3rd-party/open-router';
import { OpenRouterLLM, OpenRouterRequestOptions } from './open-router';
function loadModulesDynamically(path: string) {
const moduleDir = readdirSync(path,
{ withFileTypes: true, encoding: 'utf-8' });
const modules = moduleDir.filter((x) => x.isFile() && x.name.endsWith('.js')).map((x) => x.name);
const apiClasses: (typeof AbstractLLM<unknown>)[] = [];
for (const m of modules) {
if (m === 'index.js') {
continue;
}
try {
// if (process.env.hasOwnProperty(`LLM_DISABLE_${m.toUpperCase()}`)) {
// continue;
// }
// FIXME: Does not work with esm
// eslint-disable-next-line @typescript-eslint/no-var-requires
const mod = require(`${path}/${m}`);
for (const [_k, v] of Object.entries<Function>(mod)) {
if (v?.prototype instanceof AsyncService && (v as any).windowSize) {
apiClasses.push(v as any);
}
}
} catch (err) {
// ignore
console.warn(`Failed to load module ${m}`, err);
}
}
return apiClasses;
}
@singleton()
export class LLMManager extends AsyncService {
modelClasses: (typeof AbstractLLM<unknown>)[];
modelMap: Record<string, AbstractLLM<unknown> | undefined> = {};
logger = this.globalLogger.child({ service: this.constructor.name });
constructor(
protected globalLogger: GlobalLogger,
protected envConfig: EnvConfig,
) {
super(...arguments);
this.modelClasses = loadModulesDynamically(__dirname);
const instances: AsyncService[] = [];
for (const x of this.modelClasses) {
const instance: AbstractLLM<unknown> = container.resolve(x as any);
const names = _.uniq([x.name, x.name.toLowerCase(), ...(x.aliases || [])]);
for (const n of names) {
this.modelMap[n] = instance;
}
instances.push(instance);
}
this.dependsOn(...instances);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
async importOpenRouterModel(...openRouterModels: string[]) {
if (!openRouterModels.length) {
return;
}
if (!this.envConfig.OPENROUTER_API_KEY) {
this.logger.warn('OpenRouter API key not found, skipping OpenRouter model import');
return;
}
const openRouterClient = new OpenRouterHTTP(this.envConfig.OPENROUTER_API_KEY);
const models = (await openRouterClient.listModels()).data.data;
const modelsToImport = models.filter((x) => {
if (!openRouterModels.includes(x.id)) {
return false;
}
if (!x.architecture.output_modalities.includes('text')) {
return false;
}
return true;
});
for (const model of modelsToImport) {
this.logger.debug(`Importing OpenRouter model: ${model.id}`);
class CustomizedOpenRouterRequestOptions extends OpenRouterRequestOptions {
static override contextLength = model.context_length;
}
@injectable()
class CustomizedOpenRouterLLM extends OpenRouterLLM {
static override description = model.description;
static override aliases = [
model.id
];
static override modelOptionsType = CustomizedOpenRouterRequestOptions;
static override windowSize = CustomizedOpenRouterRequestOptions.contextLength;
static override modelName = model.id;
static override functionCalling = model.supported_parameters.includes('tools') ? 'native' : 'none';
static override functionCallingSupported = model.supported_parameters.includes('tools');
static override interleavedPromptSupported = model.architecture.input_modalities.includes('image');
}
Object.defineProperty(CustomizedOpenRouterLLM, 'name', { value: model.name });
this.modelClasses.push(CustomizedOpenRouterLLM as any);
const llmInstance = container.resolve(CustomizedOpenRouterLLM);
this.modelMap[model.id] = llmInstance;
this.dependsOn(llmInstance);
await llmInstance.serviceReady();
this.logger.debug(`Imported OpenRouter model: ${model.id}`);
}
}
getModel(name: string) {
const mdl = this.modelMap[name];
return mdl;
}
assertModel(name: string) {
const mdl = this.modelMap[name];
if (!mdl) {
throw new AssertionFailureError(`Unknown model: ${name}`);
}
return mdl;
}
hasModel(name: string) {
return Boolean(this.modelMap[name]);
}
async withModel(model: string, func: (mdl: AbstractLLM<unknown>, client: unknown) => any, options: {
maxTry?: number;
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: (client: AbstractLLM<unknown>['clients'][number]) => boolean;
} = {}): Promise<ReturnType<typeof func>> {
const instance = this.getModel(model);
if (!instance) {
throw new Error(`Model not found: ${model}`);
}
await instance.serviceReady();
if (!instance.clients.length) {
throw new Error('Model client not ready');
}
const clients = options?.filterClients ? instance.clients.filter(options.filterClients) : instance.clients;
const maxTry = options?.maxTry ?? (clients.length + 1);
let i = 0;
while (true) {
for (const client of options.reverseProviderPreference ? [...clients].reverse() : clients) {
if (typeof client !== 'object' || !client) {
throw new Error('Invalid client');
}
try {
this.logger.debug(`Calling LLM model ${instance.constructor.name} with ${(client as any).name || client.constructor.name}`);
const t0 = Date.now();
const r = await func(instance, client);
if (isReadable(r as any)) {
(r as unknown as Readable).once('end', () => {
this.logger.debug(`Streaming succeeded. LLM model ${instance.constructor.name} with ${(client as any).name || client.constructor.name} in ${Date.now() - t0}ms.`);
});
} else {
this.logger.debug(`Call succeeded. LLM model ${instance.constructor.name} with ${(client as any).name || client.constructor.name} in ${Date.now() - t0}ms.`);
}
return r;
} catch (err: any) {
if (err?.cause instanceof ApplicationError) {
this.logger.error(`Failed to run LLM model ${this.constructor.name} with ${(client as any).name || client.constructor.name}. Non recoverable error: Code ${err.code}.`, { error: err });
throw err.cause;
}
const code = err.status || (err as any).code;
if ((((client as any).Error && err instanceof (client as any).Error) || err instanceof HarmfulContentError) && (typeof code === 'number')) {
if (code !== 20 && code !== 429 && code !== 401 && !(code >= 500 && code < 600)) {
this.logger.error(`Failed to run LLM model ${instance.constructor.name} with ${(client as any).name || client.constructor.name}. Non recoverable error: Code ${code}.`, { error: err });
throw err;
}
}
if (++i >= maxTry) {
this.logger.error(`Failed to run LLM model ${instance.constructor.name} with ${(client as any).name || client.constructor.name}. No tries left.`, { error: err });
throw err;
}
const delayMs = Math.floor((options?.delayFactor ?? 1) * (1 + Math.random() * 0.4 - 0.2) * 100);
this.logger.warn(`Failed to run LLM model ${instance.constructor.name} with ${(client as any).name || client.constructor.name}, retrying after ${delayMs}ms...`, { error: err });
await delay(delayMs);
}
}
}
throw new Error('unreachable');
}
run<T, MO extends FunctionCallingAwareLLMModelOptions<T>>(model: string, options: {
prompt: string | PromptChunk[],
options?: MO,
requestOptions?: HTTPServiceRequestOptions,
filterClients?: (client: unknown) => boolean;
maxTry?: number;
delayFactor?: number;
reverseProviderPreference?: boolean;
}): Promise<DependsOnOptions<MO>> {
const mdl = this.getModel(model);
if (!mdl) {
throw new Error(`Model not found: ${model}`);
}
return mdl.exec(options.prompt, options.options, options);
}
iterRun<T, MO extends FunctionCallingAwareLLMModelOptions<T>>(model: string, options: {
prompt: string | PromptChunk[],
options?: MO,
requestOptions?: HTTPServiceRequestOptions,
filterClients?: (client: unknown) => boolean;
maxTry?: number;
delayFactor?: number;
reverseProviderPreference?: boolean;
}) {
const mdl = this.getModel(model);
if (!mdl) {
throw new Error(`Model not found: ${model}`);
}
return mdl.iterExec<MO>(options.prompt, options.options, options);
}
runStructured<U extends (typeof LLMDto), MO extends FunctionCallingAwareLLMModelOptions<unknown>>(
model: string,
expectOutputClass: U,
input: string | FunctionCallingAwareLLMMessage[],
options: {
options?: MO,
requestOptions?: HTTPServiceRequestOptions,
maxTry?: number;
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: (client: unknown) => boolean;
asyncValidate?: (parsed: InstanceType<U>, raw: string) => Promise<boolean>;
peekStream?: LLMPeakStream;
peekChannel?: string;
}) {
const mdl = this.getModel(model);
if (!mdl) {
throw new Error(`Model not found: ${model}`);
}
return mdl.structured<U, MO>(expectOutputClass, input, options.options, _.omit(options, 'options'));
}
iterStructured<U extends (typeof LLMDto), MO extends FunctionCallingAwareLLMModelOptions<unknown>>(
model: string,
expectOutputClass: U,
input: string | FunctionCallingAwareLLMMessage[],
options: {
options?: MO,
requestOptions?: HTTPServiceRequestOptions,
maxTry?: number;
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: (client: unknown) => boolean;
asyncValidate?: (parsed: InstanceType<U>, raw: string) => Promise<boolean>;
peekStream?: LLMPeakStream;
peekChannel?: string;
}) {
const mdl = this.getModel(model);
if (!mdl) {
throw new Error(`Model not found: ${model}`);
}
return mdl.iterStructured<U, MO>(expectOutputClass, input, options.options, _.omit(options, 'options'));
}
async runProfile<R extends PromptProfile, MO extends FunctionCallingAwareLLMModelOptions<unknown>>(
profile: R,
options: {
options?: MO,
maxTry?: number;
requestOptions?: HTTPServiceRequestOptions,
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: (client: unknown) => boolean;
peekStream?: LLMPeakStream;
peekChannel?: string;
} = {}): Promise<DetectFunctions<MO & Awaited<ReturnType<R['renderModelOptions']>>, NonNullable<R['modelOutput']>>> {
await this.serviceReady();
const systemPrompt = await profile.renderSystemPrompt();
const prompt = await profile.renderPrompt();
const mdlPreference = await profile.selectModel();
const modelOptions = await profile.renderModelOptions();
const finalModelOptions = _.merge(_.cloneDeep(modelOptions), options?.options);
const mdl = this.getModel(mdlPreference);
if (!mdl) {
throw new Error(`Model not found: ${mdlPreference}`);
}
const msgMixin: any = {};
let extPrompt = '';
if ((typeof prompt === 'string') && prompt) {
if (finalModelOptions?.messages) {
finalModelOptions.messages = [
{
role: 'user',
content: prompt,
},
...finalModelOptions.messages
];
} else {
extPrompt = prompt;
}
} else {
msgMixin.messages = prompt;
}
const modelOutputDto = profile.modelOutputDto;
if (isCoercibleClass(modelOutputDto) || isScalarLike(modelOutputDto)
) {
const execModelOptions = {
system: systemPrompt,
stream: options.peekStream ? (mdl.streamingSupported) : false,
...msgMixin,
...finalModelOptions,
};
const thisReflect: PromptProfileRuntimeMetadata = {
modelName: mdlPreference,
model: mdl,
prompt,
modelOptions: execModelOptions,
iterations: [],
};
profile.runtime = thisReflect;
const r = await mdl.structured(modelOutputDto, prompt, execModelOptions, {
..._.omit(options, 'options'),
asyncValidate: profile.acceptModelOutput.bind(profile),
collectRawInput: (opts) => {
thisReflect.iterations.push({ input: opts });
},
collectRawOutput: (rawOutput) => {
const lastIter = thisReflect.iterations[thisReflect.iterations.length - 1];
if (lastIter) {
lastIter.output = rawOutput;
}
}
});
return r;
}
const execModelOptions = {
...msgMixin,
...finalModelOptions,
};
const thisReflect: PromptProfileRuntimeMetadata = {
modelName: mdlPreference,
model: mdl,
prompt,
modelOptions: execModelOptions,
iterations: [{ input: execModelOptions }],
};
profile.runtime = thisReflect;
const r = await mdl.exec(extPrompt, execModelOptions, _.omit(options, 'options'));
thisReflect.iterations[0].output = r;
if (r instanceof LLMFunctionCallRequest) {
return r as any;
}
await profile.acceptModelOutput(r, r as any);
return r as any;
}
async *iterProfile<R extends PromptProfile, MO extends FunctionCallingAwareLLMModelOptions<unknown>>(
profile: R,
options: {
options?: MO,
maxTry?: number;
requestOptions?: HTTPServiceRequestOptions,
delayFactor?: number;
reverseProviderPreference?: boolean;
filterClients?: (client: unknown) => boolean;
peekStream?: LLMPeakStream;
peekChannel?: string;
} = {}): AsyncGenerator<DetectFunctions<MO & Awaited<ReturnType<R['renderModelOptions']>>, NonNullable<R['modelOutput']>>, unknown, unknown> {
await this.serviceReady();
const systemPrompt = await profile.renderSystemPrompt();
const prompt = await profile.renderPrompt();
const mdlPreference = await profile.selectModel();
const modelOptions = await profile.renderModelOptions();
const finalModelOptions = _.merge(_.cloneDeep(modelOptions), options?.options);
const mdl = this.getModel(mdlPreference);
if (!mdl) {
throw new Error(`Model not found: ${mdlPreference}`);
}
const msgMixin: any = {};
let extPrompt = '';
if ((typeof prompt === 'string') && prompt) {
if (finalModelOptions?.messages) {
finalModelOptions.messages = [
{
role: 'user',
content: prompt,
},
...finalModelOptions.messages
];
} else {
extPrompt = prompt;
}
} else {
msgMixin.messages = prompt;
}
const modelOutputDto = profile.modelOutputDto;
if (isCoercibleClass(modelOutputDto) || isScalarLike(modelOutputDto)
) {
const execModelOptions = {
system: systemPrompt,
stream: options.peekStream ? (mdl.streamingSupported) : false,
...msgMixin,
...finalModelOptions,
};
const thisReflect: PromptProfileRuntimeMetadata = {
modelName: mdlPreference,
model: mdl,
prompt,
modelOptions: execModelOptions,
iterations: [],
};
profile.runtime = thisReflect;
yield* mdl.iterStructured(modelOutputDto, prompt, execModelOptions, {
..._.omit(options, 'options'),
asyncValidate: profile.acceptModelOutput.bind(profile),
collectRawInput: (opts) => {
thisReflect.iterations.push({ input: opts });
},
collectRawOutput: (rawOutput) => {
const lastIter = thisReflect.iterations[thisReflect.iterations.length - 1];
if (lastIter) {
lastIter.output = rawOutput;
}
}
}) as any;
return;
}
const execModelOptions = {
...msgMixin,
...finalModelOptions,
};
const thisReflect: PromptProfileRuntimeMetadata = {
modelName: mdlPreference,
model: mdl,
prompt,
modelOptions: execModelOptions,
iterations: [{ input: execModelOptions }],
};
profile.runtime = thisReflect;
const r = mdl.iterExec(extPrompt, execModelOptions, {
..._.omit(options, 'options'),
collectRawOutput(rawOutput) {
thisReflect.iterations[0].output = rawOutput;
}
});
const textChunks: string[] = [];
let lastYield: string | LLMFunctionCallRequest = '';
for await (const x of r) {
lastYield = x;
if (typeof x === 'string' && x) {
textChunks.push(x);
yield textChunks.join('') as any;
} else if (x) {
yield x as any;
}
if (x instanceof LLMFunctionCallRequest) {
return;
}
}
await profile.acceptModelOutput(lastYield, lastYield as any);
return;
}
brief() {
return Object.fromEntries(this.modelClasses.map((x) => [x.aliases![0], x.description!]));
}
}
const instance = container.resolve(LLMManager);
export default instance;
src/services/common-llm/vertex-gemini.ts
+216
View File
@@ -0,0 +1,216 @@
import _ from "lodash";
import { injectable } from "tsyringe";
import { GeminiPro } from './google-gemini';
import { GoogleAuth } from 'google-auth-library';
import { VertexGeminiHTTP } from '../../3rd-party/google-gemini';
import { PromptChunk } from './misc';
import { readFile } from 'fs/promises';
import { InjectProperty } from '../registry';
import { HashManager } from 'civkit/hash';
import { DefaultBucket } from '../default-bucket';
import { maxConcurrency } from 'civkit/decorators';
import { downloadFile } from 'civkit/download';
import { mimeOf } from 'civkit/mime';
const sha256Hasher = new HashManager('sha256', 'hex');
@injectable()
export class VertexGeminiPro extends GeminiPro {
static override description = 'Vertex AI Gemini Pro 1.0 Model';
static override aliases = ['vertex-gemini-1.0-pro'];
static override modelName = 'gemini-1.0-pro';
protected gAuth = new GoogleAuth({
scopes: 'https://www.googleapis.com/auth/cloud-platform',
});
@InjectProperty(DefaultBucket)
defaultBucket!: DefaultBucket;
protected refreshCredsInterval?: ReturnType<typeof setInterval>;
constructor() {
// @ts-ignore
super(...arguments);
}
override async init() {
await this.dependencyReady();
if (process.env.NODE_ENV?.includes('dry-run') || !process.env.GCLOUD_PROJECT) {
this.emit('ready');
return;
}
if (this.refreshCredsInterval) {
clearInterval(this.refreshCredsInterval);
this.refreshCredsInterval = undefined;
}
try {
const cred = await this.refreshCreds();
this.clients = [
new VertexGeminiHTTP(cred.token, cred.project),
];
for (const x of this.clients) {
x.on('unauthorized', () => this.refreshCreds());
}
this.refreshCredsInterval = setInterval(() => {
this.refreshCreds().catch((e) => {
this.logger.error('Failed to refresh creds', e);
});
}, 60 * 20 * 1000).unref();
} catch (err) {
this.logger.warn('Failed to initialize Vertex Gemini clients', err);
}
this.emit('ready');
}
@maxConcurrency(1)
async refreshCreds() {
const ad = await this.gAuth.getApplicationDefault();
const token = (await ad.credential.getAccessToken()).token;
if (!token) {
throw new Error('Failed to get access token');
}
if (!ad.projectId) {
throw new Error('Failed to get project ID');
}
for (const x of this.clients) {
x.apiKey = token;
}
return {
project: ad.projectId,
token
};
}
override async encodePromptChunks(promptChunks: string | PromptChunk[] | null): Promise<any> {
if (!(this.constructor as typeof GeminiPro).interleavedPromptSupported) {
return super.encodePromptChunks(promptChunks);
}
if (typeof promptChunks === 'string') {
return super.encodePromptChunks(promptChunks);
}
if (!promptChunks) {
return super.encodePromptChunks(promptChunks);
}
const ps = promptChunks.map(async (x) => {
if (typeof x === 'string') {
return { text: x };
}
if (x instanceof URL) {
const tmpPath = this.tempFileManager.alloc();
await downloadFile(x.toString(), tmpPath);
const buff = await readFile(tmpPath);
const mimeVec = await mimeOf(buff);
if (x.hostname === 'storage.googleapis.com') {
const gsUrl = `gs://${x.pathname.slice(1)}`;
return {
fileData: {
mimeType: `${mimeVec.mediaType}/${mimeVec.subType}`,
fileUri: gsUrl,
}
};
}
const digest = sha256Hasher.hash(buff);
const fname = `llmTmp/${digest}`;
await this.defaultBucket.putBufferIfNotExist(fname, buff, {
'Content-Type': `${mimeVec.mediaType}/${mimeVec.subType}`
});
const r = {
fileData: {
mimeType: `${mimeVec.mediaType}/${mimeVec.subType}`,
fileUri: `gs://${this.defaultBucket.bucket}/${fname}`,
}
};
this.tempFileManager.bindPathTo(r, tmpPath);
return r;
}
if (x instanceof Blob) {
const buff = Buffer.from(await x.arrayBuffer());
const digest = sha256Hasher.hash(buff);
const fname = `llmTmp/${digest}`;
await this.defaultBucket.putBufferIfNotExist(fname, buff, {
'Content-Type': x.type
});
return {
fileData: {
mimeType: x.type,
fileUri: `gs://${this.defaultBucket.bucket}/${fname}`,
}
};
}
if (Buffer.isBuffer(x)) {
const buff = x;
const mimeVec = await mimeOf(buff);
const digest = sha256Hasher.hash(buff);
const fname = `llmTmp/${digest}`;
await this.defaultBucket.putBufferIfNotExist(fname, buff, {
'Content-Type': `${mimeVec.mediaType}/${mimeVec.subType}`,
});
return {
fileData: {
mimeType: `${mimeVec.mediaType}/${mimeVec.subType}`,
fileUri: `gs://${this.defaultBucket.bucket}/${fname}`,
}
};
}
return x;
});
return await Promise.all(ps);
}
}
export class VertexGemini25Flash extends VertexGeminiPro {
static override description = 'Vertex AI Gemini 2.5 Flash Model';
static override aliases = ['vertex-gemini-2.5-flash'];
static override modelName = 'gemini-2.5-flash';
static override interleavedPromptSupported = true;
static override jsonModeSchemaSupported = true;
static override windowSize = 1_000_000;
}
export class VertexGemini25FlashLite extends VertexGeminiPro {
static override description = 'Vertex AI Gemini 2.5 Flash Lite Model';
static override aliases = ['vertex-gemini-2.5-flash-lite'];
static override modelName = 'gemini-2.5-flash-lite';
static override interleavedPromptSupported = true;
static override jsonModeSchemaSupported = true;
static override windowSize = 1_000_000;
}
export class VertexGemini31FlashLite extends VertexGeminiPro {
static override description = 'Vertex AI Gemini 3.1 Flash Lite Preview';
static override aliases = ['vertex-gemini-3.1-flash-lite', 'vertex-gemini-3.1-flash-lite-preview'];
static override modelName = 'gemini-3.1-flash-lite-preview';
static override interleavedPromptSupported = true;
static override jsonModeSchemaSupported = true;
static override windowSize = 1_000_000;
}
export class VertexGemini31Pro extends VertexGeminiPro {
static override description = 'Vertex AI Gemini 3.1 Pro Preview';
static override aliases = ['vertex-gemini-3.1-pro', 'vertex-gemini-3.1-pro-preview'];
static override modelName = 'gemini-3.1-pro-preview';
static override interleavedPromptSupported = true;
static override jsonModeSchemaSupported = true;
static override windowSize = 1_000_000;
}
src/services/curl.ts
+643
View File
@@ -0,0 +1,643 @@
import { AsyncService } from 'civkit/async-service';
import { singleton } from 'tsyringe';
import { Blob } from 'buffer';
import { Curl, CurlCode, CurlFeature, HeaderInfo, Browser, getChromeConfig } from '@nomagick/node-libcurl-impersonate';
import { parseString as parseSetCookieString } from 'set-cookie-parser';
import { ScrappingOptions } from './puppeteer';
import { GlobalLogger } from './logger';
import { AssertionFailureError } from 'civkit/civ-rpc';
import { FancyFile } from 'civkit/fancy-file';
import { ServiceBadAttemptError, ServiceBadApproachError, TargetFileTooLargeError } from './errors';
import { TempFileManager } from '../services/temp-file';
import _ from 'lodash';
import { PassThrough, Readable } from 'stream';
import { AsyncLocalContext } from './async-context';
import { BlackHoleDetector } from './blackhole-detector';
import { humanReadableDataSize } from 'civkit/readability';
export interface CURLScrappingOptions<T = any> extends ScrappingOptions<T> {
method?: string;
body?: string | Buffer;
maxSize?: number;
}
export const REDIRECTION_CODES = new Set([301, 302, 303, 307, 308]);
@singleton()
export class CurlControl extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
chromeVersion: string = `136`;
safariVersion: string = `537.36`;
platform: string = `Linux`;
ua: string = `Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/${this.safariVersion} (KHTML, like Gecko) Chrome/${this.chromeVersion}.0.0.0 Safari/${this.safariVersion}`;
lifeCycleTrack = new WeakMap();
impersonateCfg?: ReturnType<typeof getChromeConfig>;
constructor(
protected globalLogger: GlobalLogger,
protected tempFileManager: TempFileManager,
protected asyncLocalContext: AsyncLocalContext,
protected blackHoleDetector: BlackHoleDetector,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
if (process.platform === 'darwin') {
this.platform = `macOS`;
} else if (process.platform === 'win32') {
this.platform = `Windows`;
}
this.emit('ready');
}
impersonateChrome(ua: string) {
this.chromeVersion = ua.match(/Chrome\/(\d+)/)![1];
this.safariVersion = ua.match(/AppleWebKit\/([\d\.]+)/)![1];
this.ua = ua;
this.impersonateCfg = getChromeConfig({
version: this.chromeVersion,
fingerprint: {
ja3: '772,4865-4866-4867-49195-49199-49196-49200-52393-52392-49171-49172-156-157-47-53,23-0-45-17613-65281-18-43-65037-5-51-16-10-27-13-35-11,4588-29-23-24,0',
ja4: 't13d1514h2_002f,0035,009c,009d,1301,1302,1303,c013,c014,c02b,c02c,c02f,c030,cca8,cca9_0005,000a,000b,000d,0012,0017,001b,0023,002b,002d,0033,44cd,fe0d,ff01_0401,0403,0501,0503,0601,0804,0805,0806',
akami: '1:65536;2:0;4:6291456;6:262144|15663105|0|m,a,s,p'
}
});
}
curlImpersonateHeader(curl: Curl, headers?: object) {
let uaPlatform = this.platform;
if (this.ua.includes('Windows')) {
uaPlatform = 'Windows';
} else if (this.ua.includes('Android')) {
uaPlatform = 'Android';
} else if (this.ua.includes('iPhone') || this.ua.includes('iPad') || this.ua.includes('iPod')) {
uaPlatform = 'iOS';
} else if (this.ua.includes('CrOS')) {
uaPlatform = 'Chrome OS';
} else if (this.ua.includes('Macintosh')) {
uaPlatform = 'macOS';
}
const mixinHeaders: Record<string, string> = {
'sec-ch-ua': `"Chromium";v="${this.chromeVersion}", "Google Chrome";v="${this.chromeVersion}", "Not.A/Brand";v="99"`,
'sec-ch-ua-mobile': '?0',
'sec-ch-ua-platform': `"${uaPlatform}"`,
'Upgrade-Insecure-Requests': '1',
'User-Agent': this.ua,
'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
'Sec-Fetch-Site': 'none',
'Sec-Fetch-Mode': 'navigate',
'Sec-Fetch-User': '?1',
'Sec-Fetch-Dest': 'document',
// 'Accept-Encoding': 'gzip, deflate, br, zstd',
'Accept-Language': 'en-US,en;q=0.9',
'Priority': 'u=0, i'
};
const headersCopy: Record<string, string | undefined> = { ...headers };
for (const k of Object.keys(mixinHeaders)) {
const lowerK = k.toLowerCase();
if (headersCopy[lowerK]) {
mixinHeaders[k] = headersCopy[lowerK];
delete headersCopy[lowerK];
}
}
Object.assign(mixinHeaders, headersCopy);
curl.setOpt(Curl.option.HTTPHEADER, Object.entries(mixinHeaders).flatMap(([k, v]) => {
if (Array.isArray(v) && v.length) {
return v.map((v2) => `${k}: ${v2}`);
}
return [`${k}: ${v}`];
}));
return curl;
}
urlToStream(urlToCrawl: URL, crawlOpts?: CURLScrappingOptions) {
return new Promise<{
statusCode: number,
statusText?: string,
data?: Readable,
headers: HeaderInfo[],
}>((resolve, reject) => {
const actx = this.asyncLocalContext.ctx;
let contentType = '';
const curl = Curl.impersonate(this.impersonateCfg || Browser.Chrome);
curl.enable(CurlFeature.StreamResponse);
curl.setOpt('URL', urlToCrawl.toString());
curl.setOpt(Curl.option.FOLLOWLOCATION, false);
curl.setOpt(Curl.option.SSL_VERIFYPEER, false);
curl.setOpt(Curl.option.TIMEOUT_MS, crawlOpts?.timeoutMs || 30_000);
curl.setOpt(Curl.option.CONNECTTIMEOUT_MS, 3_000);
curl.setOpt(Curl.option.LOW_SPEED_LIMIT, 32768);
curl.setOpt(Curl.option.LOW_SPEED_TIME, 5_000);
if (crawlOpts?.method) {
curl.setOpt(Curl.option.CUSTOMREQUEST, crawlOpts.method.toUpperCase());
}
if (crawlOpts?.body) {
// @ts-ignore
curl.setOpt(Curl.option.POSTFIELDS, crawlOpts.body.toString());
}
const headersToSet = { ...crawlOpts?.extraHeaders };
if (crawlOpts?.cookies?.length) {
const cookieKv: Record<string, string> = {};
for (const cookie of crawlOpts.cookies) {
cookieKv[cookie.name] = cookie.value;
}
for (const cookie of crawlOpts.cookies) {
if (cookie.maxAge && cookie.maxAge < 0) {
delete cookieKv[cookie.name];
continue;
}
if (cookie.expires && cookie.expires < new Date()) {
delete cookieKv[cookie.name];
continue;
}
if (cookie.secure && urlToCrawl.protocol !== 'https:') {
delete cookieKv[cookie.name];
continue;
}
if (cookie.domain && !urlToCrawl.hostname.endsWith(cookie.domain)) {
delete cookieKv[cookie.name];
continue;
}
if (cookie.path && !urlToCrawl.pathname.startsWith(cookie.path)) {
delete cookieKv[cookie.name];
continue;
}
}
const cookieChunks = Object.entries(cookieKv).map(([k, v]) => `${k}=${encodeURIComponent(v)}`);
headersToSet.cookie ??= cookieChunks.join('; ');
}
if (crawlOpts?.referer) {
headersToSet.referer ??= crawlOpts.referer;
}
if (crawlOpts?.overrideUserAgent) {
headersToSet['user-agent'] ??= crawlOpts.overrideUserAgent;
}
this.curlImpersonateHeader(curl, headersToSet);
if (crawlOpts?.proxyUrl) {
const proxyUrlCopy = new URL(crawlOpts.proxyUrl);
curl.setOpt(Curl.option.PROXY, proxyUrlCopy.href);
curl.setOpt(Curl.option.CONNECTTIMEOUT_MS, 6_000);
curl.setOpt(Curl.option.PROXY_CREDENTIAL_NO_REUSE, 1);
}
let curlStream: Readable | undefined;
let streamRecvCounter = 0;
curl.on('error', (err, errCode) => {
this.asyncLocalContext.bridge(actx, () => {
curl.close();
const err2 = this.digestCurlCode(errCode, err.message) ||
new AssertionFailureError(`Failed to access ${urlToCrawl.origin}: ${err.message}`);
err2.cause ??= err;
if (curlStream) {
// For some reason, manually emitting error event is required for curlStream.
if (errCode === CurlCode.CURLE_HTTP2_STREAM && streamRecvCounter) {
this.logger.warn(`Curl(${errCode}) ${urlToCrawl.origin}: ${err}, but some data was received, not emitting error to the stream.`, { err, urlToCrawl });
curlStream.push(null);
return;
}
curlStream.emit('error', err2);
curlStream.destroy(err2);
} else {
this.logger.warn(`Curl(${errCode}) ${urlToCrawl.origin}: ${err}`, { err, urlToCrawl });
}
reject(err2);
});
});
curl.setOpt(Curl.option.MAXFILESIZE, 4 * 1024 * 1024 * 1024); // 4GB
let status = -1;
let statusText: string | undefined;
let contentEncoding = '';
curl.once('end', () => {
if (curlStream) {
curlStream.once('end', () => curl.close());
return;
}
curl.close();
});
curl.on('stream', (stream, statusCode, headers) => {
status = statusCode;
for (const headerSet of (headers as HeaderInfo[])) {
for (const [k, v] of Object.entries(headerSet)) {
if (k.trim().endsWith(':')) {
Reflect.set(headerSet, k.slice(0, k.indexOf(':')), v || '');
Reflect.deleteProperty(headerSet, k);
continue;
}
if (v === undefined) {
Reflect.set(headerSet, k, '');
continue;
}
if (k.toLowerCase() === 'content-type' && typeof v === 'string') {
contentType = v.toLowerCase();
}
}
}
const lastResHeaders = headers[headers.length - 1];
statusText = (lastResHeaders as HeaderInfo).result?.reason;
for (const [k, v] of Object.entries(lastResHeaders)) {
const kl = k.toLowerCase();
if (kl === 'content-type') {
contentType = (v || '').toLowerCase();
}
if (kl === 'content-encoding') {
contentEncoding = (v || '').toLowerCase();
}
if (contentType && contentEncoding) {
break;
}
}
curlStream = stream;
stream.on('data', (buf) => {
if (buf.length > 0) {
streamRecvCounter += buf.length;
}
});
const passThrough = new PassThrough();
stream.pipe(passThrough);
stream.once('error', (err) => {
passThrough.destroy(err);
});
stream = passThrough;
stream.once('end', () => {
this.asyncLocalContext.bridge(actx, () => {
this.logger.debug(`CURL: [${statusCode}${statusText ? ` ${statusText}` : ''}] ${urlToCrawl.origin} ${humanReadableDataSize(streamRecvCounter)} ${contentType} ${contentEncoding}`, {
statusCode, statusText, traffic: streamRecvCounter, contentType, contentEncoding
});
});
});
if (
REDIRECTION_CODES.has(statusCode) &&
headers.some((h) => h.hasOwnProperty('Location') || h.hasOwnProperty('location'))
) {
stream.resume();
stream.once('end', () => {
resolve({
statusCode: status,
statusText,
data: undefined,
headers: headers as HeaderInfo[],
});
});
return;
}
resolve({
statusCode: status,
statusText,
data: stream,
headers: headers as HeaderInfo[],
});
});
curl.perform();
});
}
async urlToFile(urlToCrawl: URL, crawlOpts?: CURLScrappingOptions) {
let leftRedirection = 6;
let cookieRedirects = 0;
let opts = { ...crawlOpts };
let nextHopUrl = urlToCrawl;
const fakeHeaderInfos: HeaderInfo[] = [];
do {
const s = await this.urlToStream(nextHopUrl, opts);
const r = { ...s } as {
statusCode: number,
statusText?: string,
data?: FancyFile,
headers: HeaderInfo[],
};
const headers = r.headers[r.headers.length - 1];
const claimedSize = parseInt(headers['Content-Length'] || headers['content-length']);
if (crawlOpts?.maxSize && claimedSize > crawlOpts.maxSize) {
throw new TargetFileTooLargeError(`Target file is too large: ${humanReadableDataSize(claimedSize)} > ${humanReadableDataSize(crawlOpts.maxSize)}`);
}
if (r.data) {
const fpath = this.tempFileManager.alloc();
const fancyFile = FancyFile.auto(r.data, fpath);
this.tempFileManager.bindPathTo(fancyFile, fpath);
r.data = fancyFile;
}
if (REDIRECTION_CODES.has(r.statusCode)) {
fakeHeaderInfos.push(...r.headers);
const location: string | undefined = headers.Location || headers.location;
const setCookieHeader = headers['Set-Cookie'] || headers['set-cookie'];
if (setCookieHeader) {
const cookieAssignments = Array.isArray(setCookieHeader) ? setCookieHeader : [setCookieHeader];
const parsed = cookieAssignments.filter(Boolean).map((x) => parseSetCookieString(x, { decodeValues: true }));
if (parsed.length) {
opts.cookies = [...(opts.cookies || []), ...parsed];
}
if (!location) {
cookieRedirects += 1;
}
}
if (!location && !setCookieHeader) {
let sizeAcc = 0;
s.data!.on('data', (chunk) => {
sizeAcc += chunk.length;
if (crawlOpts?.maxSize && sizeAcc > crawlOpts.maxSize) {
s.data!.destroy(new TargetFileTooLargeError(`Max size exceeded: ${humanReadableDataSize(sizeAcc)}`));
}
});
// Follow curl behavior
return {
statusCode: r.statusCode,
data: r.data,
headers: fakeHeaderInfos.concat(r.headers),
};
}
if (!location && cookieRedirects > 1) {
throw new ServiceBadApproachError(`Failed to access ${urlToCrawl}: Browser required to solve complex cookie preconditions.`);
}
nextHopUrl = new URL(location || '', nextHopUrl);
leftRedirection -= 1;
continue;
}
let sizeAcc = 0;
s.data!.on('data', (chunk) => {
sizeAcc += chunk.length;
if (crawlOpts?.maxSize && sizeAcc > crawlOpts.maxSize) {
s.data!.destroy(new TargetFileTooLargeError(`Max size exceeded: ${humanReadableDataSize(sizeAcc)}`));
}
});
return {
statusCode: r.statusCode,
statusText: r.statusText,
data: r.data,
headers: fakeHeaderInfos.concat(r.headers),
};
} while (leftRedirection > 0);
throw new ServiceBadAttemptError(`Failed to access ${urlToCrawl}: Too many redirections.`);
}
async sideLoad(targetUrl: URL, crawlOpts?: CURLScrappingOptions) {
const curlResult = await this.urlToFile(targetUrl, crawlOpts);
this.blackHoleDetector.itWorked();
let finalURL = targetUrl;
const sideLoadOpts: CURLScrappingOptions<FancyFile>['sideLoad'] = {
impersonate: {},
proxyOrigin: {},
};
for (const headers of curlResult.headers) {
const potentialLastResult = sideLoadOpts.impersonate[finalURL.href];
sideLoadOpts.impersonate[finalURL.href] = {
status: headers.result?.code || -1,
headers: _.mergeWith(potentialLastResult?.headers || {}, _.omit(headers, 'result'), (objValue, srcValue) => {
if (Array.isArray(objValue) && Array.isArray(srcValue)) {
return [...objValue, ...srcValue];
}
return undefined;
}),
contentType: headers['Content-Type'] || headers['content-type'],
};
if (crawlOpts?.proxyUrl) {
sideLoadOpts.proxyOrigin[finalURL.origin] = crawlOpts.proxyUrl;
}
if (headers.result?.code && [301, 302, 307, 308].includes(headers.result.code)) {
const location = headers.Location || headers.location;
if (location) {
finalURL = new URL(location, finalURL);
}
}
}
const lastHeaders = curlResult.headers[curlResult.headers.length - 1];
const contentType = (lastHeaders['Content-Type'] || lastHeaders['content-type'])?.toLowerCase() || (await curlResult.data?.mimeType) || 'application/octet-stream';
const contentDisposition = lastHeaders['Content-Disposition'] || lastHeaders['content-disposition'];
const fileName = contentDisposition?.match(/filename="([^"]+)"/i)?.[1] || finalURL.pathname.split('/').pop();
if (sideLoadOpts.impersonate[finalURL.href] && (await curlResult.data?.size)) {
sideLoadOpts.impersonate[finalURL.href].body = curlResult.data;
}
// This should keep the file from being garbage collected and deleted until this asyncContext/request is done.
this.lifeCycleTrack.set(this.asyncLocalContext.ctx, curlResult.data);
return {
finalURL,
sideLoadOpts,
chain: curlResult.headers,
status: curlResult.statusCode,
statusText: curlResult.statusText,
headers: lastHeaders,
contentType,
contentDisposition,
fileName,
file: curlResult.data
};
}
async urlToBlob(urlToCrawl: URL, crawlOpts?: CURLScrappingOptions) {
let leftRedirection = 6;
let cookieRedirects = 0;
let opts = { ...crawlOpts };
let nextHopUrl = urlToCrawl;
const fakeHeaderInfos: HeaderInfo[] = [];
do {
const s = await this.urlToStream(nextHopUrl, opts);
const r = { ...s } as {
statusCode: number,
statusText?: string,
data?: Blob,
headers: HeaderInfo[],
};
const headers = r.headers[r.headers.length - 1];
const claimedSize = parseInt(headers['Content-Length'] || headers['content-length']);
if (crawlOpts?.maxSize && claimedSize > crawlOpts.maxSize) {
throw new TargetFileTooLargeError(`Target file is too large: ${humanReadableDataSize(claimedSize)} > ${humanReadableDataSize(crawlOpts.maxSize)}`);
}
if (REDIRECTION_CODES.has(r.statusCode)) {
fakeHeaderInfos.push(...r.headers);
const location: string | undefined = headers.Location || headers.location;
const setCookieHeader = headers['Set-Cookie'] || headers['set-cookie'];
if (setCookieHeader) {
const cookieAssignments = Array.isArray(setCookieHeader) ? setCookieHeader : [setCookieHeader];
const parsed = cookieAssignments.filter(Boolean).map((x) => parseSetCookieString(x, { decodeValues: true }));
if (parsed.length) {
opts.cookies = [...(opts.cookies || []), ...parsed];
}
if (!location) {
cookieRedirects += 1;
}
}
if (!location && !setCookieHeader) {
// Follow curl behavior
if (s.data) {
const chunks: Buffer[] = [];
let sizeAcc = 0;
s.data.on('data', (chunk) => {
chunks.push(chunk);
sizeAcc += chunk.length;
if (crawlOpts?.maxSize && sizeAcc > crawlOpts.maxSize) {
s.data!.destroy(new TargetFileTooLargeError(`Max size exceeded: ${humanReadableDataSize(sizeAcc)}`));
}
});
await new Promise((resolve, reject) => {
s.data!.once('end', resolve);
s.data!.once('error', reject);
});
r.data = new Blob(chunks, { type: headers['Content-Type'] || headers['content-type'] });
}
return {
statusCode: r.statusCode,
data: r.data,
headers: fakeHeaderInfos.concat(r.headers),
};
}
if (!location && cookieRedirects > 1) {
throw new ServiceBadApproachError(`Failed to access ${urlToCrawl}: Browser required to solve complex cookie preconditions.`);
}
nextHopUrl = new URL(location || '', nextHopUrl);
leftRedirection -= 1;
continue;
}
if (s.data) {
const chunks: Buffer[] = [];
let sizeAcc = 0;
s.data.on('data', (chunk) => {
chunks.push(chunk);
sizeAcc += chunk.length;
if (crawlOpts?.maxSize && sizeAcc > crawlOpts.maxSize) {
s.data!.destroy(new TargetFileTooLargeError(`Max size exceeded: ${humanReadableDataSize(sizeAcc)}`));
}
});
await new Promise((resolve, reject) => {
s.data!.once('end', resolve);
s.data!.once('error', reject);
});
r.data = new Blob(chunks, { type: headers['Content-Type'] || headers['content-type'] });
}
return {
statusCode: r.statusCode,
statusText: r.statusText,
data: r.data,
headers: fakeHeaderInfos.concat(r.headers),
};
} while (leftRedirection > 0);
throw new ServiceBadAttemptError(`Failed to access ${urlToCrawl}: Too many redirections.`);
}
async sideLoadBlob(targetUrl: URL, crawlOpts?: CURLScrappingOptions) {
const curlResult = await this.urlToBlob(targetUrl, crawlOpts);
this.blackHoleDetector.itWorked();
let finalURL = targetUrl;
const sideLoadOpts: CURLScrappingOptions<Blob>['sideLoad'] = {
impersonate: {},
proxyOrigin: {},
};
for (const headers of curlResult.headers) {
const potentialLastResult = sideLoadOpts.impersonate[finalURL.href];
sideLoadOpts.impersonate[finalURL.href] = {
status: headers.result?.code || -1,
headers: _.mergeWith(potentialLastResult?.headers || {}, _.omit(headers, 'result'), (objValue, srcValue) => {
if (Array.isArray(objValue) && Array.isArray(srcValue)) {
return [...objValue, ...srcValue];
}
return undefined;
}),
contentType: headers['Content-Type'] || headers['content-type'],
};
if (crawlOpts?.proxyUrl) {
sideLoadOpts.proxyOrigin[finalURL.origin] = crawlOpts.proxyUrl;
}
if (headers.result?.code && [301, 302, 307, 308].includes(headers.result.code)) {
const location = headers.Location || headers.location;
if (location) {
finalURL = new URL(location, finalURL);
}
}
}
const lastHeaders = curlResult.headers[curlResult.headers.length - 1];
const contentType = (lastHeaders['Content-Type'] || lastHeaders['content-type'])?.toLowerCase() || (curlResult.data?.type) || 'application/octet-stream';
const contentDisposition = lastHeaders['Content-Disposition'] || lastHeaders['content-disposition'];
const fileName = contentDisposition?.match(/filename="([^"]+)"/i)?.[1] || finalURL.pathname.split('/').pop();
if (sideLoadOpts.impersonate[finalURL.href] && (curlResult.data?.size)) {
sideLoadOpts.impersonate[finalURL.href].body = curlResult.data;
}
// This should keep the file from being garbage collected and deleted until this asyncContext/request is done.
this.lifeCycleTrack.set(this.asyncLocalContext.ctx, curlResult.data);
return {
finalURL,
sideLoadOpts,
chain: curlResult.headers,
status: curlResult.statusCode,
statusText: curlResult.statusText,
headers: lastHeaders,
contentType,
contentDisposition,
fileName,
file: curlResult.data
};
}
digestCurlCode(code: CurlCode, msg: string) {
switch (code) {
// 400 User errors
case CurlCode.CURLE_COULDNT_RESOLVE_HOST: {
return new AssertionFailureError(msg);
}
// Maybe retry but dont retry with curl again
case CurlCode.CURLE_OPERATION_TIMEDOUT:
case CurlCode.CURLE_UNSUPPORTED_PROTOCOL:
case CurlCode.CURLE_PEER_FAILED_VERIFICATION: {
return new ServiceBadApproachError(msg);
}
// Retryable errors
case CurlCode.CURLE_PROXY:
case CurlCode.CURLE_REMOTE_ACCESS_DENIED:
case CurlCode.CURLE_SEND_ERROR:
case CurlCode.CURLE_RECV_ERROR:
case CurlCode.CURLE_GOT_NOTHING:
case CurlCode.CURLE_SSL_CONNECT_ERROR:
case CurlCode.CURLE_QUIC_CONNECT_ERROR:
case CurlCode.CURLE_COULDNT_RESOLVE_PROXY:
case CurlCode.CURLE_COULDNT_CONNECT:
case CurlCode.CURLE_PARTIAL_FILE: {
return new ServiceBadAttemptError(msg);
}
default: {
return undefined;
}
}
}
}
src/services/default-bucket.ts
+63
View File
@@ -0,0 +1,63 @@
import { singleton } from 'tsyringe';
import { AbstractObjectStorageService, ObjectStorageOptions } from 'civkit/abstract/object-storage';
import { GlobalLogger } from './logger';
import { EnvConfig } from './envconfig';
import { once } from 'events';
@singleton()
export class DefaultBucket extends AbstractObjectStorageService {
logger = this.globalLogger.child({ service: this.constructor.name });
override options!: ObjectStorageOptions;
constructor(
protected globalLogger: GlobalLogger,
protected envConfig: EnvConfig,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
const bucketUrl = new URL(process.env.GCP_STORAGE_ENDPOINT || 'https://storage.googleapis.com');
this.options = {
region: this.envConfig.GCP_STORAGE_REGION || 'us-central1',
bucket: process.env.GCP_STORAGE_BUCKET || `${process.env.GCLOUD_PROJECT}.appspot.com`,
endPoint: bucketUrl.hostname,
port: Number(bucketUrl.port) || (bucketUrl.protocol === 'https:' ? 443 : 80),
useSSL: bucketUrl.protocol === 'https:',
accessKey: this.envConfig.GCP_STORAGE_ACCESS_KEY,
secretKey: this.envConfig.GCP_STORAGE_SECRET_KEY,
};
await super.init();
this.emit('ready');
}
async readSingleFile(path: string) {
const stream = await this.minioClient.getObject(this.bucket, path);
const chunks: Buffer[] = [];
stream.on('data', (chunk: Buffer) => {
chunks.push(chunk);
});
await once(stream, 'end');
return Buffer.concat(chunks);
}
async putBuffer(path: string, buffer: Buffer, metadata?: Record<string, any>) {
const r = await this.minioClient.putObject(this.bucket, path, buffer, buffer.byteLength, metadata);
return r;
}
async putBufferIfNotExist(path: string, buffer: Buffer, metadata?: Record<string, any>) {
const s = await this.minioClient.statObject(this.bucket, path).catch(() => undefined);
if (!s) {
const r = await this.minioClient.putObject(this.bucket, path, buffer, buffer.byteLength, metadata);
return r;
}
return s;
}
}
src/services/envconfig.ts
+67
View File
@@ -0,0 +1,67 @@
import { container, singleton } from 'tsyringe';
export const SPECIAL_COMBINED_ENV_KEY = 'SECRETS_COMBINED';
const CONF_ENV = [
'GCP_STORAGE_ACCESS_KEY',
'GCP_STORAGE_SECRET_KEY',
'GCP_STORAGE_REGION',
'OPENAI_API_KEY',
'ANTHROPIC_API_KEY',
'OPENROUTER_API_KEY',
'REPLICATE_API_KEY',
'GOOGLE_AI_STUDIO_API_KEY',
'THORDATA_PROXY_URL',
'THORDATA_PROXY_URL_ALT',
'THORDATA_SERP_API_KEY',
'BRIGHTDATA_PROXY_URL',
'BRIGHTDATA_ISP_PROXY_URL',
'BRIGHTDATA_SERP_API_KEY',
'SERPER_SEARCH_API_KEY',
'CLOUD_FLARE_API_KEY',
] as const;
@singleton()
export class EnvConfig {
dynamic!: Record<string, string>;
combined: Record<string, string> = {};
originalEnv: Record<string, string | undefined> = { ...process.env };
constructor() {
if (process.env[SPECIAL_COMBINED_ENV_KEY]) {
Object.assign(this.combined, JSON.parse(
Buffer.from(process.env[SPECIAL_COMBINED_ENV_KEY]!, 'base64').toString('utf-8')
));
delete process.env[SPECIAL_COMBINED_ENV_KEY];
}
// Static config
for (const x of CONF_ENV) {
const s = process.env[x] || this.combined[x] || '';
Reflect.set(this, x, s);
if (x in process.env) {
delete process.env[x];
}
}
// Dynamic config
this.dynamic = new Proxy({
get: (_target: any, prop: string) => {
return this.combined[prop] || process.env[prop] || '';
}
}, {}) as any;
}
}
// eslint-disable-next-line @typescript-eslint/no-empty-interface
export interface EnvConfig extends Record<typeof CONF_ENV[number], string> { }
const instance = container.resolve(EnvConfig);
export default instance;
src/services/errors.ts
+87
View File
@@ -0,0 +1,87 @@
import { ApplicationError, Prop, RPC_TRANSFER_PROTOCOL_META_SYMBOL, StatusCode } from 'civkit/civ-rpc';
import _ from 'lodash';
import dayjs from 'dayjs';
import utc from 'dayjs/plugin/utc';
dayjs.extend(utc);
@StatusCode(50301)
export class ServiceDisabledError extends ApplicationError { }
@StatusCode(50302)
export class ServiceCrashedError extends ApplicationError { }
@StatusCode(50303)
export class ServiceNodeResourceDrainError extends ApplicationError { }
@StatusCode(50304)
export class ServiceBadAttemptError extends ApplicationError { }
@StatusCode(50305)
export class ServiceBadApproachError extends ServiceBadAttemptError { }
@StatusCode(40104)
export class EmailUnverifiedError extends ApplicationError { }
@StatusCode(40201)
export class InsufficientCreditsError extends ApplicationError { }
@StatusCode(40202)
export class TierFeatureConstraintError extends ApplicationError { }
@StatusCode(40203)
export class InsufficientBalanceError extends ApplicationError { }
@StatusCode(40903)
export class LockConflictError extends ApplicationError { }
@StatusCode(40904)
export class BudgetExceededError extends ApplicationError { }
@StatusCode(45101)
export class HarmfulContentError extends ApplicationError { }
@StatusCode(45102)
export class SecurityCompromiseError extends ApplicationError { }
@StatusCode(45103)
export class USGovernmentConcernError extends ApplicationError { }
@StatusCode(45104)
export class ChineseGovernmentConcernError extends ApplicationError { }
@StatusCode(45105)
export class EUGovernmentConcernError extends ApplicationError { }
@StatusCode(41201)
export class BatchSizeTooLargeError extends ApplicationError { }
@StatusCode(41302)
export class TargetFileTooLargeError extends ApplicationError { }
@StatusCode(42903)
export class RateLimitTriggeredError extends ApplicationError {
@Prop({
desc: 'Retry after seconds',
})
retryAfter?: number;
@Prop({
desc: 'Retry after date',
})
retryAfterDate?: Date;
protected override get [RPC_TRANSFER_PROTOCOL_META_SYMBOL]() {
const retryAfter = this.retryAfter || this.retryAfterDate;
if (!retryAfter) {
return super[RPC_TRANSFER_PROTOCOL_META_SYMBOL];
}
return _.merge(_.cloneDeep(super[RPC_TRANSFER_PROTOCOL_META_SYMBOL]), {
headers: {
'Retry-After': `${retryAfter instanceof Date ? dayjs(retryAfter).utc().format('ddd, DD MMM YYYY HH:mm:ss [GMT]') : retryAfter}`,
}
});
}
}
src/services/finalizer.ts
+41
View File
@@ -0,0 +1,41 @@
import { AbstractFinalizerService } from 'civkit/finalizer';
import { container, singleton } from 'tsyringe';
import { isMainThread } from 'worker_threads';
import { GlobalLogger } from './logger';
const realProcessExit = process.exit;
@singleton()
export class FinalizerService extends AbstractFinalizerService {
container = container;
logger = this.globalLogger.child({ service: this.constructor.name });
override quitProcess(code?: string | number | null | undefined): never {
return realProcessExit(code);
}
constructor(protected globalLogger: GlobalLogger) {
super(...arguments);
}
override onUnhandledRejection(err: unknown, _triggeringPromise: Promise<unknown>): void {
this.logger.warn(`Unhandled promise rejection in pid ${process.pid}`, { err });
}
}
const instance = container.resolve(FinalizerService);
export const { Finalizer } = instance.decorators();
export default instance;
process.exit = ((code?: number) => {
if (isMainThread) {
instance.terminate(code);
return;
}
return realProcessExit(code);
}) as typeof process.exit;
if (isMainThread) {
instance.serviceReady();
}
src/services/geoip.ts
+134
View File
@@ -0,0 +1,134 @@
import { container, singleton } from 'tsyringe';
import fsp from 'fs/promises';
import { CityResponse, Reader } from 'maxmind';
import { Coercible, Prop } from 'civkit/coercible';
import { GlobalLogger } from './logger';
import path from 'path';
import { Threaded } from './threaded';
import { AsyncService } from 'civkit/async-service';
import { runOnce } from 'civkit/decorators';
export enum GEOIP_SUPPORTED_LANGUAGES {
EN = 'en',
ZH_CN = 'zh-CN',
JA = 'ja',
DE = 'de',
FR = 'fr',
ES = 'es',
PT_BR = 'pt-BR',
RU = 'ru',
}
export class GeoIPInfo extends Coercible {
@Prop()
code?: string;
@Prop()
name?: string;
}
export class GeoIPCountryInfo extends GeoIPInfo {
@Prop()
eu?: boolean;
}
export class GeoIPCityResponse extends Coercible {
@Prop()
continent?: GeoIPInfo;
@Prop()
country?: GeoIPCountryInfo;
@Prop({
arrayOf: GeoIPInfo
})
subdivisions?: GeoIPInfo[];
@Prop()
city?: string;
@Prop({
arrayOf: Number
})
coordinates?: [number, number, number];
@Prop()
timezone?: string;
}
@singleton()
export class GeoIPService extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
mmdbCity!: Reader<CityResponse>;
constructor(
protected globalLogger: GlobalLogger,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
@runOnce()
async _lazyload() {
const mmdpPath = path.resolve(__dirname, '..', '..', 'licensed', 'GeoLite2-City.mmdb');
const dbBuff = await fsp.readFile(mmdpPath, { flag: 'r', encoding: null });
this.mmdbCity = new Reader<CityResponse>(dbBuff);
this.logger.info(`Loaded GeoIP database, ${dbBuff.byteLength} bytes`);
}
@Threaded()
async lookupCity(ip: string, lang: GEOIP_SUPPORTED_LANGUAGES = GEOIP_SUPPORTED_LANGUAGES.EN) {
await this._lazyload();
const r = this.mmdbCity.get(ip);
if (!r) {
return undefined;
}
return GeoIPCityResponse.from({
continent: r.continent ? {
code: r.continent?.code,
name: r.continent?.names?.[lang] || r.continent?.names?.en,
} : undefined,
country: r.country ? {
code: r.country?.iso_code,
name: r.country?.names?.[lang] || r.country?.names.en,
eu: r.country?.is_in_european_union,
} : undefined,
city: r.city?.names?.[lang] || r.city?.names?.en,
subdivisions: r.subdivisions?.map((x) => ({
code: x.iso_code,
name: x.names?.[lang] || x.names?.en,
})),
coordinates: r.location ? [
r.location.latitude, r.location.longitude, r.location.accuracy_radius
] : undefined,
timezone: r.location?.time_zone,
});
}
@Threaded()
async lookupCities(ips: string[], lang: GEOIP_SUPPORTED_LANGUAGES = GEOIP_SUPPORTED_LANGUAGES.EN) {
const r = (await Promise.all(ips.map((ip) => this.lookupCity(ip, lang)))).filter(Boolean) as GeoIPCityResponse[];
return r;
}
}
const instance = container.resolve(GeoIPService);
export default instance;
src/services/ipasn.ts
+79
View File
@@ -0,0 +1,79 @@
import { container, singleton } from 'tsyringe';
import fsp from 'fs/promises';
import { AsnResponse, Reader } from 'maxmind';
import { Coercible, Prop } from 'civkit/coercible';
import { GlobalLogger } from './logger';
import path from 'path';
import { Threaded } from './threaded';
import { AsyncService } from 'civkit/async-service';
import { runOnce } from 'civkit/decorators';
export class ASNInfo extends Coercible {
@Prop()
asn?: number;
@Prop()
org?: string;
}
@singleton()
export class IPASNService extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
mmdbASN!: Reader<AsnResponse>;
constructor(
protected globalLogger: GlobalLogger,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
@runOnce()
async _lazyload() {
const mmdpPath = path.resolve(__dirname, '..', '..', 'licensed', 'geolite2-asn.mmdb');
const dbBuff = await fsp.readFile(mmdpPath, { flag: 'r', encoding: null });
this.mmdbASN = new Reader<AsnResponse>(dbBuff);
this.logger.info(`Loaded ASN database, ${dbBuff.byteLength} bytes`);
}
@Threaded()
async lookupASN(ip: string) {
await this._lazyload();
const r = this.mmdbASN.get(ip);
if (!r) {
return undefined;
}
return ASNInfo.from({
asn: r.autonomous_system_number,
org: r.autonomous_system_organization,
});
}
@Threaded()
async lookupASNs(ips: string[]) {
const r = (await Promise.all(ips.map((ip) => this.lookupASN(ip)))).filter(Boolean) as ASNInfo[];
return r;
}
}
const instance = container.resolve(IPASNService);
export default instance;
src/services/jsdom.ts
+613
View File
@@ -0,0 +1,613 @@
/// <reference lib="dom" />
import { container, singleton } from 'tsyringe';
import { GlobalLogger } from './logger';
import { ExtendedSnapshot, ImgBrief, PageSnapshot } from './puppeteer';
import { Readability } from '@mozilla/readability';
import { Threaded } from '../services/threaded';
import type { ExtraScrappingOptions } from '../api/crawler';
import { tailwindClasses } from '../utils/tailwind-classes';
import { countGPTToken } from '../utils/openai';
import { AsyncService } from 'civkit/async-service';
import { ApplicationError, AssertionFailureError } from 'civkit/civ-rpc';
import { MarkifyService } from './markify';
const pLinkedom = import('linkedom');
@singleton()
export class JSDomControl extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
linkedom!: Awaited<typeof pLinkedom>;
constructor(
protected globalLogger: GlobalLogger,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.linkedom = await pLinkedom;
this.emit('ready');
}
async narrowSnapshot(snapshot: PageSnapshot | undefined, options?: ExtraScrappingOptions) {
if (snapshot?.parsed && !options?.targetSelector && !options?.removeSelector && !options?.withIframe && !options?.withShadowDom) {
return snapshot;
}
if (!snapshot?.html) {
return snapshot;
}
try {
// SideLoad contains native objects that cannot go through thread boundaries.
return await this.actualNarrowSnapshot(snapshot, { ...options, sideLoad: undefined });
} catch (err: any) {
this.logger.warn(`Error narrowing snapshot`, { err });
if (err instanceof ApplicationError) {
throw err;
}
throw new AssertionFailureError(`Failed to process the page: ${err?.message}`);
}
}
@Threaded()
async actualNarrowSnapshot(snapshot: PageSnapshot, options?: ExtraScrappingOptions): Promise<PageSnapshot | undefined> {
const t0 = Date.now();
let sourceHTML = snapshot.html;
if (options?.withShadowDom && snapshot.shadowExpanded) {
sourceHTML = snapshot.shadowExpanded;
}
let jsdom = this.linkedom.parseHTML(sourceHTML);
if (!jsdom.window.document.documentElement) {
jsdom = this.linkedom.parseHTML(`<html><body>${sourceHTML}</body></html>`);
}
const allNodes: Node[] = [];
jsdom.window.document.querySelectorAll('svg').forEach((x) => x.innerHTML = '');
if (options?.withIframe) {
jsdom.window.document.querySelectorAll('iframe[src],frame[src]').forEach((x) => {
const origSrc = x.getAttribute('src');
if (!origSrc) {
return;
}
const src = new URL(origSrc, snapshot.rebase || snapshot.href).toString();
const thisSnapshot = snapshot.childFrames?.find((f) => f.href === src);
if (options?.withIframe === 'quoted') {
const blockquoteElem = jsdom.window.document.createElement('blockquote');
const preElem = jsdom.window.document.createElement('pre');
preElem.innerHTML = thisSnapshot?.text || '';
blockquoteElem.appendChild(preElem);
x.replaceWith(blockquoteElem);
} else if (thisSnapshot?.html) {
x.innerHTML = thisSnapshot.html;
x.querySelectorAll('script, style').forEach((s) => s.remove());
if (src) {
x.querySelectorAll('[src]').forEach((el) => {
const imgSrc = el.getAttribute('src')!;
if (URL.canParse(imgSrc, src!)) {
el.setAttribute('src', new URL(imgSrc, src!).toString());
}
});
x.querySelectorAll('[href]').forEach((el) => {
const linkHref = el.getAttribute('href')!;
if (URL.canParse(linkHref, src!)) {
el.setAttribute('href', new URL(linkHref, src!).toString());
}
});
}
}
});
}
if (Array.isArray(options?.removeSelector)) {
for (const rl of options!.removeSelector) {
jsdom.window.document.querySelectorAll(rl).forEach((x) => x.remove());
}
} else if (options?.removeSelector) {
jsdom.window.document.querySelectorAll(options.removeSelector).forEach((x) => x.remove());
}
let bewareTargetContentDoesNotExist = false;
if (Array.isArray(options?.targetSelector)) {
bewareTargetContentDoesNotExist = true;
for (const x of options!.targetSelector.map((x) => jsdom.window.document.querySelectorAll(x))) {
x.forEach((el) => {
if (!allNodes.includes(el)) {
allNodes.push(el);
}
});
}
} else if (options?.targetSelector) {
bewareTargetContentDoesNotExist = true;
jsdom.window.document.querySelectorAll(options.targetSelector).forEach((el) => {
if (!allNodes.includes(el)) {
allNodes.push(el);
}
});
} else {
allNodes.push(jsdom.window.document);
}
if (!allNodes.length) {
if (bewareTargetContentDoesNotExist) {
return undefined;
}
return snapshot;
}
const textNodes: HTMLElement[] = [];
let rootDoc: Document;
let bewareHeadHtmlUnexpectedlyPreserved = false;
if (allNodes.length === 1 && allNodes[0].nodeName === '#document' && (allNodes[0] as any).documentElement) {
rootDoc = allNodes[0] as any;
if (rootDoc.body?.innerText) {
textNodes.push(rootDoc.body);
}
} else {
rootDoc = this.linkedom.parseHTML('<html><body></body></html>').window.document;
rootDoc.head.innerHTML = jsdom.window.document.head?.innerHTML || '';
bewareHeadHtmlUnexpectedlyPreserved = true;
for (const n of allNodes) {
rootDoc.body.appendChild(n);
rootDoc.body.appendChild(rootDoc.createTextNode('\n\n'));
if ((n as HTMLElement).innerText) {
textNodes.push(n as HTMLElement);
}
}
}
const metadata: Record<string, any> = {};
const lang = rootDoc.documentElement.getAttribute('lang');
if (lang) {
metadata.lang = lang;
}
rootDoc.head?.querySelectorAll('meta[content]').forEach((el) => {
const name = el.getAttribute('name') || el.getAttribute('property');
if (!name) {
return;
}
const val = el.getAttribute('content') || '';
const curVal = Reflect.get(metadata, name);
if (Array.isArray(curVal)) {
curVal.push(val);
} else if (curVal) {
Reflect.set(metadata, name, [curVal, val]);
} else {
Reflect.set(metadata, name, val);
}
});
const baseURI = snapshot.rebase || snapshot.href;
const external = {} as { [rel: string]: { [href: string]: { [k: string]: string; }; }; };
rootDoc.head?.querySelectorAll('link[rel][href]').forEach((el) => {
const attributes: Record<string, string> = {};
el.getAttributeNames().forEach((attrName) => {
attributes[attrName] = el.getAttribute(attrName) || '';
});
const rels = attributes['rel']?.split(/\s+/g).filter(Boolean) || [];
if (!rels?.length) {
return;
}
let href = attributes['href'] || '';
if (href) {
try {
href = new URL(href, baseURI).href;
} catch (err) {
// void 0;
}
}
delete attributes.rel;
delete attributes.href;
for (const rel of rels) {
external[rel] ??= {};
external[rel][href] ??= {};
Object.assign(external[rel][href], attributes);
}
});
delete external['stylesheet'];
delete external['shortcut'];
const textChunks = textNodes.map((x) => {
const clone = x.cloneNode(true) as HTMLElement;
clone.querySelectorAll('script,style,link,svg').forEach((s) => s.remove());
return clone.innerText;
});
let parsed: any = snapshot.parsed;
if (options?.readabilityRequired && (!parsed || options?.targetSelector)) {
try {
parsed = new Readability(rootDoc.cloneNode(true) as any).parse();
} catch (err: any) {
this.logger.warn(`Failed to parse selected element`, { err });
}
}
const imgSet = new Set<string>();
const rebuiltImgs: ImgBrief[] = [];
Array.from(rootDoc.querySelectorAll('img[src],img[data-src]'))
.map((x: any) => [x.getAttribute('src'), x.getAttribute('data-src'), x.getAttribute('alt')])
.forEach(([u1, u2, alt]) => {
let absUrl: string | undefined;
if (u1) {
try {
const u1Txt = new URL(u1, snapshot.rebase || snapshot.href).toString();
imgSet.add(u1Txt);
absUrl = u1Txt;
} catch (err) {
// void 0;
}
}
if (u2) {
try {
const u2Txt = new URL(u2, snapshot.rebase || snapshot.href).toString();
imgSet.add(u2Txt);
absUrl = u2Txt;
} catch (err) {
// void 0;
}
}
if (absUrl) {
rebuiltImgs.push({
src: absUrl,
alt
});
}
});
const r = {
...snapshot,
title: snapshot.title || jsdom.window.document.title,
description: snapshot.description ||
(jsdom.window.document.head?.querySelector('meta[name$="description"][content],meta[property$="description"][content]')?.getAttribute('content') ?? ''),
parsed,
html: rootDoc.documentElement.outerHTML,
text: textChunks.join('\n'),
imgs: (snapshot.imgs || rebuiltImgs)?.filter((x) => imgSet.has(x.src)) || [],
metadata,
external,
} as PageSnapshot;
if (bewareHeadHtmlUnexpectedlyPreserved) {
rootDoc.head!.innerHTML = '';
r.html = rootDoc.documentElement.outerHTML;
}
const dt = Date.now() - t0;
if (dt > 1000) {
this.logger.warn(`Performance issue: Narrowing snapshot took ${dt}ms`, { url: snapshot.href, dt });
}
return r;
}
@Threaded()
async inferSnapshot(snapshot: PageSnapshot) {
const t0 = Date.now();
const extendedSnapshot = { ...snapshot } as ExtendedSnapshot;
let documentElement;
try {
documentElement = this.snippetToElement(snapshot.html, snapshot.href);
const dt = Date.now() - t0;
this.logger.debug(`Parsing of jsdom took ${dt}ms`, { url: snapshot.href, dt });
documentElement.querySelectorAll('svg').forEach((x) => x.innerHTML = '');
const links = Array.from(documentElement.querySelectorAll('a[href]'))
.map((x: any) => [x.textContent.replace(/\s+/g, ' ').trim(), x.getAttribute('href'),])
.map(([text, href]) => {
if (!href) {
return undefined;
}
try {
const parsed = new URL(href, snapshot.rebase || snapshot.href);
return [text, parsed.toString()] as const;
} catch (err) {
return undefined;
}
})
.filter(Boolean) as [string, string][];
extendedSnapshot.links = links;
const imgs = Array.from(documentElement.querySelectorAll('img[src],img[data-src]'))
.map((x: any) => {
let linkPreferredSrc = x.getAttribute('src') || '';
if (linkPreferredSrc.startsWith('data:')) {
const dataSrc = x.getAttribute('data-src') || '';
if (dataSrc && !dataSrc.startsWith('data:')) {
linkPreferredSrc = dataSrc;
}
}
return {
src: new URL(linkPreferredSrc, snapshot.rebase || snapshot.href).toString(),
width: parseInt(x.getAttribute('width') || '0'),
height: parseInt(x.getAttribute('height') || '0'),
alt: x.getAttribute('alt') || x.getAttribute('title'),
};
});
extendedSnapshot.imgs = imgs as any;
} catch (_err) {
void 0;
}
const dt = Date.now() - t0;
if (dt > 1000) {
this.logger.warn(`Performance issue: Inferring snapshot took ${dt}ms`, { url: snapshot.href, dt });
} else {
this.logger.debug(`Inferring snapshot took ${dt}ms`, { url: snapshot.href, dt });
}
return { documentElement, snapshot: extendedSnapshot } as const;
}
cleanRedundantEmptyLines(text: string) {
const lines = text.split(/\r?\n/g);
const mappedFlag = lines.map((line) => Boolean(line.trim()));
return lines.filter((_line, i) => mappedFlag[i] || mappedFlag[i - 1]).join('\n');
}
@Threaded()
async cleanHTMLforLMs(sourceHTML: string, ...discardSelectors: string[]): Promise<string> {
const t0 = Date.now();
let jsdom = this.linkedom.parseHTML(sourceHTML);
if (!jsdom.window.document.documentElement) {
jsdom = this.linkedom.parseHTML(`<html><body>${sourceHTML}</body></html>`);
}
for (const rl of discardSelectors) {
jsdom.window.document.querySelectorAll(rl).forEach((x) => x.remove());
}
jsdom.window.document.querySelectorAll('img[src],img[data-src]').forEach((x) => {
const src = x.getAttribute('src') || x.getAttribute('data-src');
if (src?.startsWith('data:')) {
x.setAttribute('src', 'blob:opaque');
}
x.removeAttribute('data-src');
x.removeAttribute('srcset');
});
jsdom.window.document.querySelectorAll('[class]').forEach((x) => {
const classes = x.getAttribute('class')?.split(/\s+/g) || [];
const newClasses = classes.filter((c) => !tailwindClasses.has(c));
x.setAttribute('class', newClasses.join(' '));
});
jsdom.window.document.querySelectorAll('[style]').forEach((x) => {
const style = x.getAttribute('style')?.toLocaleLowerCase() || '';
if (style.startsWith('display: none')) {
return;
}
x.removeAttribute('style');
});
const treeWalker = jsdom.window.document.createTreeWalker(
jsdom.window.document, // Start from the root document
0x80 // Only show comment nodes
);
let currentNode;
while ((currentNode = treeWalker.nextNode())) {
currentNode.parentNode?.removeChild(currentNode); // Remove each comment node
}
jsdom.window.document.querySelectorAll('*').forEach((x) => {
const attrs = x.getAttributeNames();
for (const attr of attrs) {
if (attr.startsWith('data-') || attr.startsWith('aria-')) {
x.removeAttribute(attr);
}
}
});
const final = this.cleanRedundantEmptyLines(jsdom.window.document.documentElement.outerHTML);
const dt = Date.now() - t0;
if (dt > 1000) {
this.logger.warn(`Performance issue: Cleaning HTML for LMs took ${dt}ms`, { dt });
}
return final;
}
snippetToElement(snippet?: string, url?: string) {
let parsed = this.linkedom.parseHTML(snippet || '<html><body></body></html>');
if (!parsed.window.document.documentElement) {
parsed = this.linkedom.parseHTML(`<html><body>${snippet || ''}</body></html>`);
}
// Hack for turndown gfm table plugin.
parsed.window.document.querySelectorAll('table').forEach((x) => {
Object.defineProperty(x, 'rows', { value: Array.from(x.querySelectorAll('tr')), enumerable: true });
});
Object.defineProperty(parsed.window.document.documentElement, 'cloneNode', {
value: function () { return this; },
});
if (url?.startsWith('https://jina.ai/reader')) {
const signature = parsed.window.document.createElement('p');
signature.textContent = 'Welcome home!';
parsed.window.document.body.insertBefore(signature, parsed.window.document.body.firstChild);
}
return parsed.window.document.documentElement;
}
runMarkify(markifyService: MarkifyService, html: HTMLElement) {
const t0 = Date.now();
try {
return markifyService.markify(html);
} finally {
const dt = Date.now() - t0;
if (dt > 1000) {
this.logger.warn(`Performance issue: Markify took ${dt}ms`, { dt });
}
}
}
@Threaded()
async analyzeHTMLTextLite(sourceHTML: string) {
let jsdom = this.linkedom.parseHTML(sourceHTML);
if (!jsdom.window.document.documentElement) {
jsdom = this.linkedom.parseHTML(`<html><body>${sourceHTML}</body></html>`);
}
jsdom.window.document.querySelectorAll('script,style,link,svg').forEach((s) => s.remove());
const text = jsdom.window.document.body.innerText || '';
return {
title: jsdom.window.document.title,
text,
tokens: countGPTToken(text.replaceAll(/[\s\r\n\t]+/g, ' ')),
};
}
@Threaded()
async xmlTextToSnapshot(sourceXML: string, url: URL) {
const xmlDom = new this.linkedom.DOMParser().parseFromString(sourceXML, 'text/xml');
const rootTagName = xmlDom.documentElement.nodeName.toLowerCase();
if (rootTagName === 'rss') {
const channel: Element = xmlDom.querySelector('rss channel');
if (channel) {
const snapshot = {
title: channel.querySelector('title')?.textContent || '',
href: url.href,
description: channel.querySelector('description')?.textContent || '',
text: channel.textContent || '',
html: sourceXML,
traits: ['blob'],
metadata: {
} as any,
};
const links = [] as { title?: string; href: string; description?: string; date?: string; }[];
for (const elem of channel.children) {
if (elem.tagName.toLowerCase() !== 'item') {
snapshot.metadata[elem.tagName.toLowerCase()] = elem.textContent || '';
continue;
}
links.push({
title: elem.querySelector('title')?.textContent || undefined,
href: elem.querySelector('link')?.textContent || '',
description: elem.querySelector('description')?.textContent || undefined,
date: elem.querySelector('pubDate')?.textContent || undefined,
});
}
const altHtml = `<html><head><title>${snapshot.title}</title><meta name="description" content="${snapshot.description}"></head><body>${links.map((l) => `<div><h3><a href="${l.href}">${l.title || ''}</a></h3><p>${l.description || ''}</p><a href="${l.href || ''}">${l.href || ''}</a><br /><time>${l.date || ''}</time></div>`).join('\n')}</body></html>`;
snapshot.html = altHtml;
return snapshot;
}
} else if (rootTagName === 'feed') {
const feed = xmlDom.querySelector('feed');
if (feed) {
const snapshot = {
title: xmlDom.querySelector('feed > title')?.textContent || '',
href: url.href,
description: xmlDom.querySelector('feed > subtitle')?.textContent || '',
text: xmlDom.documentElement.textContent || '',
html: sourceXML,
traits: ['blob'],
metadata: {
} as any,
};
const links = [] as { title?: string; href: string; description?: string; date?: string; }[];
for (const elem of feed.children) {
if (elem.tagName.toLowerCase() !== 'link') {
snapshot.metadata[elem.tagName.toLowerCase()] = elem.textContent || '';
continue;
}
links.push({
title: elem.querySelector('title')?.textContent || undefined,
href: elem.querySelector('link')?.getAttribute('href') || '',
description: elem.querySelector('summary')?.textContent || undefined,
date: elem.querySelector('updated')?.textContent || undefined,
});
}
const altHtml = `<html><head><title>${snapshot.title}</title><meta name="description" content="${snapshot.description}"></head><body>${links.map((l) => `<div><h3><a href="${l.href}">${l.title || ''}</a></h3><p>${l.description || ''}</p><a href="${l.href || ''}">${l.href || ''}</a><br /><time>${l.date || ''}</time></div>`).join('\n')}</body></html>`;
snapshot.html = altHtml;
return snapshot;
}
} else if (rootTagName === 'sitemapindex') {
const snapshot = {
title: 'Sitemap Index',
href: url.href,
description: '',
text: xmlDom.documentElement.textContent || '',
html: sourceXML,
traits: ['blob'],
metadata: {
} as any,
};
const links = [] as { href: string; lastmod?: string; }[];
for (const sitemap of xmlDom.querySelectorAll('sitemapindex sitemap')) {
links.push({
href: sitemap.querySelector('loc')?.textContent || '',
lastmod: sitemap.querySelector('lastmod')?.textContent || undefined,
});
}
const altHtml = `<html><head><title>${snapshot.title}</title><meta name="description" content="${snapshot.description}"></head><body>${links.map((l) => `<div><a href="${l.href}">${l.href}</a><br /><time>${l.lastmod || ''}</time></div>`).join('\n')}</body></html>`;
snapshot.html = altHtml;
return snapshot;
} else if (rootTagName === 'urlset') {
const snapshot = {
title: 'Sitemap',
href: url.href,
description: '',
text: xmlDom.documentElement.textContent || '',
html: sourceXML,
traits: ['blob'],
metadata: {
} as any,
};
const links = [] as { href: string; lastmod?: string; }[];
for (const urlElem of xmlDom.querySelectorAll('urlset url')) {
links.push({
href: urlElem.querySelector('loc')?.textContent || '',
lastmod: urlElem.querySelector('lastmod')?.textContent || undefined,
});
}
const altHtml = `<html><head><title>${snapshot.title}</title><meta name="description" content="${snapshot.description}"></head><body>${links.map((l) => `<div><a href="${l.href}">${l.href}</a><br /><time>${l.lastmod || ''}</time></div>`).join('\n')}</body></html>`;
snapshot.html = altHtml;
return snapshot;
}
return {
title: xmlDom.documentElement.tagName,
description: '',
text: xmlDom.documentElement.textContent || sourceXML,
html: sourceXML,
};
}
}
const jsdomControl = container.resolve(JSDomControl);
export default jsdomControl;
src/services/lm.ts
+152
View File
@@ -0,0 +1,152 @@
import { AsyncService } from 'civkit/async-service';
import { singleton } from 'tsyringe';
import { PageSnapshot } from './puppeteer';
import { GlobalLogger } from './logger';
import _ from 'lodash';
import { AssertionFailureError } from 'civkit/civ-rpc';
import { LLMManager } from './common-llm';
import { JSDomControl } from './jsdom';
const tripleBackTick = '```';
@singleton()
export class LmControl extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
constructor(
protected globalLogger: GlobalLogger,
protected commonLLM: LLMManager,
protected jsdomControl: JSDomControl,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
async* geminiFromBrowserSnapshot(snapshot?: PageSnapshot & {
pageshotUrl?: string,
}) {
const pageshot = snapshot?.pageshotUrl || snapshot?.pageshot;
if (!pageshot) {
throw new AssertionFailureError('Screenshot of the page is not available');
}
const html = await this.jsdomControl.cleanHTMLforLMs(snapshot.html, 'script,link,style,textarea,select>option,svg');
const it = this.commonLLM.iterRun('vertex-gemini-3.1-flash-lite', {
prompt: [
`HTML: \n${html}\n\nSCREENSHOT: \n`,
typeof pageshot === 'string' ? new URL(pageshot) : pageshot,
`Convert this webpage into a markdown source file that does not contain HTML tags, retaining the page language and visual structures.`,
],
options: {
system: 'You are ReaderLM-v7, a model that generates Markdown source files only. No HTML, notes and chit-chats allowed',
stream: true
}
});
const chunks: string[] = [];
for await (const txt of it) {
chunks.push(txt);
const output: PageSnapshot = {
...snapshot,
parsed: {
...snapshot?.parsed,
textContent: chunks.join(''),
}
};
yield output;
}
return;
}
async* readerLMMarkdownFromSnapshot(snapshot?: PageSnapshot) {
if (!snapshot) {
throw new AssertionFailureError('Snapshot of the page is not available');
}
const html = await this.jsdomControl.cleanHTMLforLMs(snapshot.html, 'script,link,style,textarea,select>option,svg');
const it = this.commonLLM.iterRun('readerlm-v2', {
prompt: `Extract the main content from the given HTML and convert it to Markdown format.\n\n${tripleBackTick}html\n${html}\n${tripleBackTick}\n`,
options: {
// system: 'You are an AI assistant developed by VENDOR_NAME',
stream: true,
modelSpecific: {
top_k: 1,
temperature: 0,
repetition_penalty: 1.13,
presence_penalty: 0.25,
frequency_penalty: 0.25,
max_tokens: 4096,
}
},
maxTry: 1,
});
const chunks: string[] = [];
for await (const txt of it) {
chunks.push(txt);
const output: PageSnapshot = {
...snapshot,
parsed: {
...snapshot?.parsed,
textContent: chunks.join(''),
}
};
yield output;
}
return;
}
async* readerLMFromSnapshot(schema?: string, instruction: string = 'Infer useful information from the HTML and present it in a structured JSON object.', snapshot?: PageSnapshot) {
if (!snapshot) {
throw new AssertionFailureError('Snapshot of the page is not available');
}
const html = await this.jsdomControl.cleanHTMLforLMs(snapshot.html, 'script,link,style,textarea,select>option,svg');
const it = this.commonLLM.iterRun('readerlm-v2', {
prompt: `${instruction}\n\n${tripleBackTick}html\n${html}\n${tripleBackTick}\n${schema ? `The JSON schema:\n${tripleBackTick}json\n${schema}\n${tripleBackTick}\n` : ''}`,
options: {
// system: 'You are an AI assistant developed by VENDOR_NAME',
stream: true,
modelSpecific: {
top_k: 1,
temperature: 0,
repetition_penalty: 1.13,
presence_penalty: 0.25,
frequency_penalty: 0.25,
max_tokens: 8192,
}
},
maxTry: 1,
});
const chunks: string[] = [];
for await (const txt of it) {
chunks.push(txt);
const output: PageSnapshot = {
...snapshot,
parsed: {
...snapshot?.parsed,
textContent: chunks.join(''),
}
};
yield output;
}
return;
}
}
src/services/logger.ts
+57
View File
@@ -0,0 +1,57 @@
import { AbstractPinoLogger } from 'civkit/pino-logger';
import { singleton, container } from 'tsyringe';
import { threadId } from 'node:worker_threads';
import { getTraceCtx } from 'civkit/async-context';
const levelToSeverityMap: { [k: string]: string | undefined; } = {
trace: 'DEFAULT',
debug: 'DEBUG',
info: 'INFO',
warn: 'WARNING',
error: 'ERROR',
fatal: 'CRITICAL',
};
@singleton()
export class GlobalLogger extends AbstractPinoLogger {
loggerOptions = {
level: 'debug',
base: {
tid: threadId,
}
};
override init(): void {
if (process.env['NODE_ENV']?.startsWith('prod')) {
super.init(process.stdout);
} else {
const PinoPretty = require('pino-pretty').PinoPretty;
super.init(PinoPretty({
singleLine: true,
colorize: true,
messageFormat(log: any, messageKey: any) {
return `${log['tid'] ? `[${log['tid']}]` : ''}[${log['service'] || 'ROOT'}] ${log[messageKey]}`;
},
}));
}
this.emit('ready');
}
override log(...args: any[]) {
const [levelObj, ...rest] = args;
const severity = levelToSeverityMap[levelObj?.level];
const traceCtx = getTraceCtx();
const patched: any= { ...levelObj, severity };
const traceId = traceCtx?.googleTraceId || traceCtx?.traceId;
if (traceId && process.env['GCLOUD_PROJECT']) {
patched['logging.googleapis.com/trace'] = `projects/${process.env['GCLOUD_PROJECT']}/traces/${traceId}`;
}
return super.log(patched, ...rest);
}
}
const instance = container.resolve(GlobalLogger);
export default instance;
src/services/markify.ts
+854
View File
@@ -0,0 +1,854 @@
import { MathMLToLaTeX } from '@nomagick/mathml-to-latex';
export interface MarkifyOptions {
headingStyle?: 'atx' | 'setext';
hr?: string;
bulletListMarker?: '-' | '+' | '*';
codeBlockStyle?: 'indented' | 'fenced';
fence?: '```' | '~~~' | undefined;
emDelimiter?: '_' | '*';
strongDelimiter?: '__' | '**';
linkStyle?: 'inlined' | 'referenced' | 'discarded';
linkReferenceStyle?: 'full' | 'collapsed' | 'shortcut';
preformattedCode?: boolean;
footnoteStyle?: 'inline' | 'document';
// Add other options as needed
baseUrl?: string;
gfm?: boolean;
}
export interface Replacer {
name: string;
replacement: (this: MarkifyService, content: string, node: Element, options?: MarkifyOptions, instance?: MarkifyService) => string;
}
export interface MarkifyRule {
filter: string | string[];
replacement: Replacer['replacement'];
}
const tableCellAlignment: Record<string, string> = {
left: ':---',
right: '---:',
center: ':---:',
};
export class MarkifyService {
protected options: MarkifyOptions;
protected rules: { [key: string]: Replacer[]; } = {};
protected keepTags: Set<string> = new Set();
links: Array<{
href: string;
// domain: string;
text: string;
title?: string;
ref: number;
}> = [];
hrefTitleMap = new Map<string, number>();
images: Array<{ src: string; alt: string; ref: number; }> = [];
protected listLevel: number = -1;
protected listStack: Array<{ tag: string; num?: number; }> = [];
protected tableStack: Array<{ tag: string; hasTh: boolean; }> = [];
refMap = new WeakMap<Element, any>();
protected ignoreDivNewline: boolean = false;
protected ignoreEmptyTextNode: boolean = true;
protected layoutByTable = false;
protected strongRegex: RegExp;
protected emRegex: RegExp;
protected trimRegex: RegExp;
private higlightRegex: RegExp;
commonMark: Record<string, Function> = {
h1: this.processHeading,
h2: this.processHeading,
h3: this.processHeading,
h4: this.processHeading,
h5: this.processHeading,
h6: this.processHeading,
p: this.processParagraph,
br: this.processBr,
hr: this.processHr,
strong: this.processStrong,
b: this.processStrong,
em: this.processEmphasis,
i: this.processEmphasis,
ul: this.processListWrapper,
ol: this.processListWrapper,
li: this.processListItem,
a: this.processAnchor,
img: this.processImage,
pre: this.processPre,
code: this.processCode,
div: this.processDivNode,
svg: this.processSVG,
title: this.processTitle,
figcaption: this.processFigcaption,
link: this.skip,
script: this.skip,
meta: this.skip,
style: this.skip,
details: this.processDivNode,
label: this.processLabel,
select: this.processSelect,
option: this.processOption,
blockquote: this.processBlockQuote,
};
gfm: Record<string, Function> = {
table: this.processTable,
thead: this.processThead,
tbody: this.processTbody,
tr: this.processTr,
th: this.processTd,
td: this.processTd,
strike: this.processStrike,
del: this.processStrike,
s: this.processStrike,
input: this.processInput,
math: this.processMath,
};
fnMap: Record<string, Function>;
constructor(options: MarkifyOptions = {}) {
this.options = Object.assign({}, options);
this.fnMap = {
...this.commonMark,
...(this.options.gfm ? this.gfm : {}),
};
this.strongRegex = /^\*\*(.*?)\*\*$/g;
this.emRegex = /^_(.*?)_$/g;
this.trimRegex = /\s+/g;
this.higlightRegex = /highlight-(?:text|source)-([a-z0-9]+)/;
if (this.options.baseUrl?.startsWith('blob:') || this.options.baseUrl?.startsWith('data:')) {
delete this.options.baseUrl;
}
this.init();
}
protected init() {
this.keepTags.clear();
this.rules = {};
this.listLevel = -1;
this.listStack = [];
this.tableStack = [];
this.links = [];
this.images = [];
this.ignoreDivNewline = false;
this.ignoreEmptyTextNode = true;
this.layoutByTable = false;
}
markify(element: HTMLElement): string {
let markdown = '';
markdown = this.processNode(element);
// Trim leading/trailing whitespace and ensure consistent newlines
const chunks = [
markdown.trim().replace(/\n{3,}/g, '\n\n'),
];
if (this.options.linkStyle === 'referenced' && this.links.length) {
chunks.push('');
switch (this.options.linkReferenceStyle) {
case 'collapsed':
case 'shortcut': {
chunks.push(this.links.map((x) => `[${x.text}]: ${x.href}${x.title ? ` "${x.title}"` : ''}`).join('\n'));
break;
}
default: {
const expposed = new Set();
chunks.push(this.links.filter((x) => {
if (expposed.has(x.ref)) {
return false;
};
expposed.add(x.ref);
return true;
}).map((x) => `[${x.ref}]: ${x.href}${x.title ? ` "${x.title}"` : ''}`).join('\n'));
break;
}
}
}
return chunks.join('\n');
}
trackLink(href: string, text: string = '', title?: string) {
const normalizedTitle = title ? title.replaceAll(/"/g, '\\"').replaceAll(/\r?\n/g, ' ').trim() : '';
const hrefTitleString = `${href}${normalizedTitle ? ` "${normalizedTitle}"` : ''}`;
const normalizedText = text.replaceAll(/\s+/g, ' ').replaceAll(/\r?\n/g, ' ').trim();
const record = { href, text: normalizedText, title: normalizedTitle, ref: -1 };
if (!this.hrefTitleMap.has(hrefTitleString)) {
this.hrefTitleMap.set(hrefTitleString, this.links.length + 1);
this.links.push(record);
}
record.ref = this.hrefTitleMap.get(hrefTitleString)!;
return record.ref;
}
protected processNode(element: Element): string {
const { nodeType, nodeName } = element;
if (nodeType === 3) {
return this.processTextNode(element as unknown as Text);
}
if (![1, 9, 11].includes(nodeType)) return '';
const tagName = nodeName.toLowerCase();
if (this.keepTags.has(tagName)) {
return element.outerHTML || '';
}
const customRuleSet = this.rules[tagName];
const fn = this.fnMap[tagName];
const r = fn ? fn.call(this, element) : this.processChildren(element);
if (customRuleSet?.length) {
return this.replaceByRules(r || '', element, customRuleSet);
}
return r;
}
protected skip(_: any) {
return '';
}
protected processDivNode(divNode: HTMLDivElement): string {
const markdown = this.processChildren(divNode);
const firstChild = divNode.firstChild;
const className = divNode.className || '';
const language = (className.match(this.higlightRegex) || [])[1];
if (firstChild?.nodeName === 'PRE' && language) {
return `\n\n${this.options.fence}${language}\n${firstChild.textContent}\n${this.options.fence}\n\n`;
}
// const trimmedMarkdown = markdown.replace(/^\n+|\n+$/g, '');
let trimmedMarkdown = markdown;
let start = 0;
let end = trimmedMarkdown.length - 1;
while (trimmedMarkdown[start] === '\n') {
start++;
}
while (trimmedMarkdown[end] === '\n') {
end--;
}
trimmedMarkdown = trimmedMarkdown.slice(start, end + 1);
if (this.ignoreDivNewline) {
return `${trimmedMarkdown} `;
}
if (trimmedMarkdown) return `\n\n${trimmedMarkdown}\n\n`;
return markdown;
}
protected processTextNode(textNode: Text): string {
let text = textNode.data || '';
if (!text.trim()) {
// return this.ignoreEmptyTextNode ? '' : ' ';
if (this.ignoreEmptyTextNode || !text.includes(' ')) return '';
return ' ';
}
const emDelimiter = this.options.emDelimiter || '_';
const strongDelimiter = this.options.strongDelimiter || '**';
// Escape markdown control characters
// if (!this.ignoreEscaping) {
// text = text.replace(/^(\*|_|`|~|>|#|-|\+|=|\!|\[|\]|\(|\))/g, '\\$1');
// }
// remove extra spaces and new lines
text = text.replace(this.trimRegex, ' ');
//handle emphasis and strong within text node.
text = text.replace(
this.strongRegex,
` ${strongDelimiter}$1${strongDelimiter} `,
);
text = text.replace(this.emRegex, ` ${emDelimiter}$1${emDelimiter} `);
return text;
}
protected replaceByRules(
text: string,
element: Element,
replacers: Replacer[],
): string {
let replacedText = text;
for (const replacer of replacers) {
replacedText = replacer.replacement.call(this, replacedText, element, this.options, this);
}
return replacedText;
}
protected processTitle(element: Element): string {
const titleText = this.processChildren(element);
if (!titleText) return '';
return `${titleText}\n`;
}
protected processHeading(element: Element): string {
const headingStyle = this.options.headingStyle || 'atx';
const headingText = this.processChildren(element).trim();
const name: string = element.tagName ?? 'h6';
const level = +name[1];
if (!headingText) return '';
if (headingStyle === 'setext' && level <= 2) {
const underline = level <= 2 ? (level === 1 ? '=' : '-') : '';
return underline
? `\n${headingText}\n${underline.repeat(headingText.length)}\n\n`
: `\n${headingText}\n\n`;
} else {
return `\n${'#'.repeat(level)} ${headingText}\n\n`;
}
}
protected processHr(element: Element): string {
const hr = this.options.hr || '* * *';
return `\n\n${hr}\n\n`;
}
protected processBr(element: Element): string {
return '\n\n';
}
protected processSelect(element: Element): string {
const markdown = this.processChildren(element);
if (!markdown) {
return '';
}
return `\n${markdown}\n\n`;
}
protected processOption(element: Element): string {
// const value = element.attribs?.value || '';
const value = element.getAttribute('value') || '';
return `${value}\n`;
}
protected processLabel(element: Element): string {
const labelText = this.processChildren(element);
if (!labelText) return '';
return `${labelText} `;
}
protected processStrong(element: Element): string {
const strongDelimiter = this.options.strongDelimiter || '**';
const text = this.processChildren(element);
if (!text) return '';
return `${strongDelimiter}${text.trim()}${strongDelimiter}`;
}
protected processEmphasis(element: Element): string {
let text = this.processChildren(element).trim();
if (!text) return '';
text = text.replace(/_/g, '\\_');
const emDelimiter = this.options.emDelimiter || '_';
return `${emDelimiter}${text}${emDelimiter}`;
}
protected processListWrapper(element: Element): string {
const tag = element.tagName?.toLowerCase() || 'ul';
// store this wrapper in the listStack
this.listStack.unshift({ tag, num: 0 });
this.listLevel++;
const markdown = this.processChildren(element);
this.listLevel--;
this.listStack.shift();
return `\n${markdown}\n`;
}
protected processListItem(element: Element): string {
const markdown: string[] = [];
let prevlistStack = this.listStack[0];
let withNewList = false;
this.ignoreDivNewline = true;
if (!prevlistStack || prevlistStack.tag === 'li') {
prevlistStack = { tag: 'ul', num: 0 };
this.listStack.unshift(prevlistStack);
this.listLevel++;
withNewList = true;
}
this.listStack.unshift({ tag: 'li' });
let listItemText = this.processChildren(element);
if (listItemText[0] === '\n') {
listItemText = listItemText.slice(1);
}
if (listItemText[listItemText.length - 1] === '\n') {
listItemText = listItemText.slice(0, -1);
}
const bulletListMarker =
prevlistStack.tag === 'ol'
? `${prevlistStack.num! + 1}.`
: this.options.bulletListMarker || '*';
if (listItemText) {
markdown.push(
`${' '.repeat(
this.listLevel,
)}${bulletListMarker} ${listItemText}\n`,
);
} else {
markdown.push('\n');
}
prevlistStack.num!++;
this.listStack.shift();
if (withNewList) {
this.listLevel--;
this.listStack.shift();
}
this.ignoreDivNewline = false;
return markdown.join('');
}
protected processAnchor(element: Element): string {
let href = element.getAttribute('href') || '';
let domain = '';
const title = (element.getAttribute('title') || '').replace(/"/g, '\\"');
let text = this.processChildren(element);
text = text.replace(/\s+/g, ' ').trim();
if (this.options.baseUrl) {
try {
const parsed = new URL(href, this.options.baseUrl);
href = parsed.href;
domain = parsed.hostname;
} catch (e) { void 0; }
}
const linkStyle = this.options.linkStyle || 'inlined';
let linkRef = undefined;
if (href) {
linkRef = this.links.length + 1;
const rec = { href, text, title, ref: linkRef, domain };
this.refMap.set(element, rec);
rec.ref = this.trackLink(href, text, title);
linkRef = rec.ref;
}
const headingRegex = /^#{1,} /;
switch (linkStyle) {
case 'referenced':
switch (this.options.linkReferenceStyle) {
case 'collapsed':
if (headingRegex.test(text)) {
const hashPart = text.slice(0, text.indexOf(' ') + 1);
const textPart = text.slice(text.indexOf(' ') + 1);
return `${hashPart}[${textPart}][]`;
}
return `[${text}][]`;
case 'shortcut':
if (headingRegex.test(text)) {
const hashPart = text.slice(0, text.indexOf(' ') + 1);
const textPart = text.slice(text.indexOf(' ') + 1);
return `${hashPart}[${textPart}]`;
}
return `[${text}]`;
case 'full':
default: {
if (headingRegex.test(text)) {
const hashPart = text.slice(0, text.indexOf(' ') + 1);
const textPart = text.slice(text.indexOf(' ') + 1);
return `${hashPart}[${textPart}][${linkRef ?? ''}]`;
}
return `[${text}][${linkRef ?? ''}]`;
}
}
case 'discarded':
return text;
case 'inlined':
default: {
if (headingRegex.test(text)) {
const hashPart = text.slice(0, text.indexOf(' ') + 1);
const textPart = text.slice(text.indexOf(' ') + 1);
return `${hashPart}[${textPart}](${href}${title ? ` "${title}"` : ''})`;
}
return `[${text}](${href}${title ? ` "${title}"` : ''})`;
}
}
}
protected processImage(element: Element): string {
const src = element.getAttribute('src') || '';
const alt = element.getAttribute('alt') || '';
const title = element.getAttribute('title') || '';
const rec = { src, alt, ref: this.images.length + 1 };
this.refMap.set(element, rec);
this.images.push(rec);
return `![${alt}](${src}${title ? ` "${title}"` : ''})`;
}
protected processPreformattedChildren(nodes?: Element[]) {
if (!nodes) return '';
const markdown: string[] = [];
nodes.forEach((node) => {
if (node.nodeType === 3) {
markdown.push((node as any).data || '');
} else if (node.tagName?.toLowerCase() === 'code') {
// If there's a <code> inside <pre>, process its text content
markdown.push(this.processCode(node));
} else {
// Process other tags as text
markdown.push(this.processNode(node));
}
});
return markdown.join('');
}
protected processCode(element: Element): string {
this.ignoreEmptyTextNode = false;
const text = this.getCodeText(element);
this.ignoreEmptyTextNode = true;
return text;
}
protected getCodeText(element: Element): string {
const classes = element.className.split(' ') || [];
const language =
classes
.find((cls) => cls.startsWith('language-') || cls.startsWith('lang-'))
?.replace(/^(language-|lang-)/, '') ?? '';
const codeBlockStyle = this.options.codeBlockStyle || 'indented';
// const code = this.processPreformattedChildren(Array.from(element.childNodes) as Element[]) || '';
const code = element.textContent || '';
if (!code) return '';
if (!language && !code.includes('\n')) {
return `\`${code.trim()}\``;
}
if (codeBlockStyle === 'fenced') {
const fence = this.options.fence || '```';
return `\n${fence}${language}\n${code.trim()}\n${fence}\n`;
} else {
return `\n ${code.trim().replace(/\n/g, '\n ')}\n`;
}
}
protected processPre(element: Element): string {
this.ignoreEmptyTextNode = false;
const preText = this.processPreformattedChildren(Array.from(element.childNodes) as Element[]);
this.ignoreEmptyTextNode = true;
return preText;
}
protected processParagraph(element: Element): string {
const text = this.processChildren(element).trim();
if (!text) return '';
return `\n${text}\n\n`;
}
protected processSVG(element: Element): string {
// Handle SVG elements if needed
return '';
}
protected processTable(element: Element): string {
this.tableStack.unshift({ tag: 'table', hasTh: false });
const markdown = this.processChildren(element);
let tableIndex = 0;
for (; tableIndex < this.tableStack.length; tableIndex++) {
if (this.tableStack[tableIndex].tag === 'table') {
break;
}
}
this.tableStack = this.tableStack.slice(tableIndex + 1);
const items: string[] = [];
if (!markdown.startsWith('\n')) {
items.unshift('\n');
}
items.push(markdown);
if (!markdown.endsWith('\n')) {
items.push('\n');
}
return items.join('');
}
protected processThead(element: Element): string {
if (element.previousSibling && element.previousSibling.textContent) {
return '\n' + this.processChildren(element);
}
return this.processChildren(element);
}
protected processTbody(element: Element): string {
return this.processChildren(element);
}
protected processTd(element: Element): string {
const items: string[] = [];
const text = this.processChildren(element);
items.push(`${this.layoutByTable ? text : text.replace(/\s+/g, ' ')}`.trim());
return items.join('');
}
protected processTr(element: Element): string {
const markdown: string[] = [];
const cellContents: string[] = [];
const alignments: string[] = [];
const children = Array.from(element.childNodes) as Element[];
let isTh = children.some((child) => child.tagName?.toLowerCase() === 'th');
if (this.tableStack[0]?.tag !== 'table') isTh = false;
let lastTable = this.tableStack.find((item) => item.tag === 'table');
this.tableStack.unshift({ tag: 'tr', hasTh: false });
const isWrappedByTable = lastTable?.tag === 'table';
if (
!isWrappedByTable ||
(!isTh && !lastTable?.hasTh)
) {
// treat as normal div
this.layoutByTable = true;
const tableText = `${this.processChildren(element)}\n`;
this.layoutByTable = false;
return tableText;
}
if (isTh && lastTable) {
lastTable.hasTh = true;
}
for (let i = 0; i < children.length; i++) {
const child = children[i];
const name = child.tagName?.toLowerCase() ?? '';
if (name === 'th' || name === 'td') {
const cellContent = (this.processChildren(child) ?? '').replace(/\s+/g, ' ').trim();
cellContents.push(cellContent);
if (name === 'th') {
alignments.push(child.getAttribute('align') || '');
}
}
}
const cellContentStr = cellContents.join(' | ');
markdown.push(`| ${cellContentStr} |\n`);
if (isTh) {
const divider = alignments
.map((align) => {
return tableCellAlignment[align] ?? '---';
})
.join(' | ');
markdown.push(`| ${divider} |\n`);
}
return markdown.join('');
}
protected processStrike(element: Element): string {
const text = this.processChildren(element);
if (!text) return '';
return `~~${text}~~`;
}
protected processInput(element: Element) {
if (element.getAttribute('type') === 'checkbox') {
return this.processCheckBox(element);
}
return '';
}
protected processCheckBox(element: Element): string {
const checked = element.getAttribute('checked') !== undefined;
return `- [${checked ? 'x' : ' '}] `;
}
protected isUnicodeSupported() {
try {
new RegExp("\\p{L}", "u");
return true;
} catch {
return false;
}
}
protected needSpace(str1: string, str2: string): boolean {
const useUnicode = this.isUnicodeSupported();
const endsWithText = useUnicode
? /[\p{L}\p{N}_]$/u.test(str1)
: /[^\s\W]$/.test(str1);
const startsWithText = useUnicode
? /^[\p{L}\p{N}_]/u.test(str2)
: /^[^\s\W]/.test(str2);
return endsWithText && startsWithText;
}
protected processChildren(element: Element): string {
let markdown = '';
const items: any[] = [];
let lastItem = '';
(Array.from(element.childNodes) as Element[] || []).forEach((child) => {
const { tagName } = child;
const cname = tagName?.toLowerCase() ?? '';
let childContent = this.processNode(child);
if (!childContent) return;
const parentTag = child.parentNode?.nodeName.toLowerCase() || '';
// first level of list
if (this.listLevel > 0 &&
(parentTag === 'ul' || parentTag === 'ol') &&
cname &&
cname !== 'li'
) {
childContent = `${' '.repeat(this.listLevel) + childContent}\n`;
}
if (['ul', 'ol'].includes(cname)) {
if (!lastItem?.endsWith('\n\n') && !childContent.startsWith('\n')) {
childContent = '\n' + childContent;
}
}
else if (this.needSpace(lastItem, childContent)) {
childContent = ' ' + childContent;
}
lastItem = childContent;
items.push(childContent);
});
markdown = items.join('');
return markdown;
}
protected processFigcaption(element: Element) {
return `\n\n${this.processChildren(element)}\n\n`;
}
protected processBlockQuote(element: Element): string {
let items: string[] = [];
(Array.from(element.childNodes) as Element[] || []).forEach((child) => {
const childContent = this.processNode(child);
if (!childContent) return;
items.push(childContent.replace(/^ /g, ''));
});
const markdown = items.join('').trim();
if (!markdown) return '';
return `\n> ${markdown.replace(/\n/g, '\n> ')}\n\n`;
}
protected processNodes(elements: Element[]): string {
const markdown: string[] = [];
for (const child of elements) {
const cmd = this.processNode(child);
markdown.push(cmd);
}
return markdown.join('');
}
protected processMath(element: Element): string {
const altText = element.getAttribute('alttext');
if (altText) {
return altText;
}
try {
const latex = MathMLToLaTeX.convert(element.outerHTML);
if (!latex.trim()) return '';
// Parent is p and no siblings, then is block. Otherwise inline.
if (element.getAttribute('display') === 'block' || (element.parentElement?.tagName.toLowerCase() === 'p' && element.parentElement.childNodes.length === 1)) {
return `\n\n$$\n${latex.trim()}\n$$\n\n`;
}
return `$${latex.trim()}$`;
} catch (e) {
return element.textContent || '';
}
}
use(addRuleFns: Array<(service: MarkifyService) => void>) {
addRuleFns.forEach((fn) => {
fn(this);
});
}
addRule(name: string, rule: MarkifyRule) {
const { filter, replacement } = rule;
const tags = Array.isArray(filter) ? filter : [filter];
tags.forEach((tag) => {
const replacers = this.rules[tag] || [];
replacers.push({
name,
replacement: replacement,
});
this.rules[tag] = replacers;
});
}
keep(tag?: string) {
if (!tag) return;
this.keepTags.add(tag.toLowerCase());
}
}
src/services/minimal-stealth.js
+603
View File
@@ -0,0 +1,603 @@
export function minimalStealth() {
/**
* A set of shared utility functions specifically for the purpose of modifying native browser APIs without leaving traces.
*
* Meant to be passed down in puppeteer and used in the context of the page (everything in here runs in NodeJS as well as a browser).
*
* Note: If for whatever reason you need to use this outside of `puppeteer-extra`:
* Just remove the `module.exports` statement at the very bottom, the rest can be copy pasted into any browser context.
*
* Alternatively take a look at the `extract-stealth-evasions` package to create a finished bundle which includes these utilities.
*
*/
const utils = {};
// const toStringRedirects = new WeakMap();
utils.init = () => {
utils.preloadCache();
// const handler = {
// apply: function (target, ctx) {
// // This fixes e.g. `HTMLMediaElement.prototype.canPlayType.toString + ""`
// if (ctx === Function.prototype.toString) {
// return utils.makeNativeString('toString');
// }
// const originalObj = toStringRedirects.get(ctx);
// // `toString` targeted at our proxied Object detected
// if (originalObj) {
// const fallback = () =>
// originalObj && originalObj.name
// ? utils.makeNativeString(originalObj.name)
// : utils.makeNativeString(ctx.name);
// // Return the toString representation of our original object if possible
// return originalObj + '' || fallback();
// }
// if (typeof ctx === 'undefined' || ctx === null) {
// return target.call(ctx);
// }
// // Check if the toString protype of the context is the same as the global prototype,
// // if not indicates that we are doing a check across different windows., e.g. the iframeWithdirect` test case
// const hasSameProto = Object.getPrototypeOf(
// Function.prototype.toString
// ).isPrototypeOf(ctx.toString); // eslint-disable-line no-prototype-builtins
// if (!hasSameProto) {
// // Pass the call on to the local Function.prototype.toString instead
// return ctx.toString();
// }
// return target.call(ctx);
// }
// };
// const toStringProxy = new Proxy(
// Function.prototype.toString,
// utils.stripProxyFromErrors(handler)
// );
// utils.replaceProperty(Function.prototype, 'toString', {
// value: toStringProxy
// });
};
/**
* Wraps a JS Proxy Handler and strips it's presence from error stacks, in case the traps throw.
*
* The presence of a JS Proxy can be revealed as it shows up in error stack traces.
*
* @param {object} handler - The JS Proxy handler to wrap
*/
utils.stripProxyFromErrors = (handler = {}) => {
const newHandler = {
setPrototypeOf: function (target, proto) {
if (proto === null)
throw new TypeError('Cannot convert object to primitive value');
if (Object.getPrototypeOf(target) === Object.getPrototypeOf(proto)) {
throw new TypeError('Cyclic __proto__ value');
}
return Reflect.setPrototypeOf(target, proto);
}
};
// We wrap each trap in the handler in a try/catch and modify the error stack if they throw
const traps = Object.getOwnPropertyNames(handler);
traps.forEach(trap => {
newHandler[trap] = function () {
try {
// Forward the call to the defined proxy handler
return handler[trap].call(this, ...(arguments || []));
} catch (err) {
// Stack traces differ per browser, we only support chromium based ones currently
if (!err || !err.stack || !err.stack.includes(`at `)) {
throw err;
}
// When something throws within one of our traps the Proxy will show up in error stacks
// An earlier implementation of this code would simply strip lines with a blacklist,
// but it makes sense to be more surgical here and only remove lines related to our Proxy.
// We try to use a known "anchor" line for that and strip it with everything above it.
// If the anchor line cannot be found for some reason we fall back to our blacklist approach.
const stripWithBlacklist = (stack, stripFirstLine = true) => {
const blacklist = [
`at Reflect.${trap} `, // e.g. Reflect.get or Reflect.apply
`at Object.${trap} `, // e.g. Object.get or Object.apply
`at Object.newHandler.<computed> [as ${trap}] ` // caused by this very wrapper :-)
];
return (
err.stack
.split('\n')
// Always remove the first (file) line in the stack (guaranteed to be our proxy)
.filter((line, index) => !(index === 1 && stripFirstLine))
// Check if the line starts with one of our blacklisted strings
.filter(line => !blacklist.some(bl => line.trim().startsWith(bl)))
.join('\n')
);
};
const stripWithAnchor = (stack, anchor) => {
const stackArr = stack.split('\n');
anchor = anchor || `at Object.newHandler.<computed> [as ${trap}] `; // Known first Proxy line in chromium
const anchorIndex = stackArr.findIndex(line =>
line.trim().startsWith(anchor)
);
if (anchorIndex === -1) {
return false; // 404, anchor not found
}
// Strip everything from the top until we reach the anchor line
// Note: We're keeping the 1st line (zero index) as it's unrelated (e.g. `TypeError`)
stackArr.splice(1, anchorIndex);
return stackArr.join('\n');
};
// Special cases due to our nested toString proxies
err.stack = err.stack.replace(
'at Object.toString (',
'at Function.toString ('
);
if ((err.stack || '').includes('at Function.toString (')) {
err.stack = stripWithBlacklist(err.stack, false);
throw err;
}
// Try using the anchor method, fallback to blacklist if necessary
err.stack = stripWithAnchor(err.stack) || stripWithBlacklist(err.stack);
throw err; // Re-throw our now sanitized error
}
};
});
return newHandler;
};
/**
* Strip error lines from stack traces until (and including) a known line the stack.
*
* @param {object} err - The error to sanitize
* @param {string} anchor - The string the anchor line starts with
*/
utils.stripErrorWithAnchor = (err, anchor) => {
const stackArr = err.stack.split('\n');
const anchorIndex = stackArr.findIndex(line => line.trim().startsWith(anchor));
if (anchorIndex === -1) {
return err; // 404, anchor not found
}
// Strip everything from the top until we reach the anchor line (remove anchor line as well)
// Note: We're keeping the 1st line (zero index) as it's unrelated (e.g. `TypeError`)
stackArr.splice(1, anchorIndex);
err.stack = stackArr.join('\n');
return err;
};
/**
* Replace the property of an object in a stealthy way.
*
* Note: You also want to work on the prototype of an object most often,
* as you'd otherwise leave traces (e.g. showing up in Object.getOwnPropertyNames(obj)).
*
* @see https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Object/defineProperty
*
* @example
* replaceProperty(WebGLRenderingContext.prototype, 'getParameter', { value: "alice" })
* // or
* replaceProperty(Object.getPrototypeOf(navigator), 'languages', { get: () => ['en-US', 'en'] })
*
* @param {object} obj - The object which has the property to replace
* @param {string} propName - The property name to replace
* @param {object} descriptorOverrides - e.g. { value: "alice" }
*/
utils.replaceProperty = (obj, propName, descriptorOverrides = {}) => {
return Object.defineProperty(obj, propName, {
// Copy over the existing descriptors (writable, enumerable, configurable, etc)
...(Object.getOwnPropertyDescriptor(obj, propName) || {}),
// Add our overrides (e.g. value, get())
...descriptorOverrides
});
};
/**
* Preload a cache of function copies and data.
*
* For a determined enough observer it would be possible to overwrite and sniff usage of functions
* we use in our internal Proxies, to combat that we use a cached copy of those functions.
*
* Note: Whenever we add a `Function.prototype.toString` proxy we should preload the cache before,
* by executing `utils.preloadCache()` before the proxy is applied (so we don't cause recursive lookups).
*
* This is evaluated once per execution context (e.g. window)
*/
utils.preloadCache = () => {
if (utils.cache) {
return;
}
utils.cache = {
// Used in our proxies
Reflect: {
get: Reflect.get.bind(Reflect),
apply: Reflect.apply.bind(Reflect)
},
// Used in `makeNativeString`
nativeToStringStr: Function.toString + '' // => `function toString() { [native code] }`
};
};
/**
* Utility function to generate a cross-browser `toString` result representing native code.
*
* There's small differences: Chromium uses a single line, whereas FF & Webkit uses multiline strings.
* To future-proof this we use an existing native toString result as the basis.
*
* The only advantage we have over the other team is that our JS runs first, hence we cache the result
* of the native toString result once, so they cannot spoof it afterwards and reveal that we're using it.
*
* @example
* makeNativeString('foobar') // => `function foobar() { [native code] }`
*
* @param {string} [name] - Optional function name
*/
utils.makeNativeString = (name = '') => {
return utils.cache.nativeToStringStr.replace('toString', name || '');
};
/**
* Helper function to modify the `toString()` result of the provided object.
*
* Note: Use `utils.redirectToString` instead when possible.
*
* There's a quirk in JS Proxies that will cause the `toString()` result to differ from the vanilla Object.
* If no string is provided we will generate a `[native code]` thing based on the name of the property object.
*
* @example
* patchToString(WebGLRenderingContext.prototype.getParameter, 'function getParameter() { [native code] }')
*
* @param {object} obj - The object for which to modify the `toString()` representation
* @param {string} str - Optional string used as a return value
*/
utils.patchToString = (obj, str = '') => {
// toStringRedirects.set(obj, str);
Object.defineProperty(obj, 'toString', {
value: ()=> str,
enumerable: false,
writable: true,
configurable: true,
});
};
/**
* Make all nested functions of an object native.
*
* @param {object} obj
*/
utils.patchToStringNested = (obj = {}) => {
return utils.execRecursively(obj, ['function'], utils.patchToString);
};
/**
* Redirect toString requests from one object to another.
*
* @param {object} proxyObj - The object that toString will be called on
* @param {object} originalObj - The object which toString result we wan to return
*/
utils.redirectToString = (proxyObj, originalObj) => {
// toStringRedirects.set(proxyObj, originalObj);
Object.defineProperty(proxyObj, 'toString', {
value: ()=> originalObj.toString(),
enumerable: false,
writable: true,
configurable: true,
});
};
/**
* All-in-one method to replace a property with a JS Proxy using the provided Proxy handler with traps.
*
* Will stealthify these aspects (strip error stack traces, redirect toString, etc).
* Note: This is meant to modify native Browser APIs and works best with prototype objects.
*
* @example
* replaceWithProxy(WebGLRenderingContext.prototype, 'getParameter', proxyHandler)
*
* @param {object} obj - The object which has the property to replace
* @param {string} propName - The name of the property to replace
* @param {object} handler - The JS Proxy handler to use
*/
utils.replaceWithProxy = (obj, propName, handler) => {
const originalObj = obj[propName];
const proxyObj = new Proxy(obj[propName], utils.stripProxyFromErrors(handler));
utils.replaceProperty(obj, propName, { value: proxyObj });
utils.redirectToString(proxyObj, originalObj);
return true;
};
/**
* All-in-one method to replace a getter with a JS Proxy using the provided Proxy handler with traps.
*
* @example
* replaceGetterWithProxy(Object.getPrototypeOf(navigator), 'vendor', proxyHandler)
*
* @param {object} obj - The object which has the property to replace
* @param {string} propName - The name of the property to replace
* @param {object} handler - The JS Proxy handler to use
*/
utils.replaceGetterWithProxy = (obj, propName, handler) => {
const fn = Object.getOwnPropertyDescriptor(obj, propName).get;
const fnStr = fn.toString(); // special getter function string
const proxyObj = new Proxy(fn, utils.stripProxyFromErrors(handler));
utils.replaceProperty(obj, propName, { get: proxyObj });
utils.patchToString(proxyObj, fnStr);
return true;
};
/**
* All-in-one method to replace a getter and/or setter. Functions get and set
* of handler have one more argument that contains the native function.
*
* @example
* replaceGetterSetter(HTMLIFrameElement.prototype, 'contentWindow', handler)
*
* @param {object} obj - The object which has the property to replace
* @param {string} propName - The name of the property to replace
* @param {object} handlerGetterSetter - The handler with get and/or set
* functions
* @see https://developer.mozilla.org/docs/Web/JavaScript/Reference/Global_Objects/Object/defineProperty#description
*/
utils.replaceGetterSetter = (obj, propName, handlerGetterSetter) => {
const ownPropertyDescriptor = Object.getOwnPropertyDescriptor(obj, propName);
const handler = { ...ownPropertyDescriptor };
if (handlerGetterSetter.get !== undefined) {
const nativeFn = ownPropertyDescriptor.get;
handler.get = function () {
return handlerGetterSetter.get.call(this, nativeFn.bind(this));
};
utils.redirectToString(handler.get, nativeFn);
}
if (handlerGetterSetter.set !== undefined) {
const nativeFn = ownPropertyDescriptor.set;
handler.set = function (newValue) {
handlerGetterSetter.set.call(this, newValue, nativeFn.bind(this));
};
utils.redirectToString(handler.set, nativeFn);
}
Object.defineProperty(obj, propName, handler);
};
/**
* All-in-one method to mock a non-existing property with a JS Proxy using the provided Proxy handler with traps.
*
* Will stealthify these aspects (strip error stack traces, redirect toString, etc).
*
* @example
* mockWithProxy(chrome.runtime, 'sendMessage', function sendMessage() {}, proxyHandler)
*
* @param {object} obj - The object which has the property to replace
* @param {string} propName - The name of the property to replace or create
* @param {object} pseudoTarget - The JS Proxy target to use as a basis
* @param {object} handler - The JS Proxy handler to use
*/
utils.mockWithProxy = (obj, propName, pseudoTarget, handler) => {
const proxyObj = new Proxy(pseudoTarget, utils.stripProxyFromErrors(handler));
utils.replaceProperty(obj, propName, { value: proxyObj });
utils.patchToString(proxyObj);
return true;
};
/**
* All-in-one method to create a new JS Proxy with stealth tweaks.
*
* This is meant to be used whenever we need a JS Proxy but don't want to replace or mock an existing known property.
*
* Will stealthify certain aspects of the Proxy (strip error stack traces, redirect toString, etc).
*
* @example
* createProxy(navigator.mimeTypes.__proto__.namedItem, proxyHandler) // => Proxy
*
* @param {object} pseudoTarget - The JS Proxy target to use as a basis
* @param {object} handler - The JS Proxy handler to use
*/
utils.createProxy = (pseudoTarget, handler) => {
const proxyObj = new Proxy(pseudoTarget, utils.stripProxyFromErrors(handler));
utils.patchToString(proxyObj);
return proxyObj;
};
/**
* Helper function to split a full path to an Object into the first part and property.
*
* @example
* splitObjPath(`HTMLMediaElement.prototype.canPlayType`)
* // => {objName: "HTMLMediaElement.prototype", propName: "canPlayType"}
*
* @param {string} objPath - The full path to an object as dot notation string
*/
utils.splitObjPath = objPath => ({
// Remove last dot entry (property) ==> `HTMLMediaElement.prototype`
objName: objPath.split('.').slice(0, -1).join('.'),
// Extract last dot entry ==> `canPlayType`
propName: objPath.split('.').slice(-1)[0]
});
/**
* Convenience method to replace a property with a JS Proxy using the provided objPath.
*
* Supports a full path (dot notation) to the object as string here, in case that makes it easier.
*
* @example
* replaceObjPathWithProxy('WebGLRenderingContext.prototype.getParameter', proxyHandler)
*
* @param {string} objPath - The full path to an object (dot notation string) to replace
* @param {object} handler - The JS Proxy handler to use
*/
utils.replaceObjPathWithProxy = (objPath, handler) => {
const { objName, propName } = utils.splitObjPath(objPath);
const obj = eval(objName); // eslint-disable-line no-eval
return utils.replaceWithProxy(obj, propName, handler);
};
/**
* Traverse nested properties of an object recursively and apply the given function on a whitelist of value types.
*
* @param {object} obj
* @param {array} typeFilter - e.g. `['function']`
* @param {Function} fn - e.g. `utils.patchToString`
*/
utils.execRecursively = (obj = {}, typeFilter = [], fn) => {
function recurse(obj) {
for (const key in obj) {
if (obj[key] === undefined) {
continue;
}
if (obj[key] && typeof obj[key] === 'object') {
recurse(obj[key]);
} else {
if (obj[key] && typeFilter.includes(typeof obj[key])) {
fn.call(this, obj[key]);
}
}
}
}
recurse(obj);
return obj;
};
/**
* Everything we run through e.g. `page.evaluate` runs in the browser context, not the NodeJS one.
* That means we cannot just use reference variables and functions from outside code, we need to pass everything as a parameter.
*
* Unfortunately the data we can pass is only allowed to be of primitive types, regular functions don't survive the built-in serialization process.
* This utility function will take an object with functions and stringify them, so we can pass them down unharmed as strings.
*
* We use this to pass down our utility functions as well as any other functions (to be able to split up code better).
*
* @see utils.materializeFns
*
* @param {object} fnObj - An object containing functions as properties
*/
utils.stringifyFns = (fnObj = { hello: () => 'world' }) => {
// Object.fromEntries() ponyfill (in 6 lines) - supported only in Node v12+, modern browsers are fine
// https://github.com/feross/fromentries
function fromEntries(iterable) {
return [...iterable].reduce((obj, [key, val]) => {
obj[key] = val;
return obj;
}, {});
}
return (Object.fromEntries || fromEntries)(
Object.entries(fnObj)
.filter(([key, value]) => typeof value === 'function')
.map(([key, value]) => [key, value.toString()]) // eslint-disable-line no-eval
);
};
/**
* Utility function to reverse the process of `utils.stringifyFns`.
* Will materialize an object with stringified functions (supports classic and fat arrow functions).
*
* @param {object} fnStrObj - An object containing stringified functions as properties
*/
utils.materializeFns = (fnStrObj = { hello: "() => 'world'" }) => {
return Object.fromEntries(
Object.entries(fnStrObj).map(([key, value]) => {
if (value.startsWith('function')) {
// some trickery is needed to make oldschool functions work :-)
return [key, eval(`() => ${value}`)()]; // eslint-disable-line no-eval
} else {
// arrow functions just work
return [key, eval(value)]; // eslint-disable-line no-eval
}
})
);
};
// Proxy handler templates for re-usability
utils.makeHandler = () => ({
// Used by simple `navigator` getter evasions
getterValue: value => ({
apply(target, ctx, args) {
// Let's fetch the value first, to trigger and escalate potential errors
// Illegal invocations like `navigator.__proto__.vendor` will throw here
utils.cache.Reflect.apply(...arguments);
return value;
}
})
});
/**
* Compare two arrays.
*
* @param {array} array1 - First array
* @param {array} array2 - Second array
*/
utils.arrayEquals = (array1, array2) => {
if (array1.length !== array2.length) {
return false;
}
for (let i = 0; i < array1.length; ++i) {
if (array1[i] !== array2[i]) {
return false;
}
}
return true;
};
/**
* Cache the method return according to its arguments.
*
* @param {Function} fn - A function that will be cached
*/
utils.memoize = fn => {
const cache = [];
return function (...args) {
if (!cache.some(c => utils.arrayEquals(c.key, args))) {
cache.push({ key: args, value: fn.apply(this, args) });
}
return cache.find(c => utils.arrayEquals(c.key, args)).value;
};
};
utils.init();
if (!('WebGL2RenderingContext' in window)) {
return;
}
const getParameterProxyHandler = {
apply: function (target, ctx, args) {
const param = (args || [])[0]
const result = utils.cache.Reflect.apply(target, ctx, args)
// UNMASKED_VENDOR_WEBGL
if (param === 37445) {
return 'Intel Inc.' // default in headless: Google Inc.
}
// UNMASKED_RENDERER_WEBGL
if (param === 37446) {
return 'Intel Iris OpenGL Engine' // default in headless: Google SwiftShader
}
return result
}
}
// There's more than one WebGL rendering context
// https://developer.mozilla.org/en-US/docs/Web/API/WebGL2RenderingContext#Browser_compatibility
// To find out the original values here: Object.getOwnPropertyDescriptors(WebGLRenderingContext.prototype.getParameter)
const addProxy = (obj, propName) => {
utils.replaceWithProxy(obj, propName, getParameterProxyHandler)
}
// For whatever weird reason loops don't play nice with Object.defineProperty, here's the next best thing:
addProxy(WebGLRenderingContext.prototype, 'getParameter')
addProxy(WebGL2RenderingContext.prototype, 'getParameter')
}
src/services/misc.ts
+126
View File
@@ -0,0 +1,126 @@
import { singleton } from 'tsyringe';
import { AsyncService } from 'civkit/async-service';
import { ParamValidationError } from 'civkit/civ-rpc';
import { isIP } from 'node:net';
import { isIPInNonPublicRange } from '../utils/ip';
import { GlobalLogger } from './logger';
import { lookup } from 'node:dns/promises';
import { Threaded } from './threaded';
import { SecurityCompromiseError } from './errors';
import { GeoIPService } from './geoip';
import _ from 'lodash';
const normalizeUrl = require('@esm2cjs/normalize-url').default;
export const privateIpNotAcceptable = Boolean(process.env['NODE_ENV']?.toLowerCase()?.includes('prod') && process.env['GCLOUD_PROJECT']);
@singleton()
export class MiscService extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
constructor(
protected globalLogger: GlobalLogger,
protected geoIpService: GeoIPService,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.emit('ready');
}
@Threaded()
async assertNormalizedUrl(input: string) {
let result: URL;
try {
result = new URL(
normalizeUrl(
input,
{
stripWWW: false,
removeTrailingSlash: false,
removeSingleSlash: false,
sortQueryParameters: false,
}
)
);
} catch (err) {
throw new ParamValidationError({
message: `${err}`,
path: 'url'
});
}
if (!['http:', 'https:', 'blob:'].includes(result.protocol)) {
throw new ParamValidationError({
message: `Invalid protocol ${result.protocol}`,
path: 'url'
});
}
const normalizedHostname = result.hostname.startsWith('[') ? result.hostname.slice(1, -1) : result.hostname;
let ips: string[] = [];
const isIp = isIP(normalizedHostname);
if (isIp) {
ips.push(normalizedHostname);
}
if (
privateIpNotAcceptable &&
(result.hostname === 'localhost') ||
(isIp && isIPInNonPublicRange(normalizedHostname))
) {
this.logger.warn(`Suspicious action: Request to localhost or non-public IP: ${normalizedHostname}`, { href: result.href });
throw new SecurityCompromiseError({
message: `Suspicious action: Request to localhost or non-public IP: ${normalizedHostname}`,
path: 'url'
});
}
if (!isIp && result.protocol !== 'blob:') {
const resolved = await lookup(result.hostname, { all: true }).catch((err) => {
if (err.code === 'ENOTFOUND') {
return Promise.reject(new ParamValidationError({
message: `Domain '${result.hostname}' could not be resolved`,
path: 'url'
}));
}
return;
});
if (resolved) {
for (const x of resolved) {
if (privateIpNotAcceptable && isIPInNonPublicRange(x.address)) {
this.logger.warn(`Suspicious action: Domain resolved to non-public IP: ${result.hostname} => ${x.address}`, { href: result.href, ip: x.address });
throw new SecurityCompromiseError({
message: `Suspicious action: Domain resolved to non-public IP: ${x.address}`,
path: 'url'
});
}
ips.push(x.address);
}
}
}
let hintCountry: string | undefined;
if (ips.length) {
const hints = await this.geoIpService.lookupCities(ips);
const board: Record<string, number> = {};
for (const x of hints) {
if (x.country?.code) {
board[x.country.code] = (board[x.country.code] || 0) + 1;
}
}
hintCountry = _.maxBy(Array.from(Object.entries(board)), 1)?.[0];
}
return {
url: result,
ips,
hintCountry,
};
}
}
src/services/pdf-extract.ts
+396
View File
@@ -0,0 +1,396 @@
import { singleton } from 'tsyringe';
import _ from 'lodash';
import { TextItem } from 'pdfjs-dist/types/src/display/api';
import { AsyncService } from 'civkit/async-service';
import { GlobalLogger } from './logger';
import dayjs from 'dayjs';
import type { PDFDocumentLoadingTask } from 'pdfjs-dist';
import path from 'path';
import { AsyncLocalContext } from './async-context';
import { CanvasService } from './canvas';
import { writeFile } from 'fs/promises';
import { pathToFileURL } from 'url';
const utc = require('dayjs/plugin/utc'); // Import the UTC plugin
dayjs.extend(utc); // Extend dayjs with the UTC plugin
const timezone = require('dayjs/plugin/timezone');
dayjs.extend(timezone);
const pPdfjs = import('pdfjs-dist/legacy/build/pdf.mjs');
const nodeCmapUrl = path.resolve(require.resolve('pdfjs-dist'), '../../cmaps') + '/';
// const standardFontsUrl = path.resolve(require.resolve('pdfjs-dist'), '../../standard_fonts') + '/';
function stdDev(numbers: number[]) {
const mean = _.mean(numbers);
const squareDiffs = numbers.map((num) => Math.pow(num - mean, 2));
const avgSquareDiff = _.mean(squareDiffs);
return Math.sqrt(avgSquareDiff);
}
function isRotatedByAtLeast35Degrees(transform?: [number, number, number, number, number, number]): boolean {
if (!transform) {
return false;
}
const [a, b, c, d, _e, _f] = transform;
// Calculate the rotation angles using arctan(b/a) and arctan(-c/d)
const angle1 = Math.atan2(b, a) * (180 / Math.PI); // from a, b
const angle2 = Math.atan2(-c, d) * (180 / Math.PI); // from c, d
// Either angle1 or angle2 can be used to determine the rotation, they should be equivalent
const rotationAngle1 = Math.abs(angle1);
const rotationAngle2 = Math.abs(angle2);
// Check if the absolute rotation angle is greater than or equal to 35 degrees
return rotationAngle1 >= 35 || rotationAngle2 >= 35;
}
export interface ExtractedPDF {
meta?: Record<string, any>;
content: string;
text: string;
cached?: boolean;
}
@singleton()
export class PDFExtractor extends AsyncService {
logger = this.globalLogger.child({ service: this.constructor.name });
pdfjs!: Awaited<typeof pPdfjs>;
cacheRetentionMs = 1000 * 3600 * 24 * 7;
constructor(
protected globalLogger: GlobalLogger,
protected asyncLocalContext: AsyncLocalContext,
protected canvasService: CanvasService,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
this.pdfjs = await pPdfjs;
this.emit('ready');
}
isDataUrl(url: string) {
return url.startsWith('data:');
}
parseDataUrl(url: string) {
const protocol = url.slice(0, url.indexOf(':'));
const contentType = url.slice(url.indexOf(':') + 1, url.indexOf(';'));
const data = url.slice(url.indexOf(',') + 1);
if (protocol !== 'data' || !data) {
throw new Error('Invalid data URL');
}
if (contentType !== 'application/pdf') {
throw new Error('Invalid data URL type');
}
return {
type: contentType,
data: data
};
}
calcArticleHeightsStats(textItems: TextItem[][]) {
const articleCharHeights: number[] = [];
for (const textItem of textItems.flat()) {
if (textItem.height) {
articleCharHeights.push(...Array(textItem.str.length).fill(textItem.height));
}
}
return {
articleAvgHeight: _.mean(articleCharHeights),
articleStdDevHeight: stdDev(articleCharHeights)
};
}
digestTextItems(textItems: TextItem[][], options?: ReturnType<typeof this.calcArticleHeightsStats>) {
const { articleAvgHeight, articleStdDevHeight } = options || this.calcArticleHeightsStats(textItems);
// const articleMedianHeight = articleCharHeights.sort()[Math.floor(articleCharHeights.length / 2)];
const mdOps: Array<{
text: string;
op?: 'new' | 'append';
mode: 'h1' | 'h2' | 'p' | 'appendix' | 'space';
}> = [];
const rawChunks: string[] = [];
let op: 'append' | 'new' = 'new';
let mode: 'h1' | 'h2' | 'p' | 'space' | 'appendix' = 'p';
for (const pageTextItems of textItems) {
const charHeights = [];
for (const textItem of pageTextItems as TextItem[]) {
if (textItem.height) {
charHeights.push(...Array(textItem.str.length).fill(textItem.height));
}
rawChunks.push(`${textItem.str || ''}${textItem.hasEOL ? '\n' : ''}`);
}
const avgHeight = _.mean(charHeights);
const stdDevHeight = stdDev(charHeights);
// const medianHeight = charHeights.sort()[Math.floor(charHeights.length / 2)];
for (const textItem of pageTextItems) {
if (textItem.height > articleAvgHeight + 3 * articleStdDevHeight) {
mode = 'h1';
} else if (textItem.height > articleAvgHeight + 2 * articleStdDevHeight) {
mode = 'h2';
} else if (textItem.height && textItem.height < avgHeight - stdDevHeight) {
mode = 'appendix';
} else if (textItem.height) {
mode = 'p';
} else {
mode = 'space';
}
if (isRotatedByAtLeast35Degrees(textItem.transform as any)) {
mode = 'appendix';
}
mdOps.push({
op,
mode,
text: textItem.str
});
if (textItem.hasEOL && !textItem.str) {
op = 'new';
} else {
op = 'append';
}
}
}
const mdChunks = [];
const appendixChunks = [];
mode = 'space';
for (const x of mdOps) {
const previousMode: string = mode;
const changeToMdChunks = [];
const isNewStart = x.mode !== 'space' && (x.op === 'new' || (previousMode === 'appendix' && x.mode !== previousMode));
if (isNewStart) {
switch (x.mode) {
case 'h1': {
changeToMdChunks.push(`\n\n# `);
mode = x.mode;
break;
}
case 'h2': {
changeToMdChunks.push(`\n\n## `);
mode = x.mode;
break;
}
case 'p': {
changeToMdChunks.push(`\n\n`);
mode = x.mode;
break;
}
case 'appendix': {
mode = x.mode;
appendixChunks.push(`\n\n`);
break;
}
default: {
break;
}
}
} else {
if (x.mode === 'appendix' && appendixChunks.length) {
const lastChunk = appendixChunks[appendixChunks.length - 1];
if (!lastChunk.match(/(\s+|-)$/) && lastChunk.length !== 1) {
appendixChunks.push(' ');
}
} else if (mdChunks.length) {
const lastChunk = mdChunks[mdChunks.length - 1];
if (!lastChunk.match(/(\s+|-)$/) && lastChunk.length !== 1) {
changeToMdChunks.push(' ');
}
}
}
if (x.text) {
if (x.mode == 'appendix') {
if (appendixChunks.length || isNewStart) {
appendixChunks.push(x.text);
} else {
changeToMdChunks.push(x.text);
}
} else {
changeToMdChunks.push(x.text);
}
}
if (isNewStart && x.mode !== 'appendix' && appendixChunks.length) {
const appendix = appendixChunks.join('').split(/\r?\n/).map((x) => x.trim()).filter(Boolean).map((x) => `> ${x}`).join('\n');
changeToMdChunks.unshift(appendix);
changeToMdChunks.unshift(`\n\n`);
appendixChunks.length = 0;
}
if (x.mode === 'space' && changeToMdChunks.length) {
changeToMdChunks.length = 1;
}
if (changeToMdChunks.length) {
mdChunks.push(...changeToMdChunks);
}
}
if (mdChunks.length) {
mdChunks[0] = mdChunks[0].trimStart();
}
return { content: mdChunks.join(''), text: rawChunks.join('') } as ExtractedPDF;
}
async extract(url: string | URL) {
let loadingTask: PDFDocumentLoadingTask;
if (typeof url === 'string' && this.isDataUrl(url)) {
const { data } = this.parseDataUrl(url);
const binary = Uint8Array.from(Buffer.from(data, 'base64'));
loadingTask = this.pdfjs.getDocument({
data: binary,
// disableFontFace: true,
useSystemFonts: true,
verbosity: 0,
cMapUrl: nodeCmapUrl,
// standardFontDataUrl: standardFontsUrl,
});
} else {
loadingTask = this.pdfjs.getDocument({
url,
// disableFontFace: true,
useSystemFonts: true,
verbosity: 0,
cMapUrl: nodeCmapUrl,
// standardFontDataUrl: standardFontsUrl,
});
}
const doc = await loadingTask.promise;
const meta = await doc.getMetadata();
const textItems: TextItem[][] = [];
for (const pg of _.range(0, doc.numPages)) {
const page = await doc.getPage(pg + 1);
const textContent = await page.getTextContent({ includeMarkedContent: true });
textItems.push((textContent.items as TextItem[]));
}
return { meta: meta.info as Record<string, any>, ...this.digestTextItems(textItems) } as ExtractedPDF;
}
async extractRendered(filePath: string, targetPath: string, pagesBeingRendered: number[] | Set<number> = [1, 2, 3]) {
const fileUrl = pathToFileURL(filePath).href;
const pageRenderSet = new Set(pagesBeingRendered);
const loadingTask = this.pdfjs.getDocument({
url: fileUrl,
// disableFontFace: true,
useSystemFonts: true,
verbosity: 0,
cMapUrl: nodeCmapUrl,
});
const doc = await loadingTask.promise;
const meta = await doc.getMetadata();
const vecs: {
page: number;
width?: number;
height?: number;
pngPath?: string;
textItems: TextItem[];
}[] = [];
for (const pg of _.range(0, doc.numPages)) {
const pn = pg + 1;
const page = await doc.getPage(pn);
const mixin: Partial<typeof vecs[number]> = {};
if (pageRenderSet.has(pn)) {
const viewport = page.getViewport({ scale: 2, });
const canvas = this.canvasService.canvas.createCanvas(viewport.width, viewport.height);
const c2d = canvas.getContext('2d');
await page.render({
// @ts-ignore
canvasContext: c2d,
viewport: viewport,
}).promise;
const pngBuffer = canvas.toBuffer('image/png');
const fPath = path.join(targetPath, `_${pn}.png`);
await writeFile(fPath, pngBuffer, { flush: true });
Object.assign(mixin, {
width: viewport.width,
height: viewport.height,
pngPath: fPath,
});
}
const textContent = await page.getTextContent({ includeMarkedContent: true });
const vec = {
...mixin,
page: pn,
textItems: textContent.items as TextItem[],
};
vecs.push(vec);
}
const textItems = vecs.map((x) => x.textItems);
const articleStats = this.calcArticleHeightsStats(textItems);
return {
...this.digestTextItems(textItems, articleStats),
meta: meta.info as Record<string, any>,
pages: vecs.map((x) => {
return {
...this.digestTextItems([x.textItems], articleStats),
...x,
};
})
};
}
parsePdfDate(pdfDate: string | undefined) {
if (!pdfDate) {
return undefined;
}
// Remove the 'D:' prefix
const cleanedDate = pdfDate.slice(2);
// Define the format without the timezone part first
const dateTimePart = cleanedDate.slice(0, 14);
const timezonePart = cleanedDate.slice(14);
// Construct the full date string in a standard format
const formattedDate = `${dateTimePart}${timezonePart.replace("'", "").replace("'", "")}`;
// Parse the date with timezone
const parsedDate = dayjs(formattedDate, "YYYYMMDDHHmmssZ");
const date = parsedDate.toDate();
if (!date.valueOf()) {
return undefined;
}
return date;
}
}
src/services/proxy-provider/brightdata.ts
+63
View File
@@ -0,0 +1,63 @@
import { randomPick } from 'civkit/random';
import { randomBytes } from 'crypto';
const TIER_1_ENGLISH_SPEAKING_COUNTRIES = ['us', 'ca', 'gb', 'au', 'nz', 'sg'];
export class BrightDataProxyProvider {
baseProxyUrl!: URL;
constructor(residential: string) {
this.baseProxyUrl = new URL(residential);
}
supports(countryCode: string) {
if (['cn', 'hk', 'mo'].includes(countryCode.toLowerCase())) {
return true;
}
if (TIER_1_ENGLISH_SPEAKING_COUNTRIES.includes(countryCode.toLowerCase())) {
return true;
}
return false;
}
async alloc(inputCountryCode?: string) {
if (!this.baseProxyUrl) {
throw new Error(`Proxy provider not available`);
}
if (inputCountryCode === 'none') {
throw new Error(`Proxy opted-out`);
}
let countryCode = inputCountryCode?.toLowerCase() || 'any';
if (countryCode === 'true' || countryCode === 'auto') {
countryCode = 'any';
}
if (countryCode === 'any') {
countryCode = randomPick(TIER_1_ENGLISH_SPEAKING_COUNTRIES);
}
const user = this.getUserName(countryCode, this.baseProxyUrl.username);
const url = new URL(this.baseProxyUrl);
url.username = user;
return url;
}
protected getUserName(countryCode: string, username: string) {
const dyNameChunks = ['brd-customer', username];
if (countryCode !== 'any') {
dyNameChunks.push(`country-${countryCode}`);
}
dyNameChunks.push(`session-${randomBytes(16).toString('hex')}`);
dyNameChunks.push('route_err-block');
return dyNameChunks.join('-');
}
}
src/services/proxy-provider/index.ts
+123
View File
@@ -0,0 +1,123 @@
import { singleton } from 'tsyringe';
import { EnvConfig } from '../envconfig';
import { ThorDataProxyProvider } from './thordata';
// import { BrightDataProxyProvider } from './brightdata';
import { AsyncService } from 'civkit/async-service';
import { randomPick } from 'civkit/random';
import { ServiceBadApproachError } from '../errors';
const TIER_1_ENGLISH_SPEAKING_COUNTRIES = ['us', 'ca', 'gb', 'au', 'nz', 'sg'];
interface ProxyProvider {
alloc(inputCountryCode?: string): Promise<URL>;
supports(countryCode: string): boolean;
}
@singleton()
export class ProxyProviderService extends AsyncService {
clients: ProxyProvider[] = [];
constructor(
protected envConfig: EnvConfig,
) {
super(...arguments);
}
override async init() {
await this.dependencyReady();
if (this.envConfig.THORDATA_PROXY_URL) {
this.clients.push(new ThorDataProxyProvider(
this.envConfig.THORDATA_PROXY_URL,
this.envConfig.THORDATA_PROXY_URL_ALT // optional
));
}
// if (this.envConfig.BRIGHTDATA_PROXY_URL) {
// this.clients.push(new BrightDataProxyProvider(this.envConfig.BRIGHTDATA_PROXY_URL));
// }
// if (this.envConfig.BRIGHTDATA_ISP_PROXY_URL) {
// this.clients.push(new BrightDataProxyProvider(this.envConfig.BRIGHTDATA_ISP_PROXY_URL));
// }
this.emit('ready');
}
supports(countryCode: string) {
const normalized = countryCode.toLowerCase();
return this.clients.some((client) => client.supports(normalized));
}
async alloc(inputCountryCode?: string) {
let countryCode = inputCountryCode?.toLowerCase() || 'any';
if (countryCode === 'true' || countryCode === 'auto') {
countryCode = 'any';
}
if (countryCode === 'any') {
countryCode = process.env.PREFERRED_PROXY_COUNTRY || randomPick(TIER_1_ENGLISH_SPEAKING_COUNTRIES);
}
const someClient = this.clients.find((client) => client.supports(countryCode));
if (!someClient) {
throw new ServiceBadApproachError(`No proxy provider found that supports country code: ${countryCode}.`);
}
return someClient.alloc(countryCode);
}
*loopClients() {
if (!this.clients.length) {
throw new ServiceBadApproachError('No proxy provider found.');
}
while (true) {
yield* this.clients;
}
}
async *iterAlloc(inputCountryCode?: string, maxAttempt = Infinity): AsyncGenerator<URL> {
let countryCode = inputCountryCode?.toLowerCase() || 'any';
if (countryCode === 'true' || countryCode === 'auto') {
countryCode = 'any';
}
if (countryCode === 'any') {
countryCode = randomPick(TIER_1_ENGLISH_SPEAKING_COUNTRIES);
}
if (!this.supports(countryCode)) {
return;
}
let i = 0;
for (const client of this.loopClients()) {
if (i > maxAttempt) {
break;
}
if (!client.supports(countryCode)) {
continue;
}
const url = await client.alloc(inputCountryCode);
yield url;
i += 1;
}
}
}
@singleton()
export class SERPProxyProviderService extends ProxyProviderService {
override async init() {
await this.dependencyReady();
if (this.envConfig.THORDATA_PROXY_URL) {
this.clients.push(new ThorDataProxyProvider(
this.envConfig.THORDATA_PROXY_URL,
this.envConfig.THORDATA_PROXY_URL_ALT // optional
));
}
this.emit('ready');
}
}

Some files were not shown because too many files have changed in this diff Show More