mirror of
https://github.com/run-llama/llama_cloud_services.git
synced 2026-07-20 00:54:09 -04:00
Compare commits
27 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 82c0022aca | |||
| efa267cb7f | |||
| 75fee8d600 | |||
| 31f54bca55 | |||
| b1ae7bb736 | |||
| 31fe12e0da | |||
| 90b0c5e295 | |||
| 79fe1930cf | |||
| ab225c3eab | |||
| 6f1de75909 | |||
| 230ed64e41 | |||
| ef126c3a93 | |||
| 51a7534733 | |||
| 4f5d2bde13 | |||
| 3d05fe5d77 | |||
| c16ca673af | |||
| 6619034bce | |||
| c56fb5d8f7 | |||
| b407a5edb5 | |||
| e6a27d17fb | |||
| 34077fd479 | |||
| 7a68ad5a7f | |||
| 74a1b6c2f2 | |||
| 9a90ae5264 | |||
| 310c1bc105 | |||
| cd20b29299 | |||
| 0cb7aeb81c |
@@ -0,0 +1,95 @@
|
||||
name: Claude Code
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
pull_request_review_comment:
|
||||
types: [created]
|
||||
issues:
|
||||
types: [opened, assigned]
|
||||
pull_request_review:
|
||||
types: [submitted]
|
||||
|
||||
jobs:
|
||||
claude:
|
||||
if: |
|
||||
(github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
|
||||
(github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
|
||||
(github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) ||
|
||||
(github.event_name == 'issues' && (contains(github.event.issue.body, '@claude') || contains(github.event.issue.title, '@claude')))
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
issues: read
|
||||
id-token: write
|
||||
steps:
|
||||
- name: Check repository access
|
||||
id: check-access
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
# Get the user who triggered the event
|
||||
case "${{ github.event_name }}" in
|
||||
"issue_comment")
|
||||
USER="${{ github.event.comment.user.login }}"
|
||||
;;
|
||||
"pull_request_review_comment")
|
||||
USER="${{ github.event.comment.user.login }}"
|
||||
;;
|
||||
"pull_request_review")
|
||||
USER="${{ github.event.review.user.login }}"
|
||||
;;
|
||||
"issues")
|
||||
USER="${{ github.event.issue.user.login }}"
|
||||
;;
|
||||
esac
|
||||
|
||||
echo "Checking repository access for user: $USER"
|
||||
|
||||
# Check if user has write access to the repository
|
||||
REPO="${{ github.repository }}"
|
||||
if gh api repos/$REPO/collaborators/$USER/permission --jq '.permission' | grep -E "(admin|write)" > /dev/null 2>&1; then
|
||||
echo "User $USER has write access to the repository"
|
||||
echo "authorized=true" >> $GITHUB_OUTPUT
|
||||
else
|
||||
echo "User $USER does not have write access to the repository"
|
||||
echo "authorized=false" >> $GITHUB_OUTPUT
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Checkout repository
|
||||
if: steps.check-access.outputs.authorized == 'true'
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 1
|
||||
|
||||
- name: Run Claude Code
|
||||
if: steps.check-access.outputs.authorized == 'true'
|
||||
id: claude
|
||||
uses: anthropics/claude-code-action@beta
|
||||
with:
|
||||
anthropic_api_key: ${{ secrets.ANTHROPIC_GITHUB_API_KEY }}
|
||||
|
||||
# Optional: Specify model (defaults to Claude Sonnet 4, uncomment for Claude Opus 4)
|
||||
# model: "claude-opus-4-20250514"
|
||||
|
||||
# Optional: Customize the trigger phrase (default: @claude)
|
||||
# trigger_phrase: "/claude"
|
||||
|
||||
# Optional: Trigger when specific user is assigned to an issue
|
||||
# assignee_trigger: "claude-bot"
|
||||
|
||||
# Optional: Allow Claude to run specific commands
|
||||
# Allow bash commands to be run, for things like running tests, linting, etc.
|
||||
allowed_tools: "Bash(rg:*),Bash(find:*),Bash(grep:*),Bash(pnpm:*),Bash(npm:*),Bash(uv:*),Bash(pip:*),Bash(pipx:*),Bash(make:*),Bash(cd:*),WebFetch"
|
||||
|
||||
# Optional: Add custom instructions for Claude to customize its behavior for your project
|
||||
# custom_instructions: |
|
||||
# Follow our coding standards
|
||||
# Ensure all new code has tests
|
||||
# Use TypeScript for new files
|
||||
|
||||
# Optional: Custom environment variables for Claude
|
||||
# claude_env: |
|
||||
# NODE_ENV: test
|
||||
@@ -28,7 +28,6 @@ jobs:
|
||||
with:
|
||||
node-version-file: "ts/llama_cloud_services/.nvmrc"
|
||||
- name: Install dependencies
|
||||
working-directory: ts/llama_cloud_services/
|
||||
run: pnpm install --no-frozen-lockfile
|
||||
- name: Run lint
|
||||
working-directory: ts/llama_cloud_services/
|
||||
|
||||
@@ -22,9 +22,12 @@ jobs:
|
||||
node-version-file: "ts/llama_cloud_services/.nvmrc"
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ts/llama_cloud_services
|
||||
run: pnpm install --no-frozen-lockfile
|
||||
|
||||
- name: Run Build
|
||||
working-directory: ts/llama_cloud_services/
|
||||
run: pnpm build
|
||||
|
||||
- name: Build tarball
|
||||
run: |
|
||||
pnpm pack
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
name: Test end-to-end - Python
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
paths:
|
||||
- "py/**"
|
||||
|
||||
env:
|
||||
UV_VERSION: "0.7.20"
|
||||
LLAMA_CLOUD_API_KEY: ${{ secrets.LLAMA_CLOUD_API_KEY }}
|
||||
|
||||
jobs:
|
||||
test_e2e:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
# You can use PyPy versions in python-version.
|
||||
# For example, pypy-2.7 and pypy-3.8
|
||||
matrix:
|
||||
python-version: ["3.12"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
with:
|
||||
version: ${{ env.UV_VERSION }}
|
||||
|
||||
- name: Set up Python
|
||||
run: uv python install ${{ matrix.python-version }} && uv python pin ${{ matrix.python-version }}
|
||||
|
||||
- name: Run Tests
|
||||
working-directory: py
|
||||
run: make e2e
|
||||
|
||||
- name: Remove virtual environment
|
||||
working-directory: py
|
||||
run: rm -rf .venv/
|
||||
@@ -12,7 +12,6 @@ on:
|
||||
|
||||
env:
|
||||
UV_VERSION: "0.7.20"
|
||||
LLAMA_CLOUD_API_KEY: ${{ secrets.LLAMA_CLOUD_API_KEY }}
|
||||
|
||||
jobs:
|
||||
test:
|
||||
@@ -36,7 +35,7 @@ jobs:
|
||||
|
||||
- name: Run Tests
|
||||
working-directory: py
|
||||
run: uv run -- pytest tests/**/test_*.py
|
||||
run: uv run pytest unit_tests/ -v
|
||||
|
||||
- name: Remove virtual environment
|
||||
working-directory: py
|
||||
|
||||
@@ -30,8 +30,13 @@ jobs:
|
||||
with:
|
||||
node-version-file: "ts/llama_cloud_services/.nvmrc"
|
||||
- name: Install dependencies
|
||||
working-directory: ts/llama_cloud_services/
|
||||
run: pnpm install --no-frozen-lockfile
|
||||
- name: Run tests
|
||||
- name: Run Build
|
||||
working-directory: ts/llama_cloud_services/
|
||||
run: pnpm build
|
||||
- name: Run Tests
|
||||
working-directory: ts/llama_cloud_services/
|
||||
run: pnpm test
|
||||
- name: Run e2e tests
|
||||
working-directory: ts/e2e-tests/
|
||||
run: pnpm test
|
||||
|
||||
@@ -15,6 +15,7 @@ repos:
|
||||
- id: end-of-file-fixer
|
||||
- id: mixed-line-ending
|
||||
- id: trailing-whitespace
|
||||
exclude: ^ts/llama_cloud_services/src/client/
|
||||
- repo: https://github.com/charliermarsh/ruff-pre-commit
|
||||
rev: v0.1.5
|
||||
|
||||
@@ -33,7 +34,7 @@ repos:
|
||||
rev: v1.0.1
|
||||
hooks:
|
||||
- id: mypy
|
||||
exclude: ^py/tests/
|
||||
exclude: ^py/tests|^py/unit_tests
|
||||
additional_dependencies:
|
||||
[
|
||||
"types-requests",
|
||||
@@ -63,13 +64,13 @@ repos:
|
||||
rev: v3.0.3
|
||||
hooks:
|
||||
- id: prettier
|
||||
exclude: ^(uv.lock|ts/llama_cloud_services/pnpm-lock.yaml)
|
||||
exclude: ^(uv.lock|ts/llama_cloud_services/pnpm-lock.yaml|ts/e2e-tests)
|
||||
- repo: https://github.com/codespell-project/codespell
|
||||
rev: v2.2.6
|
||||
hooks:
|
||||
- id: codespell
|
||||
additional_dependencies: [tomli]
|
||||
exclude: ^(uv.lock|docs|ts|examples)
|
||||
exclude: ^(uv.lock|docs|ts|examples|pnpm-lock.yaml)
|
||||
args:
|
||||
[
|
||||
"--ignore-words-list",
|
||||
@@ -86,4 +87,4 @@ repos:
|
||||
- id: toml-sort-fix
|
||||
exclude: ".*uv.lock"
|
||||
|
||||
exclude: .github/ISSUE_TEMPLATE
|
||||
exclude: ^(.github/ISSUE_TEMPLATE|ts/llama_cloud_services/src/client|pnpm-lock.yaml)
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
# LlamaCloud Services Examples - Python
|
||||
|
||||
In this folder you will find several TypeScript end-to-end applications that contain examples regarding:
|
||||
|
||||
- [LlamaParse](./parse/)
|
||||
- [LlamaCloud Index](./index/)
|
||||
|
||||
Follow the instructions in each example folder to get started!
|
||||
@@ -0,0 +1,122 @@
|
||||
# LlamaExtract Demo
|
||||
|
||||
A TypeScript demo application showcasing the power of **LlamaExract** - a structured data extraction agentic service from [LlamaCloud](https://cloud.llamaindex.ai). This demo allows you to extract structured information from scientific papers and get them into a nice markdown format.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Features](#features)
|
||||
- [Prerequisites](#prerequisites)
|
||||
- [Installation](#installation)
|
||||
- [Usage](#usage)
|
||||
- [Start the Demo](#start-the-demo)
|
||||
- [Development Mode](#development-mode)
|
||||
- [Build the Project](#build-the-project)
|
||||
- [Code Quality](#code-quality)
|
||||
- [Quick Commands Reference](#quick-commands-reference)
|
||||
- [How It Works](#how-it-works)
|
||||
- [API Dependencies](#api-dependencies)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [Common Issues](#common-issues)
|
||||
- [License](#license)
|
||||
- [Contributing](#contributing)
|
||||
|
||||
## Features
|
||||
|
||||
- 📄 **Structured Data Extraction**: Extract data from your files effortlessly, and structure them the way you want!
|
||||
- 🤖 **Markdown Rendering**: Generate markdown directly from your extracted data
|
||||
- 🎨 **Beautiful CLI**: Styled console interface with colors and ASCII art
|
||||
- ⚡ **Fast Development**: Hot reload support with watch mode
|
||||
- 🛠️ **TypeScript**: Full TypeScript support with strict type checking
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Node.js (version 18 or higher)
|
||||
- pnpm package manager
|
||||
- LlamaCloud API key
|
||||
|
||||
## Installation
|
||||
|
||||
1. Clone the repository:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/run-llama/llama_cloud_services
|
||||
cd lama_cloud_services/examples-ts/extract/
|
||||
```
|
||||
|
||||
2. Install dependencies:
|
||||
|
||||
```bash
|
||||
npm install
|
||||
```
|
||||
|
||||
3. Set up your environment variables:
|
||||
|
||||
```bash
|
||||
# Add your API key to your environment
|
||||
export LLAMA_CLOUD_API_KEY="your-llamacloud-api-key"
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Start the Demo
|
||||
|
||||
```bash
|
||||
npm run start
|
||||
```
|
||||
|
||||
The application will display a welcome screen and prompt you to enter the path to a document you'd like to process.
|
||||
|
||||
### Development Mode
|
||||
|
||||
For development with hot reload:
|
||||
|
||||
```bash
|
||||
npm run dev
|
||||
```
|
||||
|
||||
### Build the Project
|
||||
|
||||
```bash
|
||||
npm run build
|
||||
```
|
||||
|
||||
### Code Quality
|
||||
|
||||
Format code:
|
||||
|
||||
```bash
|
||||
npm run format
|
||||
```
|
||||
|
||||
Lint code:
|
||||
|
||||
```bash
|
||||
npm run lint
|
||||
```
|
||||
|
||||
## How It Works
|
||||
|
||||
1. **Document Input**: Enter the path to your document when prompted
|
||||
2. **Parsing**: LlamaExtract, based on the schema you can find [here](./src/schema.ts), processes the document and extracts structured data
|
||||
3. **Markdown Rendering**: The extracted content is rendered into beautiful markdown
|
||||
4. **Results**: View the results directly in your terminal
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
1. **Module Resolution Errors**: Ensure you're using Node.js 18+ and have all dependencies installed
|
||||
2. **API Key Issues**: Verify your LlamaCloud API key is correctly set
|
||||
3. **File Path Errors**: Use absolute paths or ensure relative paths are correct from the project root
|
||||
|
||||
## License
|
||||
|
||||
MIT License - see the [LICENSE](../../LICENSE) file for details.
|
||||
|
||||
## Contributing
|
||||
|
||||
1. Fork the repository
|
||||
2. Create a feature branch
|
||||
3. Make your changes
|
||||
4. Run `npm run format` and `npm run lint`
|
||||
5. Submit a pull request
|
||||
Binary file not shown.
@@ -0,0 +1,14 @@
|
||||
import js from "@eslint/js";
|
||||
import globals from "globals";
|
||||
import tseslint from "typescript-eslint";
|
||||
import { defineConfig } from "eslint/config";
|
||||
|
||||
export default defineConfig([
|
||||
{
|
||||
files: ["**/*.{js,mjs,cjs,ts,mts,cts}"],
|
||||
plugins: { js },
|
||||
extends: ["js/recommended"],
|
||||
languageOptions: { globals: globals.browser },
|
||||
},
|
||||
tseslint.configs.recommended,
|
||||
]);
|
||||
Generated
+3276
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,37 @@
|
||||
{
|
||||
"name": "llama-extract-demo",
|
||||
"version": "0.1.0",
|
||||
"description": "Demo for LlamaExtract in TypeScript",
|
||||
"main": "index.js",
|
||||
"scripts": {
|
||||
"test": "echo \"There are no tests\"",
|
||||
"start": "npm exec tsx src/index.ts",
|
||||
"lint": "eslint ./src/",
|
||||
"format": "prettier --write ./src/",
|
||||
"build": "tsc",
|
||||
"dev": "npm exec tsx --watch src/index.ts"
|
||||
},
|
||||
"author": "LlamaIndex",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"cli-markdown": "^3.5.1",
|
||||
"consola": "^3.4.2",
|
||||
"figlet": "^1.8.2",
|
||||
"llama-cloud-services": "file:../../ts/llama_cloud_services",
|
||||
"marked": "^15.0.12",
|
||||
"marked-terminal": "^7.3.0",
|
||||
"picocolors": "^1.1.1"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@eslint/js": "^9.32.0",
|
||||
"@types/figlet": "^1.7.0",
|
||||
"@types/marked-terminal": "^6.1.1",
|
||||
"@types/node": "^24.2.0",
|
||||
"eslint": "^9.32.0",
|
||||
"globals": "^16.3.0",
|
||||
"jiti": "^2.5.1",
|
||||
"prettier": "^3.6.2",
|
||||
"typescript": "^5.9.2",
|
||||
"typescript-eslint": "^8.39.0"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
import { LlamaExtract, ExtractConfig } from "llama-cloud-services";
|
||||
import cliMarkdown from "cli-markdown";
|
||||
import { logger } from "./logger";
|
||||
import pc from "picocolors";
|
||||
import { consoleInput, renderLogo } from "./utils";
|
||||
import { dataSchema } from "./schema";
|
||||
import { renderMarkdown, ResearchData } from "./markdown";
|
||||
|
||||
export async function main(): Promise<number> {
|
||||
const extractClient = new LlamaExtract(
|
||||
process.env.LLAMA_CLOUD_API_KEY!,
|
||||
"https://api.cloud.llamaindex.ai",
|
||||
);
|
||||
await renderLogo();
|
||||
logger.log(
|
||||
`Welcome to ${pc.bold(
|
||||
pc.magentaBright("LlamaExtract Demo✨"),
|
||||
)}, our demo for ${pc.bold(pc.green("LlamaExtract"))}, a ${pc.bold(
|
||||
pc.cyan("LlamaCloud☁️"),
|
||||
)} (https://cloud.llamaindex.ai) product!.\nIn this demo we are going to try extracting relevant information ${pc.bold(
|
||||
pc.yellowBright("from scientific papers"),
|
||||
)}. Type the path to the paper you would like to process below👇\nIf you wish to exit, just type ${pc.bold(
|
||||
pc.gray("quit"),
|
||||
)}.\n`,
|
||||
);
|
||||
while (true) {
|
||||
const userInput = await consoleInput();
|
||||
if (userInput.toLowerCase() == "quit") {
|
||||
break;
|
||||
}
|
||||
try {
|
||||
const generatedData = await extractClient.extract(
|
||||
dataSchema,
|
||||
{} as ExtractConfig,
|
||||
userInput,
|
||||
);
|
||||
const research = renderMarkdown(generatedData?.data as ResearchData); // Added await here
|
||||
logger.log(`${pc.bold(pc.cyan("Extracted information:✨"))}:\n`);
|
||||
logger.log(cliMarkdown(research));
|
||||
} catch (error) {
|
||||
logger.error(`Error processing file: ${error}`);
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
@@ -0,0 +1,172 @@
|
||||
type Author = {
|
||||
name: string;
|
||||
affiliation?: string;
|
||||
email?: string;
|
||||
};
|
||||
|
||||
type Methodology = {
|
||||
approach?: string;
|
||||
participants?: string;
|
||||
methods?: string[];
|
||||
};
|
||||
|
||||
type Result = {
|
||||
finding?: string;
|
||||
significance?: string;
|
||||
supportingData?: string;
|
||||
};
|
||||
|
||||
type Reference = {
|
||||
title: string;
|
||||
authors: string;
|
||||
year?: string;
|
||||
relevance?: string;
|
||||
};
|
||||
|
||||
type Discussion = {
|
||||
implications?: string[];
|
||||
limitations?: string[];
|
||||
futureWork?: string[];
|
||||
};
|
||||
|
||||
type Publication = {
|
||||
journal?: string;
|
||||
year: string;
|
||||
doi?: string;
|
||||
url?: string;
|
||||
};
|
||||
|
||||
export type ResearchData = {
|
||||
title: string;
|
||||
authors: Author[];
|
||||
abstract: string;
|
||||
keywords?: string[];
|
||||
mainFindings: string[];
|
||||
methodology?: Methodology;
|
||||
results?: Result[];
|
||||
discussion?: Discussion;
|
||||
references?: Reference[];
|
||||
publication?: Publication;
|
||||
};
|
||||
|
||||
export function renderMarkdown(data: ResearchData): string {
|
||||
const {
|
||||
title,
|
||||
authors,
|
||||
abstract,
|
||||
keywords,
|
||||
mainFindings,
|
||||
methodology,
|
||||
results,
|
||||
discussion,
|
||||
references,
|
||||
publication,
|
||||
} = data;
|
||||
|
||||
const md: string[] = [];
|
||||
|
||||
md.push(`# ${title}\n`);
|
||||
|
||||
// Authors
|
||||
md.push(`## Authors`);
|
||||
md.push(
|
||||
authors
|
||||
.map(
|
||||
(author) =>
|
||||
`- **${author.name}**${
|
||||
author.affiliation ? `, *${author.affiliation}*` : ""
|
||||
}${author.email ? ` (${author.email})` : ""}`,
|
||||
)
|
||||
.join("\n"),
|
||||
);
|
||||
|
||||
// Abstract
|
||||
md.push(`\n## Abstract\n${abstract}`);
|
||||
|
||||
// Keywords
|
||||
if (keywords && keywords.length > 0) {
|
||||
md.push(`\n## Keywords\n${keywords.map((k) => `- ${k}`).join("\n")}`);
|
||||
}
|
||||
|
||||
// Main Findings
|
||||
md.push(
|
||||
`\n## Main Findings\n${mainFindings.map((f) => `- ${f}`).join("\n")}`,
|
||||
);
|
||||
|
||||
// Methodology
|
||||
if (methodology) {
|
||||
md.push(`\n## Methodology`);
|
||||
if (methodology.approach) md.push(`**Approach:** ${methodology.approach}`);
|
||||
if (methodology.participants)
|
||||
md.push(`**Participants:** ${methodology.participants}`);
|
||||
if (methodology.methods?.length) {
|
||||
md.push(
|
||||
`**Methods:**\n${methodology.methods.map((m) => `- ${m}`).join("\n")}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Results
|
||||
if (results?.length) {
|
||||
md.push(`\n## Results`);
|
||||
results.forEach((result, i) => {
|
||||
md.push(`\n### Result ${i + 1}`);
|
||||
if (result.finding) md.push(`- **Finding:** ${result.finding}`);
|
||||
if (result.significance)
|
||||
md.push(`- **Significance:** ${result.significance}`);
|
||||
if (result.supportingData)
|
||||
md.push(`- **Supporting Data:** ${result.supportingData}`);
|
||||
});
|
||||
}
|
||||
|
||||
// Discussion
|
||||
if (discussion) {
|
||||
md.push(`\n## Discussion`);
|
||||
if (discussion.implications?.length) {
|
||||
md.push(
|
||||
`### Implications\n${discussion.implications
|
||||
.map((d) => `- ${d}`)
|
||||
.join("\n")}`,
|
||||
);
|
||||
}
|
||||
if (discussion.limitations?.length) {
|
||||
md.push(
|
||||
`### Limitations\n${discussion.limitations
|
||||
.map((d) => `- ${d}`)
|
||||
.join("\n")}`,
|
||||
);
|
||||
}
|
||||
if (discussion.futureWork?.length) {
|
||||
md.push(
|
||||
`### Future Work\n${discussion.futureWork
|
||||
.map((d) => `- ${d}`)
|
||||
.join("\n")}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// References
|
||||
if (references?.length) {
|
||||
md.push(`\n## References`);
|
||||
references.forEach((ref, i) => {
|
||||
md.push(
|
||||
`\n**[${i + 1}]** ${ref.title} — *${ref.authors}*${
|
||||
ref.year ? ` (${ref.year})` : ""
|
||||
}`,
|
||||
);
|
||||
if (ref.relevance) md.push(`> ${ref.relevance}`);
|
||||
});
|
||||
}
|
||||
|
||||
// Publication Info
|
||||
if (publication) {
|
||||
md.push(`\n## Publication`);
|
||||
if (publication.journal) md.push(`- **Journal:** ${publication.journal}`);
|
||||
if (publication.year) md.push(`- **Year:** ${publication.year}`);
|
||||
if (publication.doi) md.push(`- **DOI:** ${publication.doi}`);
|
||||
if (publication.url)
|
||||
md.push(`- **URL:** [${publication.url}](${publication.url})`);
|
||||
}
|
||||
|
||||
return md.join("\n");
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
export const dataSchema = {
|
||||
type: "object",
|
||||
required: ["title", "authors", "abstract", "mainFindings"],
|
||||
properties: {
|
||||
title: {
|
||||
type: "string",
|
||||
description: "The full title of the research paper",
|
||||
},
|
||||
authors: {
|
||||
type: "array",
|
||||
description: "List of all authors of the paper",
|
||||
items: {
|
||||
type: "object",
|
||||
properties: {
|
||||
name: {
|
||||
type: "string",
|
||||
description: "Full name of the author",
|
||||
},
|
||||
affiliation: {
|
||||
type: "string",
|
||||
description:
|
||||
"Institution or organization the author is affiliated with",
|
||||
},
|
||||
email: {
|
||||
type: "string",
|
||||
description: "Contact email of the author if provided",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
abstract: {
|
||||
type: "string",
|
||||
description: "Complete abstract or summary of the paper",
|
||||
},
|
||||
keywords: {
|
||||
type: "array",
|
||||
description:
|
||||
"Key terms and phrases that describe the paper's main topics",
|
||||
items: {
|
||||
type: "string",
|
||||
},
|
||||
},
|
||||
mainFindings: {
|
||||
type: "array",
|
||||
description: "Key findings, conclusions, or contributions of the paper",
|
||||
items: {
|
||||
type: "string",
|
||||
},
|
||||
},
|
||||
methodology: {
|
||||
type: "object",
|
||||
description: "Research methods and approaches used",
|
||||
properties: {
|
||||
approach: {
|
||||
type: "string",
|
||||
description: "Overall research approach or study design",
|
||||
},
|
||||
participants: {
|
||||
type: "string",
|
||||
description: "Description of study participants or data sources",
|
||||
},
|
||||
methods: {
|
||||
type: "array",
|
||||
description: "Specific methods, techniques, or tools used",
|
||||
items: {
|
||||
type: "string",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
results: {
|
||||
type: "array",
|
||||
description: "Main results and outcomes of the research",
|
||||
items: {
|
||||
type: "object",
|
||||
properties: {
|
||||
finding: {
|
||||
type: "string",
|
||||
description: "Description of the specific result or finding",
|
||||
},
|
||||
significance: {
|
||||
type: "string",
|
||||
description:
|
||||
"Statistical significance or importance of the finding",
|
||||
},
|
||||
supportingData: {
|
||||
type: "string",
|
||||
description: "Relevant statistics, measurements, or data points",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
discussion: {
|
||||
type: "object",
|
||||
properties: {
|
||||
implications: {
|
||||
type: "array",
|
||||
description: "Theoretical or practical implications of the findings",
|
||||
items: {
|
||||
type: "string",
|
||||
},
|
||||
},
|
||||
limitations: {
|
||||
type: "array",
|
||||
description: "Study limitations or constraints",
|
||||
items: {
|
||||
type: "string",
|
||||
},
|
||||
},
|
||||
futureWork: {
|
||||
type: "array",
|
||||
description: "Suggested future research directions",
|
||||
items: {
|
||||
type: "string",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
references: {
|
||||
type: "array",
|
||||
description:
|
||||
"Key papers cited that are crucial to understanding this work",
|
||||
items: {
|
||||
type: "object",
|
||||
properties: {
|
||||
title: {
|
||||
type: "string",
|
||||
description: "Title of the cited paper",
|
||||
},
|
||||
authors: {
|
||||
type: "string",
|
||||
description: "Authors of the cited paper",
|
||||
},
|
||||
year: {
|
||||
type: "string",
|
||||
description: "Publication year",
|
||||
},
|
||||
relevance: {
|
||||
type: "string",
|
||||
description: "Why this reference is important to the current paper",
|
||||
},
|
||||
},
|
||||
required: ["title", "authors"],
|
||||
},
|
||||
},
|
||||
publication: {
|
||||
type: "object",
|
||||
properties: {
|
||||
journal: {
|
||||
type: "string",
|
||||
description: "Name of the journal or conference",
|
||||
},
|
||||
year: {
|
||||
type: "string",
|
||||
description: "Year of publication",
|
||||
},
|
||||
doi: {
|
||||
type: "string",
|
||||
description: "Digital Object Identifier (DOI) of the paper",
|
||||
},
|
||||
url: {
|
||||
type: "string",
|
||||
description: "URL where the paper can be accessed",
|
||||
},
|
||||
},
|
||||
required: ["year"],
|
||||
},
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,4 @@
|
||||
declare module "cli-markdown" {
|
||||
function cliMarkdown(input: string): string;
|
||||
export default cliMarkdown;
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
import * as readline from "readline/promises";
|
||||
import figlet from "figlet";
|
||||
import pc from "picocolors";
|
||||
|
||||
export async function renderLogo(): Promise<void> {
|
||||
const logoText = figlet.textSync("Extract Demo", {
|
||||
font: "ANSI Shadow",
|
||||
horizontalLayout: "default",
|
||||
verticalLayout: "default",
|
||||
width: 100,
|
||||
whitespaceBreak: true,
|
||||
});
|
||||
|
||||
// Add some styling with picocolors
|
||||
const styledLogo = pc.bold(pc.redBright(logoText));
|
||||
|
||||
// Add some padding/margin
|
||||
console.log("\n");
|
||||
console.log(styledLogo);
|
||||
console.log(pc.gray("─".repeat(60)));
|
||||
console.log("\n");
|
||||
}
|
||||
|
||||
export async function consoleInput(): Promise<string> {
|
||||
const rl = readline.createInterface({
|
||||
input: process.stdin,
|
||||
output: process.stdout,
|
||||
});
|
||||
|
||||
const answer = await rl.question("Path to your file: ");
|
||||
rl.close();
|
||||
return answer;
|
||||
}
|
||||
@@ -40,8 +40,8 @@ A TypeScript demo application showcasing the power of **LlamaCloud Index** - a f
|
||||
1. Clone the repository:
|
||||
|
||||
```bash
|
||||
git clone <repository-url>
|
||||
cd llamaparse-demo
|
||||
git clone https://github.com/run-llama/llama_cloud_services
|
||||
cd lama_cloud_services/examples-ts/index/
|
||||
```
|
||||
|
||||
2. Install dependencies:
|
||||
@@ -120,12 +120,12 @@ pnpm run lint
|
||||
|
||||
## License
|
||||
|
||||
MIT License - see the [LICENSE](../../../LICENSE) file for details.
|
||||
MIT License - see the [LICENSE](../../LICENSE) file for details.
|
||||
|
||||
## Contributing
|
||||
|
||||
1. Fork the repository
|
||||
2. Create a feature branch
|
||||
3. Make your changes
|
||||
4. Run `pnpm format` and `pnpm lint`
|
||||
4. Run `pnpm run format` and `pnpm run lint`
|
||||
5. Submit a pull request
|
||||
@@ -42,7 +42,7 @@
|
||||
"consola": "^3.4.2",
|
||||
"dotenv": "^17.2.1",
|
||||
"figlet": "^1.8.2",
|
||||
"llama-cloud-services": "link:../../../ts/llama_cloud_services",
|
||||
"llama-cloud-services": "link:../../ts/llama_cloud_services",
|
||||
"picocolors": "^1.1.1"
|
||||
}
|
||||
}
|
||||
@@ -40,8 +40,8 @@ A TypeScript demo application showcasing the power of **LlamaParse** - an intell
|
||||
1. Clone the repository:
|
||||
|
||||
```bash
|
||||
git clone <repository-url>
|
||||
cd llamaparse-demo
|
||||
git clone https://github.com/run-llama/llama_cloud_services
|
||||
cd lama_cloud_services/examples-ts/parse/
|
||||
```
|
||||
|
||||
2. Install dependencies:
|
||||
@@ -113,12 +113,12 @@ pnpm run lint
|
||||
|
||||
## License
|
||||
|
||||
MIT License - see the [LICENSE](../../../LICENSE) file for details.
|
||||
MIT License - see the [LICENSE](../../LICENSE) file for details.
|
||||
|
||||
## Contributing
|
||||
|
||||
1. Fork the repository
|
||||
2. Create a feature branch
|
||||
3. Make your changes
|
||||
4. Run `pnpm format` and `pnpm lint`
|
||||
4. Run `pnpm run format` and `pnpm run lint`
|
||||
5. Submit a pull request
|
||||
@@ -41,7 +41,7 @@
|
||||
"ai": "^4.3.19",
|
||||
"consola": "^3.4.2",
|
||||
"figlet": "^1.8.2",
|
||||
"llama-cloud-services": "link:../../../ts/llama_cloud_services",
|
||||
"llama-cloud-services": "link:../../ts/llama_cloud_services",
|
||||
"picocolors": "^1.1.1"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
import { createConsola } from "consola";
|
||||
import type { ConsolaInstance } from "consola";
|
||||
|
||||
export const logger: ConsolaInstance = createConsola({
|
||||
formatOptions: {
|
||||
date: false,
|
||||
},
|
||||
});
|
||||
+6
-8
@@ -1,11 +1,9 @@
|
||||
# LlamaCloud Services Examples
|
||||
# LlamaCloud Services Examples - Python
|
||||
|
||||
In this folder you will find several python notebooks and two end-to-end typescript applications that contain examples regarding:
|
||||
In this folder you will find several python notebooks that contain examples regarding:
|
||||
|
||||
- [LlamaParse - Python](./parse/)
|
||||
- [LlamaParse - TypeScript](./parse-ts/)
|
||||
- [LlamaExtract - Python](./extract/)
|
||||
- [LlamaReport - Python](./report/)
|
||||
- [LlamaCloud Index - TypeScript](./index-ts/)
|
||||
- [LlamaParse](./parse/)
|
||||
- [LlamaExtract](./extract/)
|
||||
- [LlamaReport](./report/)
|
||||
|
||||
Follow the instructions of each notebook/application to get started!
|
||||
Follow the instructions in each notebook to get started!
|
||||
|
||||
+218
-27
@@ -2,14 +2,40 @@
|
||||
|
||||
LlamaExtract provides a simple API for extracting structured data from unstructured documents like PDFs, text files and images.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Quick Start](#quick-start)
|
||||
- [Supported File Types](#supported-file-types)
|
||||
- [Different Input Types](#different-input-types)
|
||||
- [Async Extraction](#async-extraction)
|
||||
- [Core Concepts](#core-concepts)
|
||||
- [Defining Schemas](#defining-schemas)
|
||||
- [Using Pydantic (Recommended)](#using-pydantic-recommended)
|
||||
- [Using JSON Schema](#using-json-schema)
|
||||
- [Important restrictions on JSON/Pydantic Schema](#important-restrictions-on-jsonpydantic-schema)
|
||||
- [Extraction Configuration](#extraction-configuration)
|
||||
- [Configuration Options](#configuration-options)
|
||||
- [Extraction Agents (Advanced)](#extraction-agents-advanced)
|
||||
- [Creating Agents](#creating-agents)
|
||||
- [Agent Batch Processing](#agent-batch-processing)
|
||||
- [Updating Agent Schemas](#updating-agent-schemas)
|
||||
- [Managing Agents](#managing-agents)
|
||||
- [When to Use Agents vs Direct Extraction](#when-to-use-agents-vs-direct-extraction)
|
||||
- [Installation](#installation)
|
||||
- [Tips & Best Practices](#tips--best-practices)
|
||||
- [Additional Resources](#additional-resources)
|
||||
|
||||
## Quick Start
|
||||
|
||||
The simplest way to get started is to use the stateless API with the extraction configuration and the file/text to extract from:
|
||||
|
||||
```python
|
||||
from llama_cloud_services import LlamaExtract
|
||||
from llama_cloud import ExtractConfig, ExtractMode
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# Initialize client
|
||||
extractor = LlamaExtract()
|
||||
extractor = LlamaExtract(api_key="YOUR_API_KEY")
|
||||
|
||||
|
||||
# Define schema using Pydantic
|
||||
@@ -19,29 +45,97 @@ class Resume(BaseModel):
|
||||
skills: list[str] = Field(description="Technical skills and technologies")
|
||||
|
||||
|
||||
# Create extraction agent
|
||||
agent = extractor.create_agent(name="resume-parser", data_schema=Resume)
|
||||
# Configure extraction settings
|
||||
config = ExtractConfig(extraction_mode=ExtractMode.FAST)
|
||||
|
||||
# Extract data from document
|
||||
result = agent.extract("resume.pdf")
|
||||
# Extract data directly from document - no agent needed!
|
||||
result = extractor.extract(Resume, config, "resume.pdf")
|
||||
print(result.data)
|
||||
```
|
||||
|
||||
### Supported File Types
|
||||
|
||||
LlamaExtract supports the following file formats:
|
||||
|
||||
- **Documents**: PDF (.pdf), Word (.docx)
|
||||
- **Text files**: Plain text (.txt), CSV (.csv), JSON (.json), HTML (.html, .htm), Markdown (.md)
|
||||
- **Images**: PNG (.png), JPEG (.jpg, .jpeg)
|
||||
|
||||
### Different Input Types
|
||||
|
||||
```python
|
||||
# From file path (string or Path)
|
||||
result = extractor.extract(Resume, config, "resume.pdf")
|
||||
|
||||
# From file handle
|
||||
with open("resume.pdf", "rb") as f:
|
||||
result = extractor.extract(Resume, config, f)
|
||||
|
||||
# From bytes with filename
|
||||
with open("resume.pdf", "rb") as f:
|
||||
file_bytes = f.read()
|
||||
from llama_cloud_services.extract import SourceText
|
||||
|
||||
result = extractor.extract(
|
||||
Resume, config, SourceText(file=file_bytes, filename="resume.pdf")
|
||||
)
|
||||
|
||||
# From text content
|
||||
text = "Name: John Doe\nEmail: john@example.com\nSkills: Python, AI"
|
||||
result = extractor.extract(Resume, config, SourceText(text_content=text))
|
||||
```
|
||||
|
||||
### Async Extraction
|
||||
|
||||
For better performance with multiple files or when integrating with async applications.
|
||||
Here `queue_extraction` will enqueue the extraction jobs and exit. Alternatively, you
|
||||
can use `aextract` to poll for the job and return the extraction results.
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
|
||||
|
||||
async def extract_resumes():
|
||||
# Async extraction
|
||||
result = await extractor.aextract(Resume, config, "resume.pdf")
|
||||
print(result.data)
|
||||
|
||||
# Queue extraction jobs (returns immediately)
|
||||
jobs = await extractor.queue_extraction(
|
||||
Resume, config, ["resume1.pdf", "resume2.pdf"]
|
||||
)
|
||||
print(f"Queued {len(jobs)} extraction jobs")
|
||||
return jobs
|
||||
|
||||
|
||||
# Run async function
|
||||
jobs = asyncio.run(extract_resumes())
|
||||
# Check job status
|
||||
for job in jobs:
|
||||
status = agent.get_extraction_job(job.id).status
|
||||
print(f"Job {job.id}: {status}")
|
||||
|
||||
# Get results when complete
|
||||
results = [agent.get_extraction_run_for_job(job.id) for job in jobs]
|
||||
```
|
||||
|
||||
## Core Concepts
|
||||
|
||||
- **Extraction Agents**: Reusable extractors configured with a specific schema and extraction settings.
|
||||
- **Data Schema**: Structure definition for the data you want to extract in the form of a JSON schema or a Pydantic model.
|
||||
- **Extraction Config**: Settings that control how extraction is performed (e.g., speed vs accuracy trade-offs).
|
||||
- **Extraction Jobs**: Asynchronous extraction tasks that can be monitored.
|
||||
- **Extraction Agents** (Advanced): Reusable extractors configured with a specific schema and extraction settings.
|
||||
|
||||
## Defining Schemas
|
||||
|
||||
Schemas can be defined using either Pydantic models or JSON Schema:
|
||||
Schemas define the structure of data you want to extract. You can use either Pydantic models or JSON Schema:
|
||||
|
||||
### Using Pydantic (Recommended)
|
||||
|
||||
```python
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional
|
||||
from llama_cloud import ExtractConfig, ExtractMode
|
||||
|
||||
|
||||
class Experience(BaseModel):
|
||||
@@ -54,6 +148,11 @@ class Experience(BaseModel):
|
||||
class Resume(BaseModel):
|
||||
name: str = Field(description="Candidate name")
|
||||
experience: List[Experience] = Field(description="Work history")
|
||||
|
||||
|
||||
# Use the schema for extraction
|
||||
config = ExtractConfig(extraction_mode=ExtractMode.FAST)
|
||||
result = extractor.extract(Resume, config, "resume.pdf")
|
||||
```
|
||||
|
||||
### Using JSON Schema
|
||||
@@ -88,7 +187,9 @@ schema = {
|
||||
},
|
||||
}
|
||||
|
||||
agent = extractor.create_agent(name="resume-parser", data_schema=schema)
|
||||
# Use the schema for extraction
|
||||
config = ExtractConfig(extraction_mode=ExtractMode.FAST)
|
||||
result = extractor.extract(schema, config, "resume.pdf")
|
||||
```
|
||||
|
||||
### Important restrictions on JSON/Pydantic Schema
|
||||
@@ -108,28 +209,100 @@ be sufficient for a wide variety of use-cases.
|
||||
your extraction workflow to fit within these constraints, e.g. by extracting subset of fields
|
||||
and later merging them together.
|
||||
|
||||
## Other Extraction APIs
|
||||
## Extraction Configuration
|
||||
|
||||
### Extraction over bytes or text
|
||||
|
||||
You can use the `SourceText` class to extract from bytes or text directly without using a file. If passing the file bytes,
|
||||
you will need to pass the filename to the `SourceText` class.
|
||||
Configure how extraction is performed using `ExtractConfig`. The schema is the most important part, but several configuration options can significantly impact the extraction process.
|
||||
|
||||
```python
|
||||
with open("resume.pdf", "rb") as f:
|
||||
file_bytes = f.read()
|
||||
result = test_agent.extract(SourceText(file=file_bytes, filename="resume.pdf"))
|
||||
```
|
||||
from llama_cloud import ExtractConfig, ExtractMode, ChunkMode, ExtractTarget
|
||||
|
||||
```python
|
||||
result = test_agent.extract(
|
||||
SourceText(text_content="Candidate Name: Jane Doe")
|
||||
# Basic configuration
|
||||
config = ExtractConfig(
|
||||
extraction_mode=ExtractMode.BALANCED, # FAST, BALANCED, MULTIMODAL, PREMIUM
|
||||
extraction_target=ExtractTarget.PER_DOC, # PER_DOC, PER_PAGE
|
||||
system_prompt="Focus on the most recent data",
|
||||
page_range="1-5,10-15", # Extract from specific pages
|
||||
)
|
||||
|
||||
# Advanced configuration
|
||||
advanced_config = ExtractConfig(
|
||||
extraction_mode=ExtractMode.MULTIMODAL,
|
||||
chunk_mode=ChunkMode.PAGE, # PAGE, SECTION
|
||||
high_resolution_mode=True, # Better OCR accuracy
|
||||
invalidate_cache=False, # Bypass cached results
|
||||
cite_sources=True, # Enable source citations
|
||||
use_reasoning=True, # Enable reasoning (not in FAST mode)
|
||||
confidence_scores=True, # MULTIMODAL/PREMIUM only
|
||||
)
|
||||
```
|
||||
|
||||
### Batch Processing
|
||||
### Key Configuration Options
|
||||
|
||||
Process multiple files asynchronously:
|
||||
**Extraction Mode**: Controls processing quality and speed
|
||||
|
||||
- `FAST`: Fastest processing, suitable for simple documents with no OCR
|
||||
- `BALANCED`: Good speed/accuracy tradeoff for text-rich documents
|
||||
- `MULTIMODAL`: For visually rich documents with text, tables, and images (recommended)
|
||||
- `PREMIUM`: Highest accuracy with OCR, complex table/header detection
|
||||
|
||||
**Extraction Target**: Defines extraction scope
|
||||
|
||||
- `PER_DOC`: Apply schema to entire document (default)
|
||||
- `PER_PAGE`: Apply schema to each page, returns array of results
|
||||
|
||||
**Advanced Options**:
|
||||
|
||||
- `system_prompt`: Additional system-level instructions
|
||||
- `page_range`: Specific pages to extract (e.g., "1,3,5-7,9")
|
||||
- `chunk_mode`: Document splitting strategy (`PAGE` or `SECTION`)
|
||||
- `high_resolution_mode`: Better OCR for small text (slower processing)
|
||||
|
||||
**Extensions** (return additional metadata):
|
||||
|
||||
- `cite_sources`: Source tracing for extracted fields
|
||||
- `use_reasoning`: Explanations for extraction decisions
|
||||
- `confidence_scores`: Quantitative confidence measures (MULTIMODAL/PREMIUM only)
|
||||
|
||||
For complete configuration options, advanced settings, and detailed examples, see the [LlamaExtract Configuration Documentation](https://docs.cloud.llamaindex.ai/llamaextract/features/options).
|
||||
|
||||
## Extraction Agents (Advanced)
|
||||
|
||||
For reusable extraction workflows, you can create extraction agents that encapsulate both schema and configuration:
|
||||
|
||||
### Creating Agents
|
||||
|
||||
```python
|
||||
from llama_cloud_services import LlamaExtract
|
||||
from llama_cloud import ExtractConfig, ExtractMode
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# Initialize client
|
||||
extractor = LlamaExtract()
|
||||
|
||||
|
||||
# Define schema
|
||||
class Resume(BaseModel):
|
||||
name: str = Field(description="Full name of candidate")
|
||||
email: str = Field(description="Email address")
|
||||
skills: list[str] = Field(description="Technical skills and technologies")
|
||||
|
||||
|
||||
# Configure extraction settings
|
||||
config = ExtractConfig(extraction_mode=ExtractMode.FAST)
|
||||
|
||||
# Create extraction agent
|
||||
agent = extractor.create_agent(
|
||||
name="resume-parser", data_schema=Resume, config=config
|
||||
)
|
||||
|
||||
# Use the agent
|
||||
result = agent.extract("resume.pdf")
|
||||
print(result.data)
|
||||
```
|
||||
|
||||
### Agent Batch Processing
|
||||
|
||||
Process multiple files with an agent:
|
||||
|
||||
```python
|
||||
# Queue multiple files for extraction
|
||||
@@ -144,7 +317,7 @@ for job in jobs:
|
||||
results = [agent.get_extraction_run_for_job(job.id) for job in jobs]
|
||||
```
|
||||
|
||||
### Updating Schemas
|
||||
### Updating Agent Schemas
|
||||
|
||||
Schemas can be modified and updated after creation:
|
||||
|
||||
@@ -169,10 +342,26 @@ agent = extractor.get_agent(name="resume-parser")
|
||||
extractor.delete_agent(agent.id)
|
||||
```
|
||||
|
||||
### When to Use Agents vs Direct Extraction
|
||||
|
||||
**Use Direct Extraction When:**
|
||||
|
||||
- One-off extractions
|
||||
- Different schemas for different documents
|
||||
- Simple workflows
|
||||
- Getting started quickly
|
||||
|
||||
**Use Extraction Agents When:**
|
||||
|
||||
- Repeated extractions with the same schema
|
||||
- Team collaboration (shared, named extractors)
|
||||
- Complex workflows requiring state management
|
||||
- Production systems with consistent extraction patterns
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install llama-extract==0.1.0
|
||||
pip install llama-cloud-services
|
||||
```
|
||||
|
||||
## Tips & Best Practices
|
||||
@@ -193,9 +382,9 @@ At the core of LlamaExtract is the schema, which defines the structure of the da
|
||||
2. **Running Extractions**:
|
||||
- Note that resetting `agent.schema` will not save the schema to the database,
|
||||
until you call `agent.save`, but it will be used for running extractions.
|
||||
- Check job status prior to accessing results. Any extraction error should be available as
|
||||
part of `job.error` or `extraction_run.error` fields for debugging.
|
||||
- Consider async operations (`queue_extraction`) for large-scale extraction once you have finalized your schema.
|
||||
- Check extraction results for any errors. Error information is available in the `result.error` field for debugging.
|
||||
- Consider async operations (`aextract` or `queue_extraction`) for large-scale extraction or when processing multiple files.
|
||||
- For repeated extractions with the same schema, consider creating an extraction agent to avoid redefining the schema each time.
|
||||
|
||||
### Hitting "The response was too long to be processed" Error
|
||||
|
||||
@@ -208,5 +397,7 @@ Another option (orthogonal to the above) is to break the document into smaller s
|
||||
|
||||
## Additional Resources
|
||||
|
||||
- [Extract Documentation](https://docs.cloud.llamaindex.ai/llamaextract/getting_started) - Details on Extract features, API and examples.
|
||||
- [Example Notebook](docs/examples-py/extract/resume_screening.ipynb) - Detailed walkthrough of resume parsing
|
||||
- [Example Application with TypeScript](./examples-ts/extract/) - End-to-end examples using LlamaExtract TypeScript client.
|
||||
- [Discord Community](https://discord.com/invite/eN6D2HQ4aX) - Get help and share feedback
|
||||
|
||||
Generated
+3289
-787
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,2 @@
|
||||
packages:
|
||||
- "ts/**"
|
||||
+9
-5
@@ -4,11 +4,15 @@ help: ## Show all Makefile targets.
|
||||
@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
|
||||
|
||||
format: ## Run code autoformatters (black).
|
||||
pre-commit install
|
||||
git ls-files | xargs pre-commit run black --files
|
||||
uv run pre-commit install
|
||||
git ls-files | xargs uv run pre-commit run black --files
|
||||
|
||||
lint: ## Run linters: pre-commit (black, ruff, codespell) and mypy
|
||||
pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
|
||||
uv run pre-commit install && git ls-files | xargs uv run pre-commit run --show-diff-on-failure --files
|
||||
|
||||
test: ## Run tests via pytest
|
||||
pytest tests
|
||||
test: ## Run unit tests via pytest
|
||||
uv run pytest -v unit_tests/
|
||||
|
||||
.PHONY: e2e
|
||||
e2e: ## Run all tests. Run with high parallelism using xdist since tests are bottlenecked bound by the slow backend parsing
|
||||
uv run pytest -v -n 32 tests/
|
||||
|
||||
@@ -37,11 +37,10 @@ Example Usage:
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
import numbers
|
||||
from llama_cloud import ExtractRun
|
||||
from llama_cloud.types.agent_data import AgentData
|
||||
from llama_cloud.types.aggregate_group import AggregateGroup
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
from pydantic import BaseModel, Field, ValidationError, model_validator, ConfigDict
|
||||
from typing import (
|
||||
Generic,
|
||||
List,
|
||||
@@ -176,6 +175,16 @@ class TypedAgentDataItems(BaseModel, Generic[AgentDataT]):
|
||||
)
|
||||
|
||||
|
||||
class FieldCitation(BaseModel):
|
||||
page: Optional[int] = Field(
|
||||
None, description="The page number that the field occurred on"
|
||||
)
|
||||
matching_text: Optional[str] = Field(
|
||||
None,
|
||||
description="The original text this field's value was derived from",
|
||||
)
|
||||
|
||||
|
||||
class ExtractedFieldMetadata(BaseModel):
|
||||
"""
|
||||
Metadata for an extracted data field, such as confidence, and citation information.
|
||||
@@ -193,14 +202,14 @@ class ExtractedFieldMetadata(BaseModel):
|
||||
None,
|
||||
description="The confidence score for the field based on the extracted text only",
|
||||
)
|
||||
page_number: Optional[int] = Field(
|
||||
None, description="The page number that the field occurred on"
|
||||
)
|
||||
matching_text: Optional[str] = Field(
|
||||
citation: Optional[List[FieldCitation]] = Field(
|
||||
None,
|
||||
description="The original text this field's value was derived from",
|
||||
description="The citation for the field, including page number and matching text",
|
||||
)
|
||||
|
||||
# Forbid unknown keys to avoid swallowing nested dicts
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
ExtractedFieldMetaDataDict = Dict[
|
||||
str, Union[ExtractedFieldMetadata, Dict[str, Any], list[Any]]
|
||||
@@ -214,10 +223,12 @@ def parse_extracted_field_metadata(
|
||||
k: _parse_extracted_field_metadata_recursive(v)
|
||||
for k, v in field_metadata.items()
|
||||
if k not in _METADATA_FIELDS_SIBLING_TO_LEAF
|
||||
and k not in _ADDITIONAL_ROOT_METADATA_FIELDS
|
||||
}
|
||||
|
||||
|
||||
_METADATA_FIELDS_SIBLING_TO_LEAF = {"reasoning"}
|
||||
_ADDITIONAL_ROOT_METADATA_FIELDS = {"error"}
|
||||
|
||||
|
||||
def _parse_extracted_field_metadata_recursive(
|
||||
@@ -238,19 +249,10 @@ def _parse_extracted_field_metadata_recursive(
|
||||
if len(indicator_fields.intersection(field_value.keys())) > 0:
|
||||
try:
|
||||
merged = {**field_value, **additional_fields}
|
||||
allowed_fields = ExtractedFieldMetadata.model_fields.keys()
|
||||
merged = {k: v for k, v in merged.items() if k in allowed_fields}
|
||||
validated = ExtractedFieldMetadata.model_validate(merged)
|
||||
|
||||
# grab the citation from the array. This is just an array for backwards compatibility.
|
||||
if "citation" in field_value and len(field_value["citation"]) > 0:
|
||||
first_citation = field_value["citation"][0]
|
||||
if "page" in first_citation and isinstance(
|
||||
first_citation["page"], numbers.Number
|
||||
):
|
||||
validated.page_number = int(first_citation["page"]) # type: ignore
|
||||
if "matching_text" in first_citation and isinstance(
|
||||
first_citation["matching_text"], str
|
||||
):
|
||||
validated.matching_text = first_citation["matching_text"]
|
||||
return validated
|
||||
except ValidationError:
|
||||
pass
|
||||
@@ -340,6 +342,28 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
|
||||
description="Additional metadata about the extracted data, such as errors, tokens, etc.",
|
||||
)
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _normalize_field_metadata_on_input(cls, value: Any) -> Any:
|
||||
# Ensure any inbound representation (including JSON round-trips)
|
||||
# gets normalized so nested dicts become ExtractedFieldMetadata where appropriate.
|
||||
if (
|
||||
isinstance(value, dict)
|
||||
and "field_metadata" in value
|
||||
and isinstance(value["field_metadata"], dict)
|
||||
):
|
||||
try:
|
||||
value = {
|
||||
**value,
|
||||
"field_metadata": parse_extracted_field_metadata(
|
||||
value["field_metadata"]
|
||||
),
|
||||
}
|
||||
except Exception:
|
||||
# Let pydantic surface detailed errors later rather than swallowing completely
|
||||
pass
|
||||
return value
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
@@ -395,11 +419,14 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
|
||||
"""
|
||||
file_id = file_id or result.file.id
|
||||
file_name = file_name or result.file.name
|
||||
job_id = result.job_id
|
||||
job_field_metadata = result.extraction_metadata.get("field_metadata", {})
|
||||
errors = job_field_metadata.get("error", None)
|
||||
if not isinstance(errors, str):
|
||||
errors = None
|
||||
|
||||
try:
|
||||
field_metadata = parse_extracted_field_metadata(
|
||||
result.extraction_metadata.get("field_metadata", {})
|
||||
)
|
||||
field_metadata = parse_extracted_field_metadata(job_field_metadata)
|
||||
except ValidationError:
|
||||
field_metadata = {}
|
||||
|
||||
@@ -408,11 +435,15 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
|
||||
return cls.create(
|
||||
data=data,
|
||||
status=status,
|
||||
field_metadata=field_metadata,
|
||||
field_metadata=job_field_metadata,
|
||||
file_id=file_id,
|
||||
file_name=file_name,
|
||||
file_hash=file_hash,
|
||||
metadata=metadata or {},
|
||||
metadata={
|
||||
**({"field_errors": errors} if errors else {}),
|
||||
"job_id": job_id,
|
||||
**(metadata or {}),
|
||||
},
|
||||
)
|
||||
except ValidationError as e:
|
||||
invalid_item = ExtractedData[Dict[str, Any]].create(
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Optional
|
||||
from pydantic import BaseModel
|
||||
from llama_cloud.client import AsyncLlamaCloud
|
||||
from llama_cloud.types import (
|
||||
ClassifierRule,
|
||||
ClassifyJobResults,
|
||||
ClassifyParsingConfiguration,
|
||||
StatusEnum,
|
||||
ClassifyJobWithStatus,
|
||||
File,
|
||||
)
|
||||
from llama_cloud.resources.classifier.client import OMIT
|
||||
from llama_cloud_services.files.client import FileClient
|
||||
from llama_cloud_services.constants import POLLING_TIMEOUT_SECONDS
|
||||
from llama_cloud_services.utils import is_terminal_status, augment_async_errors
|
||||
from llama_index.core.async_utils import DEFAULT_NUM_WORKERS, run_jobs
|
||||
|
||||
|
||||
class ClassificationOutput(BaseModel):
|
||||
file_id: str
|
||||
classification: str
|
||||
|
||||
|
||||
class ClassifyClient:
|
||||
"""
|
||||
Experimental - Client for interacting with the LlamaCloud Classifier API.
|
||||
The Classification API is currently in beta and may change in the future without notice.
|
||||
|
||||
Args:
|
||||
client: The LlamaCloud client to use.
|
||||
project_id: The project ID to use.
|
||||
organization_id: The organization ID to use.
|
||||
polling_interval: The interval to poll for job completion in seconds.
|
||||
polling_timeout: The timeout for the job to complete in seconds.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: AsyncLlamaCloud,
|
||||
project_id: Optional[str] = None,
|
||||
organization_id: Optional[str] = None,
|
||||
polling_interval: float = 1.0,
|
||||
polling_timeout: float = POLLING_TIMEOUT_SECONDS,
|
||||
):
|
||||
self.client = client
|
||||
self.project_id = project_id
|
||||
self.organization_id = organization_id
|
||||
self.polling_interval = polling_interval
|
||||
self.file_client = FileClient(client, project_id, organization_id)
|
||||
self.polling_timeout = polling_timeout
|
||||
|
||||
async def aclassify_file_ids(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_ids: list[str],
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
) -> ClassifyJobResults:
|
||||
"""
|
||||
Classify a list of files by their IDs.
|
||||
Note that even if a job fails, some of the files may have been classified successfully.
|
||||
In this case, you may want to set raise_on_error to False and check the results for successful classifications.
|
||||
|
||||
Args:
|
||||
rules: The rules to use for classification.
|
||||
file_ids: The IDs of the files to classify.
|
||||
parsing_configuration: The parsing configuration to use for classification.
|
||||
raise_on_error: Whether to raise an error if the classification job fails.
|
||||
|
||||
Returns:
|
||||
The results of the classification job.
|
||||
"""
|
||||
classify_job = await self.client.classifier.create_classify_job(
|
||||
rules=rules,
|
||||
file_ids=file_ids,
|
||||
parsing_configuration=parsing_configuration or OMIT,
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
)
|
||||
|
||||
classify_job_with_status = await self._wait_for_job_completion(classify_job.id)
|
||||
|
||||
if raise_on_error and classify_job_with_status.status == StatusEnum.ERROR:
|
||||
raise ValueError(
|
||||
f"Error classifying files under job ID {classify_job_with_status.id}"
|
||||
)
|
||||
|
||||
results = await self.client.classifier.get_classification_job_results(
|
||||
classify_job_with_status.id,
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
def classify_file_ids(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_ids: list[str],
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
) -> ClassifyJobResults:
|
||||
with augment_async_errors():
|
||||
return asyncio.run(
|
||||
self.aclassify_file_ids(
|
||||
rules, file_ids, parsing_configuration, raise_on_error
|
||||
)
|
||||
)
|
||||
|
||||
async def aclassify_file_path(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_input_path: str,
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
) -> ClassifyJobResults:
|
||||
file = await self.file_client.upload_file(file_input_path)
|
||||
return await self.aclassify_file_ids(
|
||||
rules, [file.id], parsing_configuration, raise_on_error
|
||||
)
|
||||
|
||||
def classify_file_path(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_input_path: str,
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
) -> ClassifyJobResults:
|
||||
with augment_async_errors():
|
||||
return asyncio.run(
|
||||
self.aclassify_file_path(
|
||||
rules, file_input_path, parsing_configuration, raise_on_error
|
||||
)
|
||||
)
|
||||
|
||||
async def aclassify_file_paths(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_input_paths: list[str],
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
workers: int = DEFAULT_NUM_WORKERS,
|
||||
show_progress: bool = False,
|
||||
) -> ClassifyJobResults:
|
||||
coroutines = [self.file_client.upload_file(path) for path in file_input_paths]
|
||||
files: list[File] = await run_jobs(
|
||||
coroutines,
|
||||
show_progress=show_progress,
|
||||
workers=workers,
|
||||
desc="Uploading files for classification",
|
||||
)
|
||||
return await self.aclassify_file_ids(
|
||||
rules, [file.id for file in files], parsing_configuration, raise_on_error
|
||||
)
|
||||
|
||||
def classify_file_paths(
|
||||
self,
|
||||
rules: list[ClassifierRule],
|
||||
file_input_paths: list[str],
|
||||
parsing_configuration: Optional[ClassifyParsingConfiguration] = None,
|
||||
raise_on_error: bool = True,
|
||||
) -> ClassifyJobResults:
|
||||
with augment_async_errors():
|
||||
return asyncio.run(
|
||||
self.aclassify_file_paths(
|
||||
rules, file_input_paths, parsing_configuration, raise_on_error
|
||||
)
|
||||
)
|
||||
|
||||
async def _wait_for_job_completion(self, job_id: str) -> ClassifyJobWithStatus:
|
||||
job = await self.client.classifier.get_classify_job(
|
||||
job_id, project_id=self.project_id, organization_id=self.organization_id
|
||||
)
|
||||
start_time = time.time()
|
||||
while not is_terminal_status(job.status):
|
||||
polling_duration = time.time() - start_time
|
||||
if polling_duration > self.polling_timeout:
|
||||
raise TimeoutError(
|
||||
f"Job {job_id} timed out after {polling_duration} seconds"
|
||||
)
|
||||
await asyncio.sleep(self.polling_interval)
|
||||
job = await self.client.classifier.get_classify_job(
|
||||
job_id, project_id=self.project_id, organization_id=self.organization_id
|
||||
)
|
||||
return job
|
||||
@@ -1 +1,2 @@
|
||||
EU_BASE_URL = "https://api.cloud.eu.llamaindex.ai"
|
||||
POLLING_TIMEOUT_SECONDS = 300.0
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
from io import BufferedIOBase, BufferedReader, BytesIO, TextIOWrapper
|
||||
@@ -21,6 +22,7 @@ from llama_cloud import (
|
||||
ExtractJobCreate,
|
||||
ExtractRun,
|
||||
File,
|
||||
FileData,
|
||||
ExtractMode,
|
||||
StatusEnum,
|
||||
ExtractTarget,
|
||||
@@ -31,9 +33,9 @@ from llama_cloud.client import AsyncLlamaCloud
|
||||
from llama_cloud.core.api_error import ApiError
|
||||
from llama_cloud_services.extract.utils import (
|
||||
JSONObjectType,
|
||||
augment_async_errors,
|
||||
ExperimentalWarning,
|
||||
)
|
||||
from llama_cloud_services.utils import augment_async_errors
|
||||
from llama_index.core.schema import BaseComponent
|
||||
from llama_index.core.async_utils import run_jobs
|
||||
from llama_index.core.bridge.pydantic import Field, PrivateAttr
|
||||
@@ -47,7 +49,7 @@ SchemaInput = Union[JSONObjectType, Type[BaseModel]]
|
||||
|
||||
DEFAULT_EXTRACT_CONFIG = ExtractConfig(
|
||||
extraction_target=ExtractTarget.PER_DOC,
|
||||
extraction_mode=ExtractMode.BALANCED,
|
||||
extraction_mode=ExtractMode.MULTIMODAL,
|
||||
)
|
||||
|
||||
|
||||
@@ -62,6 +64,132 @@ def _is_retryable_error(exception: BaseException) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
async def _validate_schema(
|
||||
client: AsyncLlamaCloud, data_schema: SchemaInput
|
||||
) -> JSONObjectType:
|
||||
"""Convert SchemaInput to a validated JSON schema dictionary."""
|
||||
processed_schema: JSONObjectType
|
||||
if isinstance(data_schema, dict):
|
||||
# TODO: if we expose a get_validated JSON schema method, we can use it here
|
||||
processed_schema = data_schema # type: ignore
|
||||
elif isinstance(data_schema, type) and issubclass(data_schema, BaseModel):
|
||||
processed_schema = data_schema.model_json_schema()
|
||||
else:
|
||||
raise ValueError("data_schema must be either a dictionary or a Pydantic model")
|
||||
|
||||
# Validate schema via API
|
||||
validated_schema = await client.llama_extract.validate_extraction_schema(
|
||||
data_schema=processed_schema
|
||||
)
|
||||
return validated_schema.data_schema
|
||||
|
||||
|
||||
async def _get_job_with_retry(
|
||||
client: AsyncLlamaCloud,
|
||||
job_id: str,
|
||||
max_attempts: int = 5,
|
||||
initial_wait: float = 1,
|
||||
max_wait: float = 60,
|
||||
jitter: float = 5,
|
||||
) -> ExtractJob:
|
||||
"""Get extraction job with retry logic."""
|
||||
async for attempt in AsyncRetrying(
|
||||
retry=retry_if_exception(_is_retryable_error),
|
||||
stop=stop_after_attempt(max_attempts),
|
||||
wait=wait_exponential_jitter(initial=initial_wait, max=max_wait, jitter=jitter),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
return await client.llama_extract.get_job(job_id=job_id)
|
||||
|
||||
|
||||
async def _get_run_with_retry(
|
||||
client: AsyncLlamaCloud,
|
||||
job_id: str,
|
||||
project_id: Optional[str] = None,
|
||||
organization_id: Optional[str] = None,
|
||||
max_attempts: int = 3,
|
||||
initial_wait: float = 1,
|
||||
max_wait: float = 20,
|
||||
jitter: float = 3,
|
||||
) -> ExtractRun:
|
||||
"""Get extraction run with retry logic."""
|
||||
async for attempt in AsyncRetrying(
|
||||
retry=retry_if_exception(_is_retryable_error),
|
||||
stop=stop_after_attempt(max_attempts),
|
||||
wait=wait_exponential_jitter(initial=initial_wait, max=max_wait, jitter=jitter),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
return await client.llama_extract.get_run_by_job_id(
|
||||
job_id=job_id,
|
||||
project_id=project_id,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
|
||||
|
||||
async def _wait_for_job_result(
|
||||
client: AsyncLlamaCloud,
|
||||
job_id: str,
|
||||
check_interval: int = 1,
|
||||
max_timeout: int = 2000,
|
||||
verbose: bool = False,
|
||||
project_id: Optional[str] = None,
|
||||
organization_id: Optional[str] = None,
|
||||
job_retry_attempts: int = 5,
|
||||
job_max_wait: float = 60,
|
||||
job_jitter: float = 5,
|
||||
run_retry_attempts: int = 3,
|
||||
run_max_wait: float = 20,
|
||||
run_jitter: float = 3,
|
||||
) -> Optional[ExtractRun]:
|
||||
"""Wait for and return the results of an extraction job."""
|
||||
start = time.perf_counter()
|
||||
poll_count = 0
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(check_interval)
|
||||
poll_count += 1
|
||||
job = await _get_job_with_retry(
|
||||
client,
|
||||
job_id,
|
||||
max_attempts=job_retry_attempts,
|
||||
max_wait=job_max_wait,
|
||||
jitter=job_jitter,
|
||||
)
|
||||
|
||||
if job.status == StatusEnum.SUCCESS:
|
||||
return await _get_run_with_retry(
|
||||
client,
|
||||
job_id,
|
||||
project_id,
|
||||
organization_id,
|
||||
max_attempts=run_retry_attempts,
|
||||
max_wait=run_max_wait,
|
||||
jitter=run_jitter,
|
||||
)
|
||||
elif job.status == StatusEnum.PENDING:
|
||||
end = time.perf_counter()
|
||||
if end - start > max_timeout:
|
||||
raise Exception(f"Timeout while extracting the file: {job_id}")
|
||||
if verbose and poll_count % 10 == 0:
|
||||
print(".", end="", flush=True)
|
||||
continue
|
||||
else:
|
||||
warnings.warn(
|
||||
f"Failure in job: {job_id}, status: {job.status}, error: {job.error}"
|
||||
)
|
||||
return await _get_run_with_retry(
|
||||
client,
|
||||
job_id,
|
||||
project_id,
|
||||
organization_id,
|
||||
max_attempts=run_retry_attempts,
|
||||
max_wait=run_max_wait,
|
||||
jitter=run_jitter,
|
||||
)
|
||||
|
||||
|
||||
class SourceText:
|
||||
def __init__(
|
||||
self,
|
||||
@@ -99,7 +227,7 @@ class SourceText:
|
||||
raise ValueError(f"Unsupported file type: {type(self.file)}")
|
||||
|
||||
|
||||
FileInput = Union[str, Path, BufferedIOBase, SourceText]
|
||||
FileInput = Union[str, Path, BufferedIOBase, SourceText, File]
|
||||
|
||||
|
||||
def run_in_thread(
|
||||
@@ -144,6 +272,21 @@ def _extraction_config_warning(config: ExtractConfig) -> None:
|
||||
"available in the `extraction_metadata` field for the extraction run.",
|
||||
ExperimentalWarning,
|
||||
)
|
||||
if config.use_reasoning:
|
||||
if config.extraction_mode == ExtractMode.FAST:
|
||||
raise ValueError(
|
||||
"`reasoning` is only supported with BALANCED, MULTIMODAL, or PREMIUM extraction modes."
|
||||
)
|
||||
if config.cite_sources:
|
||||
if config.extraction_mode in (ExtractMode.FAST, ExtractMode.BALANCED):
|
||||
raise ValueError(
|
||||
"`cite_sources` is only supported with MULTIMODAL or PREMIUM extraction modes."
|
||||
)
|
||||
if config.confidence_scores:
|
||||
if config.extraction_mode in (ExtractMode.FAST, ExtractMode.BALANCED):
|
||||
raise ValueError(
|
||||
"`confidence_scores` is only supported with MULTIMODAL or PREMIUM extraction modes."
|
||||
)
|
||||
|
||||
|
||||
class ExtractionAgent:
|
||||
@@ -194,22 +337,10 @@ class ExtractionAgent:
|
||||
|
||||
@data_schema.setter
|
||||
def data_schema(self, data_schema: SchemaInput) -> None:
|
||||
processed_schema: JSONObjectType
|
||||
if isinstance(data_schema, dict):
|
||||
# TODO: if we expose a get_validated JSON schema method, we can use it here
|
||||
processed_schema = data_schema # type: ignore
|
||||
elif isinstance(data_schema, type) and issubclass(data_schema, BaseModel):
|
||||
processed_schema = data_schema.model_json_schema()
|
||||
else:
|
||||
raise ValueError(
|
||||
"data_schema must be either a dictionary or a Pydantic model"
|
||||
)
|
||||
validated_schema = self._run_in_thread(
|
||||
self._client.llama_extract.validate_extraction_schema(
|
||||
data_schema=processed_schema
|
||||
)
|
||||
# Use the shared schema processing and validation function
|
||||
self._data_schema = self._run_in_thread(
|
||||
_validate_schema(self._client, data_schema)
|
||||
)
|
||||
self._data_schema = validated_schema.data_schema
|
||||
|
||||
@property
|
||||
def config(self) -> ExtractConfig:
|
||||
@@ -268,11 +399,15 @@ class ExtractionAgent:
|
||||
project_id=self._project_id, upload_file=file_contents
|
||||
)
|
||||
finally:
|
||||
if file_contents is not None and isinstance(file_contents, BufferedReader):
|
||||
if file_contents is not None and isinstance(
|
||||
file_contents, (BufferedReader, BytesIO)
|
||||
):
|
||||
file_contents.close()
|
||||
|
||||
async def _upload_file(self, file_input: FileInput) -> File:
|
||||
source_text = None
|
||||
if isinstance(file_input, File):
|
||||
return file_input
|
||||
if isinstance(file_input, SourceText):
|
||||
source_text = file_input
|
||||
elif isinstance(file_input, (str, Path)):
|
||||
@@ -296,60 +431,23 @@ class ExtractionAgent:
|
||||
|
||||
return await self.upload_file(source_text)
|
||||
|
||||
async def _get_job_with_retry(self, job_id: str) -> ExtractJob:
|
||||
"""Get job with retry logic for transient errors."""
|
||||
async for attempt in AsyncRetrying(
|
||||
retry=retry_if_exception(_is_retryable_error),
|
||||
stop=stop_after_attempt(5),
|
||||
wait=wait_exponential_jitter(initial=1, max=60, jitter=5),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
return await self._client.llama_extract.get_job(job_id=job_id)
|
||||
|
||||
async def _get_run_with_retry(self, job_id: str) -> ExtractRun:
|
||||
"""Get extraction run with retry logic for transient errors."""
|
||||
async for attempt in AsyncRetrying(
|
||||
retry=retry_if_exception(_is_retryable_error),
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential_jitter(initial=1, max=20, jitter=3),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
return await self._client.llama_extract.get_run_by_job_id(job_id=job_id)
|
||||
|
||||
async def _wait_for_job_result(self, job_id: str) -> Optional[ExtractRun]:
|
||||
"""Wait for and return the results of an extraction job."""
|
||||
start = time.perf_counter()
|
||||
tries = 0
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(self.check_interval)
|
||||
tries += 1
|
||||
|
||||
try:
|
||||
job = await self._get_job_with_retry(job_id)
|
||||
|
||||
if job.status == StatusEnum.SUCCESS:
|
||||
return await self._get_run_with_retry(job_id)
|
||||
elif job.status == StatusEnum.PENDING:
|
||||
end = time.perf_counter()
|
||||
if end - start > self.max_timeout:
|
||||
raise Exception(f"Timeout while extracting the file: {job_id}")
|
||||
if self._verbose and tries % 10 == 0:
|
||||
print(".", end="", flush=True)
|
||||
continue
|
||||
else:
|
||||
warnings.warn(
|
||||
f"Failure in job: {job_id}, status: {job.status}, error: {job.error}"
|
||||
)
|
||||
return await self._get_run_with_retry(job_id)
|
||||
|
||||
except Exception as e:
|
||||
# If we get a non-retryable error or all retries are exhausted, re-raise
|
||||
if self._verbose:
|
||||
print(f"\nError in job polling for {job_id}: {e}")
|
||||
raise e
|
||||
return await _wait_for_job_result(
|
||||
client=self._client,
|
||||
job_id=job_id,
|
||||
check_interval=self.check_interval,
|
||||
max_timeout=self.max_timeout,
|
||||
verbose=self._verbose,
|
||||
project_id=self._project_id,
|
||||
organization_id=self._organization_id,
|
||||
job_retry_attempts=5,
|
||||
job_max_wait=60,
|
||||
job_jitter=5,
|
||||
run_retry_attempts=3,
|
||||
run_max_wait=20,
|
||||
run_jitter=3,
|
||||
)
|
||||
|
||||
def save(self) -> None:
|
||||
"""Persist the extraction agent's schema and config to the database.
|
||||
@@ -437,7 +535,7 @@ class ExtractionAgent:
|
||||
|
||||
upload_tasks = [self._upload_file(file) for file in files]
|
||||
with augment_async_errors():
|
||||
uploaded_files = await run_jobs(
|
||||
uploaded_files: List[File] = await run_jobs(
|
||||
upload_tasks,
|
||||
workers=self.num_workers,
|
||||
desc="Uploading files",
|
||||
@@ -643,12 +741,14 @@ class LlamaExtract(BaseComponent):
|
||||
base_url = os.getenv("LLAMA_CLOUD_BASE_URL", None) or DEFAULT_BASE_URL
|
||||
|
||||
super().__init__(
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
api_key=api_key, # type: ignore
|
||||
base_url=base_url, # type: ignore
|
||||
check_interval=check_interval,
|
||||
max_timeout=max_timeout,
|
||||
num_workers=num_workers,
|
||||
show_progress=show_progress,
|
||||
project_id=project_id,
|
||||
organization_id=organization_id,
|
||||
verify=verify,
|
||||
httpx_timeout=httpx_timeout,
|
||||
verbose=verbose,
|
||||
@@ -702,7 +802,7 @@ class LlamaExtract(BaseComponent):
|
||||
config = DEFAULT_EXTRACT_CONFIG
|
||||
|
||||
if isinstance(data_schema, dict):
|
||||
data_schema = data_schema
|
||||
pass
|
||||
elif issubclass(data_schema, BaseModel):
|
||||
data_schema = data_schema.model_json_schema()
|
||||
else:
|
||||
@@ -803,6 +903,8 @@ class LlamaExtract(BaseComponent):
|
||||
num_workers=self.num_workers,
|
||||
show_progress=self.show_progress,
|
||||
verbose=self.verbose,
|
||||
verify=self.verify,
|
||||
httpx_timeout=self.httpx_timeout,
|
||||
)
|
||||
for agent in agents
|
||||
]
|
||||
@@ -819,6 +921,249 @@ class LlamaExtract(BaseComponent):
|
||||
)
|
||||
)
|
||||
|
||||
async def _wait_for_job_result(self, job_id: str) -> Optional[ExtractRun]:
|
||||
"""Wait for and return the results of an extraction job."""
|
||||
return await _wait_for_job_result(
|
||||
client=self._async_client,
|
||||
job_id=job_id,
|
||||
check_interval=self.check_interval,
|
||||
max_timeout=self.max_timeout,
|
||||
verbose=self.verbose,
|
||||
project_id=self._project_id,
|
||||
organization_id=self._organization_id,
|
||||
job_retry_attempts=3,
|
||||
job_max_wait=4,
|
||||
job_jitter=5,
|
||||
run_retry_attempts=3,
|
||||
run_max_wait=4,
|
||||
run_jitter=3,
|
||||
)
|
||||
|
||||
def _get_mime_type(
|
||||
self,
|
||||
filename: Optional[str] = None,
|
||||
file_path: Optional[Union[str, Path]] = None,
|
||||
) -> str:
|
||||
"""Determine MIME type for a file based on filename or path."""
|
||||
# MIME type mappings for supported formats
|
||||
MIME_TYPE_MAP = {
|
||||
# Text files
|
||||
".txt": "text/plain",
|
||||
".csv": "text/csv",
|
||||
".json": "application/json",
|
||||
".html": "text/html",
|
||||
".htm": "text/html",
|
||||
".md": "text/markdown",
|
||||
# Document files
|
||||
".pdf": "application/pdf",
|
||||
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
# Image files
|
||||
".png": "image/png",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
}
|
||||
|
||||
# Try to get extension from filename or file_path
|
||||
extension = None
|
||||
if filename:
|
||||
extension = Path(filename).suffix.lower()
|
||||
elif file_path:
|
||||
extension = Path(file_path).suffix.lower()
|
||||
|
||||
# Check if the extension is supported
|
||||
if extension and extension in MIME_TYPE_MAP:
|
||||
return MIME_TYPE_MAP[extension]
|
||||
|
||||
# If we don't have a supported extension, provide helpful error message
|
||||
supported_extensions = [ext[1:] for ext in MIME_TYPE_MAP.keys()] # Remove dots
|
||||
supported_list = ", ".join(sorted(supported_extensions))
|
||||
|
||||
if extension:
|
||||
ext_without_dot = extension[1:] # Remove the leading dot
|
||||
raise ValueError(
|
||||
f"Unsupported file type: '{ext_without_dot}'. "
|
||||
f"Supported formats are: {supported_list}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Could not determine file type. Please provide a filename with one of these supported extensions: {supported_list}"
|
||||
)
|
||||
|
||||
def _convert_file_to_file_data(
|
||||
self, file_input: FileInput
|
||||
) -> Union[FileData, str, File]:
|
||||
"""Convert FileInput to FileData or text string for stateless extraction."""
|
||||
if isinstance(file_input, File):
|
||||
return file_input
|
||||
|
||||
if isinstance(file_input, SourceText):
|
||||
if file_input.text_content is not None:
|
||||
return file_input.text_content
|
||||
elif file_input.file is not None:
|
||||
if isinstance(file_input.file, bytes):
|
||||
data = file_input.file
|
||||
filename = file_input.filename
|
||||
elif isinstance(file_input.file, (str, Path)):
|
||||
with open(file_input.file, "rb") as f:
|
||||
data = f.read()
|
||||
filename = file_input.filename or str(file_input.file)
|
||||
elif isinstance(file_input.file, (BufferedIOBase, TextIOWrapper)):
|
||||
if hasattr(file_input.file, "read"):
|
||||
content = file_input.file.read()
|
||||
if isinstance(content, str):
|
||||
data = content.encode("utf-8")
|
||||
else:
|
||||
data = content
|
||||
else:
|
||||
raise ValueError("File object must have a read method")
|
||||
filename = file_input.filename or getattr(
|
||||
file_input.file, "name", None
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported file type: {type(file_input.file)}")
|
||||
|
||||
# Encode as base64
|
||||
encoded_data = base64.b64encode(data).decode("utf-8")
|
||||
|
||||
# Determine mime type
|
||||
mime_type = self._get_mime_type(filename=filename)
|
||||
|
||||
return FileData(data=encoded_data, mime_type=mime_type)
|
||||
else:
|
||||
raise ValueError("SourceText must have either text_content or file")
|
||||
|
||||
elif isinstance(file_input, (str, Path)):
|
||||
with open(file_input, "rb") as f:
|
||||
data = f.read()
|
||||
encoded_data = base64.b64encode(data).decode("utf-8")
|
||||
mime_type = self._get_mime_type(file_path=file_input)
|
||||
return FileData(data=encoded_data, mime_type=mime_type)
|
||||
|
||||
elif isinstance(file_input, bytes):
|
||||
# For raw bytes, we can't determine the file type, so we need to raise an error
|
||||
raise ValueError(
|
||||
"Cannot determine file type from raw bytes. Please use SourceText with a filename, or provide a file path."
|
||||
)
|
||||
|
||||
elif isinstance(file_input, (BufferedIOBase, TextIOWrapper)):
|
||||
if hasattr(file_input, "read"):
|
||||
content = file_input.read()
|
||||
if isinstance(content, str):
|
||||
data = content.encode("utf-8")
|
||||
else:
|
||||
data = content
|
||||
encoded_data = base64.b64encode(data).decode("utf-8")
|
||||
|
||||
# Try to get filename from the file object
|
||||
filename = getattr(file_input, "name", None)
|
||||
mime_type = self._get_mime_type(filename=filename)
|
||||
|
||||
return FileData(data=encoded_data, mime_type=mime_type)
|
||||
else:
|
||||
raise ValueError("File object must have a read method")
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported file input type: {type(file_input)}")
|
||||
|
||||
async def queue_extraction(
|
||||
self,
|
||||
data_schema: SchemaInput,
|
||||
config: ExtractConfig,
|
||||
files: Union[FileInput, List[FileInput]],
|
||||
) -> Union[ExtractJob, List[ExtractJob]]:
|
||||
"""Queue extraction jobs using stateless extraction (no agent required).
|
||||
|
||||
Args:
|
||||
data_schema: The schema defining what data to extract
|
||||
config: The extraction configuration
|
||||
files: File(s) to extract from
|
||||
|
||||
Returns:
|
||||
ExtractJob or list of ExtractJobs
|
||||
"""
|
||||
_extraction_config_warning(config)
|
||||
processed_schema = await _validate_schema(self._async_client, data_schema)
|
||||
|
||||
if not isinstance(files, list):
|
||||
files = [files]
|
||||
|
||||
jobs = []
|
||||
for file_input in files:
|
||||
file_data_or_text = self._convert_file_to_file_data(file_input)
|
||||
|
||||
if isinstance(file_data_or_text, File):
|
||||
file_args = {"file_id": file_data_or_text.id}
|
||||
|
||||
elif isinstance(file_data_or_text, str):
|
||||
# It's text content
|
||||
file_args = {"text": file_data_or_text}
|
||||
else:
|
||||
# It's FileData
|
||||
file_args = {"file": file_data_or_text}
|
||||
|
||||
job = await self._async_client.llama_extract.extract_stateless(
|
||||
project_id=self._project_id,
|
||||
organization_id=self._organization_id,
|
||||
data_schema=processed_schema,
|
||||
config=config,
|
||||
**file_args,
|
||||
)
|
||||
jobs.append(job)
|
||||
|
||||
return jobs[0] if len(jobs) == 1 else jobs
|
||||
|
||||
async def aextract(
|
||||
self,
|
||||
data_schema: SchemaInput,
|
||||
config: ExtractConfig,
|
||||
files: Union[FileInput, List[FileInput]],
|
||||
) -> Union[ExtractRun, List[ExtractRun]]:
|
||||
"""Run stateless extraction and wait for results.
|
||||
|
||||
Args:
|
||||
data_schema: The schema defining what data to extract
|
||||
config: The extraction configuration
|
||||
files: File(s) to extract from
|
||||
|
||||
Returns:
|
||||
ExtractRun or list of ExtractRuns with the extraction results
|
||||
"""
|
||||
jobs = await self.queue_extraction(data_schema, config, files)
|
||||
|
||||
if isinstance(jobs, list):
|
||||
runs = []
|
||||
for job in jobs:
|
||||
run = await self._wait_for_job_result(job.id)
|
||||
if run is None:
|
||||
raise RuntimeError(
|
||||
f"Failed to get extraction result for job {job.id}"
|
||||
)
|
||||
runs.append(run)
|
||||
return runs
|
||||
else:
|
||||
run = await self._wait_for_job_result(jobs.id)
|
||||
if run is None:
|
||||
raise RuntimeError(f"Failed to get extraction result for job {jobs.id}")
|
||||
return run
|
||||
|
||||
def extract(
|
||||
self,
|
||||
data_schema: SchemaInput,
|
||||
config: ExtractConfig,
|
||||
files: Union[FileInput, List[FileInput]],
|
||||
) -> Union[ExtractRun, List[ExtractRun]]:
|
||||
"""Run stateless extraction and wait for results (synchronous version).
|
||||
|
||||
Args:
|
||||
data_schema: The schema defining what data to extract
|
||||
config: The extraction configuration
|
||||
files: File(s) to extract from
|
||||
|
||||
Returns:
|
||||
ExtractRun or list of ExtractRuns with the extraction results
|
||||
"""
|
||||
return self._run_in_thread(self.aextract(data_schema, config, files))
|
||||
|
||||
def __del__(self) -> None:
|
||||
"""Cleanup resources properly."""
|
||||
try:
|
||||
@@ -833,6 +1178,20 @@ if __name__ == "__main__":
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Example usage:
|
||||
#
|
||||
# # Basic usage with stateless extraction (no agent required)
|
||||
# extractor = LlamaExtract()
|
||||
# schema = {"name": {"type": "string"}, "email": {"type": "string"}}
|
||||
# config = ExtractConfig(extraction_mode=ExtractMode.FAST)
|
||||
# files = ["path/to/document.pdf"]
|
||||
#
|
||||
# # Queue extraction jobs
|
||||
# jobs = extractor.queue_extraction(schema, config, files)
|
||||
#
|
||||
# # Or run extraction and wait for results
|
||||
# results = extractor.extract(schema, config, files)
|
||||
|
||||
data_dir = Path(__file__).parent.parent / "tests" / "data"
|
||||
extractor = LlamaExtract()
|
||||
try:
|
||||
|
||||
@@ -1,12 +1,4 @@
|
||||
from typing import Any, Dict, List, Union, Generator
|
||||
from contextlib import contextmanager
|
||||
|
||||
# Asyncio error messages
|
||||
nest_asyncio_err = "cannot be called from a running event loop"
|
||||
nest_asyncio_msg = (
|
||||
"The event loop is already running. "
|
||||
"Add `import nest_asyncio; nest_asyncio.apply()` to your code to fix this issue."
|
||||
)
|
||||
from typing import Any, Dict, List, Union
|
||||
|
||||
|
||||
def is_jupyter() -> bool:
|
||||
@@ -19,17 +11,6 @@ def is_jupyter() -> bool:
|
||||
return False
|
||||
|
||||
|
||||
@contextmanager
|
||||
def augment_async_errors() -> Generator[None, None, None]:
|
||||
"""Context manager to add helpful information for errors due to nested event loops."""
|
||||
try:
|
||||
yield
|
||||
except RuntimeError as e:
|
||||
if nest_asyncio_err in str(e):
|
||||
raise RuntimeError(nest_asyncio_msg)
|
||||
raise
|
||||
|
||||
|
||||
JSONType = Union[Dict[str, Any], List[Any], str, int, float, bool, None]
|
||||
JSONObjectType = Dict[str, JSONType]
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
from .client import FileClient
|
||||
|
||||
__all__ = ["FileClient"]
|
||||
@@ -0,0 +1,97 @@
|
||||
from io import BytesIO
|
||||
from typing import BinaryIO
|
||||
import os
|
||||
from llama_cloud.client import AsyncLlamaCloud
|
||||
from llama_cloud.types import File, FileCreate
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class FileClient:
|
||||
"""
|
||||
Higher-level client for interacting with the LlamaCloud Files API.
|
||||
Uses presigned URLs for uploads by default.
|
||||
|
||||
Args:
|
||||
client: The LlamaCloud client to use.
|
||||
project_id: The project ID to use.
|
||||
organization_id: The organization ID to use.
|
||||
use_presigned_url: Whether to use presigned URLs for uploads (set to False when uploading to BYOC deployments).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: AsyncLlamaCloud,
|
||||
project_id: Optional[str] = None,
|
||||
organization_id: Optional[str] = None,
|
||||
use_presigned_url: bool = True,
|
||||
):
|
||||
self.client = client
|
||||
self.project_id = project_id
|
||||
self.organization_id = organization_id
|
||||
self.use_presigned_url = use_presigned_url
|
||||
|
||||
async def get_file(self, file_id: str) -> File:
|
||||
return await self.client.files.get_file(
|
||||
file_id, project_id=self.project_id, organization_id=self.organization_id
|
||||
)
|
||||
|
||||
async def read_file_content(self, file_id: str) -> bytes:
|
||||
presigned_url = await self.client.files.read_file_content(
|
||||
file_id,
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
)
|
||||
httpx_client = self.client._client_wrapper.httpx_client
|
||||
response = await httpx_client.get(presigned_url.url)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
|
||||
async def upload_file(
|
||||
self, file_path: str, external_file_id: Optional[str] = None
|
||||
) -> File:
|
||||
external_file_id = external_file_id or file_path
|
||||
file_size = os.path.getsize(file_path)
|
||||
with open(file_path, "rb") as file:
|
||||
return await self.upload_buffer(file, external_file_id, file_size)
|
||||
|
||||
async def upload_bytes(self, bytes: bytes, external_file_id: str) -> File:
|
||||
return await self.upload_buffer(BytesIO(bytes), external_file_id, len(bytes))
|
||||
|
||||
async def upload_buffer(
|
||||
self,
|
||||
buffer: BinaryIO,
|
||||
external_file_id: str,
|
||||
file_size: int,
|
||||
) -> File:
|
||||
if self.use_presigned_url:
|
||||
if getattr(buffer, "name", None):
|
||||
name = os.path.basename(str(getattr(buffer, "name", external_file_id)))
|
||||
else:
|
||||
name = external_file_id
|
||||
presigned_url = await self.client.files.generate_presigned_url(
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
request=FileCreate(
|
||||
name=name,
|
||||
external_file_id=external_file_id,
|
||||
file_size=file_size,
|
||||
),
|
||||
)
|
||||
httpx_client = self.client._client_wrapper.httpx_client
|
||||
upload_response = await httpx_client.put(
|
||||
presigned_url.url,
|
||||
data=buffer.read(),
|
||||
)
|
||||
upload_response.raise_for_status()
|
||||
return await self.client.files.get_file(
|
||||
presigned_url.file_id,
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
)
|
||||
else:
|
||||
return await self.client.files.upload_file(
|
||||
upload_file=buffer,
|
||||
external_file_id=external_file_id,
|
||||
project_id=self.project_id,
|
||||
organization_id=self.organization_id,
|
||||
)
|
||||
@@ -25,7 +25,7 @@ from llama_index.core.readers.base import BasePydanticReader
|
||||
from llama_index.core.readers.file.base import get_default_fs
|
||||
from llama_index.core.schema import Document
|
||||
|
||||
from llama_cloud_services.utils import check_extra_params
|
||||
from llama_cloud_services.utils import check_extra_params, check_for_updates
|
||||
from llama_cloud_services.parse.types import JobResult
|
||||
from llama_cloud_services.parse.utils import (
|
||||
SUPPORTED_FILE_TYPES,
|
||||
@@ -541,6 +541,11 @@ class LlamaParse(BasePydanticReader):
|
||||
description="Whether to use the vendor multimodal API.",
|
||||
)
|
||||
|
||||
check_for_updates: Optional[bool] = Field(
|
||||
default=False,
|
||||
description="Automatically check for Python SDK updates.",
|
||||
)
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def warn_extra_params(cls, data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
@@ -596,6 +601,16 @@ class LlamaParse(BasePydanticReader):
|
||||
|
||||
return self._aclient
|
||||
|
||||
_update_checked = False
|
||||
|
||||
async def _check_for_updates(self) -> None:
|
||||
if self.check_for_updates and not self._update_checked:
|
||||
try:
|
||||
await check_for_updates(self.aclient, quiet=False)
|
||||
self._update_checked = True
|
||||
except (ValueError, httpx.HTTPStatusError):
|
||||
pass
|
||||
|
||||
@asynccontextmanager
|
||||
async def client_context(self) -> AsyncGenerator[httpx.AsyncClient, None]:
|
||||
"""Create a context for the HTTPX client."""
|
||||
@@ -645,6 +660,7 @@ class LlamaParse(BasePydanticReader):
|
||||
fs: Optional[AbstractFileSystem] = None,
|
||||
partition_target_pages: Optional[str] = None,
|
||||
) -> str:
|
||||
await self._check_for_updates()
|
||||
files = None
|
||||
file_handle = None
|
||||
input_url = file_input if self._is_input_url(file_input) else None
|
||||
@@ -1534,6 +1550,7 @@ class LlamaParse(BasePydanticReader):
|
||||
self, json_result: List[dict], download_path: str, asset_key: str
|
||||
) -> List[dict]:
|
||||
"""Download assets (images or charts) from the parsed result."""
|
||||
await self._check_for_updates()
|
||||
# Make the download path
|
||||
if not os.path.exists(download_path):
|
||||
os.makedirs(download_path)
|
||||
@@ -1642,6 +1659,7 @@ class LlamaParse(BasePydanticReader):
|
||||
self, json_result: List[dict], download_path: str
|
||||
) -> List[dict]:
|
||||
"""Download xlsx from the parsed result."""
|
||||
await self._check_for_updates()
|
||||
# make the download path
|
||||
if not os.path.exists(download_path):
|
||||
os.makedirs(download_path)
|
||||
|
||||
@@ -52,6 +52,10 @@ class PageItem(BaseModel):
|
||||
bBox: Optional[BBox] = Field(
|
||||
default=None, description="The bounding box of the item."
|
||||
)
|
||||
html: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The HTML-formatted content of the item. Only applicable for table items when output_tables_as_HTML=True.",
|
||||
)
|
||||
|
||||
|
||||
class ImageItem(BaseModel):
|
||||
@@ -108,7 +112,7 @@ class ChartItem(BaseModel):
|
||||
class Page(BaseModel):
|
||||
"""A page of the document."""
|
||||
|
||||
page: int = Field(description="The page number.")
|
||||
page: int = Field(default=0, description="The page number.")
|
||||
text: Optional[str] = Field(default=None, description="The text of the page.")
|
||||
md: Optional[str] = Field(default=None, description="The markdown of the page.")
|
||||
images: List[ImageItem] = Field(
|
||||
@@ -149,6 +153,12 @@ class Page(BaseModel):
|
||||
noTextContent: bool = Field(
|
||||
default=False, description="Whether the page has no text content."
|
||||
)
|
||||
isAudioTranscript: bool = Field(
|
||||
default=False, description="Whether the page is an audio transcript."
|
||||
)
|
||||
durationInSeconds: Optional[float] = Field(
|
||||
default=None, description="The duration of the audio transcript in seconds."
|
||||
)
|
||||
|
||||
|
||||
class JobResult(BaseModel):
|
||||
|
||||
@@ -1,7 +1,21 @@
|
||||
import os
|
||||
import importlib.metadata
|
||||
from contextlib import contextmanager
|
||||
from typing import Generator
|
||||
import difflib
|
||||
from llama_cloud.types import StatusEnum
|
||||
import httpx
|
||||
import packaging.version
|
||||
from pydantic import BaseModel
|
||||
from typing import Any, Dict, List, Tuple, Type
|
||||
|
||||
# Asyncio error messages
|
||||
nest_asyncio_err = "cannot be called from a running event loop"
|
||||
nest_asyncio_msg = (
|
||||
"The event loop is already running. "
|
||||
"Add `import nest_asyncio; nest_asyncio.apply()` to your code to fix this issue."
|
||||
)
|
||||
|
||||
|
||||
def check_extra_params(
|
||||
model_cls: Type[BaseModel], data: Dict[str, Any]
|
||||
@@ -27,3 +41,66 @@ def check_extra_params(
|
||||
)
|
||||
|
||||
return extra_params, suggestions
|
||||
|
||||
|
||||
def is_terminal_status(status: StatusEnum) -> bool:
|
||||
"""
|
||||
Check if a status is terminal, i.e. the job is done and no more updates are expected.
|
||||
Note: this must be updated if the status enum is updated.
|
||||
|
||||
Args:
|
||||
status: The status to check
|
||||
|
||||
Returns:
|
||||
True if the status is terminal, False otherwise
|
||||
"""
|
||||
return status in {
|
||||
StatusEnum.SUCCESS,
|
||||
StatusEnum.ERROR,
|
||||
StatusEnum.CANCELLED,
|
||||
StatusEnum.PARTIAL_SUCCESS,
|
||||
}
|
||||
|
||||
|
||||
async def check_for_updates(client: httpx.AsyncClient, quiet: bool = True) -> bool:
|
||||
"""Check if an SDK update is available.
|
||||
|
||||
Args:
|
||||
client: HTTPX client to use.
|
||||
quiet: If False, update availability will also be printed to stdout.
|
||||
|
||||
Returns: True if an update is available.
|
||||
|
||||
Raises:
|
||||
ValueError: Failed to get a valid release version from PyPI.
|
||||
"""
|
||||
package_name = "llama-cloud-services"
|
||||
r = await client.get(f"https://pypi.org/pypi/{package_name}/json")
|
||||
version = r.json().get("info", {}).get("version", "")
|
||||
if not version:
|
||||
raise ValueError("Failed to fetch package info from PyPI")
|
||||
latest = packaging.version.parse(version)
|
||||
current = packaging.version.parse(importlib.metadata.version(package_name))
|
||||
if current < latest:
|
||||
if not quiet:
|
||||
msg = [
|
||||
f"\u26A0\uFE0F {package_name} is out of date",
|
||||
f"Current version: {current} | Latest: {latest}",
|
||||
"To upgrade: pip install -U --force-reinstall llama-cloud-services",
|
||||
]
|
||||
print(os.linesep.join(msg))
|
||||
return True
|
||||
elif not quiet:
|
||||
print(f"{package_name} is up to date")
|
||||
return False
|
||||
|
||||
|
||||
@contextmanager
|
||||
def augment_async_errors() -> Generator[None, None, None]:
|
||||
"""Context manager to add helpful information for errors due to nested event loops."""
|
||||
try:
|
||||
yield
|
||||
except RuntimeError as e:
|
||||
if nest_asyncio_err in str(e):
|
||||
raise RuntimeError(nest_asyncio_msg)
|
||||
raise
|
||||
|
||||
@@ -11,13 +11,13 @@ dev = [
|
||||
|
||||
[project]
|
||||
name = "llama-parse"
|
||||
version = "0.6.54"
|
||||
version = "0.6.60"
|
||||
description = "Parse files into RAG-Optimized formats."
|
||||
authors = [{name = "Logan Markewich", email = "logan@llamaindex.ai"}]
|
||||
requires-python = ">=3.9,<4.0"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
dependencies = ["llama-cloud-services>=0.6.54"]
|
||||
dependencies = ["llama-cloud-services>=0.6.60"]
|
||||
|
||||
[project.scripts]
|
||||
llama-parse = "llama_parse.cli.main:parse"
|
||||
|
||||
+7
-4
@@ -5,6 +5,7 @@ build-backend = "hatchling.build"
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pytest>=8.0.0,<9",
|
||||
"pytest-xdist>=3.6.1,<4",
|
||||
"pytest-asyncio",
|
||||
"ipykernel>=6.29.0,<7",
|
||||
"pre-commit==3.2.0",
|
||||
@@ -12,12 +13,13 @@ dev = [
|
||||
"deepdiff>=8.1.1,<9",
|
||||
"ipython>=8.12.3,<9",
|
||||
"jupyter>=1.1.1,<2",
|
||||
"mypy>=1.14.1,<2"
|
||||
"mypy>=1.14.1,<2",
|
||||
"pydantic-settings>=2.10.1"
|
||||
]
|
||||
|
||||
[project]
|
||||
name = "llama-cloud-services"
|
||||
version = "0.6.54"
|
||||
version = "0.6.61"
|
||||
description = "Tailored SDK clients for LlamaCloud services."
|
||||
authors = [{name = "Logan Markewich", email = "logan@runllama.ai"}]
|
||||
requires-python = ">=3.9,<4.0"
|
||||
@@ -25,13 +27,14 @@ readme = "README.md"
|
||||
license = "MIT"
|
||||
dependencies = [
|
||||
"llama-index-core>=0.12.0",
|
||||
"llama-cloud==0.1.35",
|
||||
"llama-cloud==0.1.39",
|
||||
"pydantic>=2.8,!=2.10",
|
||||
"click>=8.1.7,<9",
|
||||
"python-dotenv>=1.0.1,<2",
|
||||
"eval-type-backport>=0.2.0,<0.3 ; python_version < '3.10'",
|
||||
"platformdirs>=4.3.7,<5",
|
||||
"tenacity>=8.5.0, <10.0"
|
||||
"tenacity>=8.5.0, <10.0",
|
||||
"packaging>=25.0"
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
[pytest]
|
||||
asyncio_mode = auto
|
||||
@@ -37,7 +37,7 @@ LLAMA_CLOUD_BASE_URL = os.getenv("LLAMA_CLOUD_BASE_URL")
|
||||
LLAMA_DEPLOY_DEPLOYMENT_NAME = os.getenv("LLAMA_DEPLOY_DEPLOYMENT_NAME")
|
||||
|
||||
|
||||
class TestData(BaseModel):
|
||||
class ExampleData(BaseModel):
|
||||
"""Simple test data model for agent data testing"""
|
||||
|
||||
name: str
|
||||
@@ -66,13 +66,13 @@ async def test_agent_data_crud_operations():
|
||||
# Create agent data client with unique collection name
|
||||
agent_data_client = AsyncAgentDataClient(
|
||||
client=client,
|
||||
type=TestData,
|
||||
type=ExampleData,
|
||||
collection=f"test-collection-{test_id[:8]}",
|
||||
agent_url_id=LLAMA_DEPLOY_DEPLOYMENT_NAME,
|
||||
)
|
||||
|
||||
# Create test data
|
||||
test_data = TestData(name="test-item", test_id=test_id, value=42)
|
||||
test_data = ExampleData(name="test-item", test_id=test_id, value=42)
|
||||
|
||||
created_item = None
|
||||
try:
|
||||
@@ -107,7 +107,7 @@ async def test_agent_data_crud_operations():
|
||||
assert aggregate_results.items[0].count == 1
|
||||
|
||||
# Test UPDATE
|
||||
updated_data = TestData(name="updated-item", test_id=test_id, value=84)
|
||||
updated_data = ExampleData(name="updated-item", test_id=test_id, value=84)
|
||||
updated_item = await agent_data_client.update_item(
|
||||
created_item.id, updated_data
|
||||
)
|
||||
|
||||
@@ -0,0 +1,212 @@
|
||||
import os
|
||||
import pytest
|
||||
from llama_cloud.client import AsyncLlamaCloud
|
||||
from llama_cloud.types import Project, ClassifierRule, ClassifyJobResults
|
||||
from llama_cloud_services.beta.classifier.client import ClassifyClient
|
||||
from llama_cloud_services.files.client import FileClient
|
||||
from llama_cloud.errors.unprocessable_entity_error import UnprocessableEntityError
|
||||
from tests.conftest import EndToEndTestSettings
|
||||
import nest_asyncio
|
||||
|
||||
# Skip all tests if API key is not set
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not os.getenv("LLAMA_CLOUD_API_KEY"), reason="LLAMA_CLOUD_API_KEY not set"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_llama_cloud_client(
|
||||
e2e_test_settings: EndToEndTestSettings,
|
||||
) -> AsyncLlamaCloud:
|
||||
return AsyncLlamaCloud(
|
||||
token=e2e_test_settings.LLAMA_CLOUD_API_KEY.get_secret_value(),
|
||||
base_url=e2e_test_settings.LLAMA_CLOUD_BASE_URL,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def project(
|
||||
async_llama_cloud_client: AsyncLlamaCloud, e2e_test_settings: EndToEndTestSettings
|
||||
) -> Project:
|
||||
projects = await async_llama_cloud_client.projects.list_projects(
|
||||
project_name=e2e_test_settings.LLAMA_CLOUD_PROJECT_NAME,
|
||||
organization_id=e2e_test_settings.LLAMA_CLOUD_ORGANIZATION_ID,
|
||||
)
|
||||
assert len(projects) == 1
|
||||
return projects[0]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def classify_client(
|
||||
async_llama_cloud_client: AsyncLlamaCloud, project: Project
|
||||
) -> ClassifyClient:
|
||||
return ClassifyClient(
|
||||
async_llama_cloud_client,
|
||||
project_id=project.id,
|
||||
organization_id=project.organization_id,
|
||||
polling_interval=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def file_client(
|
||||
async_llama_cloud_client: AsyncLlamaCloud, project: Project
|
||||
) -> FileClient:
|
||||
return FileClient(
|
||||
async_llama_cloud_client,
|
||||
project_id=project.id,
|
||||
organization_id=project.organization_id,
|
||||
use_presigned_url=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def apply_nest_asyncio():
|
||||
nest_asyncio.apply()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def simple_pdf_file_path() -> str:
|
||||
return "tests/test_files/index/Simple PDF Slides.pdf"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def research_paper_path() -> str:
|
||||
return "tests/test_files/attention_is_all_you_need_chart.pdf"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def classification_rules() -> list[ClassifierRule]:
|
||||
return [
|
||||
ClassifierRule(
|
||||
type="number",
|
||||
description="Documents with numbers",
|
||||
),
|
||||
ClassifierRule(
|
||||
type="research_paper",
|
||||
description="Research papers",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
parameterize_sync_and_async = pytest.mark.parametrize("sync", [True, False])
|
||||
|
||||
|
||||
@parameterize_sync_and_async
|
||||
@pytest.mark.asyncio
|
||||
async def test_classify_file_ids(
|
||||
classify_client: ClassifyClient,
|
||||
file_client: FileClient,
|
||||
simple_pdf_file_path: str,
|
||||
research_paper_path: str,
|
||||
classification_rules: list[ClassifierRule],
|
||||
sync: bool,
|
||||
):
|
||||
"""Test classifying files by their IDs"""
|
||||
# Upload test files first to get their IDs
|
||||
pdf_file = await file_client.upload_file(simple_pdf_file_path)
|
||||
research_paper_file = await file_client.upload_file(research_paper_path)
|
||||
|
||||
# Classify the uploaded files
|
||||
if sync:
|
||||
results = classify_client.classify_file_ids(
|
||||
rules=classification_rules, file_ids=[pdf_file.id, research_paper_file.id]
|
||||
)
|
||||
else:
|
||||
results = await classify_client.aclassify_file_ids(
|
||||
rules=classification_rules, file_ids=[pdf_file.id, research_paper_file.id]
|
||||
)
|
||||
|
||||
assert isinstance(results, ClassifyJobResults)
|
||||
assert len(results.items) == 2
|
||||
|
||||
file_id_to_expected_type = {
|
||||
pdf_file.id: "number",
|
||||
research_paper_file.id: "research_paper",
|
||||
}
|
||||
# Verify each file got classified
|
||||
for item in results.items:
|
||||
expected_type = file_id_to_expected_type[item.file_id]
|
||||
assert item.result.type == expected_type
|
||||
|
||||
|
||||
@parameterize_sync_and_async
|
||||
@pytest.mark.asyncio
|
||||
async def test_classify_file_path(
|
||||
classify_client: ClassifyClient,
|
||||
simple_pdf_file_path: str,
|
||||
classification_rules: list[ClassifierRule],
|
||||
sync: bool,
|
||||
):
|
||||
"""Test classifying a single file by path"""
|
||||
# Classify the file
|
||||
if sync:
|
||||
results = classify_client.classify_file_path(
|
||||
rules=classification_rules, file_input_path=simple_pdf_file_path
|
||||
)
|
||||
else:
|
||||
results = await classify_client.aclassify_file_path(
|
||||
rules=classification_rules, file_input_path=simple_pdf_file_path
|
||||
)
|
||||
|
||||
assert isinstance(results, ClassifyJobResults)
|
||||
assert len(results.items) == 1
|
||||
|
||||
# Verify the file got classified
|
||||
item = results.items[0]
|
||||
assert item.result.type == "number"
|
||||
|
||||
|
||||
@parameterize_sync_and_async
|
||||
@pytest.mark.asyncio
|
||||
async def test_classify_file_paths(
|
||||
classify_client: ClassifyClient,
|
||||
file_client: FileClient,
|
||||
simple_pdf_file_path: str,
|
||||
research_paper_path: str,
|
||||
classification_rules: list[ClassifierRule],
|
||||
sync: bool,
|
||||
):
|
||||
"""Test classifying multiple files by paths"""
|
||||
# Classify all test files
|
||||
if sync:
|
||||
results = classify_client.classify_file_paths(
|
||||
rules=classification_rules,
|
||||
file_input_paths=[simple_pdf_file_path, research_paper_path],
|
||||
)
|
||||
else:
|
||||
results = await classify_client.aclassify_file_paths(
|
||||
rules=classification_rules,
|
||||
file_input_paths=[simple_pdf_file_path, research_paper_path],
|
||||
)
|
||||
|
||||
assert isinstance(results, ClassifyJobResults)
|
||||
assert len(results.items) == 2
|
||||
|
||||
file_name_to_expected_type = {
|
||||
os.path.basename(simple_pdf_file_path): "number",
|
||||
os.path.basename(research_paper_path): "research_paper",
|
||||
}
|
||||
# Verify each file got classified
|
||||
for item in results.items:
|
||||
file = await file_client.get_file(item.file_id)
|
||||
expected_type = file_name_to_expected_type[file.name]
|
||||
assert item.result.type == expected_type
|
||||
|
||||
|
||||
@parameterize_sync_and_async
|
||||
@pytest.mark.asyncio
|
||||
async def test_classify_empty_file_list(
|
||||
classify_client: ClassifyClient,
|
||||
classification_rules: list[ClassifierRule],
|
||||
sync: bool,
|
||||
):
|
||||
"""Test classifying an empty list of files"""
|
||||
# This should throw an error
|
||||
with pytest.raises(UnprocessableEntityError):
|
||||
if sync:
|
||||
classify_client.classify_file_ids(rules=classification_rules, file_ids=[])
|
||||
else:
|
||||
await classify_client.aclassify_file_ids(
|
||||
rules=classification_rules, file_ids=[]
|
||||
)
|
||||
@@ -0,0 +1,26 @@
|
||||
import pytest
|
||||
from llama_index.core.constants import DEFAULT_BASE_URL
|
||||
from pydantic import Field, SecretStr
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class EndToEndTestSettings(BaseSettings):
|
||||
LLAMA_CLOUD_BASE_URL: str = Field(
|
||||
description="The base URL of the LlamaCloud API", default=DEFAULT_BASE_URL
|
||||
)
|
||||
LLAMA_CLOUD_API_KEY: SecretStr = Field(
|
||||
description="The API key for the LlamaCloud API"
|
||||
)
|
||||
LLAMA_CLOUD_ORGANIZATION_ID: str | None = Field(
|
||||
description="The organization ID for the LlamaCloud API",
|
||||
default=None,
|
||||
)
|
||||
LLAMA_CLOUD_PROJECT_NAME: str = Field(
|
||||
description="The project name for the LlamaCloud API",
|
||||
default="framework_integration_test",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def e2e_test_settings() -> EndToEndTestSettings:
|
||||
return EndToEndTestSettings()
|
||||
@@ -6,6 +6,12 @@ from llama_cloud_services.extract import LlamaExtract
|
||||
_TEST_AGENTS_TO_CLEANUP: List[str] = []
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
"""Register custom markers for extract tests."""
|
||||
config.addinivalue_line("markers", "agent_name: custom agent name for test")
|
||||
config.addinivalue_line("markers", "agent_schema: custom agent schema for test")
|
||||
|
||||
|
||||
def pytest_sessionfinish(session, exitstatus):
|
||||
"""Hook that runs after all tests complete - cleanup agents here"""
|
||||
print(
|
||||
|
||||
@@ -23,8 +23,9 @@ LLAMA_CLOUD_API_KEY = os.getenv("LLAMA_CLOUD_API_KEY")
|
||||
LLAMA_CLOUD_BASE_URL = os.getenv("LLAMA_CLOUD_BASE_URL")
|
||||
LLAMA_CLOUD_PROJECT_ID = os.getenv("LLAMA_CLOUD_PROJECT_ID")
|
||||
|
||||
TestCase = namedtuple(
|
||||
"TestCase", ["name", "schema_path", "config", "input_file", "expected_output"]
|
||||
BenchmarkTestCase = namedtuple(
|
||||
"BenchmarkTestCase",
|
||||
["name", "schema_path", "config", "input_file", "expected_output"],
|
||||
)
|
||||
|
||||
|
||||
@@ -32,7 +33,7 @@ def get_test_cases():
|
||||
"""Get all test cases from TEST_DIR.
|
||||
|
||||
Returns:
|
||||
List[TestCase]: List of test cases
|
||||
List[BenchmarkTestCase]: List of test cases
|
||||
"""
|
||||
test_cases = []
|
||||
|
||||
@@ -73,7 +74,7 @@ def get_test_cases():
|
||||
test_name = f"{data_type}/{os.path.basename(input_file)}"
|
||||
for setting in settings:
|
||||
test_cases.append(
|
||||
TestCase(
|
||||
BenchmarkTestCase(
|
||||
name=test_name,
|
||||
schema_path=schema_path,
|
||||
input_file=input_file,
|
||||
@@ -100,7 +101,7 @@ def extractor():
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def extraction_agent(test_case: TestCase, extractor: LlamaExtract):
|
||||
def extraction_agent(test_case: BenchmarkTestCase, extractor: LlamaExtract):
|
||||
"""Fixture to create and cleanup extraction agent for each test."""
|
||||
# Create unique name with random UUID (important for CI to avoid conflicts)
|
||||
unique_id = uuid.uuid4().hex[:8]
|
||||
@@ -124,13 +125,13 @@ def extraction_agent(test_case: TestCase, extractor: LlamaExtract):
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
"CI" in os.environ,
|
||||
reason="CI environment is not suitable for benchmarking",
|
||||
"CI" in os.environ or not LLAMA_CLOUD_API_KEY,
|
||||
reason="LLAMA_CLOUD_API_KEY not set or CI environment not suitable for benchmarking",
|
||||
)
|
||||
@pytest.mark.parametrize("test_case", get_test_cases(), ids=lambda x: x.name)
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_extraction(
|
||||
test_case: TestCase, extraction_agent: ExtractionAgent
|
||||
test_case: BenchmarkTestCase, extraction_agent: ExtractionAgent
|
||||
) -> None:
|
||||
start = perf_counter()
|
||||
result = await extraction_agent._run_extraction_test(
|
||||
|
||||
@@ -4,6 +4,7 @@ from pathlib import Path
|
||||
from pydantic import BaseModel
|
||||
|
||||
from llama_cloud_services.extract import LlamaExtract, ExtractionAgent, SourceText
|
||||
from llama_cloud.types import ExtractConfig, ExtractMode, ExtractRun
|
||||
from tests.extract.util import load_test_dotenv
|
||||
from .conftest import register_agent_for_cleanup
|
||||
|
||||
@@ -21,7 +22,7 @@ pytestmark = pytest.mark.skipif(
|
||||
|
||||
|
||||
# Test data
|
||||
class TestSchema(BaseModel):
|
||||
class ExampleSchema(BaseModel):
|
||||
title: str
|
||||
summary: str
|
||||
|
||||
@@ -107,7 +108,7 @@ class TestLlamaExtract:
|
||||
assert isinstance(test_agent, ExtractionAgent)
|
||||
|
||||
@pytest.mark.agent_name("test-pydantic-schema-agent")
|
||||
@pytest.mark.agent_schema((TestSchema,))
|
||||
@pytest.mark.agent_schema((ExampleSchema,))
|
||||
def test_create_agent_with_pydantic_schema(self, test_agent):
|
||||
assert isinstance(test_agent, ExtractionAgent)
|
||||
|
||||
@@ -222,7 +223,189 @@ class TestExtractionAgent:
|
||||
|
||||
def test_delete_extraction_run(self, test_agent: ExtractionAgent):
|
||||
assert test_agent.list_extraction_runs().total == 0
|
||||
run = test_agent.extract(TEST_PDF)
|
||||
run: ExtractRun = test_agent.extract(TEST_PDF)
|
||||
test_agent.delete_extraction_run(run.id)
|
||||
runs = test_agent.list_extraction_runs()
|
||||
assert runs.total == 0
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
"CI" in os.environ, reason="Test locally; functionality is mostly duplicated."
|
||||
)
|
||||
class TestStatelessExtraction:
|
||||
"""Tests for stateless extraction methods that don't require creating an agent."""
|
||||
|
||||
@pytest.fixture
|
||||
def test_config(self):
|
||||
return ExtractConfig(extraction_mode=ExtractMode.FAST)
|
||||
|
||||
@pytest.fixture
|
||||
def test_schema_dict(self):
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"title": {"type": "string"},
|
||||
"summary": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aextract_single_file(
|
||||
self, llama_extract, test_schema_dict, test_config
|
||||
):
|
||||
"""Test async stateless extraction with a single file."""
|
||||
result = await llama_extract.aextract(test_schema_dict, test_config, TEST_PDF)
|
||||
assert result.status == "SUCCESS"
|
||||
assert result.data is not None
|
||||
assert isinstance(result.data, dict)
|
||||
assert "title" in result.data
|
||||
assert "summary" in result.data
|
||||
|
||||
def test_extract_single_file(self, llama_extract, test_schema_dict, test_config):
|
||||
"""Test synchronous stateless extraction with a single file."""
|
||||
result = llama_extract.extract(test_schema_dict, test_config, TEST_PDF)
|
||||
assert result.status == "SUCCESS"
|
||||
assert result.data is not None
|
||||
assert isinstance(result.data, dict)
|
||||
assert "title" in result.data
|
||||
assert "summary" in result.data
|
||||
|
||||
def test_extract_from_bytes_with_source_text(
|
||||
self, llama_extract, test_schema_dict, test_config
|
||||
):
|
||||
"""Test stateless extraction from bytes using SourceText with filename."""
|
||||
with open(TEST_PDF, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
source_text = SourceText(file=file_bytes, filename=TEST_PDF.name)
|
||||
result = llama_extract.extract(test_schema_dict, test_config, source_text)
|
||||
assert result.status == "SUCCESS"
|
||||
assert result.data is not None
|
||||
assert isinstance(result.data, dict)
|
||||
assert "title" in result.data
|
||||
assert "summary" in result.data
|
||||
|
||||
def test_extract_from_source_text_with_file(
|
||||
self, llama_extract, test_schema_dict, test_config
|
||||
):
|
||||
"""Test stateless extraction from SourceText with file."""
|
||||
source_text = SourceText(file=TEST_PDF, filename=TEST_PDF.name)
|
||||
result = llama_extract.extract(test_schema_dict, test_config, source_text)
|
||||
assert result.status == "SUCCESS"
|
||||
assert result.data is not None
|
||||
assert isinstance(result.data, dict)
|
||||
assert "title" in result.data
|
||||
assert "summary" in result.data
|
||||
|
||||
def test_extract_from_buffered_io(
|
||||
self, llama_extract, test_schema_dict, test_config
|
||||
):
|
||||
"""Test stateless extraction from BufferedIO file handle."""
|
||||
with open(TEST_PDF, "rb") as file_handle:
|
||||
result = llama_extract.extract(test_schema_dict, test_config, file_handle)
|
||||
assert result.status == "SUCCESS"
|
||||
assert result.data is not None
|
||||
assert isinstance(result.data, dict)
|
||||
assert "title" in result.data
|
||||
assert "summary" in result.data
|
||||
|
||||
def test_extract_from_source_text_with_text_content(
|
||||
self, llama_extract, test_schema_dict, test_config
|
||||
):
|
||||
"""Test stateless extraction from SourceText with text content."""
|
||||
TEST_TEXT = """
|
||||
# Llamas
|
||||
Llamas are social animals and live with others as a herd. Their wool is soft and
|
||||
contains only a small amount of lanolin. Llamas can learn simple tasks after a
|
||||
few repetitions. When using a pack, they can carry about 25 to 30% of their body
|
||||
weight for 8 to 13 km (5–8 miles). The name llama was adopted by European settlers
|
||||
from native Peruvians.
|
||||
"""
|
||||
source_text = SourceText(text_content=TEST_TEXT)
|
||||
result = llama_extract.extract(test_schema_dict, test_config, source_text)
|
||||
assert result.status == "SUCCESS"
|
||||
assert result.data is not None
|
||||
assert isinstance(result.data, dict)
|
||||
assert "title" in result.data
|
||||
assert "summary" in result.data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_extraction_single_file(
|
||||
self, llama_extract, test_schema_dict, test_config
|
||||
):
|
||||
"""Test queuing extraction job without waiting for completion."""
|
||||
job = await llama_extract.queue_extraction(
|
||||
test_schema_dict, test_config, TEST_PDF
|
||||
)
|
||||
assert hasattr(job, "id")
|
||||
assert hasattr(job, "status")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_multiple_files(
|
||||
self, llama_extract, test_schema_dict, test_config
|
||||
):
|
||||
"""Test stateless extraction with multiple files."""
|
||||
files = [TEST_PDF, TEST_PDF] # Using same file twice for testing
|
||||
results = await llama_extract.aextract(test_schema_dict, test_config, files)
|
||||
|
||||
assert isinstance(results, list)
|
||||
assert len(results) == 2
|
||||
|
||||
for result in results:
|
||||
assert result.status == "SUCCESS"
|
||||
assert result.data is not None
|
||||
assert isinstance(result.data, dict)
|
||||
assert "title" in result.data
|
||||
assert "summary" in result.data
|
||||
|
||||
def test_extract_with_pydantic_schema(self, llama_extract, test_config):
|
||||
"""Test stateless extraction with Pydantic schema."""
|
||||
result = llama_extract.extract(ExampleSchema, test_config, TEST_PDF)
|
||||
assert result.status == "SUCCESS"
|
||||
assert result.data is not None
|
||||
assert isinstance(result.data, dict)
|
||||
assert "title" in result.data
|
||||
assert "summary" in result.data
|
||||
|
||||
def test_extract_from_raw_bytes_raises_error(
|
||||
self, llama_extract, test_schema_dict, test_config
|
||||
):
|
||||
"""Test that raw bytes without filename raises an error."""
|
||||
with open(TEST_PDF, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Cannot determine file type from raw bytes"
|
||||
):
|
||||
llama_extract.extract(test_schema_dict, test_config, file_bytes)
|
||||
|
||||
def test_mime_type_detection(self, llama_extract):
|
||||
"""Test that MIME types are correctly detected for various file types."""
|
||||
# Test PDF
|
||||
assert llama_extract._get_mime_type(filename="test.pdf") == "application/pdf"
|
||||
|
||||
# Test DOCX
|
||||
assert (
|
||||
llama_extract._get_mime_type(filename="test.docx")
|
||||
== "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
)
|
||||
|
||||
# Test text files
|
||||
assert llama_extract._get_mime_type(filename="test.txt") == "text/plain"
|
||||
assert llama_extract._get_mime_type(filename="test.csv") == "text/csv"
|
||||
assert llama_extract._get_mime_type(filename="test.json") == "application/json"
|
||||
|
||||
# Test image files
|
||||
assert llama_extract._get_mime_type(filename="test.png") == "image/png"
|
||||
assert llama_extract._get_mime_type(filename="test.jpg") == "image/jpeg"
|
||||
|
||||
# Test file path
|
||||
from pathlib import Path
|
||||
|
||||
assert (
|
||||
llama_extract._get_mime_type(file_path=Path("test.pdf"))
|
||||
== "application/pdf"
|
||||
)
|
||||
|
||||
# Test unsupported file type
|
||||
with pytest.raises(ValueError, match="Unsupported file type: 'xyz'"):
|
||||
llama_extract._get_mime_type(filename="test.xyz")
|
||||
|
||||
@@ -19,8 +19,9 @@ LLAMA_CLOUD_API_KEY = os.getenv("LLAMA_CLOUD_API_KEY")
|
||||
LLAMA_CLOUD_BASE_URL = os.getenv("LLAMA_CLOUD_BASE_URL")
|
||||
LLAMA_CLOUD_PROJECT_ID = os.getenv("LLAMA_CLOUD_PROJECT_ID")
|
||||
|
||||
TestCase = namedtuple(
|
||||
"TestCase", ["name", "schema_path", "config", "input_file", "expected_output"]
|
||||
ExtractionTestCase = namedtuple(
|
||||
"ExtractionTestCase",
|
||||
["name", "schema_path", "config", "input_file", "expected_output"],
|
||||
)
|
||||
|
||||
|
||||
@@ -28,7 +29,7 @@ def get_test_cases():
|
||||
"""Get all test cases from TEST_DIR.
|
||||
|
||||
Returns:
|
||||
List[TestCase]: List of test cases
|
||||
List[ExtractionTestCase]: List of test cases
|
||||
"""
|
||||
test_cases = []
|
||||
|
||||
@@ -69,7 +70,7 @@ def get_test_cases():
|
||||
test_name = f"{data_type}/{os.path.basename(input_file)}"
|
||||
for setting in settings:
|
||||
test_cases.append(
|
||||
TestCase(
|
||||
ExtractionTestCase(
|
||||
name=test_name,
|
||||
schema_path=schema_path,
|
||||
input_file=input_file,
|
||||
@@ -96,7 +97,7 @@ def extractor():
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def extraction_agent(test_case: TestCase, extractor: LlamaExtract):
|
||||
def extraction_agent(test_case: ExtractionTestCase, extractor: LlamaExtract):
|
||||
"""Fixture to create and cleanup extraction agent for each test."""
|
||||
# Create unique name with random UUID (important for CI to avoid conflicts)
|
||||
unique_id = uuid.uuid4().hex[:8]
|
||||
@@ -128,7 +129,9 @@ def extraction_agent(test_case: TestCase, extractor: LlamaExtract):
|
||||
reason="LLAMA_CLOUD_API_KEY not set",
|
||||
)
|
||||
@pytest.mark.parametrize("test_case", get_test_cases(), ids=lambda x: x.name)
|
||||
def test_extraction(test_case: TestCase, extraction_agent: ExtractionAgent) -> None:
|
||||
def test_extraction(
|
||||
test_case: ExtractionTestCase, extraction_agent: ExtractionAgent
|
||||
) -> None:
|
||||
result = extraction_agent.extract(test_case.input_file).data # type: ignore
|
||||
with open(test_case.expected_output, "r") as f:
|
||||
expected = json.load(f)
|
||||
|
||||
@@ -16,7 +16,6 @@ from llama_cloud import (
|
||||
from llama_cloud.client import LlamaCloud
|
||||
from llama_index.core.bridge.pydantic import BaseModel
|
||||
from llama_index.core.constants import DEFAULT_BASE_URL
|
||||
from llama_index.core.indices.managed.base import BaseManagedIndex
|
||||
from llama_index.core.schema import Document, ImageNode
|
||||
from llama_cloud_services.index import (
|
||||
LlamaCloudIndex,
|
||||
@@ -28,8 +27,6 @@ api_key = os.environ.get("LLAMA_CLOUD_API_KEY", None)
|
||||
organization_id = os.environ.get("LLAMA_CLOUD_ORGANIZATION_ID", None)
|
||||
project_name = os.environ.get("LLAMA_CLOUD_PROJECT_NAME", "framework_integration_test")
|
||||
|
||||
print("api-key", api_key, "base-url", base_url)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def remote_file() -> Tuple[str, str]:
|
||||
@@ -93,16 +90,6 @@ def _setup_index_with_file(
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_class():
|
||||
names_of_base_classes = [b.__name__ for b in LlamaCloudIndex.__mro__]
|
||||
assert BaseManagedIndex.__name__ in names_of_base_classes
|
||||
|
||||
|
||||
def test_conflicting_index_identifiers():
|
||||
with pytest.raises(ValueError):
|
||||
LlamaCloudIndex(name="test", pipeline_id="test", index_id="test")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not base_url or not api_key, reason="No platform base url or api key set"
|
||||
)
|
||||
|
||||
@@ -202,3 +202,12 @@ async def test_get_result(markdown_parser: LlamaParse) -> None:
|
||||
result = await markdown_parser.aget_result(expected.job_id)
|
||||
assert result.job_id == expected.job_id
|
||||
assert len(result.pages) == len(expected.pages)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_audio() -> None:
|
||||
parser = LlamaParse()
|
||||
filepath = "tests/test_files/hello_world.m4a"
|
||||
|
||||
result = await parser.aparse(filepath)
|
||||
assert result.job_id is not None
|
||||
|
||||
Binary file not shown.
@@ -0,0 +1,147 @@
|
||||
import os
|
||||
import pytest
|
||||
from io import BytesIO
|
||||
from llama_cloud.client import AsyncLlamaCloud
|
||||
from llama_cloud.types import Project, File
|
||||
from llama_cloud_services.files.client import FileClient
|
||||
from tests.conftest import EndToEndTestSettings
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def llama_cloud_client(
|
||||
e2e_test_settings: EndToEndTestSettings,
|
||||
) -> AsyncLlamaCloud:
|
||||
return AsyncLlamaCloud(
|
||||
token=e2e_test_settings.LLAMA_CLOUD_API_KEY.get_secret_value(),
|
||||
base_url=e2e_test_settings.LLAMA_CLOUD_BASE_URL,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def project(
|
||||
llama_cloud_client: AsyncLlamaCloud, e2e_test_settings: EndToEndTestSettings
|
||||
) -> Project:
|
||||
projects = await llama_cloud_client.projects.list_projects(
|
||||
project_name=e2e_test_settings.LLAMA_CLOUD_PROJECT_NAME,
|
||||
organization_id=e2e_test_settings.LLAMA_CLOUD_ORGANIZATION_ID,
|
||||
)
|
||||
assert len(projects) == 1
|
||||
return projects[0]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def file_client(
|
||||
llama_cloud_client: AsyncLlamaCloud, project: Project, use_presigned_url: bool
|
||||
) -> FileClient:
|
||||
return FileClient(
|
||||
llama_cloud_client,
|
||||
project_id=project.id,
|
||||
organization_id=project.organization_id,
|
||||
use_presigned_url=use_presigned_url,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_file() -> str:
|
||||
return "tests/test_files/index/Simple PDF Slides.pdf"
|
||||
|
||||
|
||||
parametrize_use_presigned_url = pytest.mark.parametrize(
|
||||
"use_presigned_url", [True, False]
|
||||
)
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_file_from_path(file_client: FileClient, test_file: str):
|
||||
"""Test uploading a file from file path"""
|
||||
external_file_id = f"test_upload_path_{os.getpid()}"
|
||||
uploaded_file = await file_client.upload_file(test_file, external_file_id)
|
||||
|
||||
assert isinstance(uploaded_file, File)
|
||||
expected_name = os.path.basename(test_file)
|
||||
assert uploaded_file.name == expected_name
|
||||
assert uploaded_file.external_file_id == external_file_id
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_bytes(
|
||||
file_client: FileClient, test_file: str, use_presigned_url: bool
|
||||
):
|
||||
"""Test uploading a file from bytes"""
|
||||
# Read file as bytes
|
||||
with open(test_file, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
|
||||
external_file_id = f"test_upload_bytes_{os.getpid()}"
|
||||
uploaded_file = await file_client.upload_bytes(file_bytes, external_file_id)
|
||||
|
||||
assert isinstance(uploaded_file, File)
|
||||
expected_name = external_file_id if use_presigned_url else "upload"
|
||||
assert uploaded_file.name == expected_name
|
||||
assert uploaded_file.external_file_id == external_file_id
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_buffer(
|
||||
file_client: FileClient, test_file: str, use_presigned_url: bool
|
||||
):
|
||||
"""Test uploading a file from buffer"""
|
||||
# Read file as bytes and create buffer
|
||||
with open(test_file, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
|
||||
buffer = BytesIO(file_bytes)
|
||||
file_size = len(file_bytes)
|
||||
external_file_id = f"test_upload_buffer_{os.getpid()}"
|
||||
|
||||
uploaded_file = await file_client.upload_buffer(buffer, external_file_id, file_size)
|
||||
|
||||
assert isinstance(uploaded_file, File)
|
||||
expected_name = external_file_id if use_presigned_url else "upload"
|
||||
assert uploaded_file.name == expected_name
|
||||
assert uploaded_file.external_file_id == external_file_id
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_file(file_client: FileClient, test_file: str):
|
||||
"""Test retrieving a file by ID"""
|
||||
# Upload a file first
|
||||
external_file_id = f"test_get_file_{os.getpid()}"
|
||||
uploaded_file = await file_client.upload_file(test_file, external_file_id)
|
||||
|
||||
# Retrieve the file by ID
|
||||
retrieved_file = await file_client.get_file(uploaded_file.id)
|
||||
|
||||
assert isinstance(retrieved_file, File)
|
||||
assert retrieved_file == uploaded_file
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_with_default_external_id(file_client: FileClient, test_file: str):
|
||||
"""Test uploading file with default external_file_id"""
|
||||
# Upload file without specifying external_file_id
|
||||
uploaded_file = await file_client.upload_file(test_file)
|
||||
|
||||
assert isinstance(uploaded_file, File)
|
||||
assert uploaded_file.name == os.path.basename(test_file)
|
||||
assert uploaded_file.external_file_id == test_file
|
||||
|
||||
|
||||
@parametrize_use_presigned_url
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_file_content(file_client: FileClient, test_file: str):
|
||||
"""Test reading a file content"""
|
||||
# Upload a file first
|
||||
external_file_id = f"test_read_file_content_{os.getpid()}"
|
||||
uploaded_file = await file_client.upload_file(test_file, external_file_id)
|
||||
|
||||
# Read the file content
|
||||
file_content = await file_client.read_file_content(uploaded_file.id)
|
||||
with open(test_file, "rb") as f:
|
||||
expected_file_content = f.read()
|
||||
assert file_content == expected_file_content
|
||||
+119
-29
@@ -1,15 +1,18 @@
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import pytest
|
||||
from llama_cloud import ExtractRun, File
|
||||
from llama_cloud.types.agent_data import AgentData
|
||||
from llama_cloud.types.aggregate_group import AggregateGroup
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
|
||||
from llama_cloud_services.beta.agent_data.schema import (
|
||||
ExtractedData,
|
||||
ExtractedFieldMetadata,
|
||||
FieldCitation,
|
||||
InvalidExtractionData,
|
||||
TypedAgentData,
|
||||
TypedAggregateGroup,
|
||||
@@ -21,8 +24,8 @@ from llama_cloud_services.beta.agent_data.schema import (
|
||||
# Test data models
|
||||
class Person(BaseModel):
|
||||
name: str
|
||||
age: int
|
||||
email: str
|
||||
age: Optional[int] = None
|
||||
email: Optional[str] = None
|
||||
|
||||
|
||||
class Company(BaseModel):
|
||||
@@ -81,8 +84,12 @@ def test_extracted_data_create_method():
|
||||
|
||||
# Test with custom values using ExtractedFieldMetadata
|
||||
field_metadata = {
|
||||
"name": ExtractedFieldMetadata(confidence=0.99, page_number=1),
|
||||
"age": ExtractedFieldMetadata(confidence=0.85, page_number=1),
|
||||
"name": ExtractedFieldMetadata(
|
||||
confidence=0.99, citation=[FieldCitation(page=1)]
|
||||
),
|
||||
"age": ExtractedFieldMetadata(
|
||||
confidence=0.85, citation=[FieldCitation(page=1)]
|
||||
),
|
||||
}
|
||||
extracted_custom = ExtractedData.create(
|
||||
person, status="accepted", field_metadata=field_metadata
|
||||
@@ -254,14 +261,16 @@ def test_parse_extracted_field_metadata():
|
||||
# name should have parsed citation data
|
||||
assert isinstance(result["name"], ExtractedFieldMetadata)
|
||||
assert result["name"].confidence == 0.95
|
||||
assert result["name"].page_number == 1
|
||||
assert result["name"].matching_text == "John Smith"
|
||||
assert result["name"].citation == [
|
||||
FieldCitation(page=1, matching_text="John Smith")
|
||||
]
|
||||
|
||||
# age should handle float page number
|
||||
assert isinstance(result["age"], ExtractedFieldMetadata)
|
||||
assert result["age"].confidence == 0.87
|
||||
assert result["age"].page_number == 2 # Should be converted to int
|
||||
assert result["age"].matching_text == "25 years old"
|
||||
assert result["age"].citation == [
|
||||
FieldCitation(page=2, matching_text="25 years old")
|
||||
]
|
||||
|
||||
# email should handle empty citations
|
||||
assert isinstance(result["email"], ExtractedFieldMetadata)
|
||||
@@ -327,30 +336,38 @@ def test_parse_extracted_field_metadata_complex():
|
||||
reasoning="Combined key parametrics and construction from the datasheet for a structured title.",
|
||||
confidence=0.9470628580889779,
|
||||
extraction_confidence=0.9470628580889779,
|
||||
page_number=1,
|
||||
matching_text="PHE844/F844, Film, Metallized Polypropylene, Safety, 0.47 uF",
|
||||
citation=[
|
||||
FieldCitation(
|
||||
page=1,
|
||||
matching_text="PHE844/F844, Film, Metallized Polypropylene, Safety, 0.47 uF",
|
||||
)
|
||||
],
|
||||
),
|
||||
"manufacturer": ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.9997446550976602,
|
||||
extraction_confidence=0.9997446550976602,
|
||||
page_number=1,
|
||||
matching_text="YAGEO KEMET",
|
||||
citation=[FieldCitation(page=1, matching_text="YAGEO KEMET")],
|
||||
),
|
||||
"features": [
|
||||
ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.9999308195540074,
|
||||
extraction_confidence=0.9999308195540074,
|
||||
page_number=1,
|
||||
matching_text="Features</td><td>EMI Safety",
|
||||
citation=[
|
||||
FieldCitation(
|
||||
page=1,
|
||||
matching_text="Features</td><td>EMI Safety",
|
||||
)
|
||||
],
|
||||
),
|
||||
ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.8642493886452225,
|
||||
extraction_confidence=0.8642493886452225,
|
||||
page_number=1,
|
||||
matching_text="THB Performance</td><td>Yes",
|
||||
citation=[
|
||||
FieldCitation(page=1, matching_text="THB Performance</td><td>Yes")
|
||||
],
|
||||
),
|
||||
],
|
||||
"dimensions": {
|
||||
@@ -358,15 +375,13 @@ def test_parse_extracted_field_metadata_complex():
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.8986941382802304,
|
||||
extraction_confidence=0.8986941382802304,
|
||||
page_number=1,
|
||||
matching_text="L</td><td>41mm MAX",
|
||||
citation=[FieldCitation(page=1, matching_text="L</td><td>41mm MAX")],
|
||||
),
|
||||
"width": ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.9999377974447091,
|
||||
extraction_confidence=0.9999377974447091,
|
||||
page_number=1,
|
||||
matching_text="T</td><td>13mm MAX",
|
||||
citation=[FieldCitation(page=1, matching_text="T</td><td>13mm MAX")],
|
||||
),
|
||||
},
|
||||
}
|
||||
@@ -390,6 +405,7 @@ def create_file(
|
||||
|
||||
def create_extract_run(
|
||||
id: str = "extract-123",
|
||||
job_id: str = "job-123",
|
||||
data: Dict[str, Any] = {"name": "John Doe", "age": 30, "email": "john@example.com"},
|
||||
extraction_metadata: Dict[str, Any] = {
|
||||
"name": {
|
||||
@@ -408,6 +424,7 @@ def create_extract_run(
|
||||
return ExtractRun.parse_obj(
|
||||
{
|
||||
"id": id,
|
||||
"job_id": job_id,
|
||||
"data": data,
|
||||
"extraction_metadata": {
|
||||
"field_metadata": extraction_metadata,
|
||||
@@ -450,8 +467,9 @@ def test_extracted_data_from_extraction_result_success():
|
||||
# Verify field metadata was parsed
|
||||
assert isinstance(extracted.field_metadata["name"], ExtractedFieldMetadata)
|
||||
assert extracted.field_metadata["name"].confidence == 0.95
|
||||
assert extracted.field_metadata["name"].page_number == 1
|
||||
assert extracted.field_metadata["name"].matching_text == "John Doe"
|
||||
assert extracted.field_metadata["name"].citation == [
|
||||
FieldCitation(page=1, matching_text="John Doe")
|
||||
]
|
||||
|
||||
# Verify overall confidence was calculated
|
||||
expected_confidence = (0.95 + 0.87 + 0.92) / 3
|
||||
@@ -485,10 +503,9 @@ def test_extracted_data_from_extraction_result_invalid_data():
|
||||
# Create ExtractRun with data that doesn't match Person schema
|
||||
extract_run = create_extract_run(
|
||||
data={
|
||||
"name": "Valid Name",
|
||||
"missing_name": "Valid Name",
|
||||
"age": "not_a_number",
|
||||
"missing_email": True,
|
||||
}, # Invalid age, missing email
|
||||
}, # Invalid age, missing name
|
||||
extraction_metadata={
|
||||
"name": {"confidence": 0.9},
|
||||
},
|
||||
@@ -507,9 +524,8 @@ def test_extracted_data_from_extraction_result_invalid_data():
|
||||
assert isinstance(invalid_data, ExtractedData)
|
||||
assert invalid_data.status == "error"
|
||||
assert invalid_data.data == {
|
||||
"name": "Valid Name",
|
||||
"missing_name": "Valid Name",
|
||||
"age": "not_a_number",
|
||||
"missing_email": True,
|
||||
}
|
||||
assert invalid_data.file_id == "error-file"
|
||||
assert invalid_data.file_name == "bad_data.pdf"
|
||||
@@ -523,3 +539,77 @@ def test_extracted_data_from_extraction_result_invalid_data():
|
||||
assert isinstance(invalid_data.field_metadata["name"], ExtractedFieldMetadata)
|
||||
assert invalid_data.field_metadata["name"].confidence == 0.9
|
||||
assert invalid_data.overall_confidence == 0.9
|
||||
|
||||
|
||||
class Dimensions(BaseModel):
|
||||
length: Optional[str] = Field(
|
||||
None, description="Length in mm (Size, Longest Side, L)"
|
||||
)
|
||||
width: Optional[str] = Field(
|
||||
None, description="Width in mm (Breadth, Side Width, W)"
|
||||
)
|
||||
height: Optional[str] = Field(
|
||||
None, description="Height in mm (Thickness, Vertical Size, H)"
|
||||
)
|
||||
diameter: Optional[str] = Field(
|
||||
None,
|
||||
description="Diameter in mm (for radial or cylindrical types) (Outer Diameter, dt, OD, D, d<sub>t</sub>)",
|
||||
)
|
||||
lead_spacing: Optional[str] = Field(
|
||||
None, description="Lead spacing in mm (Pin Pitch, Terminal Gap, LS)"
|
||||
)
|
||||
|
||||
|
||||
class Capacitor(BaseModel):
|
||||
dimensions: Optional[Dimensions] = None
|
||||
|
||||
|
||||
def test_full_parse_nested_dimensions():
|
||||
with open(Path(__file__).parent.parent.parent / "data" / "capacitor.json") as f:
|
||||
data = json.load(f)
|
||||
result = ExtractedData.from_extraction_result(ExtractRun.parse_obj(data), Capacitor)
|
||||
expected = {
|
||||
"dimensions": {
|
||||
"diameter": ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=1.0,
|
||||
extraction_confidence=1.0,
|
||||
),
|
||||
"lead_spacing": ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.9999999031936799,
|
||||
extraction_confidence=0.9999999031936799,
|
||||
),
|
||||
"length": ExtractedFieldMetadata(
|
||||
reasoning="VERBATIM EXTRACTION",
|
||||
confidence=0.9999968039036192,
|
||||
extraction_confidence=0.9999968039036192,
|
||||
),
|
||||
}
|
||||
}
|
||||
assert result.field_metadata == expected
|
||||
parsed = ExtractedData.model_validate_json(result.model_dump_json())
|
||||
assert parsed.field_metadata == expected
|
||||
|
||||
|
||||
def test_parses_field_metadata_with_error_field():
|
||||
extract_run = create_extract_run(
|
||||
extraction_metadata={
|
||||
"name": {
|
||||
"confidence": 0.95,
|
||||
"citation": [{"page": 1, "matching_text": "John Smith"}],
|
||||
},
|
||||
"error": "This is an error",
|
||||
},
|
||||
)
|
||||
|
||||
parsed = ExtractedData.from_extraction_result(extract_run, Person)
|
||||
|
||||
assert parsed.field_metadata == {
|
||||
"name": ExtractedFieldMetadata(
|
||||
confidence=0.95,
|
||||
citation=[FieldCitation(page=1, matching_text="John Smith")],
|
||||
),
|
||||
}
|
||||
assert parsed.metadata.get("field_errors") == "This is an error"
|
||||
assert parsed.metadata.get("job_id") == "job-123"
|
||||
@@ -0,0 +1,116 @@
|
||||
{
|
||||
"id": "de058dda-6ca7-4eea-a426-da802f84f971",
|
||||
"created_at": "2025-08-13T15:45:39.286921Z",
|
||||
"updated_at": "2025-08-13T15:47:04.878069Z",
|
||||
"extraction_agent_id": "e834e99f-1f35-4748-b82f-03de4bd07ca6",
|
||||
"data_schema": {
|
||||
"additionalProperties": false,
|
||||
"properties": {
|
||||
"dimensions": {
|
||||
"anyOf": [
|
||||
{
|
||||
"additionalProperties": false,
|
||||
"properties": {
|
||||
"length": {
|
||||
"anyOf": [{ "type": "string" }, { "type": "null" }],
|
||||
"description": "Length in mm (Size, Longest Side, L)"
|
||||
},
|
||||
"width": {
|
||||
"anyOf": [{ "type": "string" }, { "type": "null" }],
|
||||
"description": "Width in mm (Breadth, Side Width, W)"
|
||||
},
|
||||
"height": {
|
||||
"anyOf": [{ "type": "string" }, { "type": "null" }],
|
||||
"description": "Height in mm (Thickness, Vertical Size, H)"
|
||||
},
|
||||
"diameter": {
|
||||
"anyOf": [{ "type": "string" }, { "type": "null" }],
|
||||
"description": "Diameter in mm (for radial or cylindrical types) (Outer Diameter, OD, D)"
|
||||
},
|
||||
"lead_spacing": {
|
||||
"anyOf": [{ "type": "string" }, { "type": "null" }],
|
||||
"description": "Lead spacing in mm (Pin Pitch, Terminal Gap, LS)"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"length",
|
||||
"width",
|
||||
"height",
|
||||
"diameter",
|
||||
"lead_spacing"
|
||||
],
|
||||
"type": "object"
|
||||
},
|
||||
{ "type": "null" }
|
||||
]
|
||||
}
|
||||
},
|
||||
"required": ["dimensions"],
|
||||
"type": "object"
|
||||
},
|
||||
"config": {
|
||||
"priority": null,
|
||||
"extraction_target": "PER_DOC",
|
||||
"extraction_mode": "PREMIUM",
|
||||
"multimodal_fast_mode": false,
|
||||
"system_prompt": "",
|
||||
"use_reasoning": true,
|
||||
"cite_sources": false,
|
||||
"confidence_scores": true,
|
||||
"chunk_mode": "PAGE",
|
||||
"high_resolution_mode": false,
|
||||
"invalidate_cache": false,
|
||||
"page_range": null
|
||||
},
|
||||
"file": {
|
||||
"id": "9233cc2b-00e3-4ddd-b426-b8b59357d4cb",
|
||||
"created_at": "2025-08-12T18:31:54.269440Z",
|
||||
"updated_at": "2025-08-13T15:45:39.064906Z",
|
||||
"name": "document (2).pdf.txt",
|
||||
"external_file_id": "document (2).pdf.txt",
|
||||
"file_size": 2562,
|
||||
"file_type": "txt",
|
||||
"project_id": "77bdc79f-fb69-49ae-a783-fcc573eec7ce",
|
||||
"last_modified_at": "2025-08-13T15:45:39Z",
|
||||
"resource_info": {
|
||||
"file_size": 2562,
|
||||
"last_modified_at": "2025-08-13T15:45:39"
|
||||
},
|
||||
"permission_info": null,
|
||||
"data_source_id": null
|
||||
},
|
||||
"status": "SUCCESS",
|
||||
"error": null,
|
||||
"job_id": "5bb8a583-366c-416c-ba55-4f5724fef9a9",
|
||||
"data": {
|
||||
"dimensions": {
|
||||
"length": "82 mm",
|
||||
"width": null,
|
||||
"height": null,
|
||||
"diameter": "35 mm",
|
||||
"lead_spacing": "6.0 mm"
|
||||
}
|
||||
},
|
||||
"extraction_metadata": {
|
||||
"field_metadata": {
|
||||
"dimensions": {
|
||||
"length": {
|
||||
"extraction_confidence": 0.9999968039036192,
|
||||
"confidence": 0.9999968039036192
|
||||
},
|
||||
"diameter": { "extraction_confidence": 1.0, "confidence": 1.0 },
|
||||
"lead_spacing": {
|
||||
"extraction_confidence": 0.9999999031936799,
|
||||
"confidence": 0.9999999031936799
|
||||
},
|
||||
"reasoning": "VERBATIM EXTRACTION"
|
||||
}
|
||||
},
|
||||
"usage": {
|
||||
"num_pages_extracted": 2,
|
||||
"num_document_tokens": 1034,
|
||||
"num_output_tokens": 3440
|
||||
}
|
||||
},
|
||||
"from_ui": false
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from llama_cloud.types import File as CloudFile
|
||||
from llama_cloud_services.extract import LlamaExtract, ExtractionAgent
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _set_dummy_env(monkeypatch):
|
||||
monkeypatch.setenv("LLAMA_CLOUD_API_KEY", "test-api-key")
|
||||
monkeypatch.setenv("LLAMA_CLOUD_BASE_URL", "https://example.test")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llama_file() -> CloudFile:
|
||||
return CloudFile(
|
||||
id="file_123",
|
||||
name="sample.pdf",
|
||||
external_file_id="ext_123",
|
||||
project_id="proj_123",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def extractor() -> LlamaExtract:
|
||||
return LlamaExtract(
|
||||
api_key=os.environ["LLAMA_CLOUD_API_KEY"],
|
||||
base_url=os.environ["LLAMA_CLOUD_BASE_URL"],
|
||||
verify=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_external_validation(monkeypatch):
|
||||
import llama_cloud_services.extract.extract as extract_mod
|
||||
|
||||
async def _noop_validate_schema(client, data_schema):
|
||||
return data_schema
|
||||
|
||||
# Disable config warnings and external schema validation
|
||||
monkeypatch.setattr(
|
||||
extract_mod, "_extraction_config_warning", lambda *_args, **_kwargs: None
|
||||
)
|
||||
monkeypatch.setattr(extract_mod, "_validate_schema", _noop_validate_schema)
|
||||
|
||||
|
||||
def test_convert_fileinput_accepts_llama_file_directly(
|
||||
extractor: LlamaExtract, llama_file: CloudFile
|
||||
):
|
||||
result = extractor._convert_file_to_file_data(llama_file)
|
||||
assert result is llama_file
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_extraction_with_llama_file_uses_file_id(
|
||||
extractor: LlamaExtract, llama_file: CloudFile, no_external_validation, monkeypatch
|
||||
):
|
||||
calls = []
|
||||
|
||||
async def fake_extract_stateless(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return SimpleNamespace(id="job_1")
|
||||
|
||||
# Patch the client's method that would normally hit the network
|
||||
monkeypatch.setattr(
|
||||
extractor._async_client.llama_extract,
|
||||
"extract_stateless",
|
||||
fake_extract_stateless,
|
||||
)
|
||||
|
||||
# Minimal schema and dummy config (warnings disabled by fixture)
|
||||
schema = {"type": "object", "properties": {}}
|
||||
dummy_config = SimpleNamespace()
|
||||
|
||||
job = await extractor.queue_extraction(schema, dummy_config, llama_file)
|
||||
|
||||
assert getattr(job, "id") == "job_1"
|
||||
assert len(calls) == 1
|
||||
kwargs = calls[0]
|
||||
assert "file_id" in kwargs and kwargs["file_id"] == llama_file.id
|
||||
assert "file" not in kwargs
|
||||
assert "text" not in kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extraction_agent_upload_file_accepts_llama_file_directly(
|
||||
llama_file: CloudFile,
|
||||
):
|
||||
# Build a minimal agent without hitting external services
|
||||
dummy_async_client = SimpleNamespace()
|
||||
dummy_agent = SimpleNamespace(id="agent_1", name="dummy", data_schema={}, config={})
|
||||
|
||||
agent = ExtractionAgent(
|
||||
client=dummy_async_client,
|
||||
agent=dummy_agent,
|
||||
project_id=None,
|
||||
organization_id=None,
|
||||
check_interval=0,
|
||||
max_timeout=0,
|
||||
num_workers=1,
|
||||
show_progress=False,
|
||||
verbose=False,
|
||||
verify=False,
|
||||
httpx_timeout=1,
|
||||
)
|
||||
|
||||
result = await agent._upload_file(llama_file)
|
||||
assert result is llama_file
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extraction_agent_aextract_accepts_llama_file(
|
||||
monkeypatch, llama_file: CloudFile
|
||||
):
|
||||
# Build a minimal agent without network
|
||||
dummy_llama_extract_iface = SimpleNamespace()
|
||||
|
||||
async def fake_run_job(**kwargs):
|
||||
# Ensure we are receiving a request with the right file_id
|
||||
request = kwargs.get("request")
|
||||
assert hasattr(request, "file_id")
|
||||
assert request.file_id == llama_file.id
|
||||
return SimpleNamespace(id="job_42")
|
||||
|
||||
dummy_llama_extract_iface.run_job = fake_run_job
|
||||
dummy_async_client = SimpleNamespace(llama_extract=dummy_llama_extract_iface)
|
||||
dummy_agent = SimpleNamespace(id="agent_1", name="dummy", data_schema={}, config={})
|
||||
|
||||
agent = ExtractionAgent(
|
||||
client=dummy_async_client,
|
||||
agent=dummy_agent,
|
||||
project_id=None,
|
||||
organization_id=None,
|
||||
check_interval=0,
|
||||
max_timeout=0,
|
||||
num_workers=1,
|
||||
show_progress=False,
|
||||
verbose=False,
|
||||
verify=False,
|
||||
httpx_timeout=1,
|
||||
)
|
||||
|
||||
# Ensure _upload_file returns the File directly and is called with our File
|
||||
calls = {}
|
||||
|
||||
async def fake_upload_file(file_input):
|
||||
calls["upload_called_with"] = file_input
|
||||
assert file_input is llama_file
|
||||
return file_input
|
||||
|
||||
monkeypatch.setattr(agent, "_upload_file", fake_upload_file)
|
||||
|
||||
# Avoid polling logic by short-circuiting result wait
|
||||
async def fake_wait(job_id: str):
|
||||
assert job_id == "job_42"
|
||||
return SimpleNamespace(id="run_42", status="SUCCESS", data={})
|
||||
|
||||
monkeypatch.setattr(agent, "_wait_for_job_result", fake_wait)
|
||||
|
||||
result = await agent.aextract(llama_file)
|
||||
|
||||
assert calls.get("upload_called_with") is llama_file
|
||||
assert getattr(result, "status") == "SUCCESS"
|
||||
assert getattr(result, "id") == "run_42"
|
||||
@@ -0,0 +1,16 @@
|
||||
import pytest
|
||||
|
||||
from llama_index.core.indices.managed.base import BaseManagedIndex
|
||||
from llama_cloud_services.index import (
|
||||
LlamaCloudIndex,
|
||||
)
|
||||
|
||||
|
||||
def test_class():
|
||||
names_of_base_classes = [b.__name__ for b in LlamaCloudIndex.__mro__]
|
||||
assert BaseManagedIndex.__name__ in names_of_base_classes
|
||||
|
||||
|
||||
def test_conflicting_index_identifiers():
|
||||
with pytest.raises(ValueError):
|
||||
LlamaCloudIndex(name="test", pipeline_id="test", index_id="test")
|
||||
@@ -1,5 +1,9 @@
|
||||
from unittest.mock import Mock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
from llama_cloud_services.utils import check_extra_params
|
||||
from llama_cloud_services.utils import check_extra_params, check_for_updates
|
||||
|
||||
|
||||
class MyModel(BaseModel):
|
||||
@@ -64,3 +68,30 @@ def test_check_extra_params_completely_invalid():
|
||||
for suggestion in suggestions:
|
||||
assert "check the documentation or update the package" in suggestion
|
||||
assert "Did you mean" not in suggestion
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_for_updates(capsys: pytest.CaptureFixture):
|
||||
"""Test update checker."""
|
||||
|
||||
mock_response = Mock()
|
||||
mock_client = Mock(spec=httpx.AsyncClient)
|
||||
mock_client.get.return_value = mock_response
|
||||
|
||||
mock_response.json.return_value = {"info": {"version": "0.0.0"}}
|
||||
assert not await check_for_updates(mock_client)
|
||||
out, err = capsys.readouterr()
|
||||
assert not out and not err
|
||||
|
||||
assert not await check_for_updates(mock_client, quiet=False)
|
||||
out, _ = capsys.readouterr()
|
||||
assert "up to date" in out
|
||||
|
||||
mock_response.json.return_value = {"info": {"version": "999.0.0"}}
|
||||
assert await check_for_updates(mock_client)
|
||||
out, err = capsys.readouterr()
|
||||
assert not out and not err
|
||||
|
||||
assert await check_for_updates(mock_client, quiet=False)
|
||||
out, _ = capsys.readouterr()
|
||||
assert "out of date" in out
|
||||
Generated
+47
-5
@@ -741,6 +741,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/36/f4/c6e662dade71f56cd2f3735141b265c3c79293c109549c1e6933b0651ffc/exceptiongroup-1.3.0-py3-none-any.whl", hash = "sha256:4d111e6e0c13d0644cad6ddaa7ed0261a0b36971f6d23e7ec9b4b9097da78a10", size = 16674, upload-time = "2025-05-10T17:42:49.33Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "execnet"
|
||||
version = "2.1.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/bb/ff/b4c0dc78fbe20c3e59c0c7334de0c27eb4001a2b2017999af398bf730817/execnet-2.1.1.tar.gz", hash = "sha256:5189b52c6121c24feae288166ab41b32549c7e2348652736540b9e6e7d4e72e3", size = 166524, upload-time = "2024-04-08T09:04:19.245Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/43/09/2aea36ff60d16dd8879bdb2f5b3ee0ba8d08cbbdcdfe870e695ce3784385/execnet-2.1.1-py3-none-any.whl", hash = "sha256:26dee51f1b80cebd6d0ca8e74dd8745419761d3bef34163928cbebbdc4749fdc", size = 40612, upload-time = "2024-04-08T09:04:17.414Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "executing"
|
||||
version = "2.2.0"
|
||||
@@ -1573,21 +1582,21 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "llama-cloud"
|
||||
version = "0.1.35"
|
||||
version = "0.1.39"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "certifi" },
|
||||
{ name = "httpx" },
|
||||
{ name = "pydantic" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/9b/72/816e6e900448e1b4a8137d90e65876b296c5264a23db6ae888bd3e6660ba/llama_cloud-0.1.35.tar.gz", hash = "sha256:200349d5d57424d7461f304cdb1355a58eea3e6ca1e6b0d75c66b2e937216983", size = 106403, upload-time = "2025-07-28T17:22:06.41Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/54/d3/058363cff521239b982638663f0cdf29aa94b6a5abaf2d593000f68a3333/llama_cloud-0.1.39.tar.gz", hash = "sha256:cd2e30540df4b1c18051f680962347838178596c875a71550d518a4d9729056f", size = 120625, upload-time = "2025-08-17T16:38:56.188Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/1d/d2/8d18a021ab757cea231428404f21fe3186bf1ebaac3f57a73c379483fd3f/llama_cloud-0.1.35-py3-none-any.whl", hash = "sha256:b7abab4423118e6f638d2f326749e7a07c6426543bea6da99b623c715b22af71", size = 303280, upload-time = "2025-07-28T17:22:04.946Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/87/39/830d7d00e6964204c01fafe1f6d186a2f365ba009861f8e9e26be1b8360e/llama_cloud-0.1.39-py3-none-any.whl", hash = "sha256:1cb741937d0b5ae853761439f6115cd3d64f4ae88a6e9109699f30662806c993", size = 317407, upload-time = "2025-08-17T16:38:54.31Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "llama-cloud-services"
|
||||
version = "0.6.54"
|
||||
version = "0.6.61"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "click", version = "8.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
|
||||
@@ -1595,6 +1604,7 @@ dependencies = [
|
||||
{ name = "eval-type-backport", marker = "python_full_version < '3.10'" },
|
||||
{ name = "llama-cloud" },
|
||||
{ name = "llama-index-core" },
|
||||
{ name = "packaging" },
|
||||
{ name = "platformdirs" },
|
||||
{ name = "pydantic" },
|
||||
{ name = "python-dotenv" },
|
||||
@@ -1611,16 +1621,19 @@ dev = [
|
||||
{ name = "jupyter" },
|
||||
{ name = "mypy" },
|
||||
{ name = "pre-commit" },
|
||||
{ name = "pydantic-settings" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-xdist" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "click", specifier = ">=8.1.7,<9" },
|
||||
{ name = "eval-type-backport", marker = "python_full_version < '3.10'", specifier = ">=0.2.0,<0.3" },
|
||||
{ name = "llama-cloud", specifier = "==0.1.35" },
|
||||
{ name = "llama-cloud", specifier = "==0.1.39" },
|
||||
{ name = "llama-index-core", specifier = ">=0.12.0" },
|
||||
{ name = "packaging", specifier = ">=25.0" },
|
||||
{ name = "platformdirs", specifier = ">=4.3.7,<5" },
|
||||
{ name = "pydantic", specifier = ">=2.8,!=2.10" },
|
||||
{ name = "python-dotenv", specifier = ">=1.0.1,<2" },
|
||||
@@ -1636,8 +1649,10 @@ dev = [
|
||||
{ name = "jupyter", specifier = ">=1.1.1,<2" },
|
||||
{ name = "mypy", specifier = ">=1.14.1,<2" },
|
||||
{ name = "pre-commit", specifier = "==3.2.0" },
|
||||
{ name = "pydantic-settings", specifier = ">=2.10.1" },
|
||||
{ name = "pytest", specifier = ">=8.0.0,<9" },
|
||||
{ name = "pytest-asyncio" },
|
||||
{ name = "pytest-xdist", specifier = ">=3.6.1,<4" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2853,6 +2868,20 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d4/29/3cade8a924a61f60ccfa10842f75eb12787e1440e2b8660ceffeb26685e7/pydantic_core-2.33.2-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:2807668ba86cb38c6817ad9bc66215ab8584d1d304030ce4f0887336f28a5e27", size = 2066661, upload-time = "2025-04-23T18:33:49.995Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pydantic-settings"
|
||||
version = "2.10.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pydantic" },
|
||||
{ name = "python-dotenv" },
|
||||
{ name = "typing-inspection" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/68/85/1ea668bbab3c50071ca613c6ab30047fb36ab0da1b92fa8f17bbc38fd36c/pydantic_settings-2.10.1.tar.gz", hash = "sha256:06f0062169818d0f5524420a360d632d5857b83cffd4d42fe29597807a1614ee", size = 172583, upload-time = "2025-06-24T13:26:46.841Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/58/f0/427018098906416f580e3cf1366d3b1abfb408a0652e9f31600c24a1903c/pydantic_settings-2.10.1-py3-none-any.whl", hash = "sha256:a60952460b99cf661dc25c29c0ef171721f98bfcb52ef8d9ea4c943d7c8cc796", size = 45235, upload-time = "2025-06-24T13:26:45.485Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pygments"
|
||||
version = "2.19.2"
|
||||
@@ -2894,6 +2923,19 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c7/9d/bf86eddabf8c6c9cb1ea9a869d6873b46f105a5d292d3a6f7071f5b07935/pytest_asyncio-1.1.0-py3-none-any.whl", hash = "sha256:5fe2d69607b0bd75c656d1211f969cadba035030156745ee09e7d71740e58ecf", size = 15157, upload-time = "2025-07-16T04:29:24.929Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest-xdist"
|
||||
version = "3.8.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "execnet" },
|
||||
{ name = "pytest" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/78/b4/439b179d1ff526791eb921115fca8e44e596a13efeda518b9d845a619450/pytest_xdist-3.8.0.tar.gz", hash = "sha256:7e578125ec9bc6050861aa93f2d59f1d8d085595d6551c2c90b6f4fad8d3a9f1", size = 88069, upload-time = "2025-07-01T13:30:59.346Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88", size = 46396, upload-time = "2025-07-01T13:30:56.632Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "python-dateutil"
|
||||
version = "2.9.0.post0"
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"name": "@llamaindex/llama-cloud-services-e2e-tests",
|
||||
"private": true,
|
||||
"version": "0.0.1",
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"test": "pnpm run test:node16 && pnpm run test:nodenext && pnpm run test:bundler",
|
||||
"build": "pnpm run build:node16 && pnpm run build:nodenext && pnpm run build:bundler",
|
||||
"build:node16": "tsc -p src/tsconfig.node16.json --outDir dist/node16",
|
||||
"test:node16": "pnpm run build:node16 && node --test dist/node16/index.e2e.js",
|
||||
"build:nodenext": "tsc -p src/tsconfig.nodenext.json --outDir dist/nodenext",
|
||||
"test:nodenext": "pnpm run build:nodenext && node --test dist/nodenext/index.e2e.js",
|
||||
"build:bundler": "tsc -p src/tsconfig.bundler.json --outDir dist/bundler",
|
||||
"test:bundler": "pnpm run build:bundler && node --test dist/bundler/index.e2e.js",
|
||||
"build:node": "tsc -p src/tsconfig.node.json --outDir dist/node",
|
||||
"test:node": "pnpm run build:node && node --test dist/node/index.e2e.js"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/node": "^24.0.13",
|
||||
"llama-cloud-services": "workspace:*",
|
||||
"typescript": "^5.9.2"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
import { ok } from "node:assert";
|
||||
import { test } from "node:test";
|
||||
|
||||
import { LlamaCloudIndex } from "llama-cloud-services";
|
||||
import { LlamaParseReader } from "llama-cloud-services";
|
||||
|
||||
test("LlamaIndex module resolution test", async (t) => {
|
||||
await t.test("works with ES module", () => {
|
||||
const index = new LlamaCloudIndex({
|
||||
name: "test-index",
|
||||
projectName: "Default",
|
||||
});
|
||||
const reader = new LlamaParseReader({
|
||||
resultType: "markdown",
|
||||
verbose: false,
|
||||
});
|
||||
ok(index !== undefined);
|
||||
ok(reader !== undefined);
|
||||
});
|
||||
|
||||
await t.test("works with dynamic imports", async () => {
|
||||
const mod = await import("llama-cloud-services"); // simulates commonjs
|
||||
ok(mod !== undefined);
|
||||
const index = new mod.LlamaCloudIndex({
|
||||
name: "test-index",
|
||||
projectName: "Default",
|
||||
});
|
||||
ok(index !== undefined);
|
||||
});
|
||||
|
||||
await t.test("all imports work", () => {
|
||||
const allImports = [
|
||||
LlamaCloudIndex,
|
||||
];
|
||||
|
||||
ok(allImports.filter(Boolean).length === allImports.length);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"module": "esnext",
|
||||
"moduleResolution": "bundler",
|
||||
"target": "esnext",
|
||||
"outDir": "dist",
|
||||
"rootDir": ".",
|
||||
"strict": true,
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true
|
||||
},
|
||||
"include": ["**/*.ts"],
|
||||
"exclude": ["node_modules", "dist"]
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"module": "esnext",
|
||||
"moduleResolution": "node",
|
||||
"target": "esnext",
|
||||
"outDir": "dist",
|
||||
"rootDir": ".",
|
||||
"strict": true,
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true
|
||||
},
|
||||
"include": ["**/*.ts"],
|
||||
"exclude": ["node_modules", "dist"]
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"module": "node16",
|
||||
"moduleResolution": "node16",
|
||||
"target": "esnext",
|
||||
"outDir": "dist",
|
||||
"rootDir": ".",
|
||||
"strict": true,
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true
|
||||
},
|
||||
"include": ["**/*.ts"],
|
||||
"exclude": ["node_modules", "dist"]
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"module": "nodenext",
|
||||
"moduleResolution": "nodenext",
|
||||
"target": "esnext",
|
||||
"outDir": "dist",
|
||||
"rootDir": ".",
|
||||
"strict": true,
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true
|
||||
},
|
||||
"include": ["**/*.ts"],
|
||||
"exclude": ["node_modules", "dist"]
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"target": "ESNext",
|
||||
"module": "NodeNext",
|
||||
"moduleResolution": "NodeNext",
|
||||
"verbatimModuleSyntax": true,
|
||||
"esModuleInterop": true,
|
||||
"forceConsistentCasingInFileNames": true,
|
||||
"noUncheckedIndexedAccess": true,
|
||||
"strictNullChecks": true,
|
||||
"exactOptionalPropertyTypes": true,
|
||||
"strict": true,
|
||||
"skipLibCheck": true,
|
||||
"rootDir": "./src",
|
||||
"outDir": "./dist/type",
|
||||
"tsBuildInfoFile": "./dist/.tsbuildinfo",
|
||||
"incremental": true,
|
||||
"lib": [
|
||||
"ES2022",
|
||||
"DOM",
|
||||
"DOM.Iterable",
|
||||
"DOM.AsyncIterable"
|
||||
],
|
||||
"types": [
|
||||
"node"
|
||||
]
|
||||
},
|
||||
"include": [
|
||||
"./src"
|
||||
],
|
||||
"exclude": [
|
||||
"node_modules"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"type": "module",
|
||||
"main": "./dist/index.cjs",
|
||||
"module": "./dist/index.js",
|
||||
"types": "./dist/index.d.ts",
|
||||
"exports": "./dist/index.js",
|
||||
"private": true
|
||||
}
|
||||
@@ -12855,6 +12855,72 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/extraction/run": {
|
||||
"post": {
|
||||
"tags": ["LlamaExtract"],
|
||||
"summary": "Extract Stateless",
|
||||
"description": "Stateless extraction endpoint that uses a default extraction agent in the user's default project. Requires data_schema, config, and either file_id, text, or base64 encoded file data.",
|
||||
"operationId": "extract_stateless_api_v1_extraction_run_post",
|
||||
"security": [
|
||||
{
|
||||
"HTTPBearer": []
|
||||
},
|
||||
{
|
||||
"HTTPBearer": []
|
||||
}
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"name": "session",
|
||||
"in": "cookie",
|
||||
"required": false,
|
||||
"schema": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Session"
|
||||
}
|
||||
}
|
||||
],
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/StatelessExtractionRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Successful Response",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ExtractJob"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"422": {
|
||||
"description": "Validation Error",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/api/v1/extraction/jobs": {
|
||||
"get": {
|
||||
"tags": ["LlamaExtract"],
|
||||
@@ -35483,6 +35549,64 @@
|
||||
"title": "WebhookConfiguration",
|
||||
"description": "Allows the user to configure webhook options for notifications and callbacks."
|
||||
},
|
||||
"StatelessExtractionRequest": {
|
||||
"type": "object",
|
||||
"required": ["data_schema"],
|
||||
"properties": {
|
||||
"data_schema": {
|
||||
"anyOf": [
|
||||
{
|
||||
"additionalProperties": {
|
||||
"anyOf": [
|
||||
{
|
||||
"additionalProperties": true,
|
||||
"type": "object"
|
||||
},
|
||||
{
|
||||
"items": {},
|
||||
"type": "array"
|
||||
},
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "number"
|
||||
},
|
||||
{
|
||||
"type": "boolean"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"type": "object"
|
||||
},
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"config": {
|
||||
"$ref": "#/components/schemas/ExtractConfig",
|
||||
"description": "The configuration parameters for the extraction agent."
|
||||
},
|
||||
"file_id": {
|
||||
"type": "string",
|
||||
"format": "uuid",
|
||||
"title": "File Id",
|
||||
"description": "ID of an uploaded file to extract from"
|
||||
}
|
||||
},
|
||||
"title": "StatelessExtractionRequest",
|
||||
"description": "Request body for stateless extraction. Must include either file_id, text, or base64."
|
||||
},
|
||||
"llama_index__core__base__llms__types__ChatMessage": {
|
||||
"properties": {
|
||||
"role": {
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "llama-cloud-services",
|
||||
"version": "0.2.0",
|
||||
"version": "0.3.3",
|
||||
"type": "module",
|
||||
"license": "MIT",
|
||||
"scripts": {
|
||||
@@ -20,7 +20,8 @@
|
||||
"./api",
|
||||
"./reader",
|
||||
"./parse",
|
||||
"./beta/agent"
|
||||
"./beta/agent",
|
||||
"./extract"
|
||||
],
|
||||
"exports": {
|
||||
"./openapi.json": "./openapi.json",
|
||||
@@ -68,6 +69,17 @@
|
||||
},
|
||||
"default": "./parse/dist/index.js"
|
||||
},
|
||||
"./extract": {
|
||||
"require": {
|
||||
"types": "./extract/dist/index.d.cts",
|
||||
"default": "./extract/dist/index.cjs"
|
||||
},
|
||||
"import": {
|
||||
"types": "./extract/dist/index.d.ts",
|
||||
"default": "./extract/dist/index.js"
|
||||
},
|
||||
"default": "./extract/dist/index.js"
|
||||
},
|
||||
".": {
|
||||
"require": {
|
||||
"types": "./dist/index.d.cts",
|
||||
@@ -113,6 +125,7 @@
|
||||
},
|
||||
"dependencies": {
|
||||
"ajv": "^8.17.1",
|
||||
"file-type": "^21.0.0",
|
||||
"p-retry": "^6.2.1",
|
||||
"zod": "^3.25.76"
|
||||
},
|
||||
|
||||
Generated
-3562
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,225 @@
|
||||
import { createClient, createConfig, type Client } from "@hey-api/client-fetch";
|
||||
import { File } from "buffer";
|
||||
import * as extract from "./extract";
|
||||
import type { ExtractAgent, ExtractConfig } from "./extract";
|
||||
import { getEnv } from "@llamaindex/env";
|
||||
import type { ExtractResult } from "./type";
|
||||
|
||||
const URLS = {
|
||||
us: "https://api.cloud.llamaindex.ai",
|
||||
eu: "https://api.cloud.eu.llamaindex.ai",
|
||||
"us-staging": "https://api.staging.llamaindex.ai",
|
||||
} as const;
|
||||
|
||||
function getUrl(baseUrl: string | undefined, region: string | undefined) {
|
||||
if (typeof baseUrl != "undefined") {
|
||||
return baseUrl;
|
||||
}
|
||||
if (typeof region === "undefined") {
|
||||
return URLS["us"];
|
||||
} else if (region === "us" || region === "eu" || region === "us-staging") {
|
||||
return URLS[region];
|
||||
} else {
|
||||
throw new Error(`Unsupported region: ${region}`);
|
||||
}
|
||||
}
|
||||
|
||||
export class LlamaExtractAgent {
|
||||
private agent: ExtractAgent;
|
||||
private client: Client;
|
||||
id: string;
|
||||
name: string;
|
||||
dataSchema: {
|
||||
[key: string]:
|
||||
| string
|
||||
| number
|
||||
| boolean
|
||||
| {
|
||||
[key: string]: unknown;
|
||||
}
|
||||
| unknown[]
|
||||
| null;
|
||||
};
|
||||
|
||||
constructor(agent: ExtractAgent, client: Client) {
|
||||
this.agent = agent;
|
||||
this.client = client;
|
||||
this.id = agent.id;
|
||||
this.name = agent.name;
|
||||
this.dataSchema = agent.data_schema;
|
||||
}
|
||||
|
||||
async extract(
|
||||
filePath: string | undefined = undefined,
|
||||
fileContent:
|
||||
| Buffer<ArrayBufferLike>
|
||||
| Uint8Array<ArrayBuffer>
|
||||
| string
|
||||
| File
|
||||
| undefined = undefined,
|
||||
fileName: string | undefined = undefined,
|
||||
project_id: string | null = null,
|
||||
organization_id: string | null = null,
|
||||
fromUi: boolean | undefined = undefined,
|
||||
pollingInterval: number = 1,
|
||||
maxPollingIterations: number = 1800,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<ExtractResult | undefined> {
|
||||
return await extract.extract(
|
||||
this.agent.id,
|
||||
filePath,
|
||||
fileContent,
|
||||
fileName,
|
||||
project_id,
|
||||
organization_id,
|
||||
this.client,
|
||||
fromUi,
|
||||
pollingInterval,
|
||||
maxPollingIterations,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
export class LlamaExtract {
|
||||
private client: Client;
|
||||
|
||||
constructor(
|
||||
apiKey: string | undefined = undefined,
|
||||
baseUrl: string | undefined = undefined,
|
||||
region: string | undefined = undefined,
|
||||
) {
|
||||
const key = apiKey ?? getEnv("LLAMA_CLOUD_API_KEY");
|
||||
if (typeof key === "undefined") {
|
||||
throw new Error(
|
||||
"No API key provided and no API key found in environment. Please pass the API key or set `LLAMA_CLOUD_API_KEY` as an environment variable.",
|
||||
);
|
||||
}
|
||||
const url = getUrl(baseUrl, region);
|
||||
this.client = createClient(
|
||||
createConfig({
|
||||
baseUrl: url,
|
||||
headers: {
|
||||
Authorization: `Bearer ${key}`,
|
||||
},
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
async createAgent(
|
||||
name: string,
|
||||
dataSchema:
|
||||
| {
|
||||
[key: string]:
|
||||
| { [key: string]: unknown }
|
||||
| Array<unknown>
|
||||
| string
|
||||
| number
|
||||
| number
|
||||
| boolean
|
||||
| null;
|
||||
}
|
||||
| string,
|
||||
config: ExtractConfig | undefined = undefined,
|
||||
project_id: string | null = null,
|
||||
organization_id: string | null = null,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<LlamaExtractAgent | undefined> {
|
||||
const agent = await extract.createAgent(
|
||||
name,
|
||||
dataSchema,
|
||||
config,
|
||||
project_id,
|
||||
organization_id,
|
||||
this.client,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
);
|
||||
if (typeof agent != "undefined") {
|
||||
return new LlamaExtractAgent(agent, this.client);
|
||||
}
|
||||
}
|
||||
|
||||
async getAgent(
|
||||
name: string | undefined = undefined,
|
||||
id: string | undefined = undefined,
|
||||
project_id: string | null = null,
|
||||
organization_id: string | null = null,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<LlamaExtractAgent | undefined> {
|
||||
const agent = await extract.getAgent(
|
||||
id,
|
||||
name,
|
||||
project_id,
|
||||
organization_id,
|
||||
this.client,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
);
|
||||
if (typeof agent != "undefined") {
|
||||
return new LlamaExtractAgent(agent, this.client);
|
||||
}
|
||||
}
|
||||
|
||||
async deleteAgent(
|
||||
id: string,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 500,
|
||||
): Promise<boolean | undefined> {
|
||||
return await extract.deleteAgent(
|
||||
id,
|
||||
this.client,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
);
|
||||
}
|
||||
|
||||
async extract(
|
||||
dataSchema:
|
||||
| {
|
||||
[key: string]:
|
||||
| { [key: string]: unknown }
|
||||
| Array<unknown>
|
||||
| string
|
||||
| number
|
||||
| number
|
||||
| boolean
|
||||
| null;
|
||||
}
|
||||
| string,
|
||||
config: ExtractConfig | undefined = undefined,
|
||||
filePath: string | undefined = undefined,
|
||||
fileContent:
|
||||
| Buffer<ArrayBufferLike>
|
||||
| Uint8Array<ArrayBuffer>
|
||||
| string
|
||||
| File
|
||||
| undefined = undefined,
|
||||
fileName: string | undefined = undefined,
|
||||
project_id: string | null = null,
|
||||
organization_id: string | null = null,
|
||||
pollingInterval: number = 1,
|
||||
maxPollingIterations: number = 1800,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<ExtractResult | undefined> {
|
||||
return await extract.extractStateless(
|
||||
dataSchema,
|
||||
config,
|
||||
filePath,
|
||||
fileContent,
|
||||
fileName,
|
||||
project_id,
|
||||
organization_id,
|
||||
this.client,
|
||||
pollingInterval,
|
||||
maxPollingIterations,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,8 @@ export type {
|
||||
AggregateAgentDataOptions,
|
||||
ComparisonOperator,
|
||||
ExtractedData,
|
||||
ExtractedFieldMetadata,
|
||||
ExtractedFieldMetadataDict,
|
||||
FilterOperation,
|
||||
SearchAgentDataOptions,
|
||||
StatusType,
|
||||
|
||||
@@ -38,8 +38,12 @@ export interface ExtractedFieldMetadata {
|
||||
confidence?: number;
|
||||
/** The confidence score for the field based on the extracted text only */
|
||||
extraction_confidence?: number;
|
||||
citation: FieldCitation[];
|
||||
}
|
||||
|
||||
export interface FieldCitation {
|
||||
/** The page number that the field occurred on */
|
||||
page_number?: number;
|
||||
page?: number;
|
||||
/** The original text this field's value was derived from */
|
||||
matching_text?: string;
|
||||
}
|
||||
|
||||
@@ -16350,7 +16350,7 @@ export const PromptConfSchema = {
|
||||
default: `
|
||||
Provide a brief explanation for how you arrived at the extracted value based on the source text provided.
|
||||
- For inferred values, explain the reasoning behind the extraction briefly.
|
||||
- For simple verbatim extraction, output 'VERBATIM EXTRACTION'.
|
||||
- For simple verbatim extraction, output 'VERBATIM EXTRACTION'.
|
||||
- When supporting data is not present in the source text, output 'INSUFFICIENT DATA' and emit blank or null values for the value__ field.
|
||||
`,
|
||||
},
|
||||
@@ -17689,7 +17689,7 @@ export const StructParseConfSchema = {
|
||||
reasoning_prompt: `
|
||||
Provide a brief explanation for how you arrived at the extracted value based on the source text provided.
|
||||
- For inferred values, explain the reasoning behind the extraction briefly.
|
||||
- For simple verbatim extraction, output 'VERBATIM EXTRACTION'.
|
||||
- For simple verbatim extraction, output 'VERBATIM EXTRACTION'.
|
||||
- When supporting data is not present in the source text, output 'INSUFFICIENT DATA' and emit blank or null values for the value__ field.
|
||||
`,
|
||||
cite_sources_prompt: {
|
||||
@@ -18644,6 +18644,66 @@ export const WebhookConfigurationSchema = {
|
||||
"Allows the user to configure webhook options for notifications and callbacks.",
|
||||
} as const;
|
||||
|
||||
export const StatelessExtractionRequestSchema = {
|
||||
type: "object",
|
||||
required: ["data_schema"],
|
||||
properties: {
|
||||
data_schema: {
|
||||
anyOf: [
|
||||
{
|
||||
additionalProperties: {
|
||||
anyOf: [
|
||||
{
|
||||
additionalProperties: true,
|
||||
type: "object",
|
||||
},
|
||||
{
|
||||
items: {},
|
||||
type: "array",
|
||||
},
|
||||
{
|
||||
type: "string",
|
||||
},
|
||||
{
|
||||
type: "integer",
|
||||
},
|
||||
{
|
||||
type: "number",
|
||||
},
|
||||
{
|
||||
type: "boolean",
|
||||
},
|
||||
{
|
||||
type: "null",
|
||||
},
|
||||
],
|
||||
},
|
||||
type: "object",
|
||||
},
|
||||
{
|
||||
type: "string",
|
||||
},
|
||||
{
|
||||
type: "null",
|
||||
},
|
||||
],
|
||||
},
|
||||
config: {
|
||||
$ref: "#/components/schemas/ExtractConfig",
|
||||
description: "The configuration parameters for the extraction agent.",
|
||||
},
|
||||
file_id: {
|
||||
type: "string",
|
||||
format: "uuid",
|
||||
title: "File Id",
|
||||
description: "ID of an uploaded file to extract from",
|
||||
},
|
||||
},
|
||||
title: "StatelessExtractionRequest",
|
||||
description:
|
||||
"Request body for stateless extraction. Must include either file_id, text, or base64.",
|
||||
} as const;
|
||||
|
||||
export const llama_index__core__base__llms__types__ChatMessageSchema = {
|
||||
properties: {
|
||||
role: {
|
||||
|
||||
@@ -452,6 +452,9 @@ import type {
|
||||
UpdateExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdPutData,
|
||||
UpdateExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdPutResponse,
|
||||
UpdateExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdPutError,
|
||||
ExtractStatelessApiV1ExtractionRunPostData,
|
||||
ExtractStatelessApiV1ExtractionRunPostResponse,
|
||||
ExtractStatelessApiV1ExtractionRunPostError,
|
||||
ListJobsApiV1ExtractionJobsGetData,
|
||||
ListJobsApiV1ExtractionJobsGetResponse,
|
||||
ListJobsApiV1ExtractionJobsGetError,
|
||||
@@ -5701,6 +5704,39 @@ export const updateExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgent
|
||||
});
|
||||
};
|
||||
|
||||
/**
|
||||
* Extract Stateless
|
||||
* Stateless extraction endpoint that uses a default extraction agent in the user's default project. Requires data_schema, config, and either file_id, text, or base64 encoded file data.
|
||||
*/
|
||||
export const extractStatelessApiV1ExtractionRunPost = <
|
||||
ThrowOnError extends boolean = false,
|
||||
>(
|
||||
options: Options<ExtractStatelessApiV1ExtractionRunPostData, ThrowOnError>,
|
||||
) => {
|
||||
return (options.client ?? _heyApiClient).post<
|
||||
ExtractStatelessApiV1ExtractionRunPostResponse,
|
||||
ExtractStatelessApiV1ExtractionRunPostError,
|
||||
ThrowOnError
|
||||
>({
|
||||
security: [
|
||||
{
|
||||
scheme: "bearer",
|
||||
type: "http",
|
||||
},
|
||||
{
|
||||
scheme: "bearer",
|
||||
type: "http",
|
||||
},
|
||||
],
|
||||
url: "/api/v1/extraction/run",
|
||||
...options,
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
...options?.headers,
|
||||
},
|
||||
});
|
||||
};
|
||||
|
||||
/**
|
||||
* List Jobs
|
||||
*/
|
||||
|
||||
@@ -8272,6 +8272,35 @@ export type WebhookConfiguration = {
|
||||
> | null;
|
||||
};
|
||||
|
||||
/**
|
||||
* Request body for stateless extraction. Must include either file_id, text, or base64.
|
||||
*/
|
||||
export type StatelessExtractionRequest = {
|
||||
data_schema:
|
||||
| {
|
||||
[key: string]:
|
||||
| {
|
||||
[key: string]: unknown;
|
||||
}
|
||||
| Array<unknown>
|
||||
| string
|
||||
| number
|
||||
| number
|
||||
| boolean
|
||||
| null;
|
||||
}
|
||||
| string
|
||||
| null;
|
||||
/**
|
||||
* The configuration parameters for the extraction agent.
|
||||
*/
|
||||
config?: ExtractConfig;
|
||||
/**
|
||||
* ID of an uploaded file to extract from
|
||||
*/
|
||||
file_id?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Chat message.
|
||||
*/
|
||||
@@ -13078,6 +13107,33 @@ export type UpdateExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentI
|
||||
export type UpdateExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdPutResponse =
|
||||
UpdateExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdPutResponses[keyof UpdateExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdPutResponses];
|
||||
|
||||
export type ExtractStatelessApiV1ExtractionRunPostData = {
|
||||
body: StatelessExtractionRequest;
|
||||
path?: never;
|
||||
query?: never;
|
||||
url: "/api/v1/extraction/run";
|
||||
};
|
||||
|
||||
export type ExtractStatelessApiV1ExtractionRunPostErrors = {
|
||||
/**
|
||||
* Validation Error
|
||||
*/
|
||||
422: HttpValidationError;
|
||||
};
|
||||
|
||||
export type ExtractStatelessApiV1ExtractionRunPostError =
|
||||
ExtractStatelessApiV1ExtractionRunPostErrors[keyof ExtractStatelessApiV1ExtractionRunPostErrors];
|
||||
|
||||
export type ExtractStatelessApiV1ExtractionRunPostResponses = {
|
||||
/**
|
||||
* Successful Response
|
||||
*/
|
||||
200: ExtractJob;
|
||||
};
|
||||
|
||||
export type ExtractStatelessApiV1ExtractionRunPostResponse =
|
||||
ExtractStatelessApiV1ExtractionRunPostResponses[keyof ExtractStatelessApiV1ExtractionRunPostResponses];
|
||||
|
||||
export type ListJobsApiV1ExtractionJobsGetData = {
|
||||
body?: never;
|
||||
path?: never;
|
||||
|
||||
@@ -3219,7 +3219,7 @@ export const zPromptConf = z.object({
|
||||
reasoning_prompt: z.string().optional().default(`
|
||||
Provide a brief explanation for how you arrived at the extracted value based on the source text provided.
|
||||
- For inferred values, explain the reasoning behind the extraction briefly.
|
||||
- For simple verbatim extraction, output 'VERBATIM EXTRACTION'.
|
||||
- For simple verbatim extraction, output 'VERBATIM EXTRACTION'.
|
||||
- When supporting data is not present in the source text, output 'INSUFFICIENT DATA' and emit blank or null values for the value__ field.
|
||||
`),
|
||||
cite_sources_prompt: z
|
||||
@@ -3472,6 +3472,12 @@ export const zUserOrganizationRoleCreate = z.object({
|
||||
role_id: z.string().uuid(),
|
||||
});
|
||||
|
||||
export const zStatelessExtractionRequest = z.object({
|
||||
data_schema: z.union([z.object({}), z.string(), z.null()]),
|
||||
config: zExtractConfig.optional(),
|
||||
file_id: z.string().uuid().optional(),
|
||||
});
|
||||
|
||||
export const zListKeysApiV1ApiKeysGetResponse = z.array(zApiKey);
|
||||
|
||||
export const zGenerateKeyApiV1ApiKeysPostResponse = zApiKey;
|
||||
@@ -3829,6 +3835,8 @@ export const zGetExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentId
|
||||
export const zUpdateExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdPutResponse =
|
||||
zExtractAgent;
|
||||
|
||||
export const zExtractStatelessApiV1ExtractionRunPostResponse = zExtractJob;
|
||||
|
||||
export const zListJobsApiV1ExtractionJobsGetResponse = z.array(zExtractJob);
|
||||
|
||||
export const zRunJobApiV1ExtractionJobsPostResponse = zExtractJob;
|
||||
|
||||
@@ -0,0 +1,651 @@
|
||||
import { emitWarning } from "process";
|
||||
import fs from "fs/promises";
|
||||
import { Blob } from "buffer";
|
||||
import * as path from "path";
|
||||
import type { ExtractResult } from "./type";
|
||||
import { randomUUID } from "@llamaindex/env";
|
||||
import { File } from "buffer";
|
||||
import {
|
||||
type Options,
|
||||
type ExtractAgentCreate,
|
||||
type ExtractConfig,
|
||||
type ExtractJobCreate,
|
||||
type ExtractAgent,
|
||||
type ExtractJob,
|
||||
type CreateExtractionAgentApiV1ExtractionExtractionAgentsPostData,
|
||||
type GetExtractionAgentByNameApiV1ExtractionExtractionAgentsByNameNameGetData,
|
||||
type GetExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdGetData,
|
||||
type RunJobApiV1ExtractionJobsPostData,
|
||||
type GetJobApiV1ExtractionJobsJobIdGetData,
|
||||
type GetJobResultApiV1ExtractionJobsJobIdResultGetData,
|
||||
StatusEnum,
|
||||
type UploadFileApiV1FilesPostData,
|
||||
type StatelessExtractionRequest,
|
||||
type ExtractStatelessApiV1ExtractionRunPostData,
|
||||
type DeleteExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdDeleteData,
|
||||
createExtractionAgentApiV1ExtractionExtractionAgentsPost,
|
||||
getExtractionAgentByNameApiV1ExtractionExtractionAgentsByNameNameGet,
|
||||
getExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdGet,
|
||||
runJobApiV1ExtractionJobsPost,
|
||||
getJobApiV1ExtractionJobsJobIdGet,
|
||||
getJobResultApiV1ExtractionJobsJobIdResultGet,
|
||||
uploadFileApiV1FilesPost,
|
||||
extractStatelessApiV1ExtractionRunPost,
|
||||
deleteExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdDelete,
|
||||
} from "./api";
|
||||
import type { Client } from "@hey-api/client-fetch";
|
||||
import { sleep } from "./utils";
|
||||
import { fileTypeFromBuffer } from "file-type";
|
||||
|
||||
type BodyUploadFileApiV1FilesPost = {
|
||||
upload_file: Blob | File;
|
||||
};
|
||||
|
||||
export async function createAgent(
|
||||
name: string,
|
||||
dataSchema:
|
||||
| {
|
||||
[key: string]:
|
||||
| { [key: string]: unknown }
|
||||
| Array<unknown>
|
||||
| string
|
||||
| number
|
||||
| number
|
||||
| boolean
|
||||
| null;
|
||||
}
|
||||
| string,
|
||||
config: ExtractConfig = {} as ExtractConfig,
|
||||
project_id: string | null = null,
|
||||
organization_id: string | null = null,
|
||||
client: Client | undefined = undefined,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<ExtractAgent | undefined> {
|
||||
const agentData = {
|
||||
name: name,
|
||||
data_schema: dataSchema,
|
||||
config: config,
|
||||
} as ExtractAgentCreate;
|
||||
const agentDataCreation = {
|
||||
body: agentData,
|
||||
query: { project_id: project_id, organization_id: organization_id },
|
||||
} as CreateExtractionAgentApiV1ExtractionExtractionAgentsPostData;
|
||||
const options =
|
||||
agentDataCreation as Options<CreateExtractionAgentApiV1ExtractionExtractionAgentsPostData>;
|
||||
if (typeof client != "undefined") {
|
||||
options.client = client;
|
||||
}
|
||||
let retries: number = 0;
|
||||
while (true) {
|
||||
if (retries > maxRetriesOnError) {
|
||||
throw new Error(
|
||||
"Error while creating the agent: Exceeded maximum number of retries, the API keeps returning errors.",
|
||||
);
|
||||
}
|
||||
const response =
|
||||
await createExtractionAgentApiV1ExtractionExtractionAgentsPost(options);
|
||||
if (!response.response.ok) {
|
||||
if ("error" in response) {
|
||||
console.log(
|
||||
`An error occurred while creating the extraction agent.\nDetails:\n\n${JSON.stringify(
|
||||
response.error,
|
||||
)}\n\nRetrying...`,
|
||||
);
|
||||
}
|
||||
retries++;
|
||||
await sleep(retryInterval * 1000);
|
||||
} else {
|
||||
return response.data as ExtractAgent;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export async function getAgent(
|
||||
id: string | undefined = undefined,
|
||||
name: string | undefined = undefined,
|
||||
project_id: string | null = null,
|
||||
organization_id: string | null = null,
|
||||
client: Client | undefined = undefined,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<ExtractAgent | undefined> {
|
||||
if (typeof id === "undefined" && typeof name === "undefined") {
|
||||
throw new Error("One of `id` and `string` must be passed.");
|
||||
} else if (typeof id != "undefined" && typeof name != "undefined") {
|
||||
emitWarning("You passed both `id` and `name`, using only id...");
|
||||
const data = {
|
||||
path: { extraction_agent_id: id },
|
||||
} as GetExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdGetData;
|
||||
const options =
|
||||
data as Options<GetExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdGetData>;
|
||||
if (typeof client != "undefined") {
|
||||
options.client = client;
|
||||
}
|
||||
let retries: number = 0;
|
||||
while (true) {
|
||||
if (retries > maxRetriesOnError) {
|
||||
throw new Error(
|
||||
"Error while getting the agent: Exceeded maximum number of retries, the API keeps returning errors.",
|
||||
);
|
||||
}
|
||||
const response =
|
||||
await getExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdGet(
|
||||
options,
|
||||
);
|
||||
if (!response.response.ok) {
|
||||
if ("error" in response) {
|
||||
console.log(
|
||||
`An error occurred while getting the extraction agent by ID.\nDetails:\n\n${JSON.stringify(
|
||||
response.error,
|
||||
)}\n\nRetrying...`,
|
||||
);
|
||||
}
|
||||
retries++;
|
||||
await sleep(retryInterval * 1000);
|
||||
} else {
|
||||
return response.data as ExtractAgent;
|
||||
}
|
||||
}
|
||||
} else if (typeof name != "undefined" && typeof id === "undefined") {
|
||||
const data = {
|
||||
path: { name: name },
|
||||
query: { organization_id: organization_id, project_id: project_id },
|
||||
} as GetExtractionAgentByNameApiV1ExtractionExtractionAgentsByNameNameGetData;
|
||||
const options =
|
||||
data as Options<GetExtractionAgentByNameApiV1ExtractionExtractionAgentsByNameNameGetData>;
|
||||
if (typeof client != "undefined") {
|
||||
options.client = client;
|
||||
}
|
||||
let retries: number = 0;
|
||||
while (true) {
|
||||
if (retries > maxRetriesOnError) {
|
||||
throw new Error(
|
||||
"Error while getting the agent: Exceeded maximum number of retries, the API keeps returning errors.",
|
||||
);
|
||||
}
|
||||
const response =
|
||||
await getExtractionAgentByNameApiV1ExtractionExtractionAgentsByNameNameGet(
|
||||
options,
|
||||
);
|
||||
if (!response.response.ok) {
|
||||
if ("error" in response) {
|
||||
console.log(
|
||||
`An error occurred while getting the extraction agent by name.\nDetails:\n\n${JSON.stringify(
|
||||
response.error,
|
||||
)}\n\nRetrying...`,
|
||||
);
|
||||
}
|
||||
retries++;
|
||||
await sleep(retryInterval * 1000);
|
||||
} else {
|
||||
return response.data as ExtractAgent;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
const data = {
|
||||
path: { extraction_agent_id: id },
|
||||
} as GetExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdGetData;
|
||||
const options =
|
||||
data as Options<GetExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdGetData>;
|
||||
if (typeof client != "undefined") {
|
||||
options.client = client;
|
||||
}
|
||||
let retries: number = 0;
|
||||
while (true) {
|
||||
if (retries > maxRetriesOnError) {
|
||||
throw new Error(
|
||||
"Error while getting the agent: Exceeded maximum number of retries, the API keeps returning errors.",
|
||||
);
|
||||
}
|
||||
const response =
|
||||
await getExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdGet(
|
||||
options,
|
||||
);
|
||||
if (!response.response.ok) {
|
||||
if (!response.response.ok) {
|
||||
if ("error" in response) {
|
||||
console.log(
|
||||
`An error occurred while getting the extraction agent by ID.\nDetails:\n\n${JSON.stringify(
|
||||
response.error,
|
||||
)}\n\nRetrying...`,
|
||||
);
|
||||
}
|
||||
retries++;
|
||||
await sleep(retryInterval * 1000);
|
||||
}
|
||||
} else {
|
||||
return response.data as ExtractAgent;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function textToFile(text: string, fileName: string | null = null) {
|
||||
return new File(
|
||||
[text],
|
||||
fileName ?? "uploadedFile_" + randomUUID().replaceAll("-", "_") + ".txt",
|
||||
);
|
||||
}
|
||||
|
||||
async function uploadFile(
|
||||
filePath: string | undefined = undefined,
|
||||
fileContent:
|
||||
| Buffer<ArrayBufferLike>
|
||||
| File
|
||||
| Uint8Array<ArrayBuffer>
|
||||
| string
|
||||
| undefined = undefined,
|
||||
fileName: string | undefined = undefined,
|
||||
project_id: string | null = null,
|
||||
organization_id: string | null = null,
|
||||
client: Client | undefined = undefined,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<string | undefined> {
|
||||
let file: File | undefined = undefined;
|
||||
if (typeof filePath === "undefined" && typeof fileContent === "undefined") {
|
||||
throw new Error(
|
||||
"One between filePath and fileContent needs to be provided",
|
||||
);
|
||||
} else if (typeof filePath != "undefined") {
|
||||
const buffer = await fs.readFile(filePath);
|
||||
const actualFileName = fileName ?? path.basename(filePath);
|
||||
const uint8Array = new Uint8Array(buffer);
|
||||
file = new File([uint8Array], actualFileName);
|
||||
} else if (typeof fileContent != "undefined") {
|
||||
if (fileContent instanceof File) {
|
||||
file = fileContent;
|
||||
} else if (fileContent instanceof Buffer) {
|
||||
const fileType = await fileTypeFromBuffer(fileContent);
|
||||
const ext = fileType?.ext ?? "pdf";
|
||||
const uint8Array = new Uint8Array(fileContent);
|
||||
file = new File(
|
||||
[uint8Array],
|
||||
fileName ??
|
||||
"uploadedFile_" + randomUUID().replaceAll("-", "_") + "." + ext,
|
||||
);
|
||||
} else if (fileContent instanceof Uint8Array) {
|
||||
const fileType = await fileTypeFromBuffer(fileContent);
|
||||
const ext = fileType?.ext ?? "pdf";
|
||||
file = new File(
|
||||
[fileContent],
|
||||
fileName ??
|
||||
"uploadedFile_" + randomUUID().replaceAll("-", "_") + "." + ext,
|
||||
);
|
||||
} else if (typeof fileContent === "string") {
|
||||
file = textToFile(fileContent, fileName);
|
||||
} else {
|
||||
throw new Error("Unsupported fileContent type");
|
||||
}
|
||||
}
|
||||
const fileToUpload = {
|
||||
upload_file: file,
|
||||
} as BodyUploadFileApiV1FilesPost;
|
||||
const uploadData = {
|
||||
body: fileToUpload,
|
||||
query: { organization_id: organization_id, project_id: project_id },
|
||||
} as UploadFileApiV1FilesPostData;
|
||||
const uploadOptions = uploadData as Options<UploadFileApiV1FilesPostData>;
|
||||
if (typeof client != "undefined") {
|
||||
uploadOptions.client = client;
|
||||
}
|
||||
let retries: number = 0;
|
||||
while (true) {
|
||||
if (retries > maxRetriesOnError) {
|
||||
throw new Error(
|
||||
"Error while processing your file: Exceeded maximum number of retries, the API keeps returning errors.",
|
||||
);
|
||||
}
|
||||
const uploadResponse = await uploadFileApiV1FilesPost(uploadOptions);
|
||||
let fileId: string | undefined = undefined;
|
||||
if (!uploadResponse.response.ok) {
|
||||
retries++;
|
||||
await sleep(retryInterval * 1000);
|
||||
}
|
||||
if (typeof uploadResponse.data != "undefined") {
|
||||
fileId = uploadResponse.data.id as string;
|
||||
return fileId;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function createExtractJob(
|
||||
options:
|
||||
| Options<RunJobApiV1ExtractionJobsPostData>
|
||||
| Options<ExtractStatelessApiV1ExtractionRunPostData>,
|
||||
stateless: boolean = false,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<string | undefined> {
|
||||
let retries: number = 0;
|
||||
while (true) {
|
||||
if (retries > maxRetriesOnError) {
|
||||
throw new Error(
|
||||
"Error while creating the extraction job: Exceeded maximum number of retries, the API keeps returning errors.",
|
||||
);
|
||||
}
|
||||
let response:
|
||||
| {
|
||||
data: ExtractJob | undefined;
|
||||
request: Request;
|
||||
response: Response;
|
||||
}
|
||||
| undefined = undefined;
|
||||
if (!stateless) {
|
||||
response = (await runJobApiV1ExtractionJobsPost(
|
||||
options as Options<RunJobApiV1ExtractionJobsPostData>,
|
||||
)) as {
|
||||
data: ExtractJob | undefined;
|
||||
request: Request;
|
||||
response: Response;
|
||||
};
|
||||
} else {
|
||||
response = (await extractStatelessApiV1ExtractionRunPost(
|
||||
options as Options<ExtractStatelessApiV1ExtractionRunPostData>,
|
||||
)) as {
|
||||
data: ExtractJob | undefined;
|
||||
request: Request;
|
||||
response: Response;
|
||||
};
|
||||
}
|
||||
if (!response.response.ok) {
|
||||
if ("error" in response) {
|
||||
console.log(
|
||||
"An error occurred: ",
|
||||
JSON.stringify(response.error),
|
||||
"\nRetrying...",
|
||||
);
|
||||
}
|
||||
retries++;
|
||||
await sleep(retryInterval * 1000);
|
||||
}
|
||||
if (typeof response.data != "undefined") {
|
||||
const jobStatus = response.data.status as StatusEnum;
|
||||
if (jobStatus == "CANCELLED") {
|
||||
retries++;
|
||||
await sleep(retryInterval * 1000);
|
||||
} else if (jobStatus == "ERROR") {
|
||||
retries++;
|
||||
await sleep(retryInterval * 1000);
|
||||
} else {
|
||||
return response.data.id as string;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function pollForJobCompletion(
|
||||
jobId: string,
|
||||
interval: number = 1,
|
||||
maxIterations: number = 1800,
|
||||
client: Client | undefined = undefined,
|
||||
): Promise<boolean> {
|
||||
let status: StatusEnum | undefined = undefined;
|
||||
const jobData = {
|
||||
path: { job_id: jobId },
|
||||
} as GetJobApiV1ExtractionJobsJobIdGetData;
|
||||
const jobOptions = jobData as Options<GetJobApiV1ExtractionJobsJobIdGetData>;
|
||||
if (typeof client != "undefined") {
|
||||
jobOptions.client = client;
|
||||
}
|
||||
let numIterations: number = 0;
|
||||
while (true) {
|
||||
if (numIterations > maxIterations) {
|
||||
return false;
|
||||
}
|
||||
const response = await getJobApiV1ExtractionJobsJobIdGet(jobOptions);
|
||||
if (!response.response.ok) {
|
||||
numIterations++;
|
||||
}
|
||||
if (typeof response.data != "undefined") {
|
||||
status = response.data.status as StatusEnum;
|
||||
if (status == StatusEnum.CANCELLED || status == StatusEnum.ERROR) {
|
||||
throw new Error("There was an error extracting data from your file.");
|
||||
} else if (status == StatusEnum.SUCCESS) {
|
||||
return true;
|
||||
} else {
|
||||
numIterations++;
|
||||
await sleep(interval * 1000);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function getJobResult(
|
||||
jobId: string,
|
||||
client: Client | undefined = undefined,
|
||||
project_id: string | null = null,
|
||||
organization_id: string | null = null,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<ExtractResult | undefined> {
|
||||
const jobData = {
|
||||
path: { job_id: jobId },
|
||||
query: { organization_id: organization_id, project_id: project_id },
|
||||
} as GetJobResultApiV1ExtractionJobsJobIdResultGetData;
|
||||
const jobOptions =
|
||||
jobData as Options<GetJobResultApiV1ExtractionJobsJobIdResultGetData>;
|
||||
if (typeof client != "undefined") {
|
||||
jobOptions.client = client;
|
||||
}
|
||||
let retries: number = 0;
|
||||
while (true) {
|
||||
if (retries > maxRetriesOnError) {
|
||||
throw new Error(
|
||||
"Error while getting the result of the extraction job: Exceeded maximum number of retries, the API keeps returning errors.",
|
||||
);
|
||||
}
|
||||
const response =
|
||||
await getJobResultApiV1ExtractionJobsJobIdResultGet(jobOptions);
|
||||
if (!response.response.ok) {
|
||||
if ("error" in response) {
|
||||
console.log(
|
||||
"An error occurred: ",
|
||||
JSON.stringify(response.error),
|
||||
"\nRetrying...",
|
||||
);
|
||||
}
|
||||
retries++;
|
||||
await sleep(retryInterval * 1000);
|
||||
}
|
||||
if (typeof response.data != "undefined") {
|
||||
return {
|
||||
data: response.data.data,
|
||||
extractionMetadata: response.data.extraction_metadata,
|
||||
} as ExtractResult;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export async function extract(
|
||||
agentId: string,
|
||||
filePath: string | undefined = undefined,
|
||||
fileContent:
|
||||
| Buffer<ArrayBufferLike>
|
||||
| File
|
||||
| Uint8Array<ArrayBuffer>
|
||||
| string
|
||||
| undefined = undefined,
|
||||
fileName: string | undefined = undefined,
|
||||
project_id: string | null = null,
|
||||
organization_id: string | null = null,
|
||||
client: Client | undefined = undefined,
|
||||
fromUi: boolean | undefined = undefined,
|
||||
pollingInterval: number = 1,
|
||||
maxPollingIterations: number = 1800,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<ExtractResult | undefined> {
|
||||
const fileId = (await uploadFile(
|
||||
filePath,
|
||||
fileContent,
|
||||
fileName,
|
||||
project_id,
|
||||
organization_id,
|
||||
client,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
)) as string;
|
||||
const extractJobCreate = {
|
||||
extraction_agent_id: agentId,
|
||||
file_id: fileId,
|
||||
} as ExtractJobCreate;
|
||||
const extractData = {
|
||||
body: extractJobCreate,
|
||||
query: { from_ui: fromUi },
|
||||
} as RunJobApiV1ExtractionJobsPostData;
|
||||
const extractOptions =
|
||||
extractData as Options<RunJobApiV1ExtractionJobsPostData>;
|
||||
if (typeof client != "undefined") {
|
||||
extractOptions.client = client;
|
||||
}
|
||||
const jobId = (await createExtractJob(
|
||||
extractOptions,
|
||||
false,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
)) as string;
|
||||
const success = await pollForJobCompletion(
|
||||
jobId,
|
||||
pollingInterval,
|
||||
maxPollingIterations,
|
||||
client,
|
||||
);
|
||||
if (!success) {
|
||||
throw new Error("Your job is taking longer than 10 minutes, timing out...");
|
||||
} else {
|
||||
return (await getJobResult(
|
||||
jobId,
|
||||
client,
|
||||
project_id,
|
||||
organization_id,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
)) as ExtractResult;
|
||||
}
|
||||
}
|
||||
|
||||
export async function extractStateless(
|
||||
dataSchema:
|
||||
| {
|
||||
[key: string]:
|
||||
| { [key: string]: unknown }
|
||||
| Array<unknown>
|
||||
| string
|
||||
| number
|
||||
| number
|
||||
| boolean
|
||||
| null;
|
||||
}
|
||||
| string,
|
||||
config: ExtractConfig = {} as ExtractConfig,
|
||||
filePath: string | undefined = undefined,
|
||||
fileContent:
|
||||
| Buffer<ArrayBufferLike>
|
||||
| File
|
||||
| Uint8Array<ArrayBuffer>
|
||||
| string
|
||||
| undefined = undefined,
|
||||
fileName: string | undefined = undefined,
|
||||
project_id: string | null = null,
|
||||
organization_id: string | null = null,
|
||||
client: Client | undefined = undefined,
|
||||
pollingInterval: number = 1,
|
||||
maxPollingIterations: number = 1800,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<ExtractResult | undefined> {
|
||||
const fileId = (await uploadFile(
|
||||
filePath,
|
||||
fileContent,
|
||||
fileName,
|
||||
project_id,
|
||||
organization_id,
|
||||
client,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
)) as string;
|
||||
const extractStatetelessCreate = {
|
||||
data_schema: dataSchema,
|
||||
file_id: fileId,
|
||||
config: config,
|
||||
} as StatelessExtractionRequest;
|
||||
const extractStatetelessData = {
|
||||
body: extractStatetelessCreate,
|
||||
} as ExtractStatelessApiV1ExtractionRunPostData;
|
||||
const extractOptions =
|
||||
extractStatetelessData as Options<ExtractStatelessApiV1ExtractionRunPostData>;
|
||||
if (typeof client != "undefined") {
|
||||
extractOptions.client = client;
|
||||
}
|
||||
const jobId = (await createExtractJob(
|
||||
extractOptions,
|
||||
true,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
)) as string;
|
||||
const success = await pollForJobCompletion(
|
||||
jobId,
|
||||
pollingInterval,
|
||||
maxPollingIterations,
|
||||
client,
|
||||
);
|
||||
if (!success) {
|
||||
throw new Error("Your job is taking longer than 10 minutes, timing out...");
|
||||
} else {
|
||||
return (await getJobResult(
|
||||
jobId,
|
||||
client,
|
||||
project_id,
|
||||
organization_id,
|
||||
maxRetriesOnError,
|
||||
retryInterval,
|
||||
)) as ExtractResult;
|
||||
}
|
||||
}
|
||||
|
||||
export async function deleteAgent(
|
||||
id: string,
|
||||
client: Client | undefined = undefined,
|
||||
maxRetriesOnError: number = 10,
|
||||
retryInterval: number = 0.5,
|
||||
): Promise<boolean | undefined> {
|
||||
const deleteData = {
|
||||
path: { extraction_agent_id: id },
|
||||
} as DeleteExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdDeleteData;
|
||||
const deleteOptions =
|
||||
deleteData as Options<DeleteExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdDeleteData>;
|
||||
if (typeof client != "undefined") {
|
||||
deleteOptions.client = client;
|
||||
}
|
||||
let retries: number = 0;
|
||||
while (true) {
|
||||
if (retries > maxRetriesOnError) {
|
||||
throw new Error(
|
||||
"Maximum number of attempts for deleting agent " +
|
||||
id +
|
||||
" reached, but the API continues to return errors.",
|
||||
);
|
||||
}
|
||||
const response =
|
||||
await deleteExtractionAgentApiV1ExtractionExtractionAgentsExtractionAgentIdDelete(
|
||||
deleteOptions,
|
||||
);
|
||||
if (!response.response.ok) {
|
||||
if ("error" in response) {
|
||||
console.log(
|
||||
`An error occurred while deleting the agent: ${JSON.stringify(
|
||||
response.error,
|
||||
)}\nRetrying...`,
|
||||
);
|
||||
}
|
||||
retries++;
|
||||
await sleep(retryInterval * 1000);
|
||||
} else {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export { type ExtractAgent, type ExtractConfig };
|
||||
@@ -6,3 +6,5 @@ export {
|
||||
} from "./LlamaCloudRetriever.js";
|
||||
export type { CloudConstructorParams } from "./type.js";
|
||||
export { LlamaParseReader } from "./reader.js";
|
||||
export { LlamaExtract, LlamaExtractAgent } from "./LlamaExtract.js";
|
||||
export type { ExtractConfig } from "./extract.js";
|
||||
|
||||
@@ -14,7 +14,8 @@ import {
|
||||
getJobTextResultApiV1ParsingJobJobIdResultTextGet,
|
||||
uploadFileApiV1ParsingUploadPost,
|
||||
} from "./api";
|
||||
import { sleep } from "./utils";
|
||||
import { sleep, getSavePath } from "./utils";
|
||||
import type { ParseResult } from "./type";
|
||||
|
||||
export type Language = ParserLanguages;
|
||||
export type ResultType = "text" | "markdown" | "json";
|
||||
@@ -594,15 +595,15 @@ export class LlamaParseReader extends FileReader {
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads data from a file and returns an array of JSON objects.
|
||||
* Loads data from a file and returns an array of ParseResult objects.
|
||||
* To be used with resultType "json".
|
||||
*
|
||||
* @param filePathOrContent - The file path or the file content as a Uint8Array.
|
||||
* @returns A Promise that resolves to an array of JSON objects.
|
||||
* @returns A Promise that resolves to an array of ParseResult objects, encapsulating the JSON array resulting from the parsing within the pages attribute.
|
||||
*/
|
||||
async loadJson(
|
||||
filePathOrContent: string | Uint8Array,
|
||||
): Promise<Record<string, any>[]> {
|
||||
): Promise<ParseResult[]> {
|
||||
let jobId;
|
||||
const isFilePath =
|
||||
typeof filePathOrContent === "string" &&
|
||||
@@ -628,7 +629,7 @@ export class LlamaParseReader extends FileReader {
|
||||
const resultJson = await this.getJobResult(jobId, "json");
|
||||
resultJson.job_id = jobId;
|
||||
resultJson.file_path = isFilePath ? filePathOrContent : undefined;
|
||||
return [resultJson];
|
||||
return [resultJson] as ParseResult[];
|
||||
} catch (e) {
|
||||
console.error(`Error while parsing the file under job id ${jobId}`, e);
|
||||
if (this.ignoreErrors) {
|
||||
@@ -639,6 +640,80 @@ export class LlamaParseReader extends FileReader {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads data from a file or file bytes (or an array of those) and returns an array of ParseResult objects.
|
||||
*
|
||||
* @param filePathOrContent - The file path or the file content as a Uint8Array, or an array of one of those two types.
|
||||
* @returns A Promise that resolves to an array of ParseResult objects, encapsulating the JSON array resulting from the parsing within the pages attribute.
|
||||
*/
|
||||
async parse(
|
||||
filePathOrContent: string | Uint8Array | string[] | Uint8Array[],
|
||||
): Promise<ParseResult[]> {
|
||||
const jsonResults: Record<string, any>[][] = [];
|
||||
if (!Array.isArray(filePathOrContent)) {
|
||||
const jsonResult = await this.loadJson(filePathOrContent);
|
||||
jsonResults.push(jsonResult);
|
||||
} else {
|
||||
for (let i = 0; i < filePathOrContent.length; i++) {
|
||||
console.log(
|
||||
`Processing file ${i + 1} of ${filePathOrContent.length}...`,
|
||||
);
|
||||
const jsonResult = await this.loadJson(
|
||||
filePathOrContent[i] as string | Uint8Array,
|
||||
);
|
||||
jsonResults.push(jsonResult);
|
||||
}
|
||||
}
|
||||
const parseResults: ParseResult[] = [];
|
||||
for (const jsonResult of jsonResults) {
|
||||
for (const result of jsonResult) {
|
||||
const parseResult = {
|
||||
pages: result.pages,
|
||||
job_metadata: result.job_metadata,
|
||||
job_id: result.job_id,
|
||||
file_path: result?.file_path ?? "",
|
||||
is_completed: true,
|
||||
} as ParseResult;
|
||||
parseResults.push(parseResult);
|
||||
}
|
||||
}
|
||||
return parseResults;
|
||||
}
|
||||
|
||||
/**
|
||||
* Downloads and saves tables from a given array of ParseResult to a specified download path.
|
||||
*
|
||||
* @param jsonResults - The array of ParseResult containing table information.
|
||||
* @param downloadPath - The path where the downloaded tables will be saved as CSV files.
|
||||
* @returns A Promise that resolves to an array of strings representing the paths to the tables.
|
||||
*/
|
||||
async getTables(
|
||||
jsonResults: ParseResult[],
|
||||
downloadPath: string,
|
||||
): Promise<string[]> {
|
||||
const tables: string[] = [];
|
||||
for (const result of jsonResults) {
|
||||
for (const page of result.pages) {
|
||||
if ("items" in page && Array.isArray(page.items)) {
|
||||
for (let i = 0; i < page.items.length; i++) {
|
||||
if (
|
||||
"type" in page.items[i] &&
|
||||
page.items[i].type === "table" &&
|
||||
"csv" in page.items[i] &&
|
||||
typeof page.items[i].csv === "string" &&
|
||||
page.items[i].csv != ""
|
||||
) {
|
||||
const savePath = getSavePath(downloadPath, i);
|
||||
await fs.writeFile(savePath, page.items[i].csv);
|
||||
tables.push(savePath);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return tables;
|
||||
}
|
||||
|
||||
/**
|
||||
* Downloads and saves images from a given JSON result to a specified download path.
|
||||
* Currently only supports resultType "json".
|
||||
@@ -648,7 +723,7 @@ export class LlamaParseReader extends FileReader {
|
||||
* @returns A Promise that resolves to an array of image objects.
|
||||
*/
|
||||
async getImages(
|
||||
jsonResult: Record<string, any>[],
|
||||
jsonResult: ParseResult[],
|
||||
downloadPath: string,
|
||||
): Promise<Record<string, any>[]> {
|
||||
try {
|
||||
|
||||
@@ -8,3 +8,54 @@ export type CloudConstructorParams = {
|
||||
projectName: string;
|
||||
organizationId?: string | undefined;
|
||||
} & ClientParams;
|
||||
|
||||
export type ExtractResult = {
|
||||
data:
|
||||
| {
|
||||
[key: string]:
|
||||
| {
|
||||
[key: string]: unknown;
|
||||
}
|
||||
| Array<unknown>
|
||||
| string
|
||||
| number
|
||||
| number
|
||||
| boolean
|
||||
| null;
|
||||
}
|
||||
| Array<{
|
||||
[key: string]:
|
||||
| {
|
||||
[key: string]: unknown;
|
||||
}
|
||||
| Array<unknown>
|
||||
| string
|
||||
| number
|
||||
| number
|
||||
| boolean
|
||||
| null;
|
||||
}>
|
||||
| null;
|
||||
extractionMetadata: {
|
||||
[key: string]:
|
||||
| {
|
||||
[key: string]: unknown;
|
||||
}
|
||||
| Array<unknown>
|
||||
| string
|
||||
| number
|
||||
| number
|
||||
| boolean
|
||||
| null;
|
||||
};
|
||||
};
|
||||
|
||||
export type ParseResult = {
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
pages: Record<string, any>[];
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
job_metadata: Record<string, any>;
|
||||
job_id: string;
|
||||
is_completed: boolean;
|
||||
file_path: string;
|
||||
};
|
||||
|
||||
@@ -6,6 +6,8 @@ import {
|
||||
import { DEFAULT_BASE_URL } from "@llamaindex/core/global";
|
||||
import { getEnv } from "@llamaindex/env";
|
||||
import type { ClientParams } from "./type.js";
|
||||
import fs from "fs";
|
||||
import path from "path";
|
||||
|
||||
function getBaseUrl(baseUrl?: string): string {
|
||||
return baseUrl ?? getEnv("LLAMA_CLOUD_BASE_URL") ?? DEFAULT_BASE_URL;
|
||||
@@ -99,3 +101,19 @@ export async function getPipelineId(
|
||||
export async function sleep(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
export function getSavePath(downloadPath: string, i: number): string {
|
||||
const now = new Date();
|
||||
const formatted =
|
||||
now.toISOString().replace(/[-:T]/g, "_").replace(/\..+/, "") +
|
||||
"_" +
|
||||
now.getMilliseconds().toString().padStart(3, "0");
|
||||
|
||||
let savePath = path.join(downloadPath, `table_${formatted}.csv`);
|
||||
|
||||
if (fs.existsSync(savePath)) {
|
||||
savePath = savePath.replace(".csv", "_") + i.toString() + ".csv";
|
||||
}
|
||||
|
||||
return savePath;
|
||||
}
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
import { describe, it, expect, beforeEach, beforeAll } from "vitest";
|
||||
import { LlamaParseReader } from "../src/reader.js";
|
||||
import { LlamaCloudIndex } from "../src/LlamaCloudIndex.js";
|
||||
import { LlamaExtract, LlamaExtractAgent } from "../src/LlamaExtract.js";
|
||||
import { Document } from "@llamaindex/core/schema";
|
||||
import { fs } from "@llamaindex/env";
|
||||
import { ExtractConfig } from "../src/api.js";
|
||||
import { ParseResult } from "../src/type.js";
|
||||
|
||||
// Integration tests that require actual API keys and files
|
||||
describe("Integration Tests", () => {
|
||||
@@ -282,6 +285,78 @@ describe("Integration Tests", () => {
|
||||
},
|
||||
60000,
|
||||
);
|
||||
|
||||
it.skipIf(skipIfNoApiKey)(
|
||||
"should parse a file and return a ParseResult array",
|
||||
async () => {
|
||||
const parseReader = new LlamaParseReader({
|
||||
apiKey: process.env.LLAMA_CLOUD_API_KEY!,
|
||||
verbose: false,
|
||||
});
|
||||
|
||||
const testContent = "Test document for JSON parsing";
|
||||
const testFilePath = "test-json.txt";
|
||||
|
||||
await fs.writeFile(testFilePath, new TextEncoder().encode(testContent));
|
||||
|
||||
try {
|
||||
const result = await parseReader.parse(testFilePath);
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(Array.isArray(result)).toBe(true);
|
||||
expect(result.length).toBeGreaterThan(0);
|
||||
expect(result[0]).toHaveProperty("job_id");
|
||||
expect(result[0]).toHaveProperty("job_metadata");
|
||||
expect(result[0]).toHaveProperty("file_path");
|
||||
expect(result[0]).toHaveProperty("pages");
|
||||
|
||||
await fs.unlink(testFilePath);
|
||||
} catch (error) {
|
||||
try {
|
||||
await fs.unlink(testFilePath);
|
||||
} catch {}
|
||||
throw error;
|
||||
}
|
||||
},
|
||||
60000,
|
||||
);
|
||||
|
||||
it.skipIf(skipIfNoApiKey)(
|
||||
"should extract tables correctly from a JSON result",
|
||||
async () => {
|
||||
const parseReader = new LlamaParseReader({
|
||||
apiKey: process.env.LLAMA_CLOUD_API_KEY!,
|
||||
verbose: false,
|
||||
});
|
||||
const pseudoJsonResult = [
|
||||
{
|
||||
pages: [
|
||||
{
|
||||
items: [
|
||||
{
|
||||
type: "table",
|
||||
csv: "Name,Age,Height (cm)\nAnna,12,140\nBob,22,175\nClaire,33,173\nDenis,44,185\n",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
job_id: "jobId",
|
||||
job_metadata: { job_id: "jobId" },
|
||||
file_path: "table.csv",
|
||||
is_completed: true,
|
||||
},
|
||||
] as ParseResult[];
|
||||
|
||||
const tmpdir = await fs.mkdtemp("tables");
|
||||
const result = await parseReader.getTables(pseudoJsonResult, tmpdir);
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(Array.isArray(result)).toBe(true);
|
||||
expect(result.length).toBeGreaterThan(0);
|
||||
expect(typeof result[0] === "string").toBe(true);
|
||||
},
|
||||
60000,
|
||||
);
|
||||
});
|
||||
|
||||
describe("LlamaCloudIndex Integration", () => {
|
||||
@@ -414,6 +489,121 @@ describe("Integration Tests", () => {
|
||||
);
|
||||
});
|
||||
|
||||
describe("LlamaExtract Integration", () => {
|
||||
it.skipIf(skipIfNoApiKey)(
|
||||
"should create agents correctly",
|
||||
async () => {
|
||||
const dataSchema = {
|
||||
properties: {
|
||||
text: {
|
||||
description: "Text from the file",
|
||||
type: "string",
|
||||
},
|
||||
},
|
||||
required: ["text"],
|
||||
type: "object",
|
||||
};
|
||||
const extractClient = new LlamaExtract(
|
||||
process.env.LLAMA_CLOUD_API_KEY!,
|
||||
"https://api.cloud.llamaindex.ai",
|
||||
);
|
||||
const agent = await extractClient.createAgent(
|
||||
"ExtractTestAgent",
|
||||
dataSchema,
|
||||
);
|
||||
expect(agent).instanceOf(LlamaExtractAgent);
|
||||
},
|
||||
60000,
|
||||
);
|
||||
it.skipIf(skipIfNoApiKey)(
|
||||
"should fetch agents correctly",
|
||||
async () => {
|
||||
const extractClient = new LlamaExtract(
|
||||
process.env.LLAMA_CLOUD_API_KEY!,
|
||||
"https://api.cloud.llamaindex.ai",
|
||||
);
|
||||
const agent = await extractClient.getAgent("ExtractTestAgent");
|
||||
expect(agent).instanceOf(LlamaExtractAgent);
|
||||
},
|
||||
60000,
|
||||
);
|
||||
it.skipIf(skipIfNoApiKey)(
|
||||
"should extract data correctly (file paths and file contents) with an agent and delete that agent",
|
||||
async () => {
|
||||
const extractClient = new LlamaExtract(
|
||||
process.env.LLAMA_CLOUD_API_KEY!,
|
||||
"https://api.cloud.llamaindex.ai",
|
||||
);
|
||||
const agent = await extractClient.getAgent("ExtractTestAgent");
|
||||
const testContent =
|
||||
"**Text to extract**: Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.";
|
||||
const testFilePath = "test-extract-agent.md";
|
||||
|
||||
await fs.writeFile(testFilePath, new TextEncoder().encode(testContent));
|
||||
const result = await agent!.extract("test-extract-agent.md");
|
||||
expect("data" in result!).toBeTruthy();
|
||||
expect("extractionMetadata" in result!).toBeTruthy();
|
||||
|
||||
const buffer = await fs.readFile("test-extract-agent.md");
|
||||
const resultBuffer = await agent!.extract(
|
||||
undefined,
|
||||
buffer,
|
||||
"test-extract-agent.md",
|
||||
);
|
||||
expect("data" in resultBuffer!).toBeTruthy();
|
||||
expect("extractionMetadata" in resultBuffer!).toBeTruthy();
|
||||
|
||||
const success = await extractClient.deleteAgent(agent!.id);
|
||||
expect(success).toBeTruthy();
|
||||
},
|
||||
60000,
|
||||
);
|
||||
it.skipIf(skipIfNoApiKey)(
|
||||
"should extract statelessly file paths and file contents",
|
||||
async () => {
|
||||
const dataSchema = {
|
||||
properties: {
|
||||
text: {
|
||||
description: "Text from the file",
|
||||
type: "string",
|
||||
},
|
||||
},
|
||||
required: ["text"],
|
||||
type: "object",
|
||||
};
|
||||
|
||||
const testContent =
|
||||
"**Text to extract**: Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.";
|
||||
const testFilePath = "test-extract.md";
|
||||
|
||||
await fs.writeFile(testFilePath, new TextEncoder().encode(testContent));
|
||||
|
||||
const extractClient = new LlamaExtract(
|
||||
process.env.LLAMA_CLOUD_API_KEY!,
|
||||
"https://api.cloud.llamaindex.ai",
|
||||
);
|
||||
const result = await extractClient.extract(
|
||||
dataSchema,
|
||||
{} as ExtractConfig,
|
||||
"test-extract.md",
|
||||
);
|
||||
expect("data" in result!).toBeTruthy();
|
||||
expect("extractionMetadata" in result!).toBeTruthy();
|
||||
|
||||
const buffer = await fs.readFile("test-extract.md");
|
||||
const resultBuffer = await extractClient.extract(
|
||||
dataSchema,
|
||||
{} as ExtractConfig,
|
||||
undefined,
|
||||
buffer,
|
||||
); // testing without passing a file name
|
||||
expect("data" in resultBuffer!).toBeTruthy();
|
||||
expect("extractionMetadata" in resultBuffer!).toBeTruthy();
|
||||
},
|
||||
60000,
|
||||
);
|
||||
});
|
||||
|
||||
describe("Error Handling Integration", () => {
|
||||
it.skipIf(skipIfNoApiKey)(
|
||||
"should handle malformed files gracefully",
|
||||
|
||||
Reference in New Issue
Block a user