mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-09-13 04:04:09 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
87a08e0240 | ||
|
|
1a2604255f | ||
|
|
94b312901b | ||
|
|
25fe95dd1a | ||
|
|
f338b66274 | ||
|
|
8b1cd47f51 | ||
|
|
48395b2f00 | ||
|
|
91ef2939c5 | ||
|
|
30d84c41ad | ||
|
|
a5c531fd29 | ||
|
|
976a20496d | ||
|
|
9bb94c2337 |
@@ -61,6 +61,7 @@ wheels/
|
|||||||
.installed.cfg
|
.installed.cfg
|
||||||
*.egg
|
*.egg
|
||||||
MANIFEST
|
MANIFEST
|
||||||
|
.python-version
|
||||||
|
|
||||||
# IDE
|
# IDE
|
||||||
.vscode/
|
.vscode/
|
||||||
|
|||||||
+60
-18
@@ -7,34 +7,76 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||||||
|
|
||||||
## [Unreleased]
|
## [Unreleased]
|
||||||
|
|
||||||
### Fixed
|
## [0.2.0] - 2026-01-10
|
||||||
- Fixed model switching bug in `TextEmbedder` where internal state was not cleared, preventing dynamic updates between `fastembed` and `sentence_transformers` (#160).
|
|
||||||
- Implemented model-intrinsic embedding dimension detection in `TextEmbedder` to ensure consistency between models and vector databases.
|
|
||||||
- Updated `set_model` to properly refresh configuration and dimensions during model switches.
|
|
||||||
|
|
||||||
### Added
|
|
||||||
- Added comprehensive unit test suite `tests/embeddings/test_model_switching.py` for verifying dynamic model transitions and dimension updates.
|
|
||||||
- Fixed `TypeError: unhashable type: 'Entity'` in `GraphAnalyzer` when processing graphs with raw `Entity` objects or dictionaries in relationships (#159).
|
|
||||||
- Robustified ID extraction across `CentralityCalculator`, `CommunityDetector`, and `ConnectivityAnalyzer` to handle various entity formats.
|
|
||||||
- Improved `Entity` class hashability and equality logic in `utils/types.py`.
|
|
||||||
- Added end-to-end integration test suite for Knowledge Graph pipeline validation (GraphBuilder -> EntityResolver -> GraphAnalyzer).
|
|
||||||
- Fixed orchestrator lazy property initialization and configuration normalization logic in `Orchestrator`.
|
|
||||||
- Resolved `AssertionError` in orchestrator tests by aligning test mocks with production component usage.
|
|
||||||
- Fixed dependency compatibility issues by pinning `protobuf==4.25.3` and `grpcio==1.67.1`.
|
|
||||||
- Added missing dependencies `GitPython` and `chardet` to `pyproject.toml`.
|
|
||||||
- Verified and aligned `FileObject.text` property usage in GraphRAG notebooks for consistent content decoding.
|
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
|
- **Amazon Neptune Support**:
|
||||||
|
- Added `AmazonNeptuneStore` providing Amazon Neptune graph database integration via Bolt protocol and OpenCypher.
|
||||||
|
- Implemented `NeptuneAuthTokenManager` extending Neo4j AuthManager for AWS IAM SigV4 signing with automatic token refresh.
|
||||||
|
- Added robust connection handling: retry logic with backoff for transient errors (signature expired, connection closed) and driver recreation.
|
||||||
|
- Added `graph-amazon-neptune` optional dependency group (boto3, neo4j).
|
||||||
|
- Comprehensive test suite covering all GraphStore interface methods.
|
||||||
|
- **Docling Integration**:
|
||||||
|
- Added `DoclingParser` in `semantica.parse` for high-fidelity document parsing using the Docling library.
|
||||||
|
- Supports multi-format parsing (PDF, DOCX, PPTX, XLSX, HTML, images) with superior table extraction and structure understanding.
|
||||||
|
- Implemented as a standalone parser supporting local execution, OCR, and multiple export formats (Markdown, HTML, JSON).
|
||||||
|
- **Robust Extraction Fallbacks**:
|
||||||
|
- Implemented comprehensive fallback chains ("ML/LLM" -> "Pattern" -> "Last Resort") across `NERExtractor`, `RelationExtractor`, and `TripletExtractor` to prevent empty result lists.
|
||||||
|
- Added "Last Resort" pattern matching in `NERExtractor` to identify capitalized words as generic entities when all other methods fail.
|
||||||
|
- Added "Last Resort" adjacency-based relation extraction in `RelationExtractor` to create weak connections between adjacent entities if no relations are found.
|
||||||
|
- Added fallback logic in `TripletExtractor` to convert relations to triplets or use rule-based extraction if standard methods fail.
|
||||||
|
- **Provenance & Tracking**:
|
||||||
|
- Added count tracking to batch processing logs in `NERExtractor`, `RelationExtractor`, and `TripletExtractor`.
|
||||||
|
- Added `batch_index` and `document_id` to the metadata of all extracted entities, relations, triplets, semantic roles, and clusters for better traceability.
|
||||||
- **Semantic Extract Improvements**:
|
- **Semantic Extract Improvements**:
|
||||||
- Introduced `auto-chunking` for long text processing in LLM extraction methods (`extract_entities_llm`, `extract_relations_llm`, `extract_triplets_llm`).
|
- Introduced `auto-chunking` for long text processing in LLM extraction methods (`extract_entities_llm`, `extract_relations_llm`, `extract_triplets_llm`).
|
||||||
- Added `silent_fail` parameter to LLM extraction methods for configurable error handling.
|
- Added `silent_fail` parameter to LLM extraction methods for configurable error handling.
|
||||||
- Implemented robust JSON parsing and automatic retry logic (3 attempts with exponential backoff) in `BaseProvider` for all LLM providers.
|
- Implemented robust JSON parsing and automatic retry logic (3 attempts with exponential backoff) in `BaseProvider` for all LLM providers.
|
||||||
- Enhanced `GroqProvider` with better diagnostics and connectivity testing.
|
- Enhanced `GroqProvider` with better diagnostics and connectivity testing.
|
||||||
- Added comprehensive entity, relation, and triplet deduplication for chunked extraction.
|
- Added comprehensive entity, relation, and triplet deduplication for chunked extraction.
|
||||||
|
- Added `semantica/semantic_extract/schemas.py` with canonical Pydantic models for consistent structured output.
|
||||||
|
- **Testing**:
|
||||||
|
- Added comprehensive robustness test suite `tests/semantic_extract/test_robustness_fallback.py` for validating extraction fallbacks and metadata propagation.
|
||||||
|
- Added comprehensive unit test suite `tests/embeddings/test_model_switching.py` for verifying dynamic model transitions and dimension updates.
|
||||||
|
- Added end-to-end integration test suite for Knowledge Graph pipeline validation (GraphBuilder -> EntityResolver -> GraphAnalyzer).
|
||||||
|
- **Other**:
|
||||||
|
- Added missing dependencies `GitPython` and `chardet` to `pyproject.toml`.
|
||||||
|
- Robustified ID extraction across `CentralityCalculator`, `CommunityDetector`, and `ConnectivityAnalyzer` to handle various entity formats.
|
||||||
|
- Improved `Entity` class hashability and equality logic in `utils/types.py`.
|
||||||
|
|
||||||
|
### Changed
|
||||||
|
- **Deduplication & Conflict Logic**:
|
||||||
|
- Removed internal deduplication logic from `NERExtractor`, `RelationExtractor`, and `TripletExtractor`.
|
||||||
|
- Removed consistency/conflict checking from `ExtractionValidator` to defer to dedicated `semantica/conflicts` module.
|
||||||
|
- Removed `_deduplicate_*` methods from `semantica/semantic_extract/methods.py`.
|
||||||
|
- **Batch Processing & Consistency**:
|
||||||
|
- Standardized batch processing across all extractors (`NERExtractor`, `RelationExtractor`, `TripletExtractor`, `SemanticNetworkExtractor`, `EventDetector`, `SemanticAnalyzer`, `CoreferenceResolver`) using a unified `extract`/`analyze`/`resolve` method pattern with progress tracking.
|
||||||
|
- Added provenance metadata (`batch_index`, `document_id`) to `SemanticNetwork` nodes/edges, `Event` objects, `SemanticRole` results, `CoreferenceChain` mentions, and `SemanticCluster` (tracking source `document_ids`).
|
||||||
|
- Updated `SemanticClusterer.cluster` and `SemanticAnalyzer.cluster_semantically` to accept list of dictionaries (with `content` and `id` keys) for better document tracking during clustering.
|
||||||
|
- Removed legacy `check_triplet_consistency` from `TripletExtractor`.
|
||||||
|
- Removed `validate_consistency` and `_check_consistency` from `ExtractionValidator`.
|
||||||
|
- **Weighted Scoring**:
|
||||||
|
- Clarified weighted confidence scoring (50% Method Confidence + 50% Type Similarity) in comments.
|
||||||
|
- Explicitly labeled "Type Similarity" as "user-provided" in code comments to remove ambiguity.
|
||||||
|
- **Refactoring**:
|
||||||
|
- Fixed orchestrator lazy property initialization and configuration normalization logic in `Orchestrator`.
|
||||||
|
- Verified and aligned `FileObject.text` property usage in GraphRAG notebooks for consistent content decoding.
|
||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
- Fixed a bug in `TripletExtractor` where the `validate_triplets` method was shadowed by an internal attribute.
|
- **Critical Fixes**:
|
||||||
- Fixed incorrect `TextSplitter` import path in the `semantic_extract.methods` module.
|
- Resolved `NameError` in `extraction_validator.py` by adding missing `Union` import.
|
||||||
|
- Resolved issues where extractors would return empty lists for valid input text when primary extraction methods failed.
|
||||||
|
- Fixed metadata initialization issue in batch processing where `batch_index` and `document_id` were occasionally missing from extracted items.
|
||||||
|
- Ensured `LLMExtraction` methods (`enhance_entities`, `enhance_relations`) return original input instead of failing or returning empty results when LLM providers are unavailable.
|
||||||
|
- **Component Fixes**:
|
||||||
|
- Fixed model switching bug in `TextEmbedder` where internal state was not cleared, preventing dynamic updates between `fastembed` and `sentence_transformers` (#160).
|
||||||
|
- Implemented model-intrinsic embedding dimension detection in `TextEmbedder` to ensure consistency between models and vector databases.
|
||||||
|
- Updated `set_model` to properly refresh configuration and dimensions during model switches.
|
||||||
|
- Fixed `TypeError: unhashable type: 'Entity'` in `GraphAnalyzer` when processing graphs with raw `Entity` objects or dictionaries in relationships (#159).
|
||||||
|
- Resolved `AssertionError` in orchestrator tests by aligning test mocks with production component usage.
|
||||||
|
- Fixed dependency compatibility issues by pinning `protobuf==4.25.3` and `grpcio==1.67.1`.
|
||||||
|
- Fixed a bug in `TripletExtractor` where the `validate_triplets` method was shadowed by an internal attribute.
|
||||||
|
- Fixed incorrect `TextSplitter` import path in the `semantic_extract.methods` module.
|
||||||
|
|
||||||
## [0.1.1] - 2026-01-05
|
## [0.1.1] - 2026-01-05
|
||||||
|
|
||||||
|
|||||||
@@ -28,7 +28,7 @@
|
|||||||
|
|
||||||
*The missing fabric between raw data and AI engineering. A comprehensive open-source framework for building semantic layers and knowledge engineering systems that transform unstructured data into AI-ready knowledge — powering Knowledge Graph-Powered RAG (GraphRAG), AI Agents, Multi-Agent Systems, and AI applications with structured semantic knowledge.*
|
*The missing fabric between raw data and AI engineering. A comprehensive open-source framework for building semantic layers and knowledge engineering systems that transform unstructured data into AI-ready knowledge — powering Knowledge Graph-Powered RAG (GraphRAG), AI Agents, Multi-Agent Systems, and AI applications with structured semantic knowledge.*
|
||||||
|
|
||||||
**100% Open Source** • **MIT Licensed** • **Latest Version: 0.1.1** • **Production Ready** • **Community Driven**
|
**100% Open Source** • **MIT Licensed** • **Latest Version: 0.2.0** • **Production Ready** • **Community Driven**
|
||||||
|
|
||||||
[**Discord**](https://discord.gg/pMHguUzG)
|
[**Discord**](https://discord.gg/pMHguUzG)
|
||||||
|
|
||||||
@@ -360,7 +360,7 @@ results = vector_store.search(query="supply chain", top_k=5)
|
|||||||
|
|
||||||
### Graph Store & Triplet Store
|
### Graph Store & Triplet Store
|
||||||
|
|
||||||
> **Neo4j, FalkorDB support** • **SPARQL queries** • **RDF triplets**
|
> **Neo4j, FalkorDB, Amazon Neptune support** • **SPARQL queries** • **RDF triplets**
|
||||||
|
|
||||||
```python
|
```python
|
||||||
from semantica.graph_store import GraphStore
|
from semantica.graph_store import GraphStore
|
||||||
@@ -370,6 +370,24 @@ from semantica.triplet_store import TripletStore
|
|||||||
graph_store = GraphStore(backend="neo4j", uri="bolt://localhost:7687", user="neo4j", password="password")
|
graph_store = GraphStore(backend="neo4j", uri="bolt://localhost:7687", user="neo4j", password="password")
|
||||||
graph_store.add_nodes([{"id": "n1", "labels": ["Person"], "properties": {"name": "Alice"}}])
|
graph_store.add_nodes([{"id": "n1", "labels": ["Person"], "properties": {"name": "Alice"}}])
|
||||||
|
|
||||||
|
# Amazon Neptune Graph Store (OpenCypher via HTTP with IAM Auth)
|
||||||
|
neptune_store = GraphStore(
|
||||||
|
backend="neptune",
|
||||||
|
endpoint="your-cluster.us-east-1.neptune.amazonaws.com",
|
||||||
|
port=8182,
|
||||||
|
region="us-east-1",
|
||||||
|
iam_auth=True, # Uses AWS credential chain (boto3, env vars, or IAM role)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Node Operations
|
||||||
|
neptune_store.add_nodes([
|
||||||
|
{"labels": ["Person"], "properties": {"id": "alice", "name": "Alice", "age": 30}},
|
||||||
|
{"labels": ["Person"], "properties": {"id": "bob", "name": "Bob", "age": 25}},
|
||||||
|
])
|
||||||
|
|
||||||
|
# Query Operations
|
||||||
|
result = neptune_store.execute_query("MATCH (p:Person) RETURN p.name, p.age")
|
||||||
|
|
||||||
# Triplet Store (Blazegraph, Jena, RDF4J)
|
# Triplet Store (Blazegraph, Jena, RDF4J)
|
||||||
triplet_store = TripletStore(backend="blazegraph", endpoint="http://localhost:9999/blazegraph")
|
triplet_store = TripletStore(backend="blazegraph", endpoint="http://localhost:9999/blazegraph")
|
||||||
triplet_store.add_triplet({"subject": "Alice", "predicate": "knows", "object": "Bob"})
|
triplet_store.add_triplet({"subject": "Alice", "predicate": "knows", "object": "Bob"})
|
||||||
|
|||||||
+3
-3
@@ -26,10 +26,10 @@ Before releasing, ensure:
|
|||||||
|
|
||||||
The project uses GitHub Actions for automated releases to PyPI.
|
The project uses GitHub Actions for automated releases to PyPI.
|
||||||
|
|
||||||
1. **Tag the commit**: Create a new git tag for the version (e.g., `v0.1.1`).
|
1. **Tag the commit**: Create a new git tag for the version (e.g., `v0.2.0`).
|
||||||
```bash
|
```bash
|
||||||
git tag -a v0.1.1 -m "Release v0.1.1"
|
git tag -a v0.2.0 -m "Release v0.2.0"
|
||||||
git push origin v0.1.1
|
git push origin v0.2.0
|
||||||
```
|
```
|
||||||
2. **GitHub Action**: The `Release` workflow will automatically trigger, build the package, create a GitHub Release, and publish to PyPI using Trusted Publishing.
|
2. **GitHub Action**: The `Release` workflow will automatically trigger, build the package, create a GitHub Release, and publish to PyPI using Trusted Publishing.
|
||||||
|
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ We actively support the following versions of Semantica with security updates:
|
|||||||
|
|
||||||
| Version | Supported |
|
| Version | Supported |
|
||||||
| ------- | ------------------ |
|
| ------- | ------------------ |
|
||||||
|
| 0.2.0 | :white_check_mark: |
|
||||||
| 0.1.1 | :white_check_mark: |
|
| 0.1.1 | :white_check_mark: |
|
||||||
| 0.1.0 | :white_check_mark: |
|
| 0.1.0 | :white_check_mark: |
|
||||||
| < 0.1.0 | :x: |
|
| < 0.1.0 | :x: |
|
||||||
|
|||||||
@@ -0,0 +1,667 @@
|
|||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"# Amazon Neptune Graph Store\n",
|
||||||
|
"\n",
|
||||||
|
"## Overview\n",
|
||||||
|
"\n",
|
||||||
|
"This notebook covers the Amazon Neptune Database integration in Semantica. Amazon Neptune is a fully managed graph database service that supports both property graphs (via OpenCypher/Gremlin) and RDF graphs (via SPARQL).\n",
|
||||||
|
"\n",
|
||||||
|
"### Key Features\n",
|
||||||
|
"\n",
|
||||||
|
"- **IAM Authentication**: Secure access using AWS SigV4 signatures via AuthManager\n",
|
||||||
|
"- **OpenCypher Support**: Query using standard OpenCypher syntax\n",
|
||||||
|
"- **Bolt Protocol**: Uses Neo4j Bolt driver for efficient binary communication\n",
|
||||||
|
"- **Native ~id Support**: Leverages Neptune's native element ID handling\n",
|
||||||
|
"- **Full CRUD Operations**: Create, read, update, delete nodes and relationships\n",
|
||||||
|
"- **Automatic Retry**: Built-in retry logic with exponential backoff for transient errors\n",
|
||||||
|
"\n",
|
||||||
|
"### Prerequisites\n",
|
||||||
|
"\n",
|
||||||
|
"- An Amazon Neptune Database cluster\n",
|
||||||
|
"- AWS credentials configured (boto3, environment variables, or IAM role)\n",
|
||||||
|
"- Network access to your Neptune cluster (VPC, security groups)\n",
|
||||||
|
"\n",
|
||||||
|
"---"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Installation\n",
|
||||||
|
"\n",
|
||||||
|
"```bash\n",
|
||||||
|
"# Install Semantica with Neptune support\n",
|
||||||
|
"pip install semantica\n",
|
||||||
|
"\n",
|
||||||
|
"# Required dependencies (installed automatically)\n",
|
||||||
|
"pip install boto3 neo4j\n",
|
||||||
|
"```"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!pip install semantica"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Configuration\n",
|
||||||
|
"\n",
|
||||||
|
"Set your Neptune cluster endpoint and AWS credentials. Replace the placeholder values with your actual configuration."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"import os\n",
|
||||||
|
"\n",
|
||||||
|
"# Neptune cluster configuration - REPLACE WITH YOUR VALUES\n",
|
||||||
|
"os.environ[\"NEPTUNE_ENDPOINT\"] = \"your-cluster.us-east-1.neptune.amazonaws.com\"\n",
|
||||||
|
"os.environ[\"NEPTUNE_PORT\"] = \"8182\"\n",
|
||||||
|
"os.environ[\"AWS_REGION\"] = \"us-east-1\"\n",
|
||||||
|
"\n",
|
||||||
|
"# AWS credentials (if using IAM Auth and not relying on IAM role or ~/.aws/credentials)\n",
|
||||||
|
"# os.environ[\"AWS_ACCESS_KEY_ID\"] = \"your-access-key-id\"\n",
|
||||||
|
"# os.environ[\"AWS_SECRET_ACCESS_KEY\"] = \"your-secret-access-key\"\n",
|
||||||
|
"# os.environ[\"AWS_SESSION_TOKEN\"] = \"your-session-token\"\n",
|
||||||
|
"\n",
|
||||||
|
"print(f\"Neptune Endpoint: {os.environ.get('NEPTUNE_ENDPOINT')}\")\n",
|
||||||
|
"print(f\"AWS Region: {os.environ.get('AWS_REGION')}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Step 1: Initialize Neptune Store\n",
|
||||||
|
"\n",
|
||||||
|
"Initialize a connection to your Amazon Neptune cluster with IAM authentication."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"import os\n",
|
||||||
|
"from semantica.graph_store import GraphStore\n",
|
||||||
|
"\n",
|
||||||
|
"# Option 1: Using GraphStore factory (recommended)\n",
|
||||||
|
"neptune_store = GraphStore(\n",
|
||||||
|
" backend=\"neptune\",\n",
|
||||||
|
" endpoint=os.environ.get(\"NEPTUNE_ENDPOINT\"),\n",
|
||||||
|
" port=int(os.environ.get(\"NEPTUNE_PORT\", 8182)),\n",
|
||||||
|
" region=os.environ.get(\"AWS_REGION\", \"us-east-1\"),\n",
|
||||||
|
" iam_auth=True,\n",
|
||||||
|
")\n",
|
||||||
|
"\n",
|
||||||
|
"# Connect to Neptune\n",
|
||||||
|
"neptune_store.connect()\n",
|
||||||
|
"print(\"Connected to Amazon Neptune!\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"### Development/Testing Without IAM Auth\n",
|
||||||
|
"\n",
|
||||||
|
"For development or testing environments where IAM authentication is not required (e.g., Neptune notebooks or VPC-only access), you can disable IAM signing:"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# For dev/test environments without IAM authentication\n",
|
||||||
|
"neptune_store_dev = GraphStore(\n",
|
||||||
|
" backend=\"neptune\",\n",
|
||||||
|
" endpoint=os.environ.get(\"NEPTUNE_ENDPOINT\"),\n",
|
||||||
|
" port=int(os.environ.get(\"NEPTUNE_PORT\", 8182)),\n",
|
||||||
|
" region=os.environ.get(\"AWS_REGION\", \"us-east-1\"),\n",
|
||||||
|
" iam_auth=False, # Disable IAM signing for dev/test\n",
|
||||||
|
")\n",
|
||||||
|
"neptune_store_dev.connect()"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"### Authentication Options\n",
|
||||||
|
"\n",
|
||||||
|
"IAM Authentication (recommended for production) automatically uses the AWS credential chain:\n",
|
||||||
|
"1. Environment variables (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)\n",
|
||||||
|
"2. AWS credentials file (~/.aws/credentials)\n",
|
||||||
|
"3. IAM role (for EC2, Lambda, ECS)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Step 2: Node Operations\n",
|
||||||
|
"\n",
|
||||||
|
"### Creating Nodes\n",
|
||||||
|
"\n",
|
||||||
|
"Nodes represent entities in your graph. Each node can have:\n",
|
||||||
|
"- **ID**: A unique identifier (custom or auto-generated UUID)\n",
|
||||||
|
"- **Labels**: Categories/types (e.g., `Person`, `Company`)\n",
|
||||||
|
"- **Properties**: Key-value pairs (e.g., `{\"name\": \"Alice\", \"age\": 30}`)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Create a single node with custom ID (id in properties)\n",
|
||||||
|
"alice = neptune_store.create_node(\n",
|
||||||
|
" labels=[\"Person\"],\n",
|
||||||
|
" properties={\"id\": \"alice\", \"name\": \"Alice\", \"age\": 30, \"role\": \"Engineer\"}\n",
|
||||||
|
")\n",
|
||||||
|
"print(f\"Created node: {alice}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Create a node with auto-generated UUID (no id in properties)\n",
|
||||||
|
"bob = neptune_store.create_node(\n",
|
||||||
|
" labels=[\"Person\"],\n",
|
||||||
|
" properties={\"name\": \"Bob\", \"age\": 25, \"role\": \"Designer\"}\n",
|
||||||
|
")\n",
|
||||||
|
"print(f\"Created node with UUID: {bob['id']}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Create a company node with auto-generated ID\n",
|
||||||
|
"acme = neptune_store.create_node(\n",
|
||||||
|
" labels=[\"Company\"],\n",
|
||||||
|
" properties={\"name\": \"Acme Corp\", \"industry\": \"Technology\", \"founded\": 2010}\n",
|
||||||
|
")\n",
|
||||||
|
"print(f\"Created company: {acme}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"### Creating Multiple Nodes (Batch)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Batch create nodes for better performance\n",
|
||||||
|
"# Include 'id' in properties for custom IDs\n",
|
||||||
|
"nodes_data = [\n",
|
||||||
|
" {\"labels\": [\"Person\"], \"properties\": {\"id\": \"charlie\", \"name\": \"Charlie\", \"age\": 35}},\n",
|
||||||
|
" {\"labels\": [\"Person\"], \"properties\": {\"id\": \"diana\", \"name\": \"Diana\", \"age\": 28}},\n",
|
||||||
|
" {\"labels\": [\"Location\"], \"properties\": {\"name\": \"San Francisco\", \"state\": \"CA\"}},\n",
|
||||||
|
"]\n",
|
||||||
|
"\n",
|
||||||
|
"created_nodes = neptune_store.create_nodes(nodes_data)\n",
|
||||||
|
"print(f\"Created {len(created_nodes)} nodes in batch\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"### Retrieving Nodes"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Get a specific node by ID\n",
|
||||||
|
"alice_node = neptune_store.get_node(node_id=\"alice\")\n",
|
||||||
|
"print(f\"Retrieved: {alice_node}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Get nodes by label\n",
|
||||||
|
"people = neptune_store.get_nodes(labels=[\"Person\"], limit=10)\n",
|
||||||
|
"print(f\"Found {len(people)} Person nodes:\")\n",
|
||||||
|
"for person in people:\n",
|
||||||
|
" print(f\" - {person.get('properties', {}).get('name')}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Get nodes by properties\n",
|
||||||
|
"engineers = neptune_store.get_nodes(\n",
|
||||||
|
" labels=[\"Person\"],\n",
|
||||||
|
" properties={\"role\": \"Engineer\"},\n",
|
||||||
|
" limit=5\n",
|
||||||
|
")\n",
|
||||||
|
"print(f\"Found {len(engineers)} engineers\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"### Updating Nodes"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Update node properties (merge mode - default)\n",
|
||||||
|
"updated_alice = neptune_store.update_node(\n",
|
||||||
|
" node_id=\"alice\",\n",
|
||||||
|
" properties={\"age\": 31, \"department\": \"AI Research\"},\n",
|
||||||
|
" merge=True\n",
|
||||||
|
")\n",
|
||||||
|
"print(f\"Updated Alice: {updated_alice}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Replace all properties (merge=False)\n",
|
||||||
|
"# WARNING: This removes properties not in the update\n",
|
||||||
|
"replaced = neptune_store.update_node(\n",
|
||||||
|
" node_id=\"charlie\",\n",
|
||||||
|
" properties={\"name\": \"Charlie\", \"age\": 36},\n",
|
||||||
|
" merge=False\n",
|
||||||
|
")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"### Deleting Nodes"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Delete a node (with detach=True to also delete relationships)\n",
|
||||||
|
"deleted = neptune_store.delete_node(node_id=\"diana\", detach=True)\n",
|
||||||
|
"print(f\"Deleted diana: {deleted}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Without detach (fails if node has relationships)\n",
|
||||||
|
"# neptune_store.delete_node(node_id=\"alice\", detach=False)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Step 3: Relationship Operations\n",
|
||||||
|
"\n",
|
||||||
|
"### Creating Relationships\n",
|
||||||
|
"\n",
|
||||||
|
"Relationships connect nodes and represent connections between entities."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Create a relationship between Alice and Acme\n",
|
||||||
|
"works_at = neptune_store.create_relationship(\n",
|
||||||
|
" start_node_id=\"alice\",\n",
|
||||||
|
" end_node_id=acme[\"id\"],\n",
|
||||||
|
" rel_type=\"WORKS_AT\",\n",
|
||||||
|
" properties={\"since\": 2020, \"position\": \"Senior Engineer\"}\n",
|
||||||
|
")\n",
|
||||||
|
"print(f\"Created relationship: {works_at}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Create a KNOWS relationship between people\n",
|
||||||
|
"knows_rel = neptune_store.create_relationship(\n",
|
||||||
|
" start_node_id=\"alice\",\n",
|
||||||
|
" end_node_id=bob[\"id\"],\n",
|
||||||
|
" rel_type=\"KNOWS\",\n",
|
||||||
|
" properties={\"since\": 2019}\n",
|
||||||
|
")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"### Retrieving Relationships"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Get all relationships for a node\n",
|
||||||
|
"alice_rels = neptune_store.get_relationships(node_id=\"alice\", direction=\"both\")\n",
|
||||||
|
"print(f\"Alice has {len(alice_rels)} relationships\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Get outgoing relationships only\n",
|
||||||
|
"outgoing = neptune_store.get_relationships(node_id=\"alice\", direction=\"out\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Filter by relationship type\n",
|
||||||
|
"works_rels = neptune_store.get_relationships(\n",
|
||||||
|
" node_id=\"alice\",\n",
|
||||||
|
" rel_type=\"WORKS_AT\",\n",
|
||||||
|
" direction=\"out\"\n",
|
||||||
|
")\n",
|
||||||
|
"print(f\"Alice's work relationships: {len(works_rels)}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"### Deleting Relationships"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Delete a specific relationship by ID\n",
|
||||||
|
"if works_at.get(\"id\"):\n",
|
||||||
|
" deleted = neptune_store.delete_relationship(rel_id=works_at[\"id\"])\n",
|
||||||
|
" print(f\"Deleted relationship: {deleted}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Step 4: OpenCypher Queries\n",
|
||||||
|
"\n",
|
||||||
|
"Amazon Neptune supports OpenCypher queries via the Bolt protocol. Execute complex graph patterns using standard Cypher syntax."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Simple query\n",
|
||||||
|
"results = neptune_store.execute_query(\n",
|
||||||
|
" \"MATCH (p:Person) RETURN p.name, p.age ORDER BY p.age\"\n",
|
||||||
|
")\n",
|
||||||
|
"print(\"People in the graph:\")\n",
|
||||||
|
"for record in results.get(\"records\", []):\n",
|
||||||
|
" print(f\" - {record.get('p.name')}: {record.get('p.age')} years old\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Using parameters (safer and more efficient)\n",
|
||||||
|
"results = neptune_store.execute_query(\n",
|
||||||
|
" \"MATCH (p:Person) WHERE p.age > $min_age RETURN p.name, p.age\",\n",
|
||||||
|
" parameters={\"min_age\": 25}\n",
|
||||||
|
")\n",
|
||||||
|
"print(f\"People over 25: {len(results.get('records', []))}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Find relationships between nodes\n",
|
||||||
|
"results = neptune_store.execute_query(\"\"\"\n",
|
||||||
|
" MATCH (p:Person)-[r:WORKS_AT]->(c:Company)\n",
|
||||||
|
" RETURN p.name as employee, c.name as company, r.since as start_year\n",
|
||||||
|
"\"\"\")\n",
|
||||||
|
"for record in results.get(\"records\", []):\n",
|
||||||
|
" print(f\"{record['employee']} works at {record['company']} since {record['start_year']}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Count and aggregate\n",
|
||||||
|
"results = neptune_store.execute_query(\"\"\"\n",
|
||||||
|
" MATCH (p:Person)\n",
|
||||||
|
" RETURN count(p) as total, avg(p.age) as avg_age, max(p.age) as max_age\n",
|
||||||
|
"\"\"\")\n",
|
||||||
|
"stats = results.get(\"records\", [{}])[0]\n",
|
||||||
|
"print(f\"Total: {stats.get('total')}, Avg Age: {stats.get('avg_age'):.1f}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Step 5: Graph Analytics\n",
|
||||||
|
"\n",
|
||||||
|
"### Get Neighbors\n",
|
||||||
|
"\n",
|
||||||
|
"Traverse the graph to find connected nodes."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Get immediate neighbors (depth=1)\n",
|
||||||
|
"neighbors = neptune_store.get_neighbors(\n",
|
||||||
|
" node_id=\"alice\",\n",
|
||||||
|
" direction=\"both\",\n",
|
||||||
|
" depth=1\n",
|
||||||
|
")\n",
|
||||||
|
"print(f\"Alice's direct neighbors: {len(neighbors)}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Get neighbors up to 2 hops away\n",
|
||||||
|
"extended = neptune_store.get_neighbors(\n",
|
||||||
|
" node_id=\"alice\",\n",
|
||||||
|
" direction=\"out\",\n",
|
||||||
|
" depth=2\n",
|
||||||
|
")\n",
|
||||||
|
"print(f\"Nodes within 2 hops: {len(extended)}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"### Shortest Path\n",
|
||||||
|
"\n",
|
||||||
|
"Find the shortest path between two nodes."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Find shortest path\n",
|
||||||
|
"path = neptune_store.shortest_path(\n",
|
||||||
|
" start_node_id=\"alice\",\n",
|
||||||
|
" end_node_id=\"charlie\",\n",
|
||||||
|
" max_depth=5\n",
|
||||||
|
")\n",
|
||||||
|
"\n",
|
||||||
|
"if path:\n",
|
||||||
|
" print(\"Path found!\")\n",
|
||||||
|
" print(f\" Length: {path.get('length')}\")\n",
|
||||||
|
" print(f\" Nodes: {len(path.get('nodes', []))}\")\n",
|
||||||
|
" print(f\" Relationships: {len(path.get('relationships', []))}\")\n",
|
||||||
|
"else:\n",
|
||||||
|
" print(\"No path found between nodes\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Step 6: Graph Statistics\n",
|
||||||
|
"\n",
|
||||||
|
"Get comprehensive statistics about your graph."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Get graph statistics\n",
|
||||||
|
"stats = neptune_store.get_stats()\n",
|
||||||
|
"\n",
|
||||||
|
"print(\"Graph Statistics:\")\n",
|
||||||
|
"print(f\" Total nodes: {stats.get('node_count', 'N/A')}\")\n",
|
||||||
|
"print(f\" Total relationships: {stats.get('relationship_count', 'N/A')}\")\n",
|
||||||
|
"\n",
|
||||||
|
"print(\"\\nNode labels:\")\n",
|
||||||
|
"for label, count in stats.get('label_counts', {}).items():\n",
|
||||||
|
" print(f\" - {label}: {count}\")\n",
|
||||||
|
"\n",
|
||||||
|
"print(\"\\nRelationship types:\")\n",
|
||||||
|
"for rel_type, count in stats.get('relationship_type_counts', {}).items():\n",
|
||||||
|
" print(f\" - {rel_type}: {count}\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Step 7: Connection Management\n",
|
||||||
|
"\n",
|
||||||
|
"Always close connections when done to free resources."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Check connection status\n",
|
||||||
|
"status = neptune_store.get_status()\n",
|
||||||
|
"print(f\"Connection status: {status}\")\n",
|
||||||
|
"\n",
|
||||||
|
"# Close the connection\n",
|
||||||
|
"neptune_store.close()\n",
|
||||||
|
"print(\"Connection closed\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Neptune-Specific Considerations\n",
|
||||||
|
"\n",
|
||||||
|
"### Native Element IDs\n",
|
||||||
|
"\n",
|
||||||
|
"Neptune uses native `~id` for element identification. Include `id` in properties to set a custom ID:\n",
|
||||||
|
"\n",
|
||||||
|
"```python\n",
|
||||||
|
"# Create a node with custom ID (include 'id' in properties)\n",
|
||||||
|
"node = neptune_store.create_node(\n",
|
||||||
|
" labels=[\"Person\"],\n",
|
||||||
|
" properties={\"id\": \"my-custom-id\", \"name\": \"Test\"}\n",
|
||||||
|
")\n",
|
||||||
|
"\n",
|
||||||
|
"# Create a node with auto-generated UUID (omit 'id' from properties)\n",
|
||||||
|
"node = neptune_store.create_node(\n",
|
||||||
|
" labels=[\"Person\"],\n",
|
||||||
|
" properties={\"name\": \"Test\"}\n",
|
||||||
|
")\n",
|
||||||
|
"\n",
|
||||||
|
"# The ID is used in id() function calls internally:\n",
|
||||||
|
"# MATCH (n) WHERE id(n) = 'my-custom-id' RETURN n\n",
|
||||||
|
"```\n",
|
||||||
|
"\n",
|
||||||
|
"### OpenCypher Considerations\n",
|
||||||
|
"\n",
|
||||||
|
"Amazon Neptune Database's OpenCypher implementation has some differences from Neo4j:\n",
|
||||||
|
"\n",
|
||||||
|
"1. **No `shortestPath()` function**: Use variable-length path patterns or `allShortestPaths()`\n",
|
||||||
|
"2. **Labels syntax**: Use `labels(n)` function to retrieve node labels\n",
|
||||||
|
"3. **Property updates**: Use `SET n += {props}` for merge behavior\n",
|
||||||
|
"\n",
|
||||||
|
"For the complete OpenCypher specification supported by Amazon Neptune Database, see the [AWS documentation](https://docs.aws.amazon.com/neptune/latest/userguide/access-graph-opencypher.html).\n",
|
||||||
|
"\n",
|
||||||
|
"### Amazon Neptune Analytics\n",
|
||||||
|
"\n",
|
||||||
|
"For analytical (OLAP) workloads such as graph algorithms, aggregations, and large-scale traversals, consider [Amazon Neptune Analytics](https://docs.aws.amazon.com/neptune-analytics/latest/userguide/what-is-neptune-analytics.html). Neptune Analytics complements Neptune Database by providing optimized performance for analytical queries while Neptune Database is optimized for transactional (OLTP) workloads.\n",
|
||||||
|
"\n",
|
||||||
|
"### Performance Tips\n",
|
||||||
|
"\n",
|
||||||
|
"1. **Use batch operations** for creating multiple nodes/relationships\n",
|
||||||
|
"2. **Use parameters** in queries to enable query caching\n",
|
||||||
|
"3. **Limit result sets** with `LIMIT` clause"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Summary\n",
|
||||||
|
"\n",
|
||||||
|
"This notebook covered the Amazon Neptune Graph Store integration:\n",
|
||||||
|
"\n",
|
||||||
|
"- **IAM Authentication**: Secure AWS SigV4 signing\n",
|
||||||
|
"- **CRUD Operations**: Full node and relationship management\n",
|
||||||
|
"- **OpenCypher Queries**: Standard graph query language\n",
|
||||||
|
"- **Graph Analytics**: Neighbors and shortest path algorithms\n",
|
||||||
|
"- **Statistics & Monitoring**: Graph metrics and status\n",
|
||||||
|
"\n",
|
||||||
|
"### Key Takeaways\n",
|
||||||
|
"\n",
|
||||||
|
"- Neptune uses native `~id` for element identification\n",
|
||||||
|
"- IAM authentication is recommended for production\n",
|
||||||
|
"- Bolt protocol provides efficient binary query interface\n",
|
||||||
|
"- Semantica abstracts Neptune-specific syntax differences\n",
|
||||||
|
"\n",
|
||||||
|
"### Next Steps\n",
|
||||||
|
"\n",
|
||||||
|
"- [Graph Store (Neo4j/FalkorDB)](09_Graph_Store.ipynb) - Compare with other backends\n",
|
||||||
|
"- [Building Knowledge Graphs](07_Building_Knowledge_Graphs.ipynb) - Build production KGs\n",
|
||||||
|
"- [Graph Analytics](10_Graph_Analytics.ipynb) - Advanced analytics algorithms"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"metadata": {
|
||||||
|
"kernelspec": {
|
||||||
|
"display_name": "Python 3",
|
||||||
|
"language": "python",
|
||||||
|
"name": "python3"
|
||||||
|
},
|
||||||
|
"language_info": {
|
||||||
|
"name": "python",
|
||||||
|
"version": "3.9.0"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 4
|
||||||
|
}
|
||||||
+5
-5
@@ -12,22 +12,22 @@ How to cite Semantica in academic papers and research.
|
|||||||
author = {Hawksight AI},
|
author = {Hawksight AI},
|
||||||
year = {2026},
|
year = {2026},
|
||||||
url = {https://github.com/Hawksight-AI/semantica},
|
url = {https://github.com/Hawksight-AI/semantica},
|
||||||
version = {0.1.1},
|
version = {0.2.0},
|
||||||
doi = {10.5281/zenodo.XXXXXXX}
|
doi = {10.5281/zenodo.XXXXXXX}
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
### APA
|
### APA
|
||||||
Hawksight AI. (2026). *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering* (Version 0.1.1) [Computer software]. https://github.com/Hawksight-AI/semantica
|
Hawksight AI. (2026). *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering* (Version 0.2.0) [Computer software]. https://github.com/Hawksight-AI/semantica
|
||||||
|
|
||||||
### MLA
|
### MLA
|
||||||
Hawksight AI. *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering*. Version 0.1.1, GitHub, 2026, https://github.com/Hawksight-AI/semantica.
|
Hawksight AI. *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering*. Version 0.2.0, GitHub, 2026, https://github.com/Hawksight-AI/semantica.
|
||||||
|
|
||||||
### Chicago
|
### Chicago
|
||||||
Hawksight AI. *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering*. Version 0.1.1. GitHub, 2026. https://github.com/Hawksight-AI/semantica.
|
Hawksight AI. *Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering*. Version 0.2.0. GitHub, 2026. https://github.com/Hawksight-AI/semantica.
|
||||||
|
|
||||||
### IEEE
|
### IEEE
|
||||||
Hawksight AI, "Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering," Version 0.1.1, GitHub, 2026. [Online]. Available: https://github.com/Hawksight-AI/semantica
|
Hawksight AI, "Semantica: An Open Source Framework for Semantic Layers and Knowledge Engineering," Version 0.2.0, GitHub, 2026. [Online]. Available: https://github.com/Hawksight-AI/semantica
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|||||||
@@ -253,7 +253,7 @@ entities = ner.extract_entities(text)
|
|||||||
# Basic relation extraction
|
# Basic relation extraction
|
||||||
rel_extractor = RelationExtractor()
|
rel_extractor = RelationExtractor()
|
||||||
relations = rel_extractor.extract(text, entities=entities)
|
relations = rel_extractor.extract(text, entities=entities)
|
||||||
# [Relation(source="Elon Musk", target="SpaceX", type="founded")]
|
# [Relation(subject="Elon Musk", predicate="founded", object="SpaceX")]
|
||||||
|
|
||||||
# With configuration
|
# With configuration
|
||||||
rel_extractor = RelationExtractor(
|
rel_extractor = RelationExtractor(
|
||||||
@@ -424,6 +424,44 @@ enhanced_entities = extractor.enhance_entities(text, entities)
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
## Batch Processing & Provenance
|
||||||
|
|
||||||
|
All extractors support batch processing for high-throughput extraction. You can pass a list of strings or a list of dictionaries (with `content` and `id` keys).
|
||||||
|
|
||||||
|
**Features:**
|
||||||
|
- **Progress Tracking**: Automatically shows a progress bar for large batches.
|
||||||
|
- **Provenance Metadata**: Each extracted item includes `batch_index` and `document_id` in its `metadata`.
|
||||||
|
|
||||||
|
```python
|
||||||
|
from semantica.semantic_extract import NERExtractor
|
||||||
|
|
||||||
|
documents = [
|
||||||
|
{"id": "doc_1", "content": "Apple Inc. was founded by Steve Jobs."},
|
||||||
|
{"id": "doc_2", "content": "Microsoft Corporation was founded by Bill Gates."}
|
||||||
|
]
|
||||||
|
|
||||||
|
extractor = NERExtractor()
|
||||||
|
batch_results = extractor.extract(documents)
|
||||||
|
|
||||||
|
for i, doc_entities in enumerate(batch_results):
|
||||||
|
print(f"Document {i} entities:")
|
||||||
|
for entity in doc_entities:
|
||||||
|
print(f" - {entity.text} ({entity.label})")
|
||||||
|
print(f" Provenance: Batch Index {entity.metadata['batch_index']}, Doc ID {entity.metadata.get('document_id')}")
|
||||||
|
```
|
||||||
|
|
||||||
|
## Robust Extraction Fallbacks
|
||||||
|
|
||||||
|
The framework implements robust fallback chains to prevent empty results when primary methods fail (e.g., due to model unavailability or obscure text).
|
||||||
|
|
||||||
|
- **NER**: `ML/LLM` -> `Pattern` -> `Last Resort` (Capitalized Words)
|
||||||
|
- **Relation**: `Primary` -> `Pattern` -> `Last Resort` (Adjacency)
|
||||||
|
- **Triplet**: `Primary` -> `Relation-to-Triplet` -> `Pattern`
|
||||||
|
|
||||||
|
This ensures that you almost always get *some* structured data, even if it requires falling back to simpler heuristics.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
## Usage Examples
|
## Usage Examples
|
||||||
|
|
||||||
```python
|
```python
|
||||||
|
|||||||
+11
-3
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|||||||
|
|
||||||
[project]
|
[project]
|
||||||
name = "semantica"
|
name = "semantica"
|
||||||
version = "0.1.1"
|
version = "0.2.0"
|
||||||
description = "🧠 Semantica - An Open Source Framework for building Semantic Layers and Knowledge Engineering "
|
description = "🧠 Semantica - An Open Source Framework for building Semantic Layers and Knowledge Engineering "
|
||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
license = {text = "MIT"}
|
license = {text = "MIT"}
|
||||||
@@ -92,6 +92,7 @@ dependencies = [
|
|||||||
"groq>=0.4.0",
|
"groq>=0.4.0",
|
||||||
"openai>=1.0.0",
|
"openai>=1.0.0",
|
||||||
"litellm>=1.0.0",
|
"litellm>=1.0.0",
|
||||||
|
"instructor>=1.0.0",
|
||||||
"click>=8.1.0",
|
"click>=8.1.0",
|
||||||
"rich>=12.5.0",
|
"rich>=12.5.0",
|
||||||
"tqdm>=4.64.0",
|
"tqdm>=4.64.0",
|
||||||
@@ -182,8 +183,11 @@ llm-deepseek = [
|
|||||||
llm-litellm = [
|
llm-litellm = [
|
||||||
"litellm>=1.0.0"
|
"litellm>=1.0.0"
|
||||||
]
|
]
|
||||||
|
llm-instructor = [
|
||||||
|
"instructor>=1.0.0"
|
||||||
|
]
|
||||||
llm-all = [
|
llm-all = [
|
||||||
"semantica[llm-openai,llm-gemini,llm-groq,llm-anthropic,llm-ollama,llm-deepseek,llm-litellm]"
|
"semantica[llm-openai,llm-gemini,llm-groq,llm-anthropic,llm-ollama,llm-deepseek,llm-litellm,llm-instructor]"
|
||||||
]
|
]
|
||||||
models-huggingface = [
|
models-huggingface = [
|
||||||
"transformers>=4.20.0",
|
"transformers>=4.20.0",
|
||||||
@@ -209,8 +213,12 @@ graph-falkordb = [
|
|||||||
"falkordb>=1.0.0",
|
"falkordb>=1.0.0",
|
||||||
"redis>=4.3.0"
|
"redis>=4.3.0"
|
||||||
]
|
]
|
||||||
|
graph-amazon-neptune = [
|
||||||
|
"boto3>=1.24.0",
|
||||||
|
"neo4j>=5.0.0"
|
||||||
|
]
|
||||||
graph-all = [
|
graph-all = [
|
||||||
"semantica[graph-neo4j,graph-falkordb]"
|
"semantica[graph-neo4j,graph-falkordb,graph-amazon-neptune]"
|
||||||
]
|
]
|
||||||
parse-docling = [
|
parse-docling = [
|
||||||
"docling>=1.0.0"
|
"docling>=1.0.0"
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ Main exports:
|
|||||||
- Config: Configuration management
|
- Config: Configuration management
|
||||||
"""
|
"""
|
||||||
|
|
||||||
__version__ = "0.1.1"
|
__version__ = "0.2.0"
|
||||||
__author__ = "Semantica Contributors"
|
__author__ = "Semantica Contributors"
|
||||||
__license__ = "MIT"
|
__license__ = "MIT"
|
||||||
|
|
||||||
|
|||||||
@@ -1,48 +1,71 @@
|
|||||||
"""
|
"""
|
||||||
Graph Store Module
|
Graph Store Module
|
||||||
|
|
||||||
This module provides comprehensive property graph database integration for the
|
This module provides comprehensive property graph database integration for
|
||||||
Semantica framework, supporting multiple graph database backends including Neo4j
|
the Semantica framework, supporting multiple graph database backends including
|
||||||
and FalkorDB for storing and querying knowledge graphs.
|
Neo4j and FalkorDB for storing and querying knowledge graphs.
|
||||||
|
|
||||||
Algorithms Used:
|
Algorithms Used:
|
||||||
|
|
||||||
Graph Store Management:
|
Graph Store Management:
|
||||||
- Store Registration: Store type detection, store factory pattern, configuration management, default store selection
|
- Store Registration: Store type detection, store factory pattern,
|
||||||
- Backend Pattern: Unified interface for multiple backends (Neo4j, FalkorDB), backend instantiation, backend-specific operation delegation
|
configuration management, default store selection
|
||||||
- Store Selection: Default store resolution, store ID lookup, store validation
|
- Backend Pattern: Unified interface for multiple backends (Neo4j,
|
||||||
|
FalkorDB), backend instantiation, backend-specific operation delegation
|
||||||
|
- Store Selection: Default store resolution, store ID lookup,
|
||||||
|
store validation
|
||||||
|
|
||||||
Node and Relationship Operations:
|
Node and Relationship Operations:
|
||||||
- Node Creation: Single node insertion, batch node insertion, property validation, label management, backend delegation
|
- Node Creation: Single node insertion, batch node insertion,
|
||||||
- Node Retrieval: Pattern matching (label/property filtering), Cypher query construction, result extraction, node reconstruction
|
property validation, label management, backend delegation
|
||||||
- Node Update: Property update, label modification, atomic update operations, conflict detection
|
- Node Retrieval: Pattern matching (label/property filtering),
|
||||||
- Node Deletion: Node matching, cascade deletion (optional), deletion operation delegation, result verification
|
Cypher query construction, result extraction, node reconstruction
|
||||||
- Relationship Creation: Single relationship insertion, batch insertion, property validation, type management
|
- Node Update: Property update, label modification, atomic update
|
||||||
|
operations, conflict detection
|
||||||
|
- Node Deletion: Node matching, cascade deletion (optional),
|
||||||
|
deletion operation delegation, result verification
|
||||||
|
- Relationship Creation: Single relationship insertion, batch insertion,
|
||||||
|
property validation, type management
|
||||||
- Relationship Retrieval: Pattern matching, path queries, traversal queries
|
- Relationship Retrieval: Pattern matching, path queries, traversal queries
|
||||||
- Relationship Update: Property update, type modification
|
- Relationship Update: Property update, type modification
|
||||||
- Relationship Deletion: Relationship matching, deletion operation delegation
|
- Relationship Deletion: Relationship matching, deletion operation
|
||||||
|
delegation
|
||||||
|
|
||||||
Graph Query Execution:
|
Graph Query Execution:
|
||||||
- Cypher Query: Full Cypher query language support for Neo4j and FalkorDB (OpenCypher)
|
- Cypher Query: Full Cypher query language support for Neo4j and
|
||||||
- Pattern Matching: Node and relationship pattern matching, variable binding, path matching
|
FalkorDB (OpenCypher)
|
||||||
- Graph Traversal: BFS/DFS traversal, shortest path algorithms, path finding
|
- Pattern Matching: Node and relationship pattern matching, variable
|
||||||
|
binding, path matching
|
||||||
|
- Graph Traversal: BFS/DFS traversal, shortest path algorithms,
|
||||||
|
path finding
|
||||||
- Aggregation: COUNT, SUM, AVG, MIN, MAX operations, GROUP BY support
|
- Aggregation: COUNT, SUM, AVG, MIN, MAX operations, GROUP BY support
|
||||||
- Query Optimization: Query caching, execution plan analysis, index utilization
|
- Query Optimization: Query caching, execution plan analysis,
|
||||||
|
index utilization
|
||||||
|
|
||||||
Graph Analytics:
|
Graph Analytics:
|
||||||
- Centrality Algorithms: Degree centrality, betweenness centrality, PageRank, closeness centrality
|
- Centrality Algorithms: Degree centrality, betweenness centrality,
|
||||||
- Community Detection: Label propagation, Louvain modularity, connected components
|
PageRank, closeness centrality
|
||||||
- Path Algorithms: Shortest path, all shortest paths, Dijkstra, A* pathfinding
|
- Community Detection: Label propagation, Louvain modularity,
|
||||||
|
connected components
|
||||||
|
- Path Algorithms: Shortest path, all shortest paths, Dijkstra,
|
||||||
|
A* pathfinding
|
||||||
- Similarity: Node similarity, Jaccard similarity, cosine similarity
|
- Similarity: Node similarity, Jaccard similarity, cosine similarity
|
||||||
|
|
||||||
Store Backends:
|
Store Backends:
|
||||||
- Neo4j Store: Official Neo4j Python driver, Bolt protocol communication, transaction support, multi-database support, APOC procedures
|
- Neo4j Store: Official Neo4j Python driver, Bolt protocol
|
||||||
- FalkorDB Store: Redis-based graph database, sparse matrix representation, linear algebra queries, OpenCypher support, ultra-fast performance
|
communication, transaction support, multi-database support,
|
||||||
|
APOC procedures
|
||||||
|
- FalkorDB Store: Redis-based graph database, sparse matrix
|
||||||
|
representation, linear algebra queries, OpenCypher support,
|
||||||
|
ultra-fast performance
|
||||||
|
|
||||||
Bulk Operations:
|
Bulk Operations:
|
||||||
- Batch Processing: Chunking algorithm (fixed-size batch creation), batch size optimization, memory management for large datasets
|
- Batch Processing: Chunking algorithm (fixed-size batch creation),
|
||||||
- Transaction Management: ACID transaction support, batch commits, rollback on failure
|
batch size optimization, memory management for large datasets
|
||||||
- Progress Tracking: Load progress calculation, elapsed time tracking, throughput calculation
|
- Transaction Management: ACID transaction support, batch commits,
|
||||||
|
rollback on failure
|
||||||
|
- Progress Tracking: Load progress calculation, elapsed time tracking,
|
||||||
|
throughput calculation
|
||||||
|
|
||||||
Key Features:
|
Key Features:
|
||||||
- Multi-backend property graph support (Neo4j, FalkorDB)
|
- Multi-backend property graph support (Neo4j, FalkorDB)
|
||||||
@@ -79,33 +102,42 @@ Convenience Functions:
|
|||||||
- list_available_methods: List registered graph store methods
|
- list_available_methods: List registered graph store methods
|
||||||
|
|
||||||
Example Usage:
|
Example Usage:
|
||||||
>>> from semantica.graph_store import GraphStore, create_node, create_relationship, execute_query
|
>>> from semantica.graph_store import GraphStore, create_node, \
|
||||||
|
... create_relationship, execute_query
|
||||||
>>> # Using convenience functions
|
>>> # Using convenience functions
|
||||||
>>> node_id = create_node(labels=["Person"], properties={"name": "Alice", "age": 30})
|
>>> node_id = create_node(labels=["Person"],
|
||||||
>>> rel_id = create_relationship(start_id=node1_id, end_id=node2_id, rel_type="KNOWS", properties={"since": 2020})
|
... properties={"name": "Alice", "age": 30})
|
||||||
>>> results = execute_query("MATCH (p:Person) WHERE p.age > 25 RETURN p.name")
|
>>> rel_id = create_relationship(start_id=node1_id, end_id=node2_id,
|
||||||
|
... rel_type="KNOWS",
|
||||||
|
... properties={"since": 2020})
|
||||||
|
>>> results = execute_query("MATCH (p:Person) WHERE p.age > 25 "
|
||||||
|
... "RETURN p.name")
|
||||||
>>> # Using classes directly
|
>>> # Using classes directly
|
||||||
>>> store = GraphStore(backend="neo4j", uri="bolt://localhost:7687")
|
>>> store = GraphStore(backend="neo4j", uri="bolt://localhost:7687")
|
||||||
>>> node_id = store.create_node(labels=["Person"], properties={"name": "Bob"})
|
>>> node_id = store.create_node(labels=["Person"],
|
||||||
|
... properties={"name": "Bob"})
|
||||||
>>> results = store.execute_query("MATCH (n) RETURN n LIMIT 10")
|
>>> results = store.execute_query("MATCH (n) RETURN n LIMIT 10")
|
||||||
|
|
||||||
Author: Semantica Contributors
|
Author: Semantica Contributors
|
||||||
License: MIT
|
License: MIT
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from .config import GraphStoreConfig, graph_store_config
|
from .amazon_neptune import (
|
||||||
from .falkordb_store import (
|
AmazonNeptuneStore,
|
||||||
FalkorDBStore,
|
NeptuneAuthTokenManager,
|
||||||
FalkorDBClient,
|
NeptuneDriver,
|
||||||
FalkorDBGraph,
|
NeptuneSession,
|
||||||
|
NeptuneTransaction,
|
||||||
)
|
)
|
||||||
|
from .config import GraphStoreConfig, graph_store_config
|
||||||
|
from .falkordb_store import FalkorDBClient, FalkorDBGraph, FalkorDBStore
|
||||||
from .graph_store import (
|
from .graph_store import (
|
||||||
|
GraphAnalytics,
|
||||||
GraphManager,
|
GraphManager,
|
||||||
GraphStore,
|
GraphStore,
|
||||||
NodeManager,
|
NodeManager,
|
||||||
QueryEngine,
|
QueryEngine,
|
||||||
RelationshipManager,
|
RelationshipManager,
|
||||||
GraphAnalytics,
|
|
||||||
)
|
)
|
||||||
from .methods import (
|
from .methods import (
|
||||||
create_node,
|
create_node,
|
||||||
@@ -125,11 +157,7 @@ from .methods import (
|
|||||||
update_node,
|
update_node,
|
||||||
update_relationship,
|
update_relationship,
|
||||||
)
|
)
|
||||||
from .neo4j_store import (
|
from .neo4j_store import Neo4jDriver, Neo4jStore, Neo4jTransaction
|
||||||
Neo4jStore,
|
|
||||||
Neo4jDriver,
|
|
||||||
Neo4jTransaction,
|
|
||||||
)
|
|
||||||
from .registry import MethodRegistry, method_registry
|
from .registry import MethodRegistry, method_registry
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
@@ -144,6 +172,12 @@ __all__ = [
|
|||||||
"Neo4jStore",
|
"Neo4jStore",
|
||||||
"Neo4jDriver",
|
"Neo4jDriver",
|
||||||
"Neo4jTransaction",
|
"Neo4jTransaction",
|
||||||
|
# Amazon Neptune
|
||||||
|
"AmazonNeptuneStore",
|
||||||
|
"NeptuneAuthTokenManager",
|
||||||
|
"NeptuneDriver",
|
||||||
|
"NeptuneSession",
|
||||||
|
"NeptuneTransaction",
|
||||||
# FalkorDB
|
# FalkorDB
|
||||||
"FalkorDBStore",
|
"FalkorDBStore",
|
||||||
"FalkorDBClient",
|
"FalkorDBClient",
|
||||||
@@ -171,4 +205,3 @@ __all__ = [
|
|||||||
"MethodRegistry",
|
"MethodRegistry",
|
||||||
"method_registry",
|
"method_registry",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -6,7 +6,8 @@ supporting multiple configuration sources including environment variables, confi
|
|||||||
and programmatic configuration.
|
and programmatic configuration.
|
||||||
|
|
||||||
Supported Configuration Sources:
|
Supported Configuration Sources:
|
||||||
- Environment variables: GRAPH_STORE_DEFAULT_BACKEND, GRAPH_STORE_NEO4J_URI, GRAPH_STORE_FALKORDB_HOST, etc.
|
- Environment variables: GRAPH_STORE_DEFAULT_BACKEND,
|
||||||
|
GRAPH_STORE_NEO4J_URI, GRAPH_STORE_FALKORDB_HOST, etc.
|
||||||
- Config files: YAML, JSON, TOML formats
|
- Config files: YAML, JSON, TOML formats
|
||||||
- Programmatic: Python API for setting graph store configurations
|
- Programmatic: Python API for setting graph store configurations
|
||||||
|
|
||||||
@@ -44,7 +45,11 @@ from ..utils.logging import get_logger
|
|||||||
|
|
||||||
|
|
||||||
class GraphStoreConfig:
|
class GraphStoreConfig:
|
||||||
"""Configuration manager for graph store module - supports .env files, environment variables, and programmatic config."""
|
"""
|
||||||
|
Configuration manager for graph store module.
|
||||||
|
|
||||||
|
Supports .env files, environment variables, and programmatic config.
|
||||||
|
"""
|
||||||
|
|
||||||
def __init__(self, config_file: Optional[str] = None):
|
def __init__(self, config_file: Optional[str] = None):
|
||||||
"""
|
"""
|
||||||
@@ -124,6 +129,15 @@ class GraphStoreConfig:
|
|||||||
"GRAPH_STORE_FALKORDB_PORT": "falkordb_port",
|
"GRAPH_STORE_FALKORDB_PORT": "falkordb_port",
|
||||||
"GRAPH_STORE_FALKORDB_PASSWORD": "falkordb_password",
|
"GRAPH_STORE_FALKORDB_PASSWORD": "falkordb_password",
|
||||||
"GRAPH_STORE_FALKORDB_GRAPH_NAME": "falkordb_graph_name",
|
"GRAPH_STORE_FALKORDB_GRAPH_NAME": "falkordb_graph_name",
|
||||||
|
# Amazon Neptune settings
|
||||||
|
"GRAPH_STORE_NEPTUNE_ENDPOINT": "neptune_endpoint",
|
||||||
|
"GRAPH_STORE_NEPTUNE_PORT": "neptune_port",
|
||||||
|
"GRAPH_STORE_NEPTUNE_REGION": "neptune_region",
|
||||||
|
"GRAPH_STORE_NEPTUNE_IAM_AUTH": "neptune_iam_auth",
|
||||||
|
"GRAPH_STORE_NEPTUNE_USE_SSL": "neptune_use_ssl",
|
||||||
|
"AWS_ACCESS_KEY_ID": "neptune_access_key",
|
||||||
|
"AWS_SECRET_ACCESS_KEY": "neptune_secret_key",
|
||||||
|
"AWS_SESSION_TOKEN": "neptune_session_token",
|
||||||
}
|
}
|
||||||
|
|
||||||
for env_var, config_key in env_mappings.items():
|
for env_var, config_key in env_mappings.items():
|
||||||
@@ -135,6 +149,7 @@ class GraphStoreConfig:
|
|||||||
"timeout",
|
"timeout",
|
||||||
"max_retries",
|
"max_retries",
|
||||||
"falkordb_port",
|
"falkordb_port",
|
||||||
|
"neptune_port",
|
||||||
]:
|
]:
|
||||||
try:
|
try:
|
||||||
self._config[config_key] = int(value)
|
self._config[config_key] = int(value)
|
||||||
@@ -142,7 +157,11 @@ class GraphStoreConfig:
|
|||||||
self.logger.warning(
|
self.logger.warning(
|
||||||
f"Invalid integer value for {env_var}: {value}"
|
f"Invalid integer value for {env_var}: {value}"
|
||||||
)
|
)
|
||||||
elif config_key in ["neo4j_encrypted"]:
|
elif config_key in [
|
||||||
|
"neo4j_encrypted",
|
||||||
|
"neptune_iam_auth",
|
||||||
|
"neptune_use_ssl",
|
||||||
|
]:
|
||||||
self._config[config_key] = value.lower() in [
|
self._config[config_key] = value.lower() in [
|
||||||
"true",
|
"true",
|
||||||
"1",
|
"1",
|
||||||
@@ -171,6 +190,15 @@ class GraphStoreConfig:
|
|||||||
"falkordb_port": 6379,
|
"falkordb_port": 6379,
|
||||||
"falkordb_password": None,
|
"falkordb_password": None,
|
||||||
"falkordb_graph_name": "default",
|
"falkordb_graph_name": "default",
|
||||||
|
# Amazon Neptune defaults
|
||||||
|
"neptune_endpoint": None,
|
||||||
|
"neptune_port": 8182,
|
||||||
|
"neptune_region": None,
|
||||||
|
"neptune_iam_auth": True,
|
||||||
|
"neptune_use_ssl": True,
|
||||||
|
"neptune_access_key": None,
|
||||||
|
"neptune_secret_key": None,
|
||||||
|
"neptune_session_token": None,
|
||||||
}
|
}
|
||||||
|
|
||||||
for key, default_value in defaults.items():
|
for key, default_value in defaults.items():
|
||||||
@@ -269,6 +297,24 @@ class GraphStoreConfig:
|
|||||||
"graph_name": self._config.get("falkordb_graph_name"),
|
"graph_name": self._config.get("falkordb_graph_name"),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
def get_neptune_config(self) -> Dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Get Amazon Neptune-specific configuration.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Neptune configuration dictionary
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
"endpoint": self._config.get("neptune_endpoint"),
|
||||||
|
"port": self._config.get("neptune_port"),
|
||||||
|
"region": self._config.get("neptune_region"),
|
||||||
|
"iam_auth": self._config.get("neptune_iam_auth"),
|
||||||
|
"use_ssl": self._config.get("neptune_use_ssl"),
|
||||||
|
"access_key": self._config.get("neptune_access_key"),
|
||||||
|
"secret_key": self._config.get("neptune_secret_key"),
|
||||||
|
"session_token": self._config.get("neptune_session_token"),
|
||||||
|
}
|
||||||
|
|
||||||
def reset(self) -> None:
|
def reset(self) -> None:
|
||||||
"""Reset configuration to defaults."""
|
"""Reset configuration to defaults."""
|
||||||
self._config.clear()
|
self._config.clear()
|
||||||
@@ -278,4 +324,3 @@ class GraphStoreConfig:
|
|||||||
|
|
||||||
# Global configuration instance
|
# Global configuration instance
|
||||||
graph_store_config = GraphStoreConfig()
|
graph_store_config = GraphStoreConfig()
|
||||||
|
|
||||||
|
|||||||
@@ -34,7 +34,7 @@ License: MIT
|
|||||||
|
|
||||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||||
|
|
||||||
from ..utils.exceptions import ProcessingError, ValidationError
|
from ..utils.exceptions import ValidationError
|
||||||
from ..utils.logging import get_logger
|
from ..utils.logging import get_logger
|
||||||
from ..utils.progress_tracker import get_progress_tracker
|
from ..utils.progress_tracker import get_progress_tracker
|
||||||
from .config import graph_store_config
|
from .config import graph_store_config
|
||||||
@@ -214,7 +214,9 @@ class RelationshipManager:
|
|||||||
Returns:
|
Returns:
|
||||||
List of relationships
|
List of relationships
|
||||||
"""
|
"""
|
||||||
return self.backend.get_relationships(node_id, rel_type, direction, limit, **options)
|
return self.backend.get_relationships(
|
||||||
|
node_id, rel_type, direction, limit, **options
|
||||||
|
)
|
||||||
|
|
||||||
def delete(
|
def delete(
|
||||||
self,
|
self,
|
||||||
@@ -290,6 +292,7 @@ class QueryEngine:
|
|||||||
) -> str:
|
) -> str:
|
||||||
"""Generate cache key for query."""
|
"""Generate cache key for query."""
|
||||||
import hashlib
|
import hashlib
|
||||||
|
|
||||||
key_str = f"{query}:{str(parameters)}"
|
key_str = f"{query}:{str(parameters)}"
|
||||||
return hashlib.md5(key_str.encode()).hexdigest()
|
return hashlib.md5(key_str.encode()).hexdigest()
|
||||||
|
|
||||||
@@ -340,7 +343,9 @@ class GraphAnalytics:
|
|||||||
Returns:
|
Returns:
|
||||||
Path information or None
|
Path information or None
|
||||||
"""
|
"""
|
||||||
return self.backend.shortest_path(start_node_id, end_node_id, rel_type, max_depth, **options)
|
return self.backend.shortest_path(
|
||||||
|
start_node_id, end_node_id, rel_type, max_depth, **options
|
||||||
|
)
|
||||||
|
|
||||||
def get_neighbors(
|
def get_neighbors(
|
||||||
self,
|
self,
|
||||||
@@ -363,7 +368,9 @@ class GraphAnalytics:
|
|||||||
Returns:
|
Returns:
|
||||||
List of neighboring nodes
|
List of neighboring nodes
|
||||||
"""
|
"""
|
||||||
return self.backend.get_neighbors(node_id, rel_type, direction, depth, **options)
|
return self.backend.get_neighbors(
|
||||||
|
node_id, rel_type, direction, depth, **options
|
||||||
|
)
|
||||||
|
|
||||||
def degree_centrality(
|
def degree_centrality(
|
||||||
self,
|
self,
|
||||||
@@ -440,7 +447,7 @@ class GraphAnalytics:
|
|||||||
Component information
|
Component information
|
||||||
"""
|
"""
|
||||||
backend_type = type(self.backend).__name__
|
backend_type = type(self.backend).__name__
|
||||||
|
|
||||||
if "Neo4j" in backend_type:
|
if "Neo4j" in backend_type:
|
||||||
query = """
|
query = """
|
||||||
CALL gds.wcc.stream({
|
CALL gds.wcc.stream({
|
||||||
@@ -452,16 +459,23 @@ class GraphAnalytics:
|
|||||||
"""
|
"""
|
||||||
params = {"label": labels[0] if labels else "*"}
|
params = {"label": labels[0] if labels else "*"}
|
||||||
result = self.backend.execute_query(query, params)
|
result = self.backend.execute_query(query, params)
|
||||||
return [{"component": r["componentId"], "nodes": r["nodes"]} for r in result]
|
return [
|
||||||
|
{"component": r["componentId"], "nodes": r["nodes"]} for r in result
|
||||||
|
]
|
||||||
|
|
||||||
elif "NetworkX" in backend_type:
|
elif "NetworkX" in backend_type:
|
||||||
import networkx as nx
|
import networkx as nx
|
||||||
|
|
||||||
G = self.backend.graph
|
G = self.backend.graph
|
||||||
components = list(nx.connected_components(G))
|
components = list(nx.connected_components(G))
|
||||||
return [{"component": i, "nodes": list(c)} for i, c in enumerate(components)]
|
return [
|
||||||
|
{"component": i, "nodes": list(c)} for i, c in enumerate(components)
|
||||||
|
]
|
||||||
|
|
||||||
else:
|
else:
|
||||||
raise NotImplementedError(f"connected_components not implemented for {backend_type}")
|
raise NotImplementedError(
|
||||||
|
f"connected_components not implemented for {backend_type}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class GraphManager:
|
class GraphManager:
|
||||||
@@ -534,7 +548,11 @@ class GraphStore:
|
|||||||
self.progress_tracker.enabled = True
|
self.progress_tracker.enabled = True
|
||||||
|
|
||||||
# Determine backend
|
# Determine backend
|
||||||
self.backend = backend or config.get("backend") or graph_store_config.get("default_backend", "neo4j")
|
self.backend = (
|
||||||
|
backend
|
||||||
|
or config.get("backend")
|
||||||
|
or graph_store_config.get("default_backend", "neo4j")
|
||||||
|
)
|
||||||
self.config = config
|
self.config = config
|
||||||
|
|
||||||
# Initialize store backend
|
# Initialize store backend
|
||||||
@@ -546,16 +564,25 @@ class GraphStore:
|
|||||||
"""Initialize the appropriate store backend based on backend type."""
|
"""Initialize the appropriate store backend based on backend type."""
|
||||||
if self.backend == "neo4j":
|
if self.backend == "neo4j":
|
||||||
from .neo4j_store import Neo4jStore
|
from .neo4j_store import Neo4jStore
|
||||||
|
|
||||||
neo4j_config = graph_store_config.get_neo4j_config()
|
neo4j_config = graph_store_config.get_neo4j_config()
|
||||||
neo4j_config.update(self.config)
|
neo4j_config.update(self.config)
|
||||||
self._store_backend = Neo4jStore(**neo4j_config)
|
self._store_backend = Neo4jStore(**neo4j_config)
|
||||||
|
|
||||||
elif self.backend == "falkordb":
|
elif self.backend == "falkordb":
|
||||||
from .falkordb_store import FalkorDBStore
|
from .falkordb_store import FalkorDBStore
|
||||||
|
|
||||||
falkordb_config = graph_store_config.get_falkordb_config()
|
falkordb_config = graph_store_config.get_falkordb_config()
|
||||||
falkordb_config.update(self.config)
|
falkordb_config.update(self.config)
|
||||||
self._store_backend = FalkorDBStore(**falkordb_config)
|
self._store_backend = FalkorDBStore(**falkordb_config)
|
||||||
|
|
||||||
|
elif self.backend == "neptune" or self.backend == "amazon_neptune":
|
||||||
|
from .amazon_neptune import AmazonNeptuneStore
|
||||||
|
|
||||||
|
neptune_config = graph_store_config.get_neptune_config()
|
||||||
|
neptune_config.update(self.config)
|
||||||
|
self._store_backend = AmazonNeptuneStore(**neptune_config)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
raise ValidationError(f"Unknown backend: {self.backend}")
|
raise ValidationError(f"Unknown backend: {self.backend}")
|
||||||
|
|
||||||
@@ -621,7 +648,9 @@ class GraphStore:
|
|||||||
**options,
|
**options,
|
||||||
) -> List[Dict[str, Any]]:
|
) -> List[Dict[str, Any]]:
|
||||||
"""Get nodes matching criteria."""
|
"""Get nodes matching criteria."""
|
||||||
return self._manager.nodes.get(labels=labels, properties=properties, limit=limit, **options)
|
return self._manager.nodes.get(
|
||||||
|
labels=labels, properties=properties, limit=limit, **options
|
||||||
|
)
|
||||||
|
|
||||||
def update_node(
|
def update_node(
|
||||||
self,
|
self,
|
||||||
@@ -665,7 +694,9 @@ class GraphStore:
|
|||||||
**options,
|
**options,
|
||||||
) -> List[Dict[str, Any]]:
|
) -> List[Dict[str, Any]]:
|
||||||
"""Get relationships."""
|
"""Get relationships."""
|
||||||
return self._manager.relationships.get(node_id, rel_type, direction, limit, **options)
|
return self._manager.relationships.get(
|
||||||
|
node_id, rel_type, direction, limit, **options
|
||||||
|
)
|
||||||
|
|
||||||
def delete_relationship(
|
def delete_relationship(
|
||||||
self,
|
self,
|
||||||
@@ -723,7 +754,9 @@ class GraphStore:
|
|||||||
node_id, rel_type, direction, actual_depth, **options
|
node_id, rel_type, direction, actual_depth, **options
|
||||||
)
|
)
|
||||||
|
|
||||||
def query(self, query: str, parameters: Optional[Dict[str, Any]] = None, **options) -> List[Dict[str, Any]]:
|
def query(
|
||||||
|
self, query: str, parameters: Optional[Dict[str, Any]] = None, **options
|
||||||
|
) -> List[Dict[str, Any]]:
|
||||||
"""
|
"""
|
||||||
Execute a query and return results (Compatibility method for ContextRetriever).
|
Execute a query and return results (Compatibility method for ContextRetriever).
|
||||||
|
|
||||||
@@ -771,42 +804,46 @@ class GraphStore:
|
|||||||
# Convert to GraphStore format (labels, properties)
|
# Convert to GraphStore format (labels, properties)
|
||||||
graph_nodes = []
|
graph_nodes = []
|
||||||
for node in nodes:
|
for node in nodes:
|
||||||
# Extract label from type
|
# Extract labels - support both 'labels' array and 'type' string
|
||||||
labels = [node.get("type", "Entity")]
|
labels = node.get("labels")
|
||||||
if isinstance(labels[0], str):
|
if not labels:
|
||||||
labels = [labels[0]] # Ensure list
|
node_type = node.get("type", "Entity")
|
||||||
|
labels = [node_type] if isinstance(node_type, str) else node_type
|
||||||
|
if isinstance(labels, str):
|
||||||
|
labels = [labels]
|
||||||
|
|
||||||
# Prepare properties
|
# Prepare properties
|
||||||
props = node.get("properties", {}).copy()
|
props = node.get("properties", {}).copy()
|
||||||
|
|
||||||
# Ensure ID is preserved
|
# Ensure ID is preserved
|
||||||
if "id" in node and "id" not in props:
|
if "id" in node and "id" not in props:
|
||||||
props["id"] = node["id"]
|
props["id"] = node["id"]
|
||||||
|
|
||||||
# Ensure content/text is preserved
|
# Ensure content/text is preserved
|
||||||
if "content" in node and "content" not in props:
|
if "content" in node and "content" not in props:
|
||||||
props["content"] = node["content"]
|
props["content"] = node["content"]
|
||||||
if "text" in node and "text" not in props:
|
if "text" in node and "text" not in props:
|
||||||
props["text"] = node["text"]
|
props["text"] = node["text"]
|
||||||
|
|
||||||
graph_nodes.append({
|
graph_nodes.append({"labels": labels, "properties": props})
|
||||||
"labels": labels,
|
|
||||||
"properties": props
|
|
||||||
})
|
|
||||||
|
|
||||||
# Use batch creation
|
# Use batch creation
|
||||||
# Note: create_nodes expects dicts with 'labels' and 'properties' keys if passed directly?
|
# Note: create_nodes expects dicts with 'labels' and 'properties'
|
||||||
|
# keys if passed directly?
|
||||||
# Let's check create_nodes signature implementation in manager.
|
# Let's check create_nodes signature implementation in manager.
|
||||||
# But here I'll assume create_nodes takes a list of such dicts or similar.
|
# But here I'll assume create_nodes takes a list of such dicts
|
||||||
|
# or similar.
|
||||||
# Actually, let's look at create_nodes wrapper in this file:
|
# Actually, let's look at create_nodes wrapper in this file:
|
||||||
# def create_nodes(self, nodes: List[Dict[str, Any]], **options)
|
# def create_nodes(self, nodes: List[Dict[str, Any]], **options)
|
||||||
# It passes to self._manager.nodes.create_batch(nodes)
|
# It passes to self._manager.nodes.create_batch(nodes)
|
||||||
|
|
||||||
# If create_batch expects specific format, I should match it.
|
# If create_batch expects specific format, I should match it.
|
||||||
# Assuming create_batch is smart enough or expects standard format.
|
# Assuming create_batch is smart enough or expects standard format.
|
||||||
# To be safe, let's look at NodeManager.create_batch if possible, but I can't easily.
|
# To be safe, let's look at NodeManager.create_batch if possible,
|
||||||
# Standard expectation: List of dicts where each dict has labels and properties.
|
# but I can't easily.
|
||||||
|
# Standard expectation: List of dicts where each dict has labels
|
||||||
|
# and properties.
|
||||||
|
|
||||||
result = self.create_nodes(graph_nodes, **options)
|
result = self.create_nodes(graph_nodes, **options)
|
||||||
return len(result)
|
return len(result)
|
||||||
|
|
||||||
@@ -827,17 +864,21 @@ class GraphStore:
|
|||||||
target_id = edge.get("target_id")
|
target_id = edge.get("target_id")
|
||||||
rel_type = edge.get("type", "RELATED_TO")
|
rel_type = edge.get("type", "RELATED_TO")
|
||||||
properties = edge.get("properties", {}).copy()
|
properties = edge.get("properties", {}).copy()
|
||||||
|
|
||||||
# Preserve weight
|
# Preserve weight
|
||||||
if "weight" in edge:
|
if "weight" in edge:
|
||||||
properties["weight"] = edge["weight"]
|
properties["weight"] = edge["weight"]
|
||||||
|
|
||||||
if source_id and target_id:
|
if source_id and target_id:
|
||||||
try:
|
try:
|
||||||
self.create_relationship(source_id, target_id, rel_type, properties, **options)
|
self.create_relationship(
|
||||||
|
source_id, target_id, rel_type, properties, **options
|
||||||
|
)
|
||||||
count += 1
|
count += 1
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.logger.warning(f"Failed to add edge {source_id}->{target_id}: {e}")
|
self.logger.warning(
|
||||||
|
f"Failed to add edge {source_id}->{target_id}: {e}"
|
||||||
|
)
|
||||||
return count
|
return count
|
||||||
|
|
||||||
def build_from_conversations(
|
def build_from_conversations(
|
||||||
@@ -872,27 +913,29 @@ class GraphStore:
|
|||||||
all_nodes = []
|
all_nodes = []
|
||||||
all_edges = []
|
all_edges = []
|
||||||
seen_nodes = set()
|
seen_nodes = set()
|
||||||
|
|
||||||
for conv in conversations:
|
for conv in conversations:
|
||||||
# Load conversation if string (file path)
|
# Load conversation if string (file path)
|
||||||
conv_data = conv
|
conv_data = conv
|
||||||
if isinstance(conv, str):
|
if isinstance(conv, str):
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from ..utils.helpers import read_json_file
|
from ..utils.helpers import read_json_file
|
||||||
|
|
||||||
conv_data = read_json_file(Path(conv))
|
conv_data = read_json_file(Path(conv))
|
||||||
|
|
||||||
nodes, edges = self._process_conversation_to_elements(
|
nodes, edges = self._process_conversation_to_elements(
|
||||||
conv_data,
|
conv_data,
|
||||||
extract_intents=extract_intents,
|
extract_intents=extract_intents,
|
||||||
extract_sentiments=extract_sentiments
|
extract_sentiments=extract_sentiments,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Add unique nodes
|
# Add unique nodes
|
||||||
for node in nodes:
|
for node in nodes:
|
||||||
if node["id"] not in seen_nodes:
|
if node["id"] not in seen_nodes:
|
||||||
all_nodes.append(node)
|
all_nodes.append(node)
|
||||||
seen_nodes.add(node["id"])
|
seen_nodes.add(node["id"])
|
||||||
|
|
||||||
all_edges.extend(edges)
|
all_edges.extend(edges)
|
||||||
|
|
||||||
if link_entities:
|
if link_entities:
|
||||||
@@ -904,13 +947,8 @@ class GraphStore:
|
|||||||
edge_count = self.add_edges(all_edges)
|
edge_count = self.add_edges(all_edges)
|
||||||
|
|
||||||
self.progress_tracker.stop_tracking(tracking_id, status="completed")
|
self.progress_tracker.stop_tracking(tracking_id, status="completed")
|
||||||
|
|
||||||
return {
|
return {"statistics": {"node_count": node_count, "edge_count": edge_count}}
|
||||||
"statistics": {
|
|
||||||
"node_count": node_count,
|
|
||||||
"edge_count": edge_count
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.progress_tracker.stop_tracking(
|
self.progress_tracker.stop_tracking(
|
||||||
@@ -930,92 +968,112 @@ class GraphStore:
|
|||||||
"""
|
"""
|
||||||
nodes = []
|
nodes = []
|
||||||
edges = []
|
edges = []
|
||||||
|
|
||||||
# Process entities
|
# Process entities
|
||||||
for entity in entities:
|
for entity in entities:
|
||||||
entity_id = entity.get("id") or entity.get("entity_id")
|
entity_id = entity.get("id") or entity.get("entity_id")
|
||||||
if entity_id:
|
if entity_id:
|
||||||
nodes.append({
|
nodes.append(
|
||||||
"id": entity_id,
|
{
|
||||||
"type": entity.get("type", "entity"),
|
"id": entity_id,
|
||||||
"properties": {
|
"type": entity.get("type", "entity"),
|
||||||
"content": entity.get("text") or entity.get("label") or entity_id,
|
"properties": {
|
||||||
**entity
|
"content": entity.get("text")
|
||||||
|
or entity.get("label")
|
||||||
|
or entity_id,
|
||||||
|
**entity,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
})
|
)
|
||||||
|
|
||||||
# Process relationships
|
# Process relationships
|
||||||
for rel in relationships:
|
for rel in relationships:
|
||||||
source = rel.get("source_id")
|
source = rel.get("source_id")
|
||||||
target = rel.get("target_id")
|
target = rel.get("target_id")
|
||||||
if source and target:
|
if source and target:
|
||||||
edges.append({
|
edges.append(
|
||||||
"source_id": source,
|
{
|
||||||
"target_id": target,
|
"source_id": source,
|
||||||
"type": rel.get("type", "related_to"),
|
"target_id": target,
|
||||||
"weight": rel.get("confidence", 1.0),
|
"type": rel.get("type", "related_to"),
|
||||||
"properties": rel
|
"weight": rel.get("confidence", 1.0),
|
||||||
})
|
"properties": rel,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
node_count = self.add_nodes(nodes)
|
node_count = self.add_nodes(nodes)
|
||||||
edge_count = self.add_edges(edges)
|
edge_count = self.add_edges(edges)
|
||||||
|
|
||||||
return {"statistics": {"node_count": node_count, "edge_count": edge_count}}
|
return {"statistics": {"node_count": node_count, "edge_count": edge_count}}
|
||||||
|
|
||||||
def _process_conversation_to_elements(self, conv_data: Dict[str, Any], **kwargs) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
|
def _process_conversation_to_elements(
|
||||||
|
self, conv_data: Dict[str, Any], **kwargs
|
||||||
|
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
|
||||||
"""Helper to process conversation into nodes and edges."""
|
"""Helper to process conversation into nodes and edges."""
|
||||||
nodes = []
|
nodes = []
|
||||||
edges = []
|
edges = []
|
||||||
|
|
||||||
conv_id = conv_data.get("id") or f"conv_{hash(str(conv_data)) % 10000}"
|
conv_id = conv_data.get("id") or f"conv_{hash(str(conv_data)) % 10000}"
|
||||||
|
|
||||||
# Conversation node
|
# Conversation node
|
||||||
nodes.append({
|
nodes.append(
|
||||||
"id": conv_id,
|
{
|
||||||
"type": "conversation",
|
"id": conv_id,
|
||||||
"properties": {
|
"type": "conversation",
|
||||||
"content": conv_data.get("content", "") or conv_data.get("summary", ""),
|
"properties": {
|
||||||
"timestamp": conv_data.get("timestamp")
|
"content": conv_data.get("content", "")
|
||||||
|
or conv_data.get("summary", ""),
|
||||||
|
"timestamp": conv_data.get("timestamp"),
|
||||||
|
},
|
||||||
}
|
}
|
||||||
})
|
)
|
||||||
|
|
||||||
name_to_id = {}
|
name_to_id = {}
|
||||||
extract_entities = kwargs.get("extract_entities", True) # Default true if not passed?
|
# Note: extract_entities option is available but not used in this
|
||||||
# Actually ContextGraph defaults to True in init, but here we are static.
|
# implementation. Default true if not passed. ContextGraph defaults
|
||||||
|
# to True in init, but here we are static.
|
||||||
# Let's assume True unless told otherwise or check config.
|
# Let's assume True unless told otherwise or check config.
|
||||||
|
|
||||||
# Extract entities
|
# Extract entities
|
||||||
for entity in conv_data.get("entities", []):
|
for entity in conv_data.get("entities", []):
|
||||||
entity_id = entity.get("id") or entity.get("entity_id")
|
entity_id = entity.get("id") or entity.get("entity_id")
|
||||||
entity_text = entity.get("text") or entity.get("label") or entity.get("name") or entity_id
|
entity_text = (
|
||||||
|
entity.get("text")
|
||||||
|
or entity.get("label")
|
||||||
|
or entity.get("name")
|
||||||
|
or entity_id
|
||||||
|
)
|
||||||
entity_type = entity.get("type", "entity")
|
entity_type = entity.get("type", "entity")
|
||||||
|
|
||||||
# Generate ID if missing
|
# Generate ID if missing
|
||||||
if not entity_id and entity_text:
|
if not entity_id and entity_text:
|
||||||
import hashlib
|
import hashlib
|
||||||
entity_hash = hashlib.md5(f"{entity_text}_{entity_type}".encode()).hexdigest()[:12]
|
|
||||||
|
entity_hash = hashlib.md5(
|
||||||
|
f"{entity_text}_{entity_type}".encode()
|
||||||
|
).hexdigest()[:12]
|
||||||
entity_id = f"{entity_type.lower()}_{entity_hash}"
|
entity_id = f"{entity_type.lower()}_{entity_hash}"
|
||||||
|
|
||||||
if entity_id:
|
if entity_id:
|
||||||
if entity_text:
|
if entity_text:
|
||||||
name_to_id[entity_text] = entity_id
|
name_to_id[entity_text] = entity_id
|
||||||
|
|
||||||
nodes.append({
|
nodes.append(
|
||||||
"id": entity_id,
|
{
|
||||||
"type": "entity", # Normalize type?
|
"id": entity_id,
|
||||||
"properties": {
|
"type": "entity", # Normalize type?
|
||||||
"content": entity_text,
|
"properties": {
|
||||||
"type": entity_type,
|
"content": entity_text,
|
||||||
**entity
|
"type": entity_type,
|
||||||
|
**entity,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
})
|
)
|
||||||
|
|
||||||
# Edge: Conversation -> Entity
|
# Edge: Conversation -> Entity
|
||||||
edges.append({
|
edges.append(
|
||||||
"source_id": conv_id,
|
{"source_id": conv_id, "target_id": entity_id, "type": "mentions"}
|
||||||
"target_id": entity_id,
|
)
|
||||||
"type": "mentions"
|
|
||||||
})
|
|
||||||
|
|
||||||
# Extract relationships
|
# Extract relationships
|
||||||
for rel in conv_data.get("relationships", []):
|
for rel in conv_data.get("relationships", []):
|
||||||
@@ -1029,43 +1087,54 @@ class GraphStore:
|
|||||||
target = name_to_id[rel.get("target")]
|
target = name_to_id[rel.get("target")]
|
||||||
|
|
||||||
if source and target:
|
if source and target:
|
||||||
edges.append({
|
edges.append(
|
||||||
"source_id": source,
|
{
|
||||||
"target_id": target,
|
"source_id": source,
|
||||||
"type": rel.get("type", "related_to"),
|
"target_id": target,
|
||||||
"weight": rel.get("confidence", 1.0),
|
"type": rel.get("type", "related_to"),
|
||||||
"properties": rel
|
"weight": rel.get("confidence", 1.0),
|
||||||
})
|
"properties": rel,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
return nodes, edges
|
return nodes, edges
|
||||||
|
|
||||||
def _link_entities_elements(self, nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
def _link_entities_elements(
|
||||||
|
self, nodes: List[Dict[str, Any]]
|
||||||
|
) -> List[Dict[str, Any]]:
|
||||||
"""Link similar entities."""
|
"""Link similar entities."""
|
||||||
edges = []
|
edges = []
|
||||||
# Lazy import to avoid circular dependency
|
# Lazy import to avoid circular dependency
|
||||||
try:
|
try:
|
||||||
from ..context.entity_linker import EntityLinker
|
from ..context.entity_linker import EntityLinker
|
||||||
linker = EntityLinker() # Use default config
|
|
||||||
|
linker = EntityLinker() # Use default config
|
||||||
except (ImportError, OSError):
|
except (ImportError, OSError):
|
||||||
return []
|
return []
|
||||||
|
|
||||||
entity_nodes = [n for n in nodes if n.get("type") == "entity"]
|
entity_nodes = [n for n in nodes if n.get("type") == "entity"]
|
||||||
for i, node1 in enumerate(entity_nodes):
|
for i, node1 in enumerate(entity_nodes):
|
||||||
content1 = node1["properties"].get("content", "")
|
content1 = node1["properties"].get("content", "")
|
||||||
if not content1: continue
|
if not content1:
|
||||||
|
continue
|
||||||
|
|
||||||
for node2 in entity_nodes[i + 1 :]:
|
for node2 in entity_nodes[i + 1 :]:
|
||||||
content2 = node2["properties"].get("content", "")
|
content2 = node2["properties"].get("content", "")
|
||||||
if not content2: continue
|
if not content2:
|
||||||
|
continue
|
||||||
similarity = linker._calculate_text_similarity(content1.lower(), content2.lower())
|
|
||||||
|
similarity = linker._calculate_text_similarity(
|
||||||
|
content1.lower(), content2.lower()
|
||||||
|
)
|
||||||
if similarity >= linker.similarity_threshold:
|
if similarity >= linker.similarity_threshold:
|
||||||
edges.append({
|
edges.append(
|
||||||
"source_id": node1["id"],
|
{
|
||||||
"target_id": node2["id"],
|
"source_id": node1["id"],
|
||||||
"type": "similar_to",
|
"target_id": node2["id"],
|
||||||
"weight": similarity
|
"type": "similar_to",
|
||||||
})
|
"weight": similarity,
|
||||||
|
}
|
||||||
|
)
|
||||||
return edges
|
return edges
|
||||||
|
|
||||||
@property
|
@property
|
||||||
@@ -1087,4 +1156,3 @@ class GraphStore:
|
|||||||
def analytics(self) -> GraphAnalytics:
|
def analytics(self) -> GraphAnalytics:
|
||||||
"""Get analytics engine."""
|
"""Get analytics engine."""
|
||||||
return self._manager.analytics
|
return self._manager.analytics
|
||||||
|
|
||||||
|
|||||||
@@ -89,6 +89,7 @@ class GraphBuilder:
|
|||||||
self.track_history = track_history
|
self.track_history = track_history
|
||||||
self.version_snapshots = version_snapshots
|
self.version_snapshots = version_snapshots
|
||||||
self.graph_store = graph_store
|
self.graph_store = graph_store
|
||||||
|
self.config = kwargs # Store additional config for extractors
|
||||||
|
|
||||||
# Initialize logging
|
# Initialize logging
|
||||||
from ..utils.logging import get_logger
|
from ..utils.logging import get_logger
|
||||||
@@ -130,6 +131,11 @@ class GraphBuilder:
|
|||||||
|
|
||||||
def _process_item(self, item: Any, all_entities: List[Any], all_relationships: List[Any], **options):
|
def _process_item(self, item: Any, all_entities: List[Any], all_relationships: List[Any], **options):
|
||||||
"""Helper to process a single item and add to entities or relationships list."""
|
"""Helper to process a single item and add to entities or relationships list."""
|
||||||
|
if isinstance(item, str):
|
||||||
|
# Treat string as text for extraction
|
||||||
|
self._extract_from_text(item, all_entities, all_relationships, **options)
|
||||||
|
return
|
||||||
|
|
||||||
if hasattr(item, "text") and (hasattr(item, "label") or hasattr(item, "type")):
|
if hasattr(item, "text") and (hasattr(item, "label") or hasattr(item, "type")):
|
||||||
# It's likely an Entity object
|
# It's likely an Entity object
|
||||||
entity_dict = {
|
entity_dict = {
|
||||||
@@ -209,30 +215,68 @@ class GraphBuilder:
|
|||||||
# If still nothing found and has 'text', try extraction
|
# If still nothing found and has 'text', try extraction
|
||||||
if not found_something and "text" in item:
|
if not found_something and "text" in item:
|
||||||
text = item["text"]
|
text = item["text"]
|
||||||
# Perform extraction if requested or if it's the only way
|
self._extract_from_text(text, all_entities, all_relationships, **options)
|
||||||
if options.get("extract", True):
|
found_something = True
|
||||||
from ..semantic_extract.ner_extractor import NERExtractor
|
|
||||||
from ..semantic_extract.triplet_extractor import TripletExtractor
|
|
||||||
|
|
||||||
ner_method = options.get("ner_method", "ml")
|
|
||||||
triplet_method = options.get("triplet_method", "pattern")
|
|
||||||
|
|
||||||
ner = NERExtractor(method=ner_method)
|
|
||||||
entities = ner.extract_entities(text)
|
|
||||||
for ent in entities:
|
|
||||||
self._process_item(ent, all_entities, all_relationships, **options)
|
|
||||||
|
|
||||||
# Only try triplets if specifically requested or if method provided
|
|
||||||
if "triplet_method" in options or options.get("extract_relations", False):
|
|
||||||
triplet = TripletExtractor(method=triplet_method)
|
|
||||||
relations = triplet.extract_triplets(text)
|
|
||||||
for rel in relations:
|
|
||||||
self._process_item(rel, all_entities, all_relationships, **options)
|
|
||||||
found_something = True
|
|
||||||
else:
|
else:
|
||||||
# Unknown type
|
# Unknown type
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
def _extract_from_text(self, text: str, all_entities: List[Any], all_relationships: List[Any], **options):
|
||||||
|
"""Helper to extract knowledge from text using configured methods."""
|
||||||
|
if not options.get("extract", True):
|
||||||
|
return
|
||||||
|
|
||||||
|
from ..semantic_extract.ner_extractor import NERExtractor
|
||||||
|
from ..semantic_extract.relation_extractor import RelationExtractor
|
||||||
|
from ..semantic_extract.triplet_extractor import TripletExtractor
|
||||||
|
|
||||||
|
# Default to LLM methods as per requirement
|
||||||
|
ner_method = options.get("ner_method", "llm")
|
||||||
|
relation_method = options.get("relation_method", "llm")
|
||||||
|
triplet_method = options.get("triplet_method", "llm")
|
||||||
|
|
||||||
|
self.logger.info(f"Extracting knowledge from text ({len(text)} chars) using {ner_method}...")
|
||||||
|
|
||||||
|
# 1. Extract Entities
|
||||||
|
ner = NERExtractor(method=ner_method, **self.config)
|
||||||
|
try:
|
||||||
|
entities = ner.extract_entities(text, **options)
|
||||||
|
extracted_count = len(entities)
|
||||||
|
self._extraction_stats["extracted_entities"] += extracted_count
|
||||||
|
self.logger.info(f"Extracted {extracted_count} entities")
|
||||||
|
for ent in entities:
|
||||||
|
self._process_item(ent, all_entities, all_relationships, **options)
|
||||||
|
except Exception as e:
|
||||||
|
self.logger.error(f"Entity extraction failed: {e}")
|
||||||
|
entities = []
|
||||||
|
|
||||||
|
# 2. Extract Relations (if requested)
|
||||||
|
if options.get("extract_relations", True):
|
||||||
|
rel_extractor = RelationExtractor(method=relation_method, **self.config)
|
||||||
|
try:
|
||||||
|
# Pass entities if available to help relation extraction
|
||||||
|
relations = rel_extractor.extract_relations(text, entities=entities, **options)
|
||||||
|
extracted_count = len(relations)
|
||||||
|
self._extraction_stats["extracted_relations"] += extracted_count
|
||||||
|
self.logger.info(f"Extracted {extracted_count} relationships")
|
||||||
|
for rel in relations:
|
||||||
|
self._process_item(rel, all_entities, all_relationships, **options)
|
||||||
|
except Exception as e:
|
||||||
|
self.logger.error(f"Relation extraction failed: {e}")
|
||||||
|
|
||||||
|
# 3. Extract Triplets (if requested)
|
||||||
|
if options.get("extract_triplets", True):
|
||||||
|
trip_extractor = TripletExtractor(method=triplet_method, **self.config)
|
||||||
|
try:
|
||||||
|
triplets = trip_extractor.extract_triplets(text, entities=entities, **options)
|
||||||
|
extracted_count = len(triplets)
|
||||||
|
self._extraction_stats["extracted_triplets"] += extracted_count
|
||||||
|
self.logger.info(f"Extracted {extracted_count} triplets")
|
||||||
|
for trip in triplets:
|
||||||
|
self._process_item(trip, all_entities, all_relationships, **options)
|
||||||
|
except Exception as e:
|
||||||
|
self.logger.error(f"Triplet extraction failed: {e}")
|
||||||
|
|
||||||
def build(
|
def build(
|
||||||
self,
|
self,
|
||||||
sources: Union[List[Any], Any],
|
sources: Union[List[Any], Any],
|
||||||
@@ -305,6 +349,14 @@ class GraphBuilder:
|
|||||||
|
|
||||||
# Track graph building
|
# Track graph building
|
||||||
build_start_time = time.time()
|
build_start_time = time.time()
|
||||||
|
|
||||||
|
# Initialize extraction statistics for traceability
|
||||||
|
self._extraction_stats = {
|
||||||
|
"extracted_entities": 0,
|
||||||
|
"extracted_relations": 0,
|
||||||
|
"extracted_triplets": 0
|
||||||
|
}
|
||||||
|
|
||||||
tracking_id = self.progress_tracker.start_tracking(
|
tracking_id = self.progress_tracker.start_tracking(
|
||||||
module="kg",
|
module="kg",
|
||||||
submodule="GraphBuilder",
|
submodule="GraphBuilder",
|
||||||
@@ -514,7 +566,7 @@ class GraphBuilder:
|
|||||||
resolution_start = time.time()
|
resolution_start = time.time()
|
||||||
resolved_entities = resolver_to_use.resolve_entities(all_entities)
|
resolved_entities = resolver_to_use.resolve_entities(all_entities)
|
||||||
resolution_time = time.time() - resolution_start
|
resolution_time = time.time() - resolution_start
|
||||||
print(f"✅ Resolved to {len(resolved_entities)} unique entities ({resolution_time:.2f}s)")
|
print(f"[DONE] Resolved to {len(resolved_entities)} unique entities ({resolution_time:.2f}s)")
|
||||||
self.logger.info(
|
self.logger.info(
|
||||||
f"Entity resolution complete: {len(all_entities)} -> {len(resolved_entities)} unique entities"
|
f"Entity resolution complete: {len(all_entities)} -> {len(resolved_entities)} unique entities"
|
||||||
)
|
)
|
||||||
@@ -534,7 +586,7 @@ class GraphBuilder:
|
|||||||
},
|
},
|
||||||
}
|
}
|
||||||
structure_time = time.time() - structure_start
|
structure_time = time.time() - structure_start
|
||||||
print(f"✅ Graph structure built ({structure_time:.2f}s)")
|
print(f"[DONE] Graph structure built ({structure_time:.2f}s)")
|
||||||
|
|
||||||
# Persist to GraphStore if available
|
# Persist to GraphStore if available
|
||||||
if self.graph_store:
|
if self.graph_store:
|
||||||
@@ -567,7 +619,7 @@ class GraphBuilder:
|
|||||||
edge_time = time.time() - edge_start
|
edge_time = time.time() - edge_start
|
||||||
total_store_time = time.time() - store_start
|
total_store_time = time.time() - store_start
|
||||||
print(f" Added {edge_count} edges ({edge_time:.2f}s)")
|
print(f" Added {edge_count} edges ({edge_time:.2f}s)")
|
||||||
print(f"✅ GraphStore persistence complete ({total_store_time:.2f}s total)")
|
print(f"[DONE] GraphStore persistence complete ({total_store_time:.2f}s total)")
|
||||||
self.logger.info(f"Persisted {node_count} nodes and {edge_count} edges")
|
self.logger.info(f"Persisted {node_count} nodes and {edge_count} edges")
|
||||||
|
|
||||||
# Detect and resolve conflicts if conflict detector is available
|
# Detect and resolve conflicts if conflict detector is available
|
||||||
@@ -604,7 +656,14 @@ class GraphBuilder:
|
|||||||
|
|
||||||
# Print final summary with timing
|
# Print final summary with timing
|
||||||
print(f"\n{'='*60}")
|
print(f"\n{'='*60}")
|
||||||
print(f"✅ Knowledge Graph Build Complete")
|
print(f"[INFO] Extraction Statistics")
|
||||||
|
print(f" Extracted Entities: {self._extraction_stats['extracted_entities']}")
|
||||||
|
print(f" Extracted Relationships: {self._extraction_stats['extracted_relations']}")
|
||||||
|
print(f" Extracted Triplets: {self._extraction_stats['extracted_triplets']}")
|
||||||
|
print(f"{'='*60}")
|
||||||
|
|
||||||
|
print(f"\n{'='*60}")
|
||||||
|
print(f"[DONE] Knowledge Graph Build Complete")
|
||||||
print(f" Entities: {len(resolved_entities)}")
|
print(f" Entities: {len(resolved_entities)}")
|
||||||
print(f" Relationships: {len(all_relationships)}")
|
print(f" Relationships: {len(all_relationships)}")
|
||||||
print(f" Total time: {total_build_time:.2f}s")
|
print(f" Total time: {total_build_time:.2f}s")
|
||||||
|
|||||||
@@ -15,6 +15,8 @@ Key Features:
|
|||||||
- Semantic network construction
|
- Semantic network construction
|
||||||
- LLM-based extraction enhancement
|
- LLM-based extraction enhancement
|
||||||
- Extraction validation and quality assessment
|
- Extraction validation and quality assessment
|
||||||
|
- Batch processing with provenance tracking (batch_index, document_id)
|
||||||
|
- Robust fallback mechanisms (ML -> Pattern -> Last Resort)
|
||||||
|
|
||||||
Main Classes:
|
Main Classes:
|
||||||
- NamedEntityRecognizer: Main NER coordinator (confidence_threshold, merge_overlapping)
|
- NamedEntityRecognizer: Main NER coordinator (confidence_threshold, merge_overlapping)
|
||||||
|
|||||||
@@ -86,6 +86,7 @@ class CoreferenceChain:
|
|||||||
mentions: List[Mention]
|
mentions: List[Mention]
|
||||||
representative: Mention
|
representative: Mention
|
||||||
entity_type: Optional[str] = None
|
entity_type: Optional[str] = None
|
||||||
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
class CoreferenceResolver:
|
class CoreferenceResolver:
|
||||||
@@ -121,12 +122,18 @@ class CoreferenceResolver:
|
|||||||
)
|
)
|
||||||
self.chain_builder = CoreferenceChainBuilder(**self.config.get("chain", {}))
|
self.chain_builder = CoreferenceChainBuilder(**self.config.get("chain", {}))
|
||||||
|
|
||||||
def resolve_coreferences(self, text: str, **options) -> List[CoreferenceChain]:
|
def resolve_coreferences(
|
||||||
|
self,
|
||||||
|
text: str,
|
||||||
|
entities: Optional[List[Entity]] = None,
|
||||||
|
**options
|
||||||
|
) -> List[CoreferenceChain]:
|
||||||
"""
|
"""
|
||||||
Resolve coreferences in text.
|
Resolve coreferences in text.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
text: Input text
|
text: Input text
|
||||||
|
entities: List of entities (optional)
|
||||||
**options: Resolution options
|
**options: Resolution options
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
@@ -139,6 +146,8 @@ class CoreferenceResolver:
|
|||||||
)
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
from .ner_extractor import NERExtractor
|
||||||
|
|
||||||
total_steps = 4 # Extract mentions, resolve pronouns, detect coreferences, build chains
|
total_steps = 4 # Extract mentions, resolve pronouns, detect coreferences, build chains
|
||||||
current_step = 0
|
current_step = 0
|
||||||
|
|
||||||
@@ -151,8 +160,38 @@ class CoreferenceResolver:
|
|||||||
total=total_steps,
|
total=total_steps,
|
||||||
message=f"Extracting mentions... ({current_step}/{total_steps}, remaining: {remaining_steps} steps)"
|
message=f"Extracting mentions... ({current_step}/{total_steps}, remaining: {remaining_steps} steps)"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Extract pronouns
|
||||||
mentions = self._extract_mentions(text)
|
mentions = self._extract_mentions(text)
|
||||||
|
|
||||||
|
# Add entities as mentions
|
||||||
|
if entities is None:
|
||||||
|
# Extract entities if not provided
|
||||||
|
ner_config = self.config.get("ner", {})
|
||||||
|
if "ner_method" in self.config:
|
||||||
|
ner_config["method"] = self.config["ner_method"]
|
||||||
|
ner = NERExtractor(
|
||||||
|
**ner_config,
|
||||||
|
**{
|
||||||
|
k: v
|
||||||
|
for k, v in self.config.items()
|
||||||
|
if k not in ["ner", "relation", "chain", "entity", "pronoun"]
|
||||||
|
},
|
||||||
|
)
|
||||||
|
entities = ner.extract_entities(text, **options)
|
||||||
|
|
||||||
|
if entities:
|
||||||
|
for entity in entities:
|
||||||
|
mentions.append(
|
||||||
|
Mention(
|
||||||
|
text=entity.text,
|
||||||
|
start_char=entity.start_char,
|
||||||
|
end_char=entity.end_char,
|
||||||
|
mention_type="entity",
|
||||||
|
metadata={"entity_label": entity.label, "confidence": entity.confidence},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
# Step 2: Resolve pronouns
|
# Step 2: Resolve pronouns
|
||||||
current_step += 1
|
current_step += 1
|
||||||
remaining_steps = total_steps - current_step
|
remaining_steps = total_steps - current_step
|
||||||
@@ -203,18 +242,120 @@ class CoreferenceResolver:
|
|||||||
)
|
)
|
||||||
raise
|
raise
|
||||||
|
|
||||||
def resolve(self, text: str, **options) -> List[CoreferenceChain]:
|
def resolve(
|
||||||
|
self,
|
||||||
|
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||||
|
entities: Optional[Union[List[Entity], List[List[Entity]]]] = None,
|
||||||
|
pipeline_id: Optional[str] = None,
|
||||||
|
**kwargs
|
||||||
|
) -> Union[List[CoreferenceChain], List[List[CoreferenceChain]]]:
|
||||||
"""
|
"""
|
||||||
Resolve coreferences in text (alias for resolve_coreferences).
|
Resolve coreferences in text or list of documents.
|
||||||
|
Handles batch processing with progress tracking.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
text: Input text
|
text: Input text or list of documents
|
||||||
**options: Resolution options
|
entities: List of entities or list of list of entities (optional)
|
||||||
|
pipeline_id: Optional pipeline ID for progress tracking
|
||||||
|
**kwargs: Resolution options
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
list: List of coreference chains
|
Union[List[CoreferenceChain], List[List[CoreferenceChain]]]: Resolved coreference chains
|
||||||
"""
|
"""
|
||||||
return self.resolve_coreferences(text, **options)
|
if isinstance(text, list):
|
||||||
|
# Handle batch resolution with progress tracking
|
||||||
|
tracking_id = self.progress_tracker.start_tracking(
|
||||||
|
module="semantic_extract",
|
||||||
|
submodule="CoreferenceResolver",
|
||||||
|
message=f"Batch resolving coreferences from {len(text)} documents",
|
||||||
|
pipeline_id=pipeline_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
results = []
|
||||||
|
total_items = len(text)
|
||||||
|
total_chains_count = 0
|
||||||
|
|
||||||
|
# Determine update interval
|
||||||
|
if total_items <= 10:
|
||||||
|
update_interval = 1
|
||||||
|
else:
|
||||||
|
update_interval = max(1, min(10, total_items // 100))
|
||||||
|
|
||||||
|
# Initial progress update
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=0,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Starting batch resolution... 0/{total_items} (remaining: {total_items})"
|
||||||
|
)
|
||||||
|
|
||||||
|
for idx, item in enumerate(text):
|
||||||
|
# Prepare arguments for single item
|
||||||
|
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||||
|
|
||||||
|
doc_entities = None
|
||||||
|
if entities and idx < len(entities):
|
||||||
|
doc_entities = entities[idx]
|
||||||
|
|
||||||
|
# Resolve
|
||||||
|
chains = self.resolve_coreferences(doc_text, entities=doc_entities, **kwargs)
|
||||||
|
|
||||||
|
# Add provenance metadata
|
||||||
|
for chain in chains:
|
||||||
|
# Update chain metadata
|
||||||
|
if chain.metadata is None:
|
||||||
|
chain.metadata = {}
|
||||||
|
chain.metadata["batch_index"] = idx
|
||||||
|
if isinstance(item, dict) and "id" in item:
|
||||||
|
chain.metadata["document_id"] = item["id"]
|
||||||
|
|
||||||
|
# Update mentions metadata
|
||||||
|
for mention in chain.mentions:
|
||||||
|
if mention.metadata is None:
|
||||||
|
mention.metadata = {}
|
||||||
|
mention.metadata["batch_index"] = idx
|
||||||
|
if isinstance(item, dict) and "id" in item:
|
||||||
|
mention.metadata["document_id"] = item["id"]
|
||||||
|
|
||||||
|
# Update representative metadata
|
||||||
|
if chain.representative:
|
||||||
|
if chain.representative.metadata is None:
|
||||||
|
chain.representative.metadata = {}
|
||||||
|
chain.representative.metadata["batch_index"] = idx
|
||||||
|
if isinstance(item, dict) and "id" in item:
|
||||||
|
chain.representative.metadata["document_id"] = item["id"]
|
||||||
|
|
||||||
|
results.append(chains)
|
||||||
|
total_chains_count += len(chains)
|
||||||
|
|
||||||
|
# Update progress
|
||||||
|
if (idx + 1) % update_interval == 0 or (idx + 1) == total_items:
|
||||||
|
remaining = total_items - (idx + 1)
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=idx + 1,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Processing... {idx + 1}/{total_items} (remaining: {remaining}) - Resolved {total_chains_count} chains"
|
||||||
|
)
|
||||||
|
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id,
|
||||||
|
status="completed",
|
||||||
|
message=f"Batch resolution completed. Processed {len(results)} documents, resolved {total_chains_count} chains.",
|
||||||
|
)
|
||||||
|
return results
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id, status="failed", message=str(e)
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
else:
|
||||||
|
# Single item
|
||||||
|
return self.resolve_coreferences(text, entities=entities, **kwargs)
|
||||||
|
|
||||||
def _extract_mentions(self, text: str) -> List[Mention]:
|
def _extract_mentions(self, text: str) -> List[Mention]:
|
||||||
"""Extract all mentions from text."""
|
"""Extract all mentions from text."""
|
||||||
mentions = []
|
mentions = []
|
||||||
@@ -346,15 +487,49 @@ class PronounResolver:
|
|||||||
if m.mention_type == "entity" or m.mention_type == "nominal"
|
if m.mention_type == "entity" or m.mention_type == "nominal"
|
||||||
]
|
]
|
||||||
|
|
||||||
# Simple resolution: find closest preceding entity
|
# Simple resolution: find closest preceding entity with compatible type
|
||||||
|
pronoun_types = {
|
||||||
|
"he": ["PERSON"],
|
||||||
|
"him": ["PERSON"],
|
||||||
|
"his": ["PERSON"],
|
||||||
|
"she": ["PERSON"],
|
||||||
|
"her": ["PERSON"],
|
||||||
|
"it": ["ORG", "GPE", "LOC", "PRODUCT", "EVENT", "FAC", "WORK_OF_ART", "LAW", "LANGUAGE", "DATE", "TIME", "PERCENT", "MONEY", "QUANTITY", "ORDINAL", "CARDINAL"],
|
||||||
|
"its": ["ORG", "GPE", "LOC", "PRODUCT", "EVENT", "FAC", "WORK_OF_ART", "LAW", "LANGUAGE", "DATE", "TIME", "PERCENT", "MONEY", "QUANTITY", "ORDINAL", "CARDINAL"],
|
||||||
|
"they": ["ORG", "GPE", "PERSON", "NORP"], # Can be groups of people or organizations
|
||||||
|
"them": ["ORG", "GPE", "PERSON", "NORP"],
|
||||||
|
"their": ["ORG", "GPE", "PERSON", "NORP"],
|
||||||
|
}
|
||||||
|
|
||||||
for pronoun in pronouns:
|
for pronoun in pronouns:
|
||||||
# Find preceding entities
|
# Find preceding entities
|
||||||
preceding = [e for e in entities if e.end_char < pronoun.start_char]
|
preceding = [e for e in entities if e.end_char < pronoun.start_char]
|
||||||
|
|
||||||
if preceding:
|
if preceding:
|
||||||
# Take closest
|
pronoun_lower = pronoun.text.lower()
|
||||||
antecedent = max(preceding, key=lambda e: e.end_char)
|
compatible_types = pronoun_types.get(pronoun_lower)
|
||||||
|
|
||||||
|
antecedent = None
|
||||||
|
|
||||||
|
if compatible_types:
|
||||||
|
# Filter by type
|
||||||
|
compatible = [
|
||||||
|
e for e in preceding
|
||||||
|
if e.metadata and e.metadata.get("entity_label") in compatible_types
|
||||||
|
]
|
||||||
|
if compatible:
|
||||||
|
# Take closest compatible
|
||||||
|
antecedent = max(compatible, key=lambda e: e.end_char)
|
||||||
|
|
||||||
|
# Fallback to closest if no compatible found or pronoun type unknown
|
||||||
|
if antecedent is None:
|
||||||
|
antecedent = max(preceding, key=lambda e: e.end_char)
|
||||||
|
|
||||||
resolutions.append((pronoun.text, antecedent.text))
|
resolutions.append((pronoun.text, antecedent.text))
|
||||||
|
|
||||||
|
# Update pronoun metadata and link to antecedent
|
||||||
|
pronoun.entity_id = antecedent.text
|
||||||
|
pronoun.metadata["antecedent_text"] = antecedent.text
|
||||||
|
|
||||||
return resolutions
|
return resolutions
|
||||||
|
|
||||||
@@ -431,32 +606,59 @@ class CoreferenceChainBuilder:
|
|||||||
list: List of coreference chains
|
list: List of coreference chains
|
||||||
"""
|
"""
|
||||||
chains = []
|
chains = []
|
||||||
|
processed_indices = set()
|
||||||
|
|
||||||
# Simple implementation: group by text similarity
|
for i, mention in enumerate(mentions):
|
||||||
processed = set()
|
if i in processed_indices:
|
||||||
|
|
||||||
for mention in mentions:
|
|
||||||
if mention.text.lower() in processed:
|
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Find similar mentions
|
# Start a new group
|
||||||
similar = [
|
group = [mention]
|
||||||
m
|
processed_indices.add(i)
|
||||||
for m in mentions
|
|
||||||
if m.text.lower() == mention.text.lower()
|
|
||||||
or self._similar_mentions(mention.text, m.text)
|
|
||||||
]
|
|
||||||
|
|
||||||
if len(similar) > 1:
|
# Find related mentions
|
||||||
processed.add(mention.text.lower())
|
for j, other in enumerate(mentions):
|
||||||
|
if j in processed_indices:
|
||||||
|
continue
|
||||||
|
|
||||||
# Representative is first (leftmost) mention
|
is_related = False
|
||||||
representative = min(similar, key=lambda m: m.start_char)
|
|
||||||
|
# 1. Text similarity
|
||||||
|
if (
|
||||||
|
other.text.lower() == mention.text.lower()
|
||||||
|
or self._similar_mentions(mention.text, other.text)
|
||||||
|
):
|
||||||
|
is_related = True
|
||||||
|
|
||||||
|
# 2. Pronoun resolution (entity_id matches text or entity_id matches entity_id)
|
||||||
|
elif mention.entity_id and (
|
||||||
|
mention.entity_id == other.text
|
||||||
|
or mention.entity_id == other.entity_id
|
||||||
|
):
|
||||||
|
is_related = True
|
||||||
|
elif other.entity_id and (
|
||||||
|
other.entity_id == mention.text
|
||||||
|
or other.entity_id == mention.entity_id
|
||||||
|
):
|
||||||
|
is_related = True
|
||||||
|
|
||||||
|
if is_related:
|
||||||
|
group.append(other)
|
||||||
|
processed_indices.add(j)
|
||||||
|
|
||||||
|
if len(group) > 1:
|
||||||
|
# Representative is first (leftmost) mention, or prefer entity over pronoun
|
||||||
|
# Prefer entity mention as representative
|
||||||
|
entities = [m for m in group if m.mention_type != "pronoun"]
|
||||||
|
if entities:
|
||||||
|
representative = min(entities, key=lambda m: m.start_char)
|
||||||
|
else:
|
||||||
|
representative = min(group, key=lambda m: m.start_char)
|
||||||
|
|
||||||
chain = CoreferenceChain(
|
chain = CoreferenceChain(
|
||||||
mentions=similar,
|
mentions=group,
|
||||||
representative=representative,
|
representative=representative,
|
||||||
entity_type=similar[0].metadata.get("entity_label"),
|
entity_type=representative.metadata.get("entity_label"),
|
||||||
)
|
)
|
||||||
chains.append(chain)
|
chains.append(chain)
|
||||||
|
|
||||||
|
|||||||
@@ -147,6 +147,97 @@ class EventDetector:
|
|||||||
"meeting": r"met|meeting|conference|summit",
|
"meeting": r"met|meeting|conference|summit",
|
||||||
}
|
}
|
||||||
|
|
||||||
|
def extract(
|
||||||
|
self,
|
||||||
|
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||||
|
pipeline_id: Optional[str] = None,
|
||||||
|
**kwargs
|
||||||
|
) -> Union[List[Event], List[List[Event]]]:
|
||||||
|
"""
|
||||||
|
Detect events in text or list of documents.
|
||||||
|
Handles batch processing with progress tracking.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: Input text or list of documents
|
||||||
|
pipeline_id: Optional pipeline ID for progress tracking
|
||||||
|
**kwargs: Detection options
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Union[List[Event], List[List[Event]]]: Detected events
|
||||||
|
"""
|
||||||
|
if isinstance(text, list):
|
||||||
|
# Handle batch detection with progress tracking
|
||||||
|
tracking_id = self.progress_tracker.start_tracking(
|
||||||
|
module="semantic_extract",
|
||||||
|
submodule="EventDetector",
|
||||||
|
message=f"Batch detecting events from {len(text)} documents",
|
||||||
|
pipeline_id=pipeline_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
results = []
|
||||||
|
total_items = len(text)
|
||||||
|
total_events_count = 0
|
||||||
|
|
||||||
|
# Determine update interval
|
||||||
|
if total_items <= 10:
|
||||||
|
update_interval = 1
|
||||||
|
else:
|
||||||
|
update_interval = max(1, min(10, total_items // 100))
|
||||||
|
|
||||||
|
# Initial progress update
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=0,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Starting batch detection... 0/{total_items} (remaining: {total_items})"
|
||||||
|
)
|
||||||
|
|
||||||
|
for idx, item in enumerate(text):
|
||||||
|
# Prepare arguments for single item
|
||||||
|
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||||
|
|
||||||
|
# Detect
|
||||||
|
events = self.detect_events(doc_text, **kwargs)
|
||||||
|
|
||||||
|
# Add provenance metadata
|
||||||
|
for event in events:
|
||||||
|
if event.metadata is None:
|
||||||
|
event.metadata = {}
|
||||||
|
event.metadata["batch_index"] = idx
|
||||||
|
if isinstance(item, dict) and "id" in item:
|
||||||
|
event.metadata["document_id"] = item["id"]
|
||||||
|
|
||||||
|
results.append(events)
|
||||||
|
total_events_count += len(events)
|
||||||
|
|
||||||
|
# Update progress
|
||||||
|
if (idx + 1) % update_interval == 0 or (idx + 1) == total_items:
|
||||||
|
remaining = total_items - (idx + 1)
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=idx + 1,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Processing... {idx + 1}/{total_items} (remaining: {remaining}) - Detected {total_events_count} events"
|
||||||
|
)
|
||||||
|
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id,
|
||||||
|
status="completed",
|
||||||
|
message=f"Batch detection completed. Processed {len(results)} documents, detected {total_events_count} events.",
|
||||||
|
)
|
||||||
|
return results
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id, status="failed", message=str(e)
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
else:
|
||||||
|
# Single item
|
||||||
|
return self.detect_events(text, **kwargs)
|
||||||
|
|
||||||
def detect_events(self, text: str, **options) -> List[Event]:
|
def detect_events(self, text: str, **options) -> List[Event]:
|
||||||
"""
|
"""
|
||||||
Detect events in text content.
|
Detect events in text content.
|
||||||
|
|||||||
@@ -12,8 +12,8 @@ Supported Methods (for future extensibility):
|
|||||||
|
|
||||||
Algorithms Used:
|
Algorithms Used:
|
||||||
- Confidence Thresholding: Statistical threshold-based filtering
|
- Confidence Thresholding: Statistical threshold-based filtering
|
||||||
- Duplicate Detection: Set-based and similarity-based deduplication
|
- Duplicate Detection: (Removed - handled by external module)
|
||||||
- Consistency Checking: Rule-based and graph-based consistency validation
|
- Consistency Checking: (Removed - handled by external module)
|
||||||
- Quality Scoring: Weighted scoring algorithms for extraction quality
|
- Quality Scoring: Weighted scoring algorithms for extraction quality
|
||||||
- Validation Metrics: Precision, recall, F1-score calculations
|
- Validation Metrics: Precision, recall, F1-score calculations
|
||||||
- Boundary Validation: Character position and text boundary checking
|
- Boundary Validation: Character position and text boundary checking
|
||||||
@@ -22,7 +22,6 @@ Key Features:
|
|||||||
- Entity validation with confidence checking
|
- Entity validation with confidence checking
|
||||||
- Relation validation and consistency checking
|
- Relation validation and consistency checking
|
||||||
- Quality scoring and metrics calculation
|
- Quality scoring and metrics calculation
|
||||||
- Duplicate detection
|
|
||||||
- Confidence-based filtering
|
- Confidence-based filtering
|
||||||
- Validation result reporting
|
- Validation result reporting
|
||||||
- Method parameter support for future method-specific validation
|
- Method parameter support for future method-specific validation
|
||||||
@@ -47,8 +46,10 @@ Author: Semantica Contributors
|
|||||||
License: MIT
|
License: MIT
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
from typing import List, Dict, Any, Optional, Set, Tuple, Union
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Any, Dict, List, Optional
|
from datetime import datetime
|
||||||
|
import re
|
||||||
|
|
||||||
from ..utils.exceptions import ProcessingError
|
from ..utils.exceptions import ProcessingError
|
||||||
from ..utils.logging import get_logger
|
from ..utils.logging import get_logger
|
||||||
@@ -66,6 +67,7 @@ class ValidationResult:
|
|||||||
errors: List[str] = field(default_factory=list)
|
errors: List[str] = field(default_factory=list)
|
||||||
warnings: List[str] = field(default_factory=list)
|
warnings: List[str] = field(default_factory=list)
|
||||||
metrics: Dict[str, Any] = field(default_factory=dict)
|
metrics: Dict[str, Any] = field(default_factory=dict)
|
||||||
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
class ExtractionValidator:
|
class ExtractionValidator:
|
||||||
@@ -79,7 +81,6 @@ class ExtractionValidator:
|
|||||||
method: Validation method (for future extensibility, currently unused)
|
method: Validation method (for future extensibility, currently unused)
|
||||||
**config: Configuration options:
|
**config: Configuration options:
|
||||||
- min_confidence: Minimum confidence threshold (default: 0.5)
|
- min_confidence: Minimum confidence threshold (default: 0.5)
|
||||||
- validate_consistency: Check consistency (default: True)
|
|
||||||
"""
|
"""
|
||||||
self.logger = get_logger("extraction_validator")
|
self.logger = get_logger("extraction_validator")
|
||||||
self.config = config
|
self.config = config
|
||||||
@@ -90,19 +91,30 @@ class ExtractionValidator:
|
|||||||
|
|
||||||
self.method = method # Reserved for future method-based validation
|
self.method = method # Reserved for future method-based validation
|
||||||
self.min_confidence = config.get("min_confidence", 0.5)
|
self.min_confidence = config.get("min_confidence", 0.5)
|
||||||
self.validate_consistency = config.get("validate_consistency", True)
|
|
||||||
|
|
||||||
def validate_entities(self, entities: List[Entity], **options) -> ValidationResult:
|
def validate_entities(self, entities: Union[List[Entity], List[List[Entity]]], **options) -> Union[ValidationResult, List[ValidationResult]]:
|
||||||
"""
|
"""
|
||||||
Validate extracted entities.
|
Validate extracted entities.
|
||||||
|
Handles both single list and batch list of entities.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
entities: List of entities
|
entities: List of entities or list of list of entities
|
||||||
**options: Validation options
|
**options: Validation options
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
ValidationResult: Validation result
|
ValidationResult or List[ValidationResult]: Validation result(s)
|
||||||
"""
|
"""
|
||||||
|
# Handle batch validation
|
||||||
|
if entities and isinstance(entities, list) and len(entities) > 0 and isinstance(entities[0], list):
|
||||||
|
results = []
|
||||||
|
for idx, batch_entities in enumerate(entities):
|
||||||
|
res = self.validate_entities(batch_entities, **options)
|
||||||
|
# Ensure metadata has batch index
|
||||||
|
if "batch_index" not in res.metadata:
|
||||||
|
res.metadata["batch_index"] = idx
|
||||||
|
results.append(res)
|
||||||
|
return results
|
||||||
|
|
||||||
tracking_id = self.progress_tracker.start_tracking(
|
tracking_id = self.progress_tracker.start_tracking(
|
||||||
module="semantic_extract",
|
module="semantic_extract",
|
||||||
submodule="ExtractionValidator",
|
submodule="ExtractionValidator",
|
||||||
@@ -126,14 +138,6 @@ class ExtractionValidator:
|
|||||||
f"{len(low_confidence)} entities below confidence threshold"
|
f"{len(low_confidence)} entities below confidence threshold"
|
||||||
)
|
)
|
||||||
|
|
||||||
# Check for duplicates
|
|
||||||
self.progress_tracker.update_tracking(
|
|
||||||
tracking_id, message="Checking for duplicates..."
|
|
||||||
)
|
|
||||||
entity_texts = [e.text.lower() for e in entities]
|
|
||||||
duplicates = len(entity_texts) - len(set(entity_texts))
|
|
||||||
if duplicates > 0:
|
|
||||||
warnings.append(f"{duplicates} duplicate entities found")
|
|
||||||
|
|
||||||
# Check for empty entities
|
# Check for empty entities
|
||||||
empty_entities = [e for e in entities if not e.text.strip()]
|
empty_entities = [e for e in entities if not e.text.strip()]
|
||||||
@@ -148,8 +152,7 @@ class ExtractionValidator:
|
|||||||
[e for e in entities if min_confidence <= e.confidence < 0.8]
|
[e for e in entities if min_confidence <= e.confidence < 0.8]
|
||||||
),
|
),
|
||||||
"low_confidence": len(low_confidence),
|
"low_confidence": len(low_confidence),
|
||||||
"unique_entities": len(set(entity_texts)),
|
"unique_entities": len(set(e.text for e in entities)),
|
||||||
"duplicates": duplicates,
|
|
||||||
"entity_types": len(set(e.label for e in entities)),
|
"entity_types": len(set(e.label for e in entities)),
|
||||||
"average_confidence": sum(e.confidence for e in entities)
|
"average_confidence": sum(e.confidence for e in entities)
|
||||||
/ len(entities)
|
/ len(entities)
|
||||||
@@ -162,12 +165,23 @@ class ExtractionValidator:
|
|||||||
|
|
||||||
valid = len(errors) == 0
|
valid = len(errors) == 0
|
||||||
|
|
||||||
|
# Collect metadata from entities
|
||||||
|
metadata = {}
|
||||||
|
if entities:
|
||||||
|
first = entities[0]
|
||||||
|
if hasattr(first, "metadata") and first.metadata:
|
||||||
|
if "batch_index" in first.metadata:
|
||||||
|
metadata["batch_index"] = first.metadata["batch_index"]
|
||||||
|
if "document_id" in first.metadata:
|
||||||
|
metadata["document_id"] = first.metadata["document_id"]
|
||||||
|
|
||||||
result = ValidationResult(
|
result = ValidationResult(
|
||||||
valid=valid,
|
valid=valid,
|
||||||
score=score,
|
score=score,
|
||||||
errors=errors,
|
errors=errors,
|
||||||
warnings=warnings,
|
warnings=warnings,
|
||||||
metrics=metrics,
|
metrics=metrics,
|
||||||
|
metadata=metadata,
|
||||||
)
|
)
|
||||||
|
|
||||||
self.progress_tracker.stop_tracking(
|
self.progress_tracker.stop_tracking(
|
||||||
@@ -184,18 +198,30 @@ class ExtractionValidator:
|
|||||||
raise
|
raise
|
||||||
|
|
||||||
def validate_relations(
|
def validate_relations(
|
||||||
self, relations: List[Relation], **options
|
self, relations: Union[List[Relation], List[List[Relation]]], **options
|
||||||
) -> ValidationResult:
|
) -> Union[ValidationResult, List[ValidationResult]]:
|
||||||
"""
|
"""
|
||||||
Validate extracted relations.
|
Validate extracted relations.
|
||||||
|
Handles both single list and batch list of relations.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
relations: List of relations
|
relations: List of relations or list of list of relations
|
||||||
**options: Validation options
|
**options: Validation options
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
ValidationResult: Validation result
|
ValidationResult or List[ValidationResult]: Validation result(s)
|
||||||
"""
|
"""
|
||||||
|
# Handle batch validation
|
||||||
|
if relations and isinstance(relations, list) and len(relations) > 0 and isinstance(relations[0], list):
|
||||||
|
results = []
|
||||||
|
for idx, batch_relations in enumerate(relations):
|
||||||
|
res = self.validate_relations(batch_relations, **options)
|
||||||
|
# Ensure metadata has batch index
|
||||||
|
if "batch_index" not in res.metadata:
|
||||||
|
res.metadata["batch_index"] = idx
|
||||||
|
results.append(res)
|
||||||
|
return results
|
||||||
|
|
||||||
errors = []
|
errors = []
|
||||||
warnings = []
|
warnings = []
|
||||||
metrics = {}
|
metrics = {}
|
||||||
@@ -218,12 +244,6 @@ class ExtractionValidator:
|
|||||||
if invalid_relations:
|
if invalid_relations:
|
||||||
errors.append(f"{len(invalid_relations)} invalid relations found")
|
errors.append(f"{len(invalid_relations)} invalid relations found")
|
||||||
|
|
||||||
# Check consistency
|
|
||||||
if self.validate_consistency:
|
|
||||||
consistency_issues = self._check_consistency(relations)
|
|
||||||
if consistency_issues:
|
|
||||||
warnings.append(f"{len(consistency_issues)} consistency issues found")
|
|
||||||
|
|
||||||
# Calculate metrics
|
# Calculate metrics
|
||||||
metrics = {
|
metrics = {
|
||||||
"total_relations": len(relations),
|
"total_relations": len(relations),
|
||||||
@@ -244,31 +264,25 @@ class ExtractionValidator:
|
|||||||
|
|
||||||
valid = len(errors) == 0
|
valid = len(errors) == 0
|
||||||
|
|
||||||
|
# Collect metadata from relations
|
||||||
|
metadata = {}
|
||||||
|
if relations:
|
||||||
|
first = relations[0]
|
||||||
|
if hasattr(first, "metadata") and first.metadata:
|
||||||
|
if "batch_index" in first.metadata:
|
||||||
|
metadata["batch_index"] = first.metadata["batch_index"]
|
||||||
|
if "document_id" in first.metadata:
|
||||||
|
metadata["document_id"] = first.metadata["document_id"]
|
||||||
|
|
||||||
return ValidationResult(
|
return ValidationResult(
|
||||||
valid=valid, score=score, errors=errors, warnings=warnings, metrics=metrics
|
valid=valid,
|
||||||
|
score=score,
|
||||||
|
errors=errors,
|
||||||
|
warnings=warnings,
|
||||||
|
metrics=metrics,
|
||||||
|
metadata=metadata
|
||||||
)
|
)
|
||||||
|
|
||||||
def _check_consistency(self, relations: List[Relation]) -> List[str]:
|
|
||||||
"""Check consistency of relations."""
|
|
||||||
issues = []
|
|
||||||
|
|
||||||
# Check for contradictory relations
|
|
||||||
relation_pairs = {}
|
|
||||||
for relation in relations:
|
|
||||||
key = (relation.subject.text, relation.object.text)
|
|
||||||
if key not in relation_pairs:
|
|
||||||
relation_pairs[key] = []
|
|
||||||
relation_pairs[key].append(relation.predicate)
|
|
||||||
|
|
||||||
# Find contradictions (e.g., "founded_by" and "founded" for same pair)
|
|
||||||
for key, predicates in relation_pairs.items():
|
|
||||||
if len(set(predicates)) > 1:
|
|
||||||
# Check for obvious contradictions
|
|
||||||
if "founded_by" in predicates and "founded" in predicates:
|
|
||||||
issues.append(f"Contradictory relations for {key}")
|
|
||||||
|
|
||||||
return issues
|
|
||||||
|
|
||||||
def _calculate_entity_score(
|
def _calculate_entity_score(
|
||||||
self, entities: List[Entity], metrics: Dict[str, Any]
|
self, entities: List[Entity], metrics: Dict[str, Any]
|
||||||
) -> float:
|
) -> float:
|
||||||
|
|||||||
@@ -289,7 +289,14 @@ Return the enhanced relation list in JSON format."""
|
|||||||
) -> List[Entity]:
|
) -> List[Entity]:
|
||||||
"""Parse LLM response for entities."""
|
"""Parse LLM response for entities."""
|
||||||
# Simplified parsing - in practice would parse JSON
|
# Simplified parsing - in practice would parse JSON
|
||||||
# For now, return original entities
|
# For now, return original entities with updated metadata
|
||||||
|
for entity in original_entities:
|
||||||
|
if entity.metadata is None:
|
||||||
|
entity.metadata = {}
|
||||||
|
entity.metadata.update({
|
||||||
|
"enhanced_by": self.provider_name,
|
||||||
|
"model": self.model
|
||||||
|
})
|
||||||
return original_entities
|
return original_entities
|
||||||
|
|
||||||
def _parse_relation_response(
|
def _parse_relation_response(
|
||||||
@@ -297,7 +304,14 @@ Return the enhanced relation list in JSON format."""
|
|||||||
) -> List[Relation]:
|
) -> List[Relation]:
|
||||||
"""Parse LLM response for relations."""
|
"""Parse LLM response for relations."""
|
||||||
# Simplified parsing - in practice would parse JSON
|
# Simplified parsing - in practice would parse JSON
|
||||||
# For now, return original relations
|
# For now, return original relations with updated metadata
|
||||||
|
for relation in original_relations:
|
||||||
|
if relation.metadata is None:
|
||||||
|
relation.metadata = {}
|
||||||
|
relation.metadata.update({
|
||||||
|
"enhanced_by": self.provider_name,
|
||||||
|
"model": self.model
|
||||||
|
})
|
||||||
return original_relations
|
return original_relations
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -106,6 +106,7 @@ License: MIT
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import re
|
import re
|
||||||
|
import difflib
|
||||||
from typing import Any, Dict, List, Optional, Union
|
from typing import Any, Dict, List, Optional, Union
|
||||||
|
|
||||||
from ..utils.exceptions import ProcessingError
|
from ..utils.exceptions import ProcessingError
|
||||||
@@ -116,6 +117,12 @@ from .registry import method_registry
|
|||||||
from .relation_extractor import Relation
|
from .relation_extractor import Relation
|
||||||
from .triplet_extractor import Triplet
|
from .triplet_extractor import Triplet
|
||||||
|
|
||||||
|
try:
|
||||||
|
from .schemas import EntitiesResponse, RelationsResponse, TripletsResponse
|
||||||
|
SCHEMAS_AVAILABLE = True
|
||||||
|
except ImportError:
|
||||||
|
SCHEMAS_AVAILABLE = False
|
||||||
|
|
||||||
logger = get_logger("methods")
|
logger = get_logger("methods")
|
||||||
|
|
||||||
# Try to import spaCy
|
# Try to import spaCy
|
||||||
@@ -124,6 +131,275 @@ from ..utils.helpers import safe_import
|
|||||||
spacy, SPACY_AVAILABLE = safe_import("spacy")
|
spacy, SPACY_AVAILABLE = safe_import("spacy")
|
||||||
|
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# Scoring Helper Functions
|
||||||
|
# ============================================================================
|
||||||
|
|
||||||
|
# Global cache for spacy model and text embedder to avoid reloading
|
||||||
|
_nlp_cache = None
|
||||||
|
_embedder_cache = None
|
||||||
|
|
||||||
|
def get_text_embedder():
|
||||||
|
"""
|
||||||
|
Get or load the TextEmbedder model for high-accuracy semantic similarity.
|
||||||
|
"""
|
||||||
|
global _embedder_cache
|
||||||
|
if _embedder_cache:
|
||||||
|
return _embedder_cache
|
||||||
|
|
||||||
|
try:
|
||||||
|
from ..embeddings.text_embedder import TextEmbedder
|
||||||
|
# Use a lightweight but effective model for speed/accuracy balance
|
||||||
|
# BAAI/bge-small-en-v1.5 is excellent for semantic similarity
|
||||||
|
# Enable caching within the embedder if supported, or use our own
|
||||||
|
_embedder_cache = TextEmbedder(model_name="BAAI/bge-small-en-v1.5", normalize=True)
|
||||||
|
logger.info("Loaded TextEmbedder for high-accuracy similarity")
|
||||||
|
return _embedder_cache
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Failed to load TextEmbedder: {e}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
def get_nlp_model():
|
||||||
|
"""
|
||||||
|
Get or load a spaCy model for similarity calculations.
|
||||||
|
Prioritizes larger models for better vectors.
|
||||||
|
"""
|
||||||
|
global _nlp_cache
|
||||||
|
if _nlp_cache:
|
||||||
|
return _nlp_cache
|
||||||
|
|
||||||
|
if not SPACY_AVAILABLE:
|
||||||
|
return None
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Prefer larger models for vectors
|
||||||
|
# Note: 'en_core_web_lg' has true vectors. 'sm' only has context tensors.
|
||||||
|
for model_name in ["en_core_web_lg", "en_core_web_md", "en_core_web_sm"]:
|
||||||
|
if spacy.util.is_package(model_name):
|
||||||
|
try:
|
||||||
|
# Disable parser/ner for speed if we only need vectors
|
||||||
|
_nlp_cache = spacy.load(model_name, disable=["parser", "ner", "lemmatizer"])
|
||||||
|
logger.info(f"Loaded spaCy model for similarity: {model_name}")
|
||||||
|
return _nlp_cache
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Try loading generic if specific ones fail
|
||||||
|
try:
|
||||||
|
_nlp_cache = spacy.load("en_core_web_sm", disable=["parser", "ner", "lemmatizer"])
|
||||||
|
return _nlp_cache
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Failed to load spaCy model for similarity: {e}")
|
||||||
|
pass
|
||||||
|
return None
|
||||||
|
|
||||||
|
def calculate_similarity(text: str, candidates: List[str]) -> float:
|
||||||
|
"""
|
||||||
|
Calculate the maximum similarity between text and a list of candidates.
|
||||||
|
Uses a hybrid approach: Exact -> Substring -> Vectors -> Fuzzy.
|
||||||
|
"""
|
||||||
|
if not candidates:
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
if not text:
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
text_lower = text.lower().strip()
|
||||||
|
if not text_lower:
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
best_score = 0.0
|
||||||
|
|
||||||
|
# 1. Exact Match (Fastest)
|
||||||
|
candidates_lower = [c.lower().strip() for c in candidates if c]
|
||||||
|
if text_lower in candidates_lower:
|
||||||
|
return 1.0
|
||||||
|
|
||||||
|
# 1b. Common Synonyms (Fast Heuristic)
|
||||||
|
# Map common NER labels and Relations to user-friendly types
|
||||||
|
synonyms = {
|
||||||
|
# Entity Types
|
||||||
|
"person": ["people", "human", "name", "individual", "artist", "actor", "author", "politician"],
|
||||||
|
"org": ["company", "organization", "business", "institution", "agency", "brand", "corporation"],
|
||||||
|
"organization": ["company", "business", "institution", "agency", "brand", "corporation"],
|
||||||
|
"gpe": ["location", "place", "city", "country", "state", "nation", "region"],
|
||||||
|
"loc": ["location", "place", "region", "area"],
|
||||||
|
"date": ["time", "year", "day", "month", "period", "duration"],
|
||||||
|
"money": ["cost", "price", "value", "currency", "amount"],
|
||||||
|
"product": ["item", "object", "commodity", "goods", "device", "tool", "vehicle", "software", "app"],
|
||||||
|
"event": ["incident", "occasion", "activity", "happening", "ceremony"],
|
||||||
|
"drug": ["medication", "medicine", "pharmaceutical", "chemical", "treatment", "therapy"],
|
||||||
|
"chemical": ["drug", "substance", "compound", "element"],
|
||||||
|
"disease": ["condition", "illness", "sickness", "disorder", "syndrome", "ailment"],
|
||||||
|
|
||||||
|
# Relation Types
|
||||||
|
"founded_by": ["founder", "creator", "established_by", "started_by", "originator"],
|
||||||
|
"acquired": ["bought", "purchased", "acquisition", "takeover", "ownership", "merged_with"],
|
||||||
|
"subsidiary_of": ["owned_by", "parent_company", "part_of", "division_of", "unit_of"],
|
||||||
|
"works_for": ["employee_of", "employed_by", "staff_of", "team_member", "employs", "hired_by"],
|
||||||
|
"located_in": ["based_in", "headquartered_in", "situated_in", "found_in", "operates_in"],
|
||||||
|
"ceo_of": ["leader_of", "head_of", "director_of", "president_of", "chief_executive", "managed_by"],
|
||||||
|
"invested_in": ["funded", "financed", "backed", "shareholder_of", "venture_capital"],
|
||||||
|
"partner_with": ["collaborate_with", "joint_venture", "alliance", "deal_with", "partnership"],
|
||||||
|
"competitor_of": ["rival", "competes_with", "opponent", "nemesis"],
|
||||||
|
"manufacturer_of": ["producer_of", "maker_of", "creator_of", "builder_of"],
|
||||||
|
"treats": ["cures", "heals", "remedy_for", "used_for", "prescribed_for"],
|
||||||
|
"causes": ["leads_to", "results_in", "triggers", "produces", "creates"],
|
||||||
|
"diagnosed_with": ["suffers_from", "has_condition", "patient_of", "victim_of"],
|
||||||
|
}
|
||||||
|
|
||||||
|
if text_lower in synonyms:
|
||||||
|
for syn in synonyms[text_lower]:
|
||||||
|
if syn in candidates_lower:
|
||||||
|
return 0.95
|
||||||
|
# Also check reverse: if candidate is in synonyms of text
|
||||||
|
|
||||||
|
# Check if any candidate is a synonym of the text
|
||||||
|
for cand in candidates_lower:
|
||||||
|
if cand in synonyms:
|
||||||
|
if text_lower in synonyms[cand]:
|
||||||
|
return 0.95
|
||||||
|
|
||||||
|
# 2. Substring Match (Fast)
|
||||||
|
# Give a boost if one is contained in the other, but penalize by length difference
|
||||||
|
for cand in candidates_lower:
|
||||||
|
if text_lower == cand:
|
||||||
|
return 1.0
|
||||||
|
if text_lower in cand or cand in text_lower:
|
||||||
|
# Calculate length ratio
|
||||||
|
ratio = min(len(text_lower), len(cand)) / max(len(text_lower), len(cand))
|
||||||
|
# Base score 0.85 for containment, adjusted by ratio
|
||||||
|
# e.g. "Apple" in "Apple Inc" -> 0.85 * (5/9) ~= 0.47 (too low?)
|
||||||
|
# Let's be more generous for containment
|
||||||
|
score = 0.9 * ratio + 0.1 # Boost slightly
|
||||||
|
if score > best_score:
|
||||||
|
best_score = score
|
||||||
|
|
||||||
|
# 3. Text Embeddings (High Accuracy Semantic)
|
||||||
|
# This is the most accurate method for diverse/unknown domains
|
||||||
|
embedder = get_text_embedder()
|
||||||
|
embedding_score = 0.0
|
||||||
|
|
||||||
|
if embedder:
|
||||||
|
try:
|
||||||
|
# Embed text and candidates
|
||||||
|
# Batch embedding is faster and scalable without caching
|
||||||
|
all_texts = [text] + candidates
|
||||||
|
embeddings = list(embedder.embed_batch(all_texts))
|
||||||
|
|
||||||
|
if embeddings and len(embeddings) > 1:
|
||||||
|
text_emb = embeddings[0]
|
||||||
|
cand_embs = embeddings[1:]
|
||||||
|
|
||||||
|
# Calculate cosine similarity manually or via numpy
|
||||||
|
import numpy as np
|
||||||
|
text_norm = np.linalg.norm(text_emb)
|
||||||
|
|
||||||
|
if text_norm > 0:
|
||||||
|
for cand_emb in cand_embs:
|
||||||
|
cand_norm = np.linalg.norm(cand_emb)
|
||||||
|
if cand_norm > 0:
|
||||||
|
sim = np.dot(text_emb, cand_emb) / (text_norm * cand_norm)
|
||||||
|
if sim > embedding_score:
|
||||||
|
embedding_score = sim
|
||||||
|
except Exception as e:
|
||||||
|
logger.debug(f"Embedding calculation failed: {e}")
|
||||||
|
pass
|
||||||
|
|
||||||
|
if embedding_score > best_score:
|
||||||
|
best_score = embedding_score
|
||||||
|
|
||||||
|
# 4. Vector Similarity (Legacy/Fallback)
|
||||||
|
# Only use if we haven't found a good match yet and embeddings failed/unavailable
|
||||||
|
if best_score < 0.9:
|
||||||
|
nlp = get_nlp_model()
|
||||||
|
vector_score = 0.0
|
||||||
|
if nlp and nlp.vocab.vectors.shape[0] > 0:
|
||||||
|
try:
|
||||||
|
# Only use vectors if the word is in vocab or we have a good model
|
||||||
|
doc = nlp(text)
|
||||||
|
if doc.vector_norm:
|
||||||
|
for candidate in candidates:
|
||||||
|
cand_doc = nlp(candidate)
|
||||||
|
if cand_doc.vector_norm:
|
||||||
|
score = doc.similarity(cand_doc)
|
||||||
|
if score > vector_score:
|
||||||
|
vector_score = score
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
if vector_score > best_score:
|
||||||
|
best_score = vector_score
|
||||||
|
|
||||||
|
# 4. Fuzzy Match (Fallback/Refinement)
|
||||||
|
# If vector score is low (e.g. OOV words), fuzzy match might be better
|
||||||
|
# But difflib is slow for many candidates.
|
||||||
|
# Only run if we don't have a very high score yet
|
||||||
|
if best_score < 0.9:
|
||||||
|
for cand in candidates_lower:
|
||||||
|
# Quick check for common characters
|
||||||
|
if not cand: continue
|
||||||
|
|
||||||
|
# SequenceMatcher
|
||||||
|
score = difflib.SequenceMatcher(None, text_lower, cand).ratio()
|
||||||
|
if score > best_score:
|
||||||
|
best_score = score
|
||||||
|
|
||||||
|
return float(best_score)
|
||||||
|
|
||||||
|
def calculate_weighted_confidence(
|
||||||
|
item_type: str,
|
||||||
|
original_confidence: float,
|
||||||
|
valid_types: Optional[List[str]] = None,
|
||||||
|
item_text: Optional[str] = None,
|
||||||
|
weight_method: float = 0.5,
|
||||||
|
weight_similarity: float = 0.5
|
||||||
|
) -> float:
|
||||||
|
"""
|
||||||
|
Calculate weighted confidence score using both Label and Content similarity.
|
||||||
|
Final Score = (weight_method * original_confidence) + (weight_similarity * max(label_sim, content_sim))
|
||||||
|
|
||||||
|
Args:
|
||||||
|
item_type: The extracted type/label/predicate (e.g., "PERSON", "founded_by")
|
||||||
|
original_confidence: The confidence score from the extraction method (0.0-1.0)
|
||||||
|
valid_types: List of valid/preferred types provided by user
|
||||||
|
item_text: The actual text content extracted (e.g., "Steve Jobs", "acquired")
|
||||||
|
weight_method: Weight for the original method confidence (default 0.5)
|
||||||
|
weight_similarity: Weight for the similarity score (default 0.5)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
float: Weighted confidence score (0.0-1.0)
|
||||||
|
"""
|
||||||
|
if not valid_types:
|
||||||
|
return original_confidence
|
||||||
|
|
||||||
|
# Similarity 1: Label vs Valid Types (e.g., "PERSON" vs "Artist")
|
||||||
|
label_similarity = calculate_similarity(item_type, valid_types)
|
||||||
|
|
||||||
|
# Similarity 2: Content vs Valid Types (e.g., "Picasso" vs "Artist")
|
||||||
|
content_similarity = 0.0
|
||||||
|
if item_text:
|
||||||
|
content_similarity = calculate_similarity(item_text, valid_types)
|
||||||
|
|
||||||
|
# Take the best similarity match
|
||||||
|
best_similarity = max(label_similarity, content_similarity)
|
||||||
|
|
||||||
|
# Normalize weights
|
||||||
|
total_weight = weight_method + weight_similarity
|
||||||
|
if total_weight <= 0:
|
||||||
|
return original_confidence
|
||||||
|
|
||||||
|
w_m = weight_method / total_weight
|
||||||
|
w_s = weight_similarity / total_weight
|
||||||
|
|
||||||
|
final_score = (w_m * original_confidence) + (w_s * best_similarity)
|
||||||
|
|
||||||
|
return max(0.0, min(1.0, final_score))
|
||||||
|
|
||||||
|
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
# Entity Extraction Methods
|
# Entity Extraction Methods
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
@@ -134,8 +410,8 @@ def extract_entities_pattern(text: str, **kwargs) -> List[Entity]:
|
|||||||
entities = []
|
entities = []
|
||||||
|
|
||||||
patterns = {
|
patterns = {
|
||||||
"PERSON": r"\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)\b",
|
"ORG": r"\b([A-Z][a-zA-Z]+(?:\s+[A-Z][a-zA-Z]+)*\s+(?:Inc|Corp|LLC|Ltd|Company)(?:\.|\b))",
|
||||||
"ORG": r"\b([A-Z][a-zA-Z]+(?:\s+[A-Z][a-zA-Z]+)*\s+(?:Inc|Corp|LLC|Ltd|Company))\b",
|
"PERSON": r"\b([A-Z][a-z]+(?:\s+(?!Inc|Corp|LLC|Ltd|Company)[A-Z][a-z]+)+)\b",
|
||||||
"GPE": r"\b([A-Z][a-z]+\s*(?:City|State|Country|Nation))\b",
|
"GPE": r"\b([A-Z][a-z]+\s*(?:City|State|Country|Nation))\b",
|
||||||
"DATE": r"\b(\d{1,2}[/-]\d{1,2}[/-]\d{2,4}|\d{4})\b",
|
"DATE": r"\b(\d{1,2}[/-]\d{1,2}[/-]\d{2,4}|\d{4})\b",
|
||||||
}
|
}
|
||||||
@@ -310,6 +586,7 @@ def extract_entities_llm(
|
|||||||
model: Optional[str] = None,
|
model: Optional[str] = None,
|
||||||
silent_fail: bool = False,
|
silent_fail: bool = False,
|
||||||
max_text_length: Optional[int] = None,
|
max_text_length: Optional[int] = None,
|
||||||
|
structured_output_mode: str = "typed",
|
||||||
**kwargs,
|
**kwargs,
|
||||||
) -> List[Entity]:
|
) -> List[Entity]:
|
||||||
"""
|
"""
|
||||||
@@ -394,26 +671,55 @@ If an entity doesn't fit any of the preferred types, use the most appropriate ty
|
|||||||
entity_types_instruction = """Entity types should be one of: PERSON, ORG, GPE, DATE, EVENT, PRODUCT, CONCEPT, or related types.
|
entity_types_instruction = """Entity types should be one of: PERSON, ORG, GPE, DATE, EVENT, PRODUCT, CONCEPT, or related types.
|
||||||
Use the most appropriate type for each entity, including variations or synonyms if they better match the context."""
|
Use the most appropriate type for each entity, including variations or synonyms if they better match the context."""
|
||||||
|
|
||||||
prompt = f"""Extract named entities from the following text.
|
if not SCHEMAS_AVAILABLE:
|
||||||
Return ONLY a valid JSON list of objects with the following structure:
|
raise ImportError("Pydantic schemas not available. Install pydantic/instructor to use LLM extraction.")
|
||||||
[
|
|
||||||
{{"text": "entity name", "label": "ENTITY_TYPE", "start": 0, "end": 10, "confidence": 0.9}}
|
|
||||||
]
|
|
||||||
|
|
||||||
{entity_types_instruction}
|
|
||||||
Do not include any conversational filler, explanations, or markdown formatting outside the JSON block.
|
|
||||||
|
|
||||||
Text: {text}"""
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# 4. EXTRACTION WITH RETRY (handled by generate_structured)
|
prompt = f"""Extract named entities from the provided text.
|
||||||
result = llm.generate_structured(prompt)
|
Return the result as a JSON object with an "entities" key containing the list of entities.
|
||||||
entities = _parse_entity_result(result, provider, model)
|
Each entity should have 'text', 'label', and 'confidence' fields.
|
||||||
|
|
||||||
|
IMPORTANT:
|
||||||
|
- Return a FLAT LIST of entities.
|
||||||
|
- DO NOT group entities by type.
|
||||||
|
- The output structure must exactly match: {{ "entities": [ {{ "text": "...", "label": "...", "confidence": ... }}, ... ] }}
|
||||||
|
|
||||||
|
Example output (JSON format only):
|
||||||
|
{{
|
||||||
|
"entities": [
|
||||||
|
{{"text": "Entity Name", "label": "CATEGORY", "confidence": 0.95}},
|
||||||
|
{{"text": "Another Entity", "label": "OTHER_CATEGORY", "confidence": 0.90}}
|
||||||
|
]
|
||||||
|
}}
|
||||||
|
|
||||||
|
Instructions:
|
||||||
|
1. Extract entities ONLY from the text provided below.
|
||||||
|
2. Do not include any entities from the example above.
|
||||||
|
3. {entity_types_instruction}
|
||||||
|
|
||||||
|
Text to extract from:
|
||||||
|
{text}"""
|
||||||
|
|
||||||
if not entities:
|
# Use typed generation with Pydantic schema
|
||||||
logger.warning(f"No entities extracted using {provider}/{model} from text preview: {text[:100]}...")
|
result_obj = llm.generate_typed(prompt, schema=EntitiesResponse)
|
||||||
|
|
||||||
logger.info(f"Successfully extracted {len(entities)} entities using {provider}/{model}")
|
# Convert back to internal Entity format
|
||||||
|
entities = []
|
||||||
|
for e_out in result_obj.entities:
|
||||||
|
entities.append(Entity(
|
||||||
|
text=e_out.text,
|
||||||
|
label=e_out.label,
|
||||||
|
start_char=e_out.start if hasattr(e_out, "start") else 0, # Schema might not force these
|
||||||
|
end_char=e_out.end if hasattr(e_out, "end") else 0,
|
||||||
|
confidence=e_out.confidence,
|
||||||
|
metadata={
|
||||||
|
"provider": provider,
|
||||||
|
"model": model,
|
||||||
|
"extraction_method": "llm_typed",
|
||||||
|
}
|
||||||
|
))
|
||||||
|
|
||||||
|
logger.info(f"Successfully extracted {len(entities)} entities using {provider}/{model} (typed)")
|
||||||
return entities
|
return entities
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -473,6 +779,7 @@ def _extract_entities_chunked(
|
|||||||
model: Optional[str],
|
model: Optional[str],
|
||||||
silent_fail: bool,
|
silent_fail: bool,
|
||||||
max_text_length: int,
|
max_text_length: int,
|
||||||
|
structured_output_mode: str = "typed",
|
||||||
**kwargs
|
**kwargs
|
||||||
) -> List[Entity]:
|
) -> List[Entity]:
|
||||||
"""Internal helper to extract entities from long text by chunking."""
|
"""Internal helper to extract entities from long text by chunking."""
|
||||||
@@ -496,6 +803,7 @@ def _extract_entities_chunked(
|
|||||||
model=model,
|
model=model,
|
||||||
silent_fail=False, # We want to know if a chunk fails
|
silent_fail=False, # We want to know if a chunk fails
|
||||||
max_text_length=len(chunk.text) + 1,
|
max_text_length=len(chunk.text) + 1,
|
||||||
|
structured_output_mode=structured_output_mode,
|
||||||
**kwargs
|
**kwargs
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -506,22 +814,10 @@ def _extract_entities_chunked(
|
|||||||
|
|
||||||
all_entities.extend(chunk_entities)
|
all_entities.extend(chunk_entities)
|
||||||
|
|
||||||
return _deduplicate_entities(all_entities)
|
return all_entities
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def _deduplicate_entities(entities: List[Entity]) -> List[Entity]:
|
|
||||||
"""Remove duplicate entities, keeping those with higher confidence or more metadata."""
|
|
||||||
if not entities:
|
|
||||||
return []
|
|
||||||
|
|
||||||
# Sort by text, start_char, and confidence
|
|
||||||
unique_entities = {}
|
|
||||||
for ent in entities:
|
|
||||||
key = (ent.text.lower(), ent.start_char, ent.end_char, ent.label)
|
|
||||||
if key not in unique_entities or ent.confidence > unique_entities[key].confidence:
|
|
||||||
unique_entities[key] = ent
|
|
||||||
|
|
||||||
return sorted(list(unique_entities.values()), key=lambda e: e.start_char)
|
|
||||||
|
|
||||||
|
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
@@ -559,17 +855,17 @@ def extract_relations_pattern(
|
|||||||
|
|
||||||
relation_patterns = {
|
relation_patterns = {
|
||||||
"founded_by": [
|
"founded_by": [
|
||||||
fr"(?P<subject>{subject_pat})\s+(?:was\s+)?founded\s+by\s+(?P<object>{ent_pat})",
|
fr"(?P<subject>{subject_pat})(?:[.,])?\s+(?:was\s+)?founded\s+by\s+(?P<object>{ent_pat})",
|
||||||
fr"(?P<object>{ent_pat})\s+founded\s+(?P<subject>{ent_pat})",
|
fr"(?P<object>{ent_pat})(?:[.,])?\s+founded\s+(?P<subject>{ent_pat})",
|
||||||
fr"(?P<subject>{subject_pat})\s+(?:was\s+)?established\s+by\s+(?P<object>{ent_pat})",
|
fr"(?P<subject>{subject_pat})(?:[.,])?\s+(?:was\s+)?established\s+by\s+(?P<object>{ent_pat})",
|
||||||
fr"(?P<object>{ent_pat})\s+established\s+(?P<subject>{ent_pat})",
|
fr"(?P<object>{ent_pat})(?:[.,])?\s+established\s+(?P<subject>{ent_pat})",
|
||||||
fr"(?P<subject>{subject_pat})\s+(?:was\s+)?created\s+by\s+(?P<object>{ent_pat})",
|
fr"(?P<subject>{subject_pat})(?:[.,])?\s+(?:was\s+)?created\s+by\s+(?P<object>{ent_pat})",
|
||||||
fr"(?P<object>{ent_pat})\s+created\s+(?P<subject>{ent_pat})",
|
fr"(?P<object>{ent_pat})(?:[.,])?\s+created\s+(?P<subject>{ent_pat})",
|
||||||
fr"(?P<subject>{subject_pat})\s+(?:was\s+)?started\s+by\s+(?P<object>{ent_pat})",
|
fr"(?P<subject>{subject_pat})(?:[.,])?\s+(?:was\s+)?started\s+by\s+(?P<object>{ent_pat})",
|
||||||
fr"(?P<object>{ent_pat})\s+started\s+(?P<subject>{ent_pat})",
|
fr"(?P<object>{ent_pat})(?:[.,])?\s+started\s+(?P<subject>{ent_pat})",
|
||||||
fr"(?P<subject>{subject_pat})\s+(?:was\s+)?co-founded\s+by\s+(?P<object>{ent_pat})",
|
fr"(?P<subject>{subject_pat})(?:[.,])?\s+(?:was\s+)?co-founded\s+by\s+(?P<object>{ent_pat})",
|
||||||
fr"(?P<object>{ent_pat})\s+co-founded\s+(?P<subject>{ent_pat})",
|
fr"(?P<object>{ent_pat})(?:[.,])?\s+co-founded\s+(?P<subject>{ent_pat})",
|
||||||
fr"(?P<object>{ent_pat})\s+is\s+(?:the\s+)?founder\s+of\s+(?P<subject>{ent_pat})",
|
fr"(?P<object>{ent_pat})(?:[.,])?\s+is\s+(?:the\s+)?founder\s+of\s+(?P<subject>{ent_pat})",
|
||||||
],
|
],
|
||||||
"located_in": [
|
"located_in": [
|
||||||
fr"(?P<subject>{subject_pat})\s+is\s+located\s+in\s+(?P<object>{ent_pat})",
|
fr"(?P<subject>{subject_pat})\s+is\s+located\s+in\s+(?P<object>{ent_pat})",
|
||||||
@@ -606,26 +902,16 @@ def extract_relations_pattern(
|
|||||||
}
|
}
|
||||||
|
|
||||||
entity_map = {e.text.lower(): e for e in entities}
|
entity_map = {e.text.lower(): e for e in entities}
|
||||||
# DEBUG: Print entities in map
|
|
||||||
print(f"DEBUG: Entity map keys: {list(entity_map.keys())}")
|
|
||||||
|
|
||||||
for relation_type, patterns in relation_patterns.items():
|
for relation_type, patterns in relation_patterns.items():
|
||||||
for pattern in patterns:
|
for pattern in patterns:
|
||||||
# DEBUG: Print pattern being tried
|
|
||||||
# print(f"DEBUG: Trying pattern: {pattern}")
|
|
||||||
for match in re.finditer(pattern, text, re.IGNORECASE):
|
for match in re.finditer(pattern, text, re.IGNORECASE):
|
||||||
subject_text = match.group("subject").strip()
|
subject_text = match.group("subject").strip()
|
||||||
object_text = match.group("object").strip()
|
object_text = match.group("object").strip()
|
||||||
|
|
||||||
# DEBUG: Print match details
|
|
||||||
print(f"DEBUG: Match found! Subject='{subject_text}', Object='{object_text}'")
|
|
||||||
|
|
||||||
subject_entity = entity_map.get(subject_text.lower())
|
subject_entity = entity_map.get(subject_text.lower())
|
||||||
object_entity = entity_map.get(object_text.lower())
|
object_entity = entity_map.get(object_text.lower())
|
||||||
|
|
||||||
# DEBUG: Print lookup results
|
|
||||||
print(f"DEBUG: Subject Entity found: {subject_entity is not None}, Object Entity found: {object_entity is not None}")
|
|
||||||
|
|
||||||
if subject_entity and object_entity:
|
if subject_entity and object_entity:
|
||||||
start = max(0, match.start() - 50)
|
start = max(0, match.start() - 50)
|
||||||
end = min(len(text), match.end() + 50)
|
end = min(len(text), match.end() + 50)
|
||||||
@@ -725,6 +1011,127 @@ def extract_relations_cooccurrence(
|
|||||||
return relations
|
return relations
|
||||||
|
|
||||||
|
|
||||||
|
def extract_relations_similarity(
|
||||||
|
text: str, entities: List[Entity], relation_types: Optional[List[str]] = None, **kwargs
|
||||||
|
) -> List[Relation]:
|
||||||
|
"""
|
||||||
|
Similarity-based relation extraction.
|
||||||
|
Uses semantic similarity to match the context between entities to provided relation types.
|
||||||
|
"""
|
||||||
|
if not entities:
|
||||||
|
return []
|
||||||
|
|
||||||
|
# If no relation types provided, we can't do similarity matching against types
|
||||||
|
if not relation_types:
|
||||||
|
# Fallback to co-occurrence if no types to match against
|
||||||
|
logger.warning("No relation types provided for similarity matching. Falling back to co-occurrence.")
|
||||||
|
return extract_relations_cooccurrence(text, entities, **kwargs)
|
||||||
|
|
||||||
|
relations = []
|
||||||
|
|
||||||
|
# Try to load spaCy model with vectors
|
||||||
|
nlp = None
|
||||||
|
if SPACY_AVAILABLE:
|
||||||
|
try:
|
||||||
|
# Prefer larger models for vectors
|
||||||
|
for model_name in ["en_core_web_lg", "en_core_web_md", "en_core_web_sm"]:
|
||||||
|
if spacy.util.is_package(model_name):
|
||||||
|
nlp = spacy.load(model_name)
|
||||||
|
break
|
||||||
|
if not nlp:
|
||||||
|
# Try loading what we have
|
||||||
|
try:
|
||||||
|
nlp = spacy.load("en_core_web_sm")
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Pre-compute relation type vectors if possible
|
||||||
|
relation_vectors = {}
|
||||||
|
has_vectors = False
|
||||||
|
if nlp:
|
||||||
|
# Check if model has vectors
|
||||||
|
if nlp.vocab.vectors.shape[0] > 0:
|
||||||
|
has_vectors = True
|
||||||
|
for rt in relation_types:
|
||||||
|
relation_vectors[rt] = nlp(rt)
|
||||||
|
|
||||||
|
for entity1 in entities:
|
||||||
|
for entity2 in entities:
|
||||||
|
if entity1 == entity2:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Check distance
|
||||||
|
distance = abs(entity1.end_char - entity2.start_char)
|
||||||
|
# Only consider entities reasonably close (e.g., within same sentence or clause)
|
||||||
|
if distance > 100:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Ensure correct order for extracting text between
|
||||||
|
if entity1.end_char < entity2.start_char:
|
||||||
|
start_pos = entity1.end_char
|
||||||
|
end_pos = entity2.start_char
|
||||||
|
else:
|
||||||
|
start_pos = entity2.end_char
|
||||||
|
end_pos = entity1.start_char
|
||||||
|
|
||||||
|
between_text = text[start_pos:end_pos].strip()
|
||||||
|
|
||||||
|
if not between_text:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Calculate similarity
|
||||||
|
best_type = None
|
||||||
|
best_score = 0.0
|
||||||
|
|
||||||
|
if has_vectors and relation_vectors:
|
||||||
|
# Vector similarity
|
||||||
|
doc = nlp(between_text)
|
||||||
|
if doc.vector_norm:
|
||||||
|
for rt, vec in relation_vectors.items():
|
||||||
|
if vec.vector_norm:
|
||||||
|
sim = doc.similarity(vec)
|
||||||
|
if sim > best_score:
|
||||||
|
best_score = sim
|
||||||
|
best_type = rt
|
||||||
|
else:
|
||||||
|
# String similarity / Keyword matching
|
||||||
|
from difflib import SequenceMatcher
|
||||||
|
for rt in relation_types:
|
||||||
|
# Check for direct keyword presence (strong signal)
|
||||||
|
if rt.lower() in between_text.lower():
|
||||||
|
score = 1.0
|
||||||
|
else:
|
||||||
|
# Fuzzy match
|
||||||
|
score = SequenceMatcher(None, rt.lower(), between_text.lower()).ratio()
|
||||||
|
|
||||||
|
if score > best_score:
|
||||||
|
best_score = score
|
||||||
|
best_type = rt
|
||||||
|
|
||||||
|
# Threshold
|
||||||
|
threshold = kwargs.get("similarity_threshold", 0.4 if has_vectors else 0.6)
|
||||||
|
|
||||||
|
if best_type and best_score >= threshold:
|
||||||
|
relations.append(
|
||||||
|
Relation(
|
||||||
|
subject=entity1,
|
||||||
|
predicate=best_type,
|
||||||
|
object=entity2,
|
||||||
|
confidence=float(best_score),
|
||||||
|
context=text[max(0, min(entity1.start_char, entity2.start_char) - 20) : min(len(text), max(entity1.end_char, entity2.end_char) + 20)],
|
||||||
|
metadata={
|
||||||
|
"extraction_method": "similarity",
|
||||||
|
"similarity_score": float(best_score),
|
||||||
|
"between_text": between_text
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
return relations
|
||||||
|
|
||||||
|
|
||||||
def extract_relations_dependency(
|
def extract_relations_dependency(
|
||||||
text: str, entities: List[Entity], model: str = "en_core_web_sm", **kwargs
|
text: str, entities: List[Entity], model: str = "en_core_web_sm", **kwargs
|
||||||
) -> List[Relation]:
|
) -> List[Relation]:
|
||||||
@@ -882,6 +1289,7 @@ def extract_relations_llm(
|
|||||||
model: Optional[str] = None,
|
model: Optional[str] = None,
|
||||||
silent_fail: bool = False,
|
silent_fail: bool = False,
|
||||||
max_text_length: Optional[int] = None,
|
max_text_length: Optional[int] = None,
|
||||||
|
structured_output_mode: str = "typed",
|
||||||
**kwargs,
|
**kwargs,
|
||||||
) -> List[Relation]:
|
) -> List[Relation]:
|
||||||
"""
|
"""
|
||||||
@@ -954,7 +1362,8 @@ def extract_relations_llm(
|
|||||||
logger.info(f"Text length ({len(text)}) exceeds limit for relations. Chunking...")
|
logger.info(f"Text length ({len(text)}) exceeds limit for relations. Chunking...")
|
||||||
return _extract_relations_chunked(
|
return _extract_relations_chunked(
|
||||||
text, entities, provider=provider, model=model,
|
text, entities, provider=provider, model=model,
|
||||||
silent_fail=silent_fail, max_text_length=max_text_length, **kwargs
|
silent_fail=silent_fail, max_text_length=max_text_length,
|
||||||
|
**kwargs
|
||||||
)
|
)
|
||||||
|
|
||||||
entities_str = ", ".join([f"{e.text} ({e.label})" for e in entities])
|
entities_str = ", ".join([f"{e.text} ({e.label})" for e in entities])
|
||||||
@@ -972,20 +1381,63 @@ If a relation doesn't fit any of the preferred types, use the most appropriate t
|
|||||||
Extract meaningful relationships between entities. Use appropriate relation types that accurately describe how entities are connected.
|
Extract meaningful relationships between entities. Use appropriate relation types that accurately describe how entities are connected.
|
||||||
Common relation types include: related_to, part_of, located_in, created_by, uses, depends_on, interacts_with, and similar variations."""
|
Common relation types include: related_to, part_of, located_in, created_by, uses, depends_on, interacts_with, and similar variations."""
|
||||||
|
|
||||||
prompt = f"""Extract relations between entities from the following text.
|
if not SCHEMAS_AVAILABLE:
|
||||||
|
raise ImportError("Pydantic schemas not available. Install pydantic/instructor to use LLM extraction.")
|
||||||
|
|
||||||
Text: {text}
|
prompt = f"""Extract relations between entities from the provided text.
|
||||||
Entities: {entities_str}{relation_types_instruction}
|
Return the result as a JSON object with a "relations" key containing the list of relations.
|
||||||
|
Each relation must have 'subject', 'predicate', and 'object' fields.
|
||||||
|
|
||||||
Return JSON format: [{{"subject": "...", "predicate": "...", "object": "...", "confidence": 0.9}}]
|
Example output (JSON format only):
|
||||||
Extract all meaningful relationships between the entities, using the most appropriate relation type for each relationship."""
|
{{
|
||||||
|
"relations": [
|
||||||
|
{{"subject": "Entity A", "predicate": "related_to", "object": "Entity B", "confidence": 0.95}},
|
||||||
|
{{"subject": "Subject Entity", "predicate": "action_verb", "object": "Object Entity", "confidence": 0.90}}
|
||||||
|
]
|
||||||
|
}}
|
||||||
|
|
||||||
|
Instructions:
|
||||||
|
1. Extract relations ONLY from the text provided below.
|
||||||
|
2. Do not include any relations from the example above.
|
||||||
|
3. Use the provided entities list as a reference for subjects and objects.
|
||||||
|
4. {relation_types_instruction}
|
||||||
|
|
||||||
|
Text to extract from:
|
||||||
|
{text}
|
||||||
|
Entities found in text: {entities_str}"""
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# 4. EXTRACTION WITH RETRY
|
# Use typed generation with Pydantic schema
|
||||||
result = llm.generate_structured(prompt)
|
result_obj = llm.generate_typed(prompt, schema=RelationsResponse)
|
||||||
relations = _parse_relation_result(result, entities, text, provider, model)
|
|
||||||
|
|
||||||
logger.info(f"Successfully extracted {len(relations)} relations using {provider}/{model}")
|
# Convert back to internal Relation format
|
||||||
|
relations = []
|
||||||
|
for r_out in result_obj.relations:
|
||||||
|
# Find matching entities
|
||||||
|
subject_entity = next(
|
||||||
|
(e for e in entities if e.text.lower() == r_out.subject.lower()),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
object_entity = next(
|
||||||
|
(e for e in entities if e.text.lower() == r_out.object.lower()),
|
||||||
|
None
|
||||||
|
)
|
||||||
|
|
||||||
|
if subject_entity and object_entity:
|
||||||
|
relations.append(Relation(
|
||||||
|
subject=subject_entity,
|
||||||
|
predicate=r_out.predicate,
|
||||||
|
object=object_entity,
|
||||||
|
confidence=r_out.confidence,
|
||||||
|
context=text, # Simplified context
|
||||||
|
metadata={
|
||||||
|
"provider": provider,
|
||||||
|
"model": model,
|
||||||
|
"extraction_method": "llm_typed"
|
||||||
|
}
|
||||||
|
))
|
||||||
|
|
||||||
|
logger.info(f"Successfully extracted {len(relations)} relations using {provider}/{model} (typed)")
|
||||||
return relations
|
return relations
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -1067,6 +1519,7 @@ def _extract_relations_chunked(
|
|||||||
model: Optional[str],
|
model: Optional[str],
|
||||||
silent_fail: bool,
|
silent_fail: bool,
|
||||||
max_text_length: int,
|
max_text_length: int,
|
||||||
|
structured_output_mode: str = "typed",
|
||||||
**kwargs
|
**kwargs
|
||||||
) -> List[Relation]:
|
) -> List[Relation]:
|
||||||
"""Internal helper to extract relations from long text by chunking."""
|
"""Internal helper to extract relations from long text by chunking."""
|
||||||
@@ -1099,29 +1552,15 @@ def _extract_relations_chunked(
|
|||||||
model=model,
|
model=model,
|
||||||
silent_fail=False,
|
silent_fail=False,
|
||||||
max_text_length=len(chunk.text) + 1,
|
max_text_length=len(chunk.text) + 1,
|
||||||
|
structured_output_mode=structured_output_mode,
|
||||||
**kwargs
|
**kwargs
|
||||||
)
|
)
|
||||||
all_relations.extend(chunk_rels)
|
all_relations.extend(chunk_rels)
|
||||||
|
|
||||||
return _deduplicate_relations(all_relations)
|
return all_relations
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def _deduplicate_relations(relations: List[Relation]) -> List[Relation]:
|
|
||||||
"""Remove duplicate relations."""
|
|
||||||
if not relations:
|
|
||||||
return []
|
|
||||||
|
|
||||||
unique_rels = {}
|
|
||||||
for rel in relations:
|
|
||||||
key = (
|
|
||||||
rel.subject.text.lower(),
|
|
||||||
rel.predicate.lower(),
|
|
||||||
rel.object.text.lower()
|
|
||||||
)
|
|
||||||
if key not in unique_rels or rel.confidence > unique_rels[key].confidence:
|
|
||||||
unique_rels[key] = rel
|
|
||||||
|
|
||||||
return list(unique_rels.values())
|
|
||||||
|
|
||||||
|
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
@@ -1248,6 +1687,7 @@ def extract_triplets_llm(
|
|||||||
model: Optional[str] = None,
|
model: Optional[str] = None,
|
||||||
silent_fail: bool = False,
|
silent_fail: bool = False,
|
||||||
max_text_length: Optional[int] = None,
|
max_text_length: Optional[int] = None,
|
||||||
|
structured_output_mode: str = "typed",
|
||||||
**kwargs,
|
**kwargs,
|
||||||
) -> List[Triplet]:
|
) -> List[Triplet]:
|
||||||
"""
|
"""
|
||||||
@@ -1314,21 +1754,67 @@ def extract_triplets_llm(
|
|||||||
logger.info(f"Text length ({len(text)}) exceeds limit for triplets. Chunking...")
|
logger.info(f"Text length ({len(text)}) exceeds limit for triplets. Chunking...")
|
||||||
return _extract_triplets_chunked(
|
return _extract_triplets_chunked(
|
||||||
text, provider=provider, model=model,
|
text, provider=provider, model=model,
|
||||||
silent_fail=silent_fail, max_text_length=max_text_length, **kwargs
|
silent_fail=silent_fail, max_text_length=max_text_length,
|
||||||
|
**kwargs
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Use custom triplet types if provided
|
||||||
|
triplet_types = kwargs.get("triplet_types")
|
||||||
|
if triplet_types:
|
||||||
|
triplet_types_str = ", ".join(triplet_types)
|
||||||
|
triplet_types_instruction = f"""
|
||||||
|
Preferred triplet predicates: {triplet_types_str}.
|
||||||
|
You may also use related or similar predicates if they better capture the relationship (e.g., variations, synonyms, or domain-specific predicates).
|
||||||
|
If a predicate doesn't fit any of the preferred types, use the most appropriate type from the preferred list or a closely related type that accurately describes the relationship."""
|
||||||
|
else:
|
||||||
|
triplet_types_instruction = """
|
||||||
|
Extract meaningful triplets (subject-predicate-object). Use appropriate predicates that accurately describe the relationship.
|
||||||
|
Common predicates include: is_a, part_of, has_property, related_to, caused_by, etc."""
|
||||||
|
|
||||||
prompt = f"""Extract RDF triplets (subject-predicate-object) from the following text.
|
if not SCHEMAS_AVAILABLE:
|
||||||
|
raise ImportError("Pydantic schemas not available. Install pydantic/instructor to use LLM extraction.")
|
||||||
|
|
||||||
Text: {text}
|
prompt = f"""Extract RDF triplets (subject-predicate-object) from the provided text.
|
||||||
|
Return the result as a JSON object with a "triplets" key containing the list of triplets.
|
||||||
|
Each triplet must have 'subject', 'predicate', and 'object' fields.
|
||||||
|
|
||||||
Return JSON format: [{{"subject": "...", "predicate": "...", "object": "...", "confidence": 0.9}}]"""
|
Example output (JSON format only):
|
||||||
|
{{
|
||||||
|
"triplets": [
|
||||||
|
{{"subject": "Subject", "predicate": "predicate_relation", "object": "Object", "confidence": 0.99}},
|
||||||
|
{{"subject": "Concept A", "predicate": "is_a", "object": "Concept B", "confidence": 0.95}}
|
||||||
|
]
|
||||||
|
}}
|
||||||
|
|
||||||
|
Instructions:
|
||||||
|
1. Extract triplets ONLY from the text provided below.
|
||||||
|
2. Do not include any triplets from the example above.
|
||||||
|
3. Ensure subjects and objects are substrings from the text.
|
||||||
|
4. {triplet_types_instruction}
|
||||||
|
|
||||||
|
Text to extract from:
|
||||||
|
{text}"""
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# 4. EXTRACTION WITH RETRY
|
# Use typed generation with Pydantic schema
|
||||||
result = llm.generate_structured(prompt)
|
result_obj = llm.generate_typed(prompt, schema=TripletsResponse)
|
||||||
triplets = _parse_triplet_result(result, provider, model)
|
|
||||||
|
|
||||||
logger.info(f"Successfully extracted {len(triplets)} triplets using {provider}/{model}")
|
# Convert back to internal Triplet format
|
||||||
|
triplets = []
|
||||||
|
for t_out in result_obj.triplets:
|
||||||
|
triplets.append(Triplet(
|
||||||
|
subject=t_out.subject,
|
||||||
|
predicate=t_out.predicate,
|
||||||
|
object=t_out.object,
|
||||||
|
confidence=t_out.confidence,
|
||||||
|
metadata={
|
||||||
|
"provider": provider,
|
||||||
|
"model": model,
|
||||||
|
"extraction_method": "llm_typed"
|
||||||
|
}
|
||||||
|
))
|
||||||
|
|
||||||
|
logger.info(f"Successfully extracted {len(triplets)} triplets using {provider}/{model} (typed)")
|
||||||
return triplets
|
return triplets
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -1389,6 +1875,7 @@ def _extract_triplets_chunked(
|
|||||||
model: Optional[str],
|
model: Optional[str],
|
||||||
silent_fail: bool,
|
silent_fail: bool,
|
||||||
max_text_length: int,
|
max_text_length: int,
|
||||||
|
structured_output_mode: str = "typed",
|
||||||
**kwargs
|
**kwargs
|
||||||
) -> List[Triplet]:
|
) -> List[Triplet]:
|
||||||
"""Internal helper to extract triplets from long text by chunking."""
|
"""Internal helper to extract triplets from long text by chunking."""
|
||||||
@@ -1411,25 +1898,14 @@ def _extract_triplets_chunked(
|
|||||||
model=model,
|
model=model,
|
||||||
silent_fail=False,
|
silent_fail=False,
|
||||||
max_text_length=len(chunk.text) + 1,
|
max_text_length=len(chunk.text) + 1,
|
||||||
|
structured_output_mode=structured_output_mode,
|
||||||
**kwargs
|
**kwargs
|
||||||
)
|
)
|
||||||
all_triplets.extend(chunk_triplets)
|
all_triplets.extend(chunk_triplets)
|
||||||
|
|
||||||
return _deduplicate_triplets(all_triplets)
|
return all_triplets
|
||||||
|
|
||||||
|
|
||||||
def _deduplicate_triplets(triplets: List[Triplet]) -> List[Triplet]:
|
|
||||||
"""Remove duplicate triplets."""
|
|
||||||
if not triplets:
|
|
||||||
return []
|
|
||||||
|
|
||||||
unique_triplets = {}
|
|
||||||
for t in triplets:
|
|
||||||
key = (t.subject.lower(), t.predicate.lower(), t.object.lower())
|
|
||||||
if key not in unique_triplets or t.confidence > unique_triplets[key].confidence:
|
|
||||||
unique_triplets[key] = t
|
|
||||||
|
|
||||||
return list(unique_triplets.values())
|
|
||||||
|
|
||||||
|
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
@@ -1476,6 +1952,7 @@ def get_relation_method(method_name: str):
|
|||||||
"pattern": extract_relations_pattern,
|
"pattern": extract_relations_pattern,
|
||||||
"regex": extract_relations_regex,
|
"regex": extract_relations_regex,
|
||||||
"cooccurrence": extract_relations_cooccurrence,
|
"cooccurrence": extract_relations_cooccurrence,
|
||||||
|
"similarity": extract_relations_similarity,
|
||||||
"dependency": extract_relations_dependency,
|
"dependency": extract_relations_dependency,
|
||||||
"ml": extract_relations_dependency, # Alias for dependency
|
"ml": extract_relations_dependency, # Alias for dependency
|
||||||
"spacy": extract_relations_dependency, # Alias for dependency
|
"spacy": extract_relations_dependency, # Alias for dependency
|
||||||
|
|||||||
@@ -20,7 +20,12 @@ Algorithms Used:
|
|||||||
- Transformer Models: BERT, RoBERTa, DistilBERT for token classification
|
- Transformer Models: BERT, RoBERTa, DistilBERT for token classification
|
||||||
- Large Language Models: GPT, Claude, Gemini for zero-shot/few-shot extraction
|
- Large Language Models: GPT, Claude, Gemini for zero-shot/few-shot extraction
|
||||||
- Ensemble Voting: Majority voting and confidence-weighted aggregation
|
- Ensemble Voting: Majority voting and confidence-weighted aggregation
|
||||||
- Deduplication: Set-based and similarity-based entity deduplication
|
- Weighted Confidence Scoring:
|
||||||
|
* Formula: Score = (0.5 * Method_Confidence) + (0.5 * Type_Similarity_Score)
|
||||||
|
* Method_Confidence: Confidence score from the extraction algorithm
|
||||||
|
* Type_Similarity_Score: Semantic match with user-provided entity types (Exact=1.0, Synonym=0.95, Embedding=Cosine_Sim)
|
||||||
|
- Hybrid Similarity Matching: Exact -> Synonym -> Substring -> Semantic Embedding (Batch Optimized)
|
||||||
|
- Last Resort Fallback: Capitalized word heuristic when all other methods fail
|
||||||
|
|
||||||
Key Features:
|
Key Features:
|
||||||
- Multiple extraction methods:
|
- Multiple extraction methods:
|
||||||
@@ -31,8 +36,9 @@ Key Features:
|
|||||||
* HuggingFace: Custom HuggingFace NER models
|
* HuggingFace: Custom HuggingFace NER models
|
||||||
* LLM-based: Large language model extraction
|
* LLM-based: Large language model extraction
|
||||||
- Fallback chain support: Try methods in order until one succeeds
|
- Fallback chain support: Try methods in order until one succeeds
|
||||||
|
- Robust Fallbacks: Prevents empty results via ML -> Pattern -> Last Resort chain
|
||||||
- Ensemble voting: Combine results from multiple methods
|
- Ensemble voting: Combine results from multiple methods
|
||||||
- Post-processing: Entity boundary validation and deduplication
|
- Post-processing: Entity boundary validation
|
||||||
- Multiple entity type support (PERSON, ORG, GPE, DATE, etc.)
|
- Multiple entity type support (PERSON, ORG, GPE, DATE, etc.)
|
||||||
- Confidence scoring and filtering
|
- Confidence scoring and filtering
|
||||||
- Batch processing capabilities
|
- Batch processing capabilities
|
||||||
@@ -90,7 +96,12 @@ class Entity:
|
|||||||
class NERExtractor:
|
class NERExtractor:
|
||||||
"""Named Entity Recognition extractor."""
|
"""Named Entity Recognition extractor."""
|
||||||
|
|
||||||
def __init__(self, method: Union[str, List[str]] = "ml", **config):
|
def __init__(
|
||||||
|
self,
|
||||||
|
method: Union[str, List[str]] = "ml",
|
||||||
|
entity_types: Optional[List[str]] = None,
|
||||||
|
**config
|
||||||
|
):
|
||||||
"""
|
"""
|
||||||
Initialize NER extractor.
|
Initialize NER extractor.
|
||||||
|
|
||||||
@@ -103,6 +114,8 @@ class NERExtractor:
|
|||||||
- "huggingface": HuggingFace model
|
- "huggingface": HuggingFace model
|
||||||
- "llm": LLM-based extraction
|
- "llm": LLM-based extraction
|
||||||
- List of methods for fallback chain
|
- List of methods for fallback chain
|
||||||
|
entity_types: List of entity types to extract (e.g., ["PERSON", "ORG"]).
|
||||||
|
If provided, extraction methods will try to limit/focus on these types.
|
||||||
**config: Configuration options:
|
**config: Configuration options:
|
||||||
- model: Model name (for ML/HuggingFace methods)
|
- model: Model name (for ML/HuggingFace methods)
|
||||||
- huggingface_model: HuggingFace model name
|
- huggingface_model: HuggingFace model name
|
||||||
@@ -115,6 +128,7 @@ class NERExtractor:
|
|||||||
"""
|
"""
|
||||||
self.logger = get_logger("ner_extractor")
|
self.logger = get_logger("ner_extractor")
|
||||||
self.config = config
|
self.config = config
|
||||||
|
self.entity_types = entity_types
|
||||||
|
|
||||||
# Method configuration
|
# Method configuration
|
||||||
self.method = method if isinstance(method, list) else [method]
|
self.method = method if isinstance(method, list) else [method]
|
||||||
@@ -166,6 +180,7 @@ class NERExtractor:
|
|||||||
try:
|
try:
|
||||||
results = []
|
results = []
|
||||||
total_items = len(text)
|
total_items = len(text)
|
||||||
|
total_entities_count = 0
|
||||||
# Update more frequently: every 1% or at least every 10 items, but always update for small datasets
|
# Update more frequently: every 1% or at least every 10 items, but always update for small datasets
|
||||||
if total_items <= 10:
|
if total_items <= 10:
|
||||||
update_interval = 1 # Update every item for small datasets
|
update_interval = 1 # Update every item for small datasets
|
||||||
@@ -183,16 +198,28 @@ class NERExtractor:
|
|||||||
|
|
||||||
for idx, item in enumerate(text, 1):
|
for idx, item in enumerate(text, 1):
|
||||||
try:
|
try:
|
||||||
|
current_entities = []
|
||||||
if isinstance(item, dict) and "content" in item:
|
if isinstance(item, dict) and "content" in item:
|
||||||
results.append(self.extract_entities(item["content"], **kwargs))
|
current_entities = self.extract_entities(item["content"], **kwargs)
|
||||||
elif isinstance(item, str):
|
elif isinstance(item, str):
|
||||||
results.append(self.extract_entities(item, **kwargs))
|
current_entities = self.extract_entities(item, **kwargs)
|
||||||
else:
|
else:
|
||||||
# Try converting to string
|
# Try converting to string
|
||||||
try:
|
try:
|
||||||
results.append(self.extract_entities(str(item), **kwargs))
|
current_entities = self.extract_entities(str(item), **kwargs)
|
||||||
except Exception:
|
except Exception:
|
||||||
results.append([])
|
current_entities = []
|
||||||
|
|
||||||
|
# Add provenance metadata
|
||||||
|
for ent in current_entities:
|
||||||
|
if ent.metadata is None:
|
||||||
|
ent.metadata = {}
|
||||||
|
ent.metadata["batch_index"] = idx - 1
|
||||||
|
if isinstance(item, dict) and "id" in item:
|
||||||
|
ent.metadata["document_id"] = item["id"]
|
||||||
|
|
||||||
|
results.append(current_entities)
|
||||||
|
total_entities_count += len(current_entities)
|
||||||
except Exception:
|
except Exception:
|
||||||
results.append([])
|
results.append([])
|
||||||
|
|
||||||
@@ -209,13 +236,13 @@ class NERExtractor:
|
|||||||
tracking_id,
|
tracking_id,
|
||||||
processed=idx,
|
processed=idx,
|
||||||
total=total_items,
|
total=total_items,
|
||||||
message=f"Processing documents... {idx}/{total_items} (remaining: {remaining})"
|
message=f"Processing documents... {idx}/{total_items} (remaining: {remaining}) - Extracted {total_entities_count} entities so far"
|
||||||
)
|
)
|
||||||
|
|
||||||
self.progress_tracker.stop_tracking(
|
self.progress_tracker.stop_tracking(
|
||||||
tracking_id,
|
tracking_id,
|
||||||
status="completed",
|
status="completed",
|
||||||
message=f"Extracted entities from {len(results)} documents",
|
message=f"Batch extraction completed. Processed {len(results)} documents, extracted {total_entities_count} entities.",
|
||||||
)
|
)
|
||||||
return results
|
return results
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -260,10 +287,12 @@ class NERExtractor:
|
|||||||
methods = [methods]
|
methods = [methods]
|
||||||
|
|
||||||
min_confidence = options.get("min_confidence", self.min_confidence)
|
min_confidence = options.get("min_confidence", self.min_confidence)
|
||||||
entity_types = options.get("entity_types")
|
entity_types = options.get("entity_types", self.entity_types)
|
||||||
|
|
||||||
# Merge config with options
|
# Merge config with options
|
||||||
all_options = {**self.config, **options}
|
all_options = {**self.config, **options}
|
||||||
|
if entity_types:
|
||||||
|
all_options["entity_types"] = entity_types
|
||||||
|
|
||||||
# Try each method in order (fallback chain)
|
# Try each method in order (fallback chain)
|
||||||
all_entities = []
|
all_entities = []
|
||||||
@@ -300,29 +329,36 @@ class NERExtractor:
|
|||||||
api_key = os.getenv(env_key)
|
api_key = os.getenv(env_key)
|
||||||
if api_key:
|
if api_key:
|
||||||
method_options["api_key"] = api_key
|
method_options["api_key"] = api_key
|
||||||
# Pass entity_types to LLM method so it can use them in the prompt
|
|
||||||
if entity_types:
|
|
||||||
method_options["entity_types"] = entity_types
|
|
||||||
|
|
||||||
entities = method_func(text, **method_options)
|
entities = method_func(text, **method_options)
|
||||||
|
|
||||||
# Filter by confidence and entity types
|
# Apply weighted scoring if entity_types are provided
|
||||||
filtered = [e for e in entities if e.confidence >= min_confidence]
|
|
||||||
if entity_types:
|
if entity_types:
|
||||||
# Case-insensitive and flexible matching for entity types
|
try:
|
||||||
entity_types_lower = {et.lower() for et in entity_types}
|
from .methods import calculate_weighted_confidence
|
||||||
filtered = [
|
for e in entities:
|
||||||
e for e in filtered
|
e.confidence = calculate_weighted_confidence(
|
||||||
if e.label.lower() in entity_types_lower
|
item_type=e.label,
|
||||||
or any(et.lower() in e.label.lower() or e.label.lower() in et.lower()
|
original_confidence=e.confidence,
|
||||||
for et in entity_types)
|
valid_types=entity_types,
|
||||||
]
|
item_text=e.text
|
||||||
|
)
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Filter by confidence
|
||||||
|
filtered = [e for e in entities if e.confidence >= min_confidence]
|
||||||
|
|
||||||
if filtered:
|
if filtered:
|
||||||
all_entities.append((method_name, filtered))
|
all_entities.append((method_name, filtered))
|
||||||
|
|
||||||
# If not using ensemble, return first successful result
|
# If not using ensemble, return first successful result
|
||||||
if not self.ensemble_voting:
|
if not self.ensemble_voting:
|
||||||
|
# Ensure default metadata
|
||||||
|
for e in filtered:
|
||||||
|
if e.metadata is None: e.metadata = {}
|
||||||
|
if "batch_index" not in e.metadata: e.metadata["batch_index"] = 0
|
||||||
|
|
||||||
self.progress_tracker.stop_tracking(
|
self.progress_tracker.stop_tracking(
|
||||||
tracking_id,
|
tracking_id,
|
||||||
status="completed",
|
status="completed",
|
||||||
@@ -342,7 +378,8 @@ class NERExtractor:
|
|||||||
elif all_entities:
|
elif all_entities:
|
||||||
entities = all_entities[0][1] # Use first successful method
|
entities = all_entities[0][1] # Use first successful method
|
||||||
else:
|
else:
|
||||||
entities = []
|
# Fallback to pattern-based extraction if all models fail
|
||||||
|
entities = self._extract_fallback(text)
|
||||||
|
|
||||||
# Post-processing if enabled
|
# Post-processing if enabled
|
||||||
if self.post_process and entities:
|
if self.post_process and entities:
|
||||||
@@ -390,19 +427,12 @@ class NERExtractor:
|
|||||||
def _post_process_entities(self, entities: List[Entity], text: str) -> List[Entity]:
|
def _post_process_entities(self, entities: List[Entity], text: str) -> List[Entity]:
|
||||||
"""Post-process entities for refinement."""
|
"""Post-process entities for refinement."""
|
||||||
processed = []
|
processed = []
|
||||||
seen = set()
|
|
||||||
|
|
||||||
for entity in entities:
|
for entity in entities:
|
||||||
# Check boundaries
|
# Check boundaries
|
||||||
if entity.start_char < 0 or entity.end_char > len(text):
|
if entity.start_char < 0 or entity.end_char > len(text):
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Check for duplicates
|
|
||||||
key = (entity.text.lower(), entity.label, entity.start_char)
|
|
||||||
if key in seen:
|
|
||||||
continue
|
|
||||||
seen.add(key)
|
|
||||||
|
|
||||||
# Validate entity text matches
|
# Validate entity text matches
|
||||||
actual_text = text[entity.start_char : entity.end_char]
|
actual_text = text[entity.start_char : entity.end_char]
|
||||||
if actual_text.lower() != entity.text.lower():
|
if actual_text.lower() != entity.text.lower():
|
||||||
@@ -457,6 +487,7 @@ class NERExtractor:
|
|||||||
def _extract_fallback(self, text: str) -> List[Entity]:
|
def _extract_fallback(self, text: str) -> List[Entity]:
|
||||||
"""Fallback entity extraction using simple patterns."""
|
"""Fallback entity extraction using simple patterns."""
|
||||||
entities = []
|
entities = []
|
||||||
|
import re
|
||||||
|
|
||||||
# Simple patterns for common entity types
|
# Simple patterns for common entity types
|
||||||
patterns = {
|
patterns = {
|
||||||
@@ -466,20 +497,49 @@ class NERExtractor:
|
|||||||
"DATE": r"\b(\d{1,2}[/-]\d{1,2}[/-]\d{2,4}|\d{4})\b",
|
"DATE": r"\b(\d{1,2}[/-]\d{1,2}[/-]\d{2,4}|\d{4})\b",
|
||||||
}
|
}
|
||||||
|
|
||||||
import re
|
# Track covered ranges to avoid overlaps
|
||||||
|
covered_ranges = set()
|
||||||
|
|
||||||
for label, pattern in patterns.items():
|
for label, pattern in patterns.items():
|
||||||
for match in re.finditer(pattern, text):
|
for match in re.finditer(pattern, text):
|
||||||
entities.append(
|
start, end = match.start(), match.end()
|
||||||
Entity(
|
# Check overlap
|
||||||
text=match.group(1),
|
is_overlap = any(r_start < end and r_end > start for r_start, r_end in covered_ranges)
|
||||||
label=label,
|
if not is_overlap:
|
||||||
start_char=match.start(),
|
# Use group 1 if available, else group 0
|
||||||
end_char=match.end(),
|
text_val = match.group(1) if match.lastindex and match.lastindex >= 1 else match.group(0)
|
||||||
confidence=0.7, # Lower confidence for pattern-based
|
|
||||||
metadata={"extraction_method": "pattern"},
|
entities.append(
|
||||||
|
Entity(
|
||||||
|
text=text_val,
|
||||||
|
label=label,
|
||||||
|
start_char=start,
|
||||||
|
end_char=end,
|
||||||
|
confidence=0.7, # Lower confidence for pattern-based
|
||||||
|
metadata={"extraction_method": "pattern"},
|
||||||
|
)
|
||||||
)
|
)
|
||||||
)
|
covered_ranges.add((start, end))
|
||||||
|
|
||||||
|
# Last Resort: If no entities found, try single capitalized words as generic entities
|
||||||
|
if not entities:
|
||||||
|
# Match any capitalized word of length > 2
|
||||||
|
cap_pattern = r"\b[A-Z][a-z]{2,}\b"
|
||||||
|
for match in re.finditer(cap_pattern, text):
|
||||||
|
start, end = match.start(), match.end()
|
||||||
|
is_overlap = any(r_start < end and r_end > start for r_start, r_end in covered_ranges)
|
||||||
|
if not is_overlap:
|
||||||
|
entities.append(
|
||||||
|
Entity(
|
||||||
|
text=match.group(0),
|
||||||
|
label="UNKNOWN",
|
||||||
|
start_char=start,
|
||||||
|
end_char=end,
|
||||||
|
confidence=0.5,
|
||||||
|
metadata={"extraction_method": "last_resort_pattern"},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
covered_ranges.add((start, end))
|
||||||
|
|
||||||
return entities
|
return entities
|
||||||
|
|
||||||
@@ -494,7 +554,7 @@ class NERExtractor:
|
|||||||
Returns:
|
Returns:
|
||||||
list: List of entity lists for each text
|
list: List of entity lists for each text
|
||||||
"""
|
"""
|
||||||
return [self.extract_entities(text, **options) for text in texts]
|
return self.extract(texts, **options)
|
||||||
|
|
||||||
def classify_entities(self, entities: List[Entity]) -> Dict[str, List[Entity]]:
|
def classify_entities(self, entities: List[Entity]) -> Dict[str, List[Entity]]:
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -71,7 +71,19 @@ License: MIT
|
|||||||
|
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
from typing import Any, Dict, List, Optional, Union
|
import time
|
||||||
|
from typing import Any, Dict, List, Optional, Union, Type
|
||||||
|
|
||||||
|
try:
|
||||||
|
from pydantic import BaseModel, ValidationError
|
||||||
|
except ImportError:
|
||||||
|
BaseModel = Any
|
||||||
|
ValidationError = Exception
|
||||||
|
|
||||||
|
try:
|
||||||
|
import instructor
|
||||||
|
except ImportError:
|
||||||
|
instructor = None
|
||||||
|
|
||||||
from ..utils.exceptions import ProcessingError
|
from ..utils.exceptions import ProcessingError
|
||||||
from ..utils.logging import get_logger
|
from ..utils.logging import get_logger
|
||||||
@@ -211,6 +223,210 @@ class BaseProvider:
|
|||||||
raise ProcessingError(f"Failed to generate structured output: {last_error}")
|
raise ProcessingError(f"Failed to generate structured output: {last_error}")
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
def generate_typed(
|
||||||
|
self,
|
||||||
|
prompt: str,
|
||||||
|
schema: Type[BaseModel],
|
||||||
|
max_retries: int = 3,
|
||||||
|
**kwargs
|
||||||
|
) -> BaseModel:
|
||||||
|
"""
|
||||||
|
Generate structured output validated against a Pydantic schema.
|
||||||
|
Uses instructor if available and supported for the provider, otherwise falls back to a repair loop.
|
||||||
|
"""
|
||||||
|
provider_name = self.__class__.__name__
|
||||||
|
|
||||||
|
# Try using instructor first if available
|
||||||
|
if instructor:
|
||||||
|
try:
|
||||||
|
client = None
|
||||||
|
mode = instructor.Mode.TOOLS # Default mode
|
||||||
|
|
||||||
|
if provider_name == "OpenAIProvider" and self.client:
|
||||||
|
client = instructor.from_openai(self.client)
|
||||||
|
elif provider_name == "AnthropicProvider" and self.client:
|
||||||
|
client = instructor.from_anthropic(self.client)
|
||||||
|
elif provider_name == "GeminiProvider" and self.client:
|
||||||
|
client = instructor.from_gemini(
|
||||||
|
self.client,
|
||||||
|
mode=instructor.Mode.GEMINI_JSON
|
||||||
|
)
|
||||||
|
elif provider_name == "GroqProvider" and self.client:
|
||||||
|
# Try using from_groq if available (newer instructor versions)
|
||||||
|
if hasattr(instructor, "from_groq"):
|
||||||
|
client = instructor.from_groq(self.client, mode=instructor.Mode.JSON)
|
||||||
|
else:
|
||||||
|
# Fallback: Create OpenAI client pointing to Groq
|
||||||
|
# This avoids the "Client should be an instance of openai.OpenAI" warning
|
||||||
|
try:
|
||||||
|
from openai import OpenAI
|
||||||
|
groq_client = OpenAI(
|
||||||
|
base_url="https://api.groq.com/openai/v1",
|
||||||
|
api_key=self.client.api_key,
|
||||||
|
)
|
||||||
|
client = instructor.from_openai(groq_client, mode=instructor.Mode.JSON)
|
||||||
|
except Exception:
|
||||||
|
# Last resort: try passing the groq client directly
|
||||||
|
client = instructor.from_openai(self.client, mode=instructor.Mode.JSON)
|
||||||
|
elif provider_name == "OllamaProvider":
|
||||||
|
# Create OpenAI-compatible client for Ollama
|
||||||
|
try:
|
||||||
|
from openai import OpenAI
|
||||||
|
# Ollama typically runs on localhost:11434/v1
|
||||||
|
base_url = getattr(self, "base_url", "http://localhost:11434")
|
||||||
|
if not base_url.endswith("/v1"):
|
||||||
|
base_url = f"{base_url.rstrip('/')}/v1"
|
||||||
|
|
||||||
|
ollama_client = OpenAI(
|
||||||
|
base_url=base_url,
|
||||||
|
api_key="ollama", # required but unused
|
||||||
|
)
|
||||||
|
client = instructor.from_openai(ollama_client, mode=instructor.Mode.JSON)
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
elif provider_name == "DeepSeekProvider" and self.client:
|
||||||
|
# DeepSeek is OpenAI compatible
|
||||||
|
# We need to wrap the underlying client if it exposes the OpenAI interface
|
||||||
|
# or create a new OpenAI client if self.client is a deepseek.Client (which might be just a wrapper)
|
||||||
|
# Assuming deepseek.Client is compatible or we can use OpenAI client
|
||||||
|
try:
|
||||||
|
# DeepSeek usually works with standard OpenAI client
|
||||||
|
# If self.client is deepseek.Client, check if we can wrap it
|
||||||
|
# Otherwise create a new OpenAI client
|
||||||
|
from openai import OpenAI
|
||||||
|
if isinstance(self.client, OpenAI):
|
||||||
|
client = instructor.from_openai(self.client, mode=instructor.Mode.JSON)
|
||||||
|
else:
|
||||||
|
# Try creating fresh client
|
||||||
|
ds_client = OpenAI(
|
||||||
|
api_key=self.api_key,
|
||||||
|
base_url="https://api.deepseek.com"
|
||||||
|
)
|
||||||
|
client = instructor.from_openai(ds_client, mode=instructor.Mode.JSON)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
if client:
|
||||||
|
# Map generate arguments to client arguments
|
||||||
|
# Instructor standardizes on chat.completions.create for OpenAI/Groq/Anthropic/Gemini
|
||||||
|
|
||||||
|
create_kwargs = {
|
||||||
|
"model": kwargs.get("model", self.model),
|
||||||
|
"messages": [{"role": "user", "content": prompt}],
|
||||||
|
"response_model": schema,
|
||||||
|
"max_retries": max_retries,
|
||||||
|
"temperature": kwargs.get("temperature", 0.1), # Low temp for structured
|
||||||
|
}
|
||||||
|
|
||||||
|
# Add provider-specific params
|
||||||
|
if provider_name == "GroqProvider":
|
||||||
|
create_kwargs["response_format"] = {"type": "json_object"}
|
||||||
|
|
||||||
|
response = client.chat.completions.create(**create_kwargs)
|
||||||
|
return response
|
||||||
|
except Exception as e:
|
||||||
|
self.logger.warning(f"Instructor generation failed ({e}), falling back to manual repair loop.")
|
||||||
|
|
||||||
|
# Fallback: Manual repair loop
|
||||||
|
last_error = None
|
||||||
|
current_prompt = prompt
|
||||||
|
|
||||||
|
for attempt in range(max_retries):
|
||||||
|
try:
|
||||||
|
# 1. Generate JSON
|
||||||
|
# We use generate_structured to get the dict/list
|
||||||
|
json_result = self.generate_structured(current_prompt, max_retries=1, **kwargs)
|
||||||
|
|
||||||
|
# 2. Validate with Schema
|
||||||
|
# If the result is a list and schema expects a wrapper, or vice versa, we might need adjustment
|
||||||
|
# But we assume the prompt asks for the correct structure matching the schema.
|
||||||
|
|
||||||
|
# Special handling if schema is a wrapper but result is a list
|
||||||
|
if isinstance(json_result, list) and hasattr(schema, "entities") and "entities" in schema.model_fields:
|
||||||
|
# Auto-wrap for entities
|
||||||
|
json_result = {"entities": json_result}
|
||||||
|
|
||||||
|
# Handle categorized dictionary input (e.g. {"PERSON": ["Name"], "ORG": ["Corp"]})
|
||||||
|
elif isinstance(json_result, dict) and hasattr(schema, "entities") and "entities" in schema.model_fields:
|
||||||
|
# Check if it's NOT already in the correct format (i.e., missing "entities" key)
|
||||||
|
if "entities" not in json_result:
|
||||||
|
# Check if values are lists, suggesting categorized output
|
||||||
|
is_categorized = any(isinstance(v, list) for v in json_result.values())
|
||||||
|
if is_categorized:
|
||||||
|
flat_entities = []
|
||||||
|
for label, items in json_result.items():
|
||||||
|
if isinstance(items, list):
|
||||||
|
for item in items:
|
||||||
|
if isinstance(item, str):
|
||||||
|
flat_entities.append({"text": item, "label": label})
|
||||||
|
elif isinstance(item, dict):
|
||||||
|
# If it's already a dict but nested under label
|
||||||
|
item["label"] = label
|
||||||
|
flat_entities.append(item)
|
||||||
|
json_result = {"entities": flat_entities}
|
||||||
|
|
||||||
|
# Handle categorized dictionary input for relations (e.g. {"founded_by": [{"subject":..., "object":...}]})
|
||||||
|
elif isinstance(json_result, dict) and hasattr(schema, "relations") and "relations" in schema.model_fields:
|
||||||
|
if "relations" not in json_result:
|
||||||
|
is_categorized = any(isinstance(v, list) for v in json_result.values())
|
||||||
|
if is_categorized:
|
||||||
|
flat_relations = []
|
||||||
|
for label, items in json_result.items():
|
||||||
|
if isinstance(items, list):
|
||||||
|
for item in items:
|
||||||
|
if isinstance(item, dict):
|
||||||
|
# If predicate is missing, use the key as predicate
|
||||||
|
if "predicate" not in item:
|
||||||
|
item["predicate"] = label
|
||||||
|
flat_relations.append(item)
|
||||||
|
json_result = {"relations": flat_relations}
|
||||||
|
|
||||||
|
# Handle categorized dictionary input for triplets
|
||||||
|
elif isinstance(json_result, dict) and hasattr(schema, "triplets") and "triplets" in schema.model_fields:
|
||||||
|
if "triplets" not in json_result:
|
||||||
|
is_categorized = any(isinstance(v, list) for v in json_result.values())
|
||||||
|
if is_categorized:
|
||||||
|
flat_triplets = []
|
||||||
|
for label, items in json_result.items():
|
||||||
|
if isinstance(items, list):
|
||||||
|
for item in items:
|
||||||
|
if isinstance(item, dict):
|
||||||
|
flat_triplets.append(item)
|
||||||
|
json_result = {"triplets": flat_triplets}
|
||||||
|
|
||||||
|
elif isinstance(json_result, list) and hasattr(schema, "relations") and "relations" in schema.model_fields:
|
||||||
|
json_result = {"relations": json_result}
|
||||||
|
elif isinstance(json_result, list) and hasattr(schema, "triplets") and "triplets" in schema.model_fields:
|
||||||
|
json_result = {"triplets": json_result}
|
||||||
|
|
||||||
|
validated = schema.model_validate(json_result)
|
||||||
|
return validated
|
||||||
|
|
||||||
|
except ValidationError as e:
|
||||||
|
last_error = e
|
||||||
|
error_summary = str(e)
|
||||||
|
# Simplify error summary for the LLM
|
||||||
|
# (You could parse e.errors() for a better message)
|
||||||
|
|
||||||
|
if attempt < max_retries - 1:
|
||||||
|
wait_time = (attempt + 1) * 1
|
||||||
|
self.logger.warning(f"Schema validation failed (attempt {attempt + 1}): {e}. Retrying with error feedback...")
|
||||||
|
|
||||||
|
# Update prompt with error info
|
||||||
|
current_prompt = f"{prompt}\n\nPrevious response was invalid JSON or didn't match schema:\n{error_summary}\n\nPlease fix the errors and return valid JSON matching the schema."
|
||||||
|
time.sleep(wait_time)
|
||||||
|
else:
|
||||||
|
self.logger.error(f"Typed generation failed validation: {e}")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
last_error = e
|
||||||
|
if attempt < max_retries - 1:
|
||||||
|
time.sleep(1)
|
||||||
|
else:
|
||||||
|
self.logger.error(f"Typed generation failed: {e}")
|
||||||
|
|
||||||
|
raise ProcessingError(f"Failed to generate typed output after {max_retries} attempts: {last_error}")
|
||||||
|
|
||||||
class OpenAIProvider(BaseProvider):
|
class OpenAIProvider(BaseProvider):
|
||||||
"""OpenAI provider implementation."""
|
"""OpenAI provider implementation."""
|
||||||
|
|
||||||
@@ -333,7 +549,7 @@ class GroqProvider(BaseProvider):
|
|||||||
"""Groq provider implementation."""
|
"""Groq provider implementation."""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self, api_key: Optional[str] = None, model: str = "llama2-70b-4096", **kwargs
|
self, api_key: Optional[str] = None, model: str = "llama-3.3-70b-versatile", **kwargs
|
||||||
):
|
):
|
||||||
"""Initialize Groq provider."""
|
"""Initialize Groq provider."""
|
||||||
super().__init__(**kwargs)
|
super().__init__(**kwargs)
|
||||||
@@ -417,11 +633,16 @@ class GroqProvider(BaseProvider):
|
|||||||
if not self.client:
|
if not self.client:
|
||||||
raise ProcessingError("Groq client not initialized.")
|
raise ProcessingError("Groq client not initialized.")
|
||||||
|
|
||||||
json_prompt = f"{prompt}\n\nReturn the response as valid JSON only."
|
# Groq requires 'json' in the prompt for json_object mode
|
||||||
|
json_prompt = prompt
|
||||||
|
if "json" not in prompt.lower():
|
||||||
|
json_prompt = f"{prompt}\n\nReturn the response as valid JSON only."
|
||||||
|
|
||||||
response = self.client.chat.completions.create(
|
response = self.client.chat.completions.create(
|
||||||
model=kwargs.get("model", self.model),
|
model=kwargs.get("model", self.model),
|
||||||
messages=[{"role": "user", "content": json_prompt}],
|
messages=[{"role": "user", "content": json_prompt}],
|
||||||
temperature=kwargs.get("temperature", 0.3),
|
temperature=kwargs.get("temperature", 0.3),
|
||||||
|
response_format={"type": "json_object"},
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
return self._parse_json(response.choices[0].message.content)
|
return self._parse_json(response.choices[0].message.content)
|
||||||
|
|||||||
@@ -20,6 +20,12 @@ Algorithms Used:
|
|||||||
- Sequence Classification: Transformer-based relation classification models
|
- Sequence Classification: Transformer-based relation classification models
|
||||||
- Large Language Models: GPT, Claude, Gemini for relation extraction
|
- Large Language Models: GPT, Claude, Gemini for relation extraction
|
||||||
- Context Window Analysis: Sliding window and context extraction algorithms
|
- Context Window Analysis: Sliding window and context extraction algorithms
|
||||||
|
- Weighted Confidence Scoring:
|
||||||
|
* Formula: Score = (0.5 * Method_Confidence) + (0.5 * Type_Similarity_Score)
|
||||||
|
* Method_Confidence: Confidence score from the extraction algorithm
|
||||||
|
* Type_Similarity_Score: Semantic match with user-provided relation types (Exact=1.0, Synonym=0.95, Embedding=Cosine_Sim)
|
||||||
|
- Hybrid Similarity Matching: Exact -> Synonym -> Substring -> Semantic Embedding (Batch Optimized)
|
||||||
|
- Last Resort Fallback: Adjacency-based heuristic when all other methods fail
|
||||||
|
|
||||||
Key Features:
|
Key Features:
|
||||||
- Multiple extraction methods:
|
- Multiple extraction methods:
|
||||||
@@ -30,6 +36,7 @@ Key Features:
|
|||||||
* HuggingFace: Custom HuggingFace relation models
|
* HuggingFace: Custom HuggingFace relation models
|
||||||
* LLM-based: LLM-powered relation extraction
|
* LLM-based: LLM-powered relation extraction
|
||||||
- Fallback chain support: Try methods in order until one succeeds
|
- Fallback chain support: Try methods in order until one succeeds
|
||||||
|
- Robust Fallbacks: Prevents empty results via Primary -> Pattern -> Last Resort chain
|
||||||
- Multiple relation types (founded_by, located_in, works_for, born_in, etc.)
|
- Multiple relation types (founded_by, located_in, works_for, born_in, etc.)
|
||||||
- Relation classification and grouping
|
- Relation classification and grouping
|
||||||
- Relation validation and consistency checking
|
- Relation validation and consistency checking
|
||||||
@@ -197,6 +204,7 @@ class RelationExtractor:
|
|||||||
results = []
|
results = []
|
||||||
# Ensure lists are same length
|
# Ensure lists are same length
|
||||||
min_len = min(len(text), len(entities))
|
min_len = min(len(text), len(entities))
|
||||||
|
total_relations_count = 0
|
||||||
# Update more frequently: every 1% or at least every 10 items, but always update for small datasets
|
# Update more frequently: every 1% or at least every 10 items, but always update for small datasets
|
||||||
if min_len <= 10:
|
if min_len <= 10:
|
||||||
update_interval = 1 # Update every item for small datasets
|
update_interval = 1 # Update every item for small datasets
|
||||||
@@ -228,7 +236,18 @@ class RelationExtractor:
|
|||||||
if not isinstance(ent_item, list):
|
if not isinstance(ent_item, list):
|
||||||
ent_item = [] # Should not happen if entities is List[List[Entity]]
|
ent_item = [] # Should not happen if entities is List[List[Entity]]
|
||||||
|
|
||||||
results.append(self.extract_relations(doc_text, ent_item, **kwargs))
|
current_relations = self.extract_relations(doc_text, ent_item, **kwargs)
|
||||||
|
|
||||||
|
# Add provenance metadata
|
||||||
|
for rel in current_relations:
|
||||||
|
if rel.metadata is None:
|
||||||
|
rel.metadata = {}
|
||||||
|
rel.metadata["batch_index"] = i
|
||||||
|
if isinstance(doc_item, dict) and "id" in doc_item:
|
||||||
|
rel.metadata["document_id"] = doc_item["id"]
|
||||||
|
|
||||||
|
results.append(current_relations)
|
||||||
|
total_relations_count += len(current_relations)
|
||||||
|
|
||||||
remaining = min_len - (i + 1)
|
remaining = min_len - (i + 1)
|
||||||
# Update progress: always update for small datasets, or at intervals for large ones
|
# Update progress: always update for small datasets, or at intervals for large ones
|
||||||
@@ -243,13 +262,13 @@ class RelationExtractor:
|
|||||||
tracking_id,
|
tracking_id,
|
||||||
processed=i + 1,
|
processed=i + 1,
|
||||||
total=min_len,
|
total=min_len,
|
||||||
message=f"Processing documents... {i + 1}/{min_len} (remaining: {remaining})"
|
message=f"Processing documents... {i + 1}/{min_len} (remaining: {remaining}) - Extracted {total_relations_count} relations so far"
|
||||||
)
|
)
|
||||||
|
|
||||||
self.progress_tracker.stop_tracking(
|
self.progress_tracker.stop_tracking(
|
||||||
tracking_id,
|
tracking_id,
|
||||||
status="completed",
|
status="completed",
|
||||||
message=f"Extracted relations from {len(results)} documents",
|
message=f"Batch extraction completed. Processed {len(results)} documents, extracted {total_relations_count} relations.",
|
||||||
)
|
)
|
||||||
return results
|
return results
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -322,6 +341,11 @@ class RelationExtractor:
|
|||||||
|
|
||||||
# Prepare method-specific options
|
# Prepare method-specific options
|
||||||
method_options = all_options.copy()
|
method_options = all_options.copy()
|
||||||
|
|
||||||
|
# Pass relation_types to all methods so they can use them (e.g. for similarity matching)
|
||||||
|
if relation_types:
|
||||||
|
method_options["relation_types"] = relation_types
|
||||||
|
|
||||||
if method_name == "huggingface":
|
if method_name == "huggingface":
|
||||||
method_options["model"] = all_options.get(
|
method_options["model"] = all_options.get(
|
||||||
"huggingface_model", all_options.get("model")
|
"huggingface_model", all_options.get("model")
|
||||||
@@ -345,9 +369,6 @@ class RelationExtractor:
|
|||||||
api_key = os.getenv(env_key)
|
api_key = os.getenv(env_key)
|
||||||
if api_key:
|
if api_key:
|
||||||
method_options["api_key"] = api_key
|
method_options["api_key"] = api_key
|
||||||
# Pass relation_types to LLM method so it can use them in the prompt
|
|
||||||
if relation_types:
|
|
||||||
method_options["relation_types"] = relation_types
|
|
||||||
elif method_name == "dependency":
|
elif method_name == "dependency":
|
||||||
method_options["model"] = all_options.get(
|
method_options["model"] = all_options.get(
|
||||||
"model", "en_core_web_sm"
|
"model", "en_core_web_sm"
|
||||||
@@ -366,6 +387,20 @@ class RelationExtractor:
|
|||||||
import sys
|
import sys
|
||||||
print(f" [RelationExtractor] Extracted {len(relations)} relations", flush=True, file=sys.stdout)
|
print(f" [RelationExtractor] Extracted {len(relations)} relations", flush=True, file=sys.stdout)
|
||||||
|
|
||||||
|
# Apply weighted scoring if relation_types are provided
|
||||||
|
if relation_types:
|
||||||
|
try:
|
||||||
|
from .methods import calculate_weighted_confidence
|
||||||
|
for r in relations:
|
||||||
|
r.confidence = calculate_weighted_confidence(
|
||||||
|
item_type=r.predicate,
|
||||||
|
original_confidence=r.confidence,
|
||||||
|
valid_types=relation_types,
|
||||||
|
item_text=r.predicate # For relations, the predicate IS the text usually
|
||||||
|
)
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
|
||||||
# Filter by confidence
|
# Filter by confidence
|
||||||
filtered = [r for r in relations if r.confidence >= min_confidence]
|
filtered = [r for r in relations if r.confidence >= min_confidence]
|
||||||
|
|
||||||
@@ -392,7 +427,12 @@ class RelationExtractor:
|
|||||||
if all_relations:
|
if all_relations:
|
||||||
relations = all_relations[0][1] # Use first successful method
|
relations = all_relations[0][1] # Use first successful method
|
||||||
else:
|
else:
|
||||||
relations = []
|
# Fallback to pattern-based extraction if all models fail
|
||||||
|
relations = self._extract_with_patterns(text, entities)
|
||||||
|
|
||||||
|
# Last resort: if patterns also fail but we have entities, force some relations
|
||||||
|
if not relations and entities and len(entities) >= 2:
|
||||||
|
relations = self._extract_last_resort_relations(text, entities)
|
||||||
|
|
||||||
# Validate if enabled
|
# Validate if enabled
|
||||||
if validate:
|
if validate:
|
||||||
@@ -411,6 +451,38 @@ class RelationExtractor:
|
|||||||
)
|
)
|
||||||
raise
|
raise
|
||||||
|
|
||||||
|
def _extract_last_resort_relations(self, text: str, entities: List[Entity]) -> List[Relation]:
|
||||||
|
"""Last resort relation extraction based on simple adjacency."""
|
||||||
|
relations = []
|
||||||
|
# Connect adjacent entities
|
||||||
|
for i in range(len(entities) - 1):
|
||||||
|
e1 = entities[i]
|
||||||
|
e2 = entities[i+1]
|
||||||
|
|
||||||
|
# Create a weak relation
|
||||||
|
start_idx = min(e1.end_char, e2.start_char)
|
||||||
|
end_idx = max(e1.end_char, e2.start_char)
|
||||||
|
# Ensure context isn't too large or invalid
|
||||||
|
if start_idx < 0: start_idx = 0
|
||||||
|
if end_idx > len(text): end_idx = len(text)
|
||||||
|
|
||||||
|
# Expand context a bit
|
||||||
|
ctx_start = max(0, start_idx - 20)
|
||||||
|
ctx_end = min(len(text), end_idx + 20)
|
||||||
|
|
||||||
|
context = text[ctx_start:ctx_end]
|
||||||
|
|
||||||
|
rel = Relation(
|
||||||
|
subject=e1,
|
||||||
|
predicate="related_to",
|
||||||
|
object=e2,
|
||||||
|
confidence=0.3,
|
||||||
|
context=context,
|
||||||
|
metadata={"extraction_method": "last_resort_adjacency"}
|
||||||
|
)
|
||||||
|
relations.append(rel)
|
||||||
|
return relations
|
||||||
|
|
||||||
def _extract_with_patterns(
|
def _extract_with_patterns(
|
||||||
self, text: str, entities: List[Entity]
|
self, text: str, entities: List[Entity]
|
||||||
) -> List[Relation]:
|
) -> List[Relation]:
|
||||||
|
|||||||
@@ -0,0 +1,127 @@
|
|||||||
|
from typing import List, Optional
|
||||||
|
from pydantic import BaseModel, Field, field_validator, model_validator, ConfigDict
|
||||||
|
|
||||||
|
class EntityOut(BaseModel):
|
||||||
|
"""Canonical schema for entity extraction output."""
|
||||||
|
model_config = ConfigDict(populate_by_name=True, extra="ignore")
|
||||||
|
|
||||||
|
text: str = Field(..., description="The text content of the entity")
|
||||||
|
label: str = Field(..., description="The type or label of the entity (e.g., PERSON, ORG)")
|
||||||
|
start: int = Field(0, description="Start character index", alias="start_char")
|
||||||
|
end: int = Field(0, description="End character index", alias="end_char")
|
||||||
|
confidence: float = Field(0.9, description="Confidence score between 0 and 1")
|
||||||
|
metadata: dict = Field(default_factory=dict, description="Additional metadata including provenance")
|
||||||
|
|
||||||
|
@field_validator("text", mode="before")
|
||||||
|
@classmethod
|
||||||
|
def clean_text(cls, v):
|
||||||
|
if isinstance(v, str):
|
||||||
|
return v.strip()
|
||||||
|
return str(v)
|
||||||
|
|
||||||
|
@field_validator("confidence", mode="before")
|
||||||
|
@classmethod
|
||||||
|
def normalize_confidence(cls, v):
|
||||||
|
if isinstance(v, str):
|
||||||
|
try:
|
||||||
|
v = float(v)
|
||||||
|
except ValueError:
|
||||||
|
return 0.9
|
||||||
|
if isinstance(v, (int, float)):
|
||||||
|
return max(0.0, min(1.0, float(v)))
|
||||||
|
return 0.9
|
||||||
|
|
||||||
|
@model_validator(mode="before")
|
||||||
|
@classmethod
|
||||||
|
def handle_aliases(cls, data):
|
||||||
|
if isinstance(data, dict):
|
||||||
|
# Handle 'type' as alias for 'label'
|
||||||
|
if "label" not in data and "type" in data:
|
||||||
|
data["label"] = data["type"]
|
||||||
|
# Handle 'value' or 'span' as alias for 'text'
|
||||||
|
if "text" not in data:
|
||||||
|
if "value" in data:
|
||||||
|
data["text"] = data["value"]
|
||||||
|
elif "span" in data:
|
||||||
|
data["text"] = data["span"]
|
||||||
|
return data
|
||||||
|
|
||||||
|
class RelationOut(BaseModel):
|
||||||
|
"""Canonical schema for relation extraction output."""
|
||||||
|
model_config = ConfigDict(populate_by_name=True, extra="ignore")
|
||||||
|
|
||||||
|
subject: str = Field(..., description="Source entity text")
|
||||||
|
object: str = Field(..., description="Target entity text")
|
||||||
|
predicate: str = Field(..., description="Relation type or predicate")
|
||||||
|
confidence: float = Field(0.9, description="Confidence score between 0 and 1")
|
||||||
|
metadata: dict = Field(default_factory=dict, description="Additional metadata including provenance")
|
||||||
|
|
||||||
|
@model_validator(mode="before")
|
||||||
|
@classmethod
|
||||||
|
def handle_aliases(cls, data):
|
||||||
|
if isinstance(data, dict):
|
||||||
|
if "subject" not in data and "source" in data:
|
||||||
|
data["subject"] = data["source"]
|
||||||
|
if "object" not in data and "target" in data:
|
||||||
|
data["object"] = data["target"]
|
||||||
|
if "predicate" not in data and "label" in data:
|
||||||
|
data["predicate"] = data["label"]
|
||||||
|
return data
|
||||||
|
|
||||||
|
@property
|
||||||
|
def source(self):
|
||||||
|
return self.subject
|
||||||
|
|
||||||
|
@property
|
||||||
|
def target(self):
|
||||||
|
return self.object
|
||||||
|
|
||||||
|
@property
|
||||||
|
def label(self):
|
||||||
|
return self.predicate
|
||||||
|
|
||||||
|
@field_validator("confidence", mode="before")
|
||||||
|
@classmethod
|
||||||
|
def normalize_confidence(cls, v):
|
||||||
|
if isinstance(v, str):
|
||||||
|
try:
|
||||||
|
v = float(v)
|
||||||
|
except ValueError:
|
||||||
|
return 0.9
|
||||||
|
if isinstance(v, (int, float)):
|
||||||
|
return max(0.0, min(1.0, float(v)))
|
||||||
|
return 0.9
|
||||||
|
|
||||||
|
class TripletOut(BaseModel):
|
||||||
|
"""Canonical schema for triplet extraction output."""
|
||||||
|
model_config = ConfigDict(populate_by_name=True, extra="ignore")
|
||||||
|
|
||||||
|
subject: str = Field(..., description="Subject of the triplet")
|
||||||
|
predicate: str = Field(..., description="Predicate or relation")
|
||||||
|
object: str = Field(..., description="Object of the triplet")
|
||||||
|
confidence: float = Field(0.9, description="Confidence score between 0 and 1")
|
||||||
|
metadata: dict = Field(default_factory=dict, description="Additional metadata including provenance")
|
||||||
|
|
||||||
|
@field_validator("confidence", mode="before")
|
||||||
|
@classmethod
|
||||||
|
def normalize_confidence(cls, v):
|
||||||
|
if isinstance(v, str):
|
||||||
|
try:
|
||||||
|
v = float(v)
|
||||||
|
except ValueError:
|
||||||
|
return 0.9
|
||||||
|
if isinstance(v, (int, float)):
|
||||||
|
return max(0.0, min(1.0, float(v)))
|
||||||
|
return 0.9
|
||||||
|
|
||||||
|
class EntitiesResponse(BaseModel):
|
||||||
|
"""Wrapper for list of entities."""
|
||||||
|
entities: List[EntityOut] = Field(default_factory=list)
|
||||||
|
|
||||||
|
class RelationsResponse(BaseModel):
|
||||||
|
"""Wrapper for list of relations."""
|
||||||
|
relations: List[RelationOut] = Field(default_factory=list)
|
||||||
|
|
||||||
|
class TripletsResponse(BaseModel):
|
||||||
|
"""Wrapper for list of triplets."""
|
||||||
|
triplets: List[TripletOut] = Field(default_factory=list)
|
||||||
@@ -71,6 +71,7 @@ class SemanticRole:
|
|||||||
start_char: int
|
start_char: int
|
||||||
end_char: int
|
end_char: int
|
||||||
confidence: float = 1.0
|
confidence: float = 1.0
|
||||||
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
@@ -81,6 +82,7 @@ class SemanticCluster:
|
|||||||
cluster_id: int
|
cluster_id: int
|
||||||
centroid: Optional[str] = None
|
centroid: Optional[str] = None
|
||||||
similarity_score: float = 0.0
|
similarity_score: float = 0.0
|
||||||
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
class SemanticAnalyzer:
|
class SemanticAnalyzer:
|
||||||
@@ -118,6 +120,103 @@ class SemanticAnalyzer:
|
|||||||
self.role_labeler = RoleLabeler(**self.config.get("role", {}))
|
self.role_labeler = RoleLabeler(**self.config.get("role", {}))
|
||||||
self.semantic_clusterer = SemanticClusterer(**self.config.get("clustering", {}))
|
self.semantic_clusterer = SemanticClusterer(**self.config.get("clustering", {}))
|
||||||
|
|
||||||
|
def analyze(
|
||||||
|
self,
|
||||||
|
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||||
|
pipeline_id: Optional[str] = None,
|
||||||
|
**kwargs
|
||||||
|
) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
|
||||||
|
"""
|
||||||
|
Perform semantic analysis on text or list of documents.
|
||||||
|
Handles batch processing with progress tracking.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: Input text or list of documents
|
||||||
|
pipeline_id: Optional pipeline ID for progress tracking
|
||||||
|
**kwargs: Analysis options
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Union[Dict[str, Any], List[Dict[str, Any]]]: Analysis results
|
||||||
|
"""
|
||||||
|
if isinstance(text, list):
|
||||||
|
# Handle batch analysis with progress tracking
|
||||||
|
tracking_id = self.progress_tracker.start_tracking(
|
||||||
|
module="semantic_extract",
|
||||||
|
submodule="SemanticAnalyzer",
|
||||||
|
message=f"Batch analyzing {len(text)} documents",
|
||||||
|
pipeline_id=pipeline_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
results = []
|
||||||
|
total_items = len(text)
|
||||||
|
|
||||||
|
# Determine update interval
|
||||||
|
if total_items <= 10:
|
||||||
|
update_interval = 1
|
||||||
|
else:
|
||||||
|
update_interval = max(1, min(10, total_items // 100))
|
||||||
|
|
||||||
|
# Initial progress update
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=0,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Starting batch analysis... 0/{total_items} (remaining: {total_items})"
|
||||||
|
)
|
||||||
|
|
||||||
|
for idx, item in enumerate(text):
|
||||||
|
# Prepare arguments for single item
|
||||||
|
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||||
|
|
||||||
|
# Analyze
|
||||||
|
analysis = self.analyze_semantics(doc_text, **kwargs)
|
||||||
|
|
||||||
|
# Add provenance metadata
|
||||||
|
analysis["batch_index"] = idx
|
||||||
|
if isinstance(item, dict) and "id" in item:
|
||||||
|
analysis["document_id"] = item["id"]
|
||||||
|
|
||||||
|
# Also inject into semantic roles if present
|
||||||
|
if "semantic_roles" in analysis:
|
||||||
|
for role in analysis["semantic_roles"]:
|
||||||
|
# role is a dict here because analyze_semantics converts it
|
||||||
|
if "metadata" not in role:
|
||||||
|
role["metadata"] = {}
|
||||||
|
|
||||||
|
role["metadata"]["batch_index"] = idx
|
||||||
|
if isinstance(item, dict) and "id" in item:
|
||||||
|
role["metadata"]["document_id"] = item["id"]
|
||||||
|
|
||||||
|
results.append(analysis)
|
||||||
|
|
||||||
|
# Update progress
|
||||||
|
if (idx + 1) % update_interval == 0 or (idx + 1) == total_items:
|
||||||
|
remaining = total_items - (idx + 1)
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=idx + 1,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Processing... {idx + 1}/{total_items} (remaining: {remaining})"
|
||||||
|
)
|
||||||
|
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id,
|
||||||
|
status="completed",
|
||||||
|
message=f"Batch analysis completed. Processed {len(results)} documents.",
|
||||||
|
)
|
||||||
|
return results
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id, status="failed", message=str(e)
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
else:
|
||||||
|
# Single item
|
||||||
|
return self.analyze_semantics(text, **kwargs)
|
||||||
|
|
||||||
def analyze_semantics(self, text: str, **options) -> Dict[str, Any]:
|
def analyze_semantics(self, text: str, **options) -> Dict[str, Any]:
|
||||||
"""
|
"""
|
||||||
Perform comprehensive semantic analysis.
|
Perform comprehensive semantic analysis.
|
||||||
@@ -200,13 +299,13 @@ class SemanticAnalyzer:
|
|||||||
return self.label_semantic_roles(text, **options)
|
return self.label_semantic_roles(text, **options)
|
||||||
|
|
||||||
def cluster_semantically(
|
def cluster_semantically(
|
||||||
self, texts: List[str], **options
|
self, texts: Union[List[str], List[Dict[str, Any]]], **options
|
||||||
) -> List[SemanticCluster]:
|
) -> List[SemanticCluster]:
|
||||||
"""
|
"""
|
||||||
Perform semantic clustering of texts.
|
Perform semantic clustering of texts.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
texts: List of texts to cluster
|
texts: List of texts or documents to cluster
|
||||||
**options: Clustering options
|
**options: Clustering options
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
@@ -372,12 +471,14 @@ class SemanticClusterer:
|
|||||||
if not self.progress_tracker.enabled:
|
if not self.progress_tracker.enabled:
|
||||||
self.progress_tracker.enabled = True
|
self.progress_tracker.enabled = True
|
||||||
|
|
||||||
def cluster(self, texts: List[str], **options) -> List[SemanticCluster]:
|
def cluster(
|
||||||
|
self, texts: Union[List[str], List[Dict[str, Any]]], **options
|
||||||
|
) -> List[SemanticCluster]:
|
||||||
"""
|
"""
|
||||||
Perform semantic clustering of texts.
|
Perform semantic clustering of texts.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
texts: List of texts to cluster
|
texts: List of texts or documents (dict with 'content' and 'id') to cluster
|
||||||
**options: Clustering options:
|
**options: Clustering options:
|
||||||
- num_clusters: Number of clusters (default: auto)
|
- num_clusters: Number of clusters (default: auto)
|
||||||
- similarity_threshold: Minimum similarity for clustering
|
- similarity_threshold: Minimum similarity for clustering
|
||||||
@@ -388,11 +489,27 @@ class SemanticClusterer:
|
|||||||
if not texts:
|
if not texts:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
# Extract content and IDs if input is list of dicts
|
||||||
|
processed_texts = []
|
||||||
|
doc_ids = []
|
||||||
|
|
||||||
|
for item in texts:
|
||||||
|
if isinstance(item, dict):
|
||||||
|
content = item.get("content", str(item))
|
||||||
|
processed_texts.append(content)
|
||||||
|
if "id" in item:
|
||||||
|
doc_ids.append(item["id"])
|
||||||
|
else:
|
||||||
|
doc_ids.append(None)
|
||||||
|
else:
|
||||||
|
processed_texts.append(str(item))
|
||||||
|
doc_ids.append(None)
|
||||||
|
|
||||||
# Track clustering
|
# Track clustering
|
||||||
tracking_id = self.progress_tracker.start_tracking(
|
tracking_id = self.progress_tracker.start_tracking(
|
||||||
module="semantic_extract",
|
module="semantic_extract",
|
||||||
submodule="SemanticClusterer",
|
submodule="SemanticClusterer",
|
||||||
message=f"Clustering {len(texts)} texts",
|
message=f"Clustering {len(processed_texts)} texts",
|
||||||
)
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -402,7 +519,7 @@ class SemanticClusterer:
|
|||||||
clusters = []
|
clusters = []
|
||||||
assigned = set()
|
assigned = set()
|
||||||
|
|
||||||
total_texts = len(texts)
|
total_texts = len(processed_texts)
|
||||||
if total_texts <= 10:
|
if total_texts <= 10:
|
||||||
update_interval = 1 # Update every item for small datasets
|
update_interval = 1 # Update every item for small datasets
|
||||||
else:
|
else:
|
||||||
@@ -418,22 +535,28 @@ class SemanticClusterer:
|
|||||||
)
|
)
|
||||||
|
|
||||||
cluster_id = 0
|
cluster_id = 0
|
||||||
for i, text1 in enumerate(texts):
|
for i, text1 in enumerate(processed_texts):
|
||||||
if i in assigned:
|
if i in assigned:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
cluster_texts = [text1]
|
cluster_texts = [text1]
|
||||||
|
cluster_doc_ids = []
|
||||||
|
if doc_ids[i] is not None:
|
||||||
|
cluster_doc_ids.append(doc_ids[i])
|
||||||
|
|
||||||
assigned.add(i)
|
assigned.add(i)
|
||||||
|
|
||||||
# Find similar texts
|
# Find similar texts
|
||||||
remaining_texts = len(texts) - (i + 1)
|
remaining_texts = len(processed_texts) - (i + 1)
|
||||||
for j, text2 in enumerate(texts[i + 1 :], start=i + 1):
|
for j, text2 in enumerate(processed_texts[i + 1 :], start=i + 1):
|
||||||
if j in assigned:
|
if j in assigned:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
similarity = similarity_analyzer.calculate_similarity(text1, text2)
|
similarity = similarity_analyzer.calculate_similarity(text1, text2)
|
||||||
if similarity >= similarity_threshold:
|
if similarity >= similarity_threshold:
|
||||||
cluster_texts.append(text2)
|
cluster_texts.append(text2)
|
||||||
|
if doc_ids[j] is not None:
|
||||||
|
cluster_doc_ids.append(doc_ids[j])
|
||||||
assigned.add(j)
|
assigned.add(j)
|
||||||
|
|
||||||
# Create cluster
|
# Create cluster
|
||||||
@@ -443,6 +566,11 @@ class SemanticClusterer:
|
|||||||
centroid=cluster_texts[0], # Use first as centroid
|
centroid=cluster_texts[0], # Use first as centroid
|
||||||
similarity_score=similarity_threshold,
|
similarity_score=similarity_threshold,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Add provenance metadata
|
||||||
|
if cluster_doc_ids:
|
||||||
|
cluster.metadata["document_ids"] = cluster_doc_ids
|
||||||
|
|
||||||
clusters.append(cluster)
|
clusters.append(cluster)
|
||||||
cluster_id += 1
|
cluster_id += 1
|
||||||
|
|
||||||
|
|||||||
@@ -36,6 +36,41 @@ print(f"Relations: {relations}")
|
|||||||
print(f"Extracted {len(entities)} entities and {len(relations)} relations")
|
print(f"Extracted {len(entities)} entities and {len(relations)} relations")
|
||||||
```
|
```
|
||||||
|
|
||||||
|
## Batch Processing & Provenance
|
||||||
|
|
||||||
|
All extractors support batch processing for high-throughput extraction. You can pass a list of strings or a list of dictionaries (with `content` and `id` keys).
|
||||||
|
|
||||||
|
**Features:**
|
||||||
|
- **Progress Tracking**: Automatically shows a progress bar for large batches.
|
||||||
|
- **Provenance Metadata**: Each extracted item includes `batch_index` and `document_id` in its `metadata`.
|
||||||
|
|
||||||
|
```python
|
||||||
|
from semantica.semantic_extract import NERExtractor
|
||||||
|
|
||||||
|
documents = [
|
||||||
|
{"id": "doc_1", "content": "Apple Inc. was founded by Steve Jobs."},
|
||||||
|
{"id": "doc_2", "content": "Microsoft Corporation was founded by Bill Gates."}
|
||||||
|
]
|
||||||
|
|
||||||
|
extractor = NERExtractor()
|
||||||
|
batch_results = extractor.extract(documents)
|
||||||
|
|
||||||
|
for i, doc_entities in enumerate(batch_results):
|
||||||
|
print(f"Document {i} entities:")
|
||||||
|
for entity in doc_entities:
|
||||||
|
print(f" - {entity.text} ({entity.label})")
|
||||||
|
print(f" Provenance: Batch Index {entity.metadata['batch_index']}, Doc ID {entity.metadata.get('document_id')}")
|
||||||
|
```
|
||||||
|
|
||||||
|
## Robust Extraction Fallbacks
|
||||||
|
|
||||||
|
The framework implements robust fallback chains to prevent empty results when primary methods fail (e.g., due to model unavailability or obscure text).
|
||||||
|
|
||||||
|
- **NER**: `ML/LLM` -> `Pattern` -> `Last Resort` (Capitalized Words)
|
||||||
|
- **Relation**: `Primary` -> `Pattern` -> `Last Resort` (Adjacency)
|
||||||
|
- **Triplet**: `Primary` -> `Relation-to-Triplet` -> `Pattern`
|
||||||
|
|
||||||
|
This ensures that you almost always get *some* structured data, even if it requires falling back to simpler heuristics.
|
||||||
|
|
||||||
## Entity Extraction
|
## Entity Extraction
|
||||||
|
|
||||||
|
|||||||
@@ -149,6 +149,120 @@ class SemanticNetworkExtractor:
|
|||||||
self.config["ner_method"] = method
|
self.config["ner_method"] = method
|
||||||
self.config["relation_method"] = method
|
self.config["relation_method"] = method
|
||||||
|
|
||||||
|
def extract(
|
||||||
|
self,
|
||||||
|
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||||
|
entities: Optional[Union[List[Entity], List[List[Entity]]]] = None,
|
||||||
|
relations: Optional[Union[List[Relation], List[List[Relation]]]] = None,
|
||||||
|
pipeline_id: Optional[str] = None,
|
||||||
|
**kwargs
|
||||||
|
) -> Union[SemanticNetwork, List[SemanticNetwork]]:
|
||||||
|
"""
|
||||||
|
Extract semantic network from text or list of documents.
|
||||||
|
Handles batch processing with progress tracking.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: Input text or list of documents
|
||||||
|
entities: Optional pre-extracted entities (single list or list of lists)
|
||||||
|
relations: Optional pre-extracted relations (single list or list of lists)
|
||||||
|
pipeline_id: Optional pipeline ID for progress tracking
|
||||||
|
**kwargs: Extraction options
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Union[SemanticNetwork, List[SemanticNetwork]]: Extracted semantic network(s)
|
||||||
|
"""
|
||||||
|
if isinstance(text, list):
|
||||||
|
# Handle batch extraction with progress tracking
|
||||||
|
tracking_id = self.progress_tracker.start_tracking(
|
||||||
|
module="semantic_extract",
|
||||||
|
submodule="SemanticNetworkExtractor",
|
||||||
|
message=f"Batch extracting semantic networks from {len(text)} documents",
|
||||||
|
pipeline_id=pipeline_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
results = []
|
||||||
|
total_items = len(text)
|
||||||
|
|
||||||
|
# Determine update interval
|
||||||
|
if total_items <= 10:
|
||||||
|
update_interval = 1
|
||||||
|
else:
|
||||||
|
update_interval = max(1, min(10, total_items // 100))
|
||||||
|
|
||||||
|
# Initial progress update
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=0,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Starting batch extraction... 0/{total_items} (remaining: {total_items})"
|
||||||
|
)
|
||||||
|
|
||||||
|
for idx, item in enumerate(text):
|
||||||
|
# Prepare arguments for single item
|
||||||
|
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||||
|
|
||||||
|
doc_entities = None
|
||||||
|
if entities and isinstance(entities, list) and idx < len(entities):
|
||||||
|
doc_entities = entities[idx]
|
||||||
|
|
||||||
|
doc_relations = None
|
||||||
|
if relations and isinstance(relations, list) and idx < len(relations):
|
||||||
|
doc_relations = relations[idx]
|
||||||
|
|
||||||
|
# Extract
|
||||||
|
network = self.extract_network(
|
||||||
|
doc_text,
|
||||||
|
entities=doc_entities,
|
||||||
|
relations=doc_relations,
|
||||||
|
**kwargs
|
||||||
|
)
|
||||||
|
|
||||||
|
# Add provenance metadata to nodes and edges
|
||||||
|
batch_meta = {"batch_index": idx}
|
||||||
|
if isinstance(item, dict) and "id" in item:
|
||||||
|
batch_meta["document_id"] = item["id"]
|
||||||
|
|
||||||
|
# Update network metadata
|
||||||
|
network.metadata.update(batch_meta)
|
||||||
|
|
||||||
|
# Update nodes metadata
|
||||||
|
for node in network.nodes:
|
||||||
|
node.metadata.update(batch_meta)
|
||||||
|
|
||||||
|
# Update edges metadata
|
||||||
|
for edge in network.edges:
|
||||||
|
edge.metadata.update(batch_meta)
|
||||||
|
|
||||||
|
results.append(network)
|
||||||
|
|
||||||
|
# Update progress
|
||||||
|
if (idx + 1) % update_interval == 0 or (idx + 1) == total_items:
|
||||||
|
remaining = total_items - (idx + 1)
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=idx + 1,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Processing... {idx + 1}/{total_items} (remaining: {remaining})"
|
||||||
|
)
|
||||||
|
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id,
|
||||||
|
status="completed",
|
||||||
|
message=f"Batch extraction completed. Processed {len(results)} documents.",
|
||||||
|
)
|
||||||
|
return results
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id, status="failed", message=str(e)
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
else:
|
||||||
|
# Single item
|
||||||
|
return self.extract_network(text, entities=entities, relations=relations, **kwargs)
|
||||||
|
|
||||||
def extract_network(
|
def extract_network(
|
||||||
self,
|
self,
|
||||||
text: str,
|
text: str,
|
||||||
|
|||||||
@@ -18,6 +18,12 @@ Algorithms Used:
|
|||||||
- Large Language Models: GPT, Claude, Gemini for structured triplet extraction
|
- Large Language Models: GPT, Claude, Gemini for structured triplet extraction
|
||||||
- RDF Serialization: Graph serialization algorithms (Turtle, N-Triples, JSON-LD)
|
- RDF Serialization: Graph serialization algorithms (Turtle, N-Triples, JSON-LD)
|
||||||
- URI Normalization: String normalization and URI formatting algorithms
|
- URI Normalization: String normalization and URI formatting algorithms
|
||||||
|
- Weighted Confidence Scoring:
|
||||||
|
* Formula: Score = (0.5 * Method_Confidence) + (0.5 * Type_Similarity_Score)
|
||||||
|
* Method_Confidence: Confidence score from the extraction algorithm
|
||||||
|
* Type_Similarity_Score: Semantic match with user-provided triplet types (Exact=1.0, Synonym=0.95, Embedding=Cosine_Sim)
|
||||||
|
- Hybrid Similarity Matching: Exact -> Synonym -> Substring -> Semantic Embedding (Batch Optimized)
|
||||||
|
- Last Resort Fallback: Relation-to-Triplet conversion when all other methods fail
|
||||||
|
|
||||||
Key Features:
|
Key Features:
|
||||||
- Multiple extraction methods:
|
- Multiple extraction methods:
|
||||||
@@ -26,6 +32,7 @@ Key Features:
|
|||||||
* HuggingFace: Custom HuggingFace triplet models
|
* HuggingFace: Custom HuggingFace triplet models
|
||||||
* LLM-based: LLM-powered triplet extraction
|
* LLM-based: LLM-powered triplet extraction
|
||||||
- Fallback chain support: Try methods in order until one succeeds
|
- Fallback chain support: Try methods in order until one succeeds
|
||||||
|
- Robust Fallbacks: Prevents empty results via Primary -> Relation-to-Triplet -> Pattern chain
|
||||||
- RDF triplet generation from entities and relations
|
- RDF triplet generation from entities and relations
|
||||||
- Subject-predicate-object extraction
|
- Subject-predicate-object extraction
|
||||||
- Triplet validation and quality checking
|
- Triplet validation and quality checking
|
||||||
@@ -99,6 +106,7 @@ class TripletExtractor:
|
|||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
method: Union[str, List[str]] = "pattern",
|
method: Union[str, List[str]] = "pattern",
|
||||||
|
triplet_types: Optional[List[str]] = None,
|
||||||
include_temporal: bool = False,
|
include_temporal: bool = False,
|
||||||
include_provenance: bool = False,
|
include_provenance: bool = False,
|
||||||
config=None,
|
config=None,
|
||||||
@@ -114,6 +122,7 @@ class TripletExtractor:
|
|||||||
- "huggingface": HuggingFace model
|
- "huggingface": HuggingFace model
|
||||||
- "llm": LLM-based extraction
|
- "llm": LLM-based extraction
|
||||||
- List of methods for fallback chain
|
- List of methods for fallback chain
|
||||||
|
triplet_types: Specific triplet types/predicates to extract (e.g., ["foundedBy", "locatedIn"])
|
||||||
include_temporal: Whether to include temporal information in triplets
|
include_temporal: Whether to include temporal information in triplets
|
||||||
include_provenance: Whether to track source sentences for provenance
|
include_provenance: Whether to track source sentences for provenance
|
||||||
config: Legacy config dict (deprecated, use kwargs)
|
config: Legacy config dict (deprecated, use kwargs)
|
||||||
@@ -135,6 +144,7 @@ class TripletExtractor:
|
|||||||
self.progress_tracker.enabled = True
|
self.progress_tracker.enabled = True
|
||||||
|
|
||||||
# Store parameters
|
# Store parameters
|
||||||
|
self.triplet_types = triplet_types
|
||||||
self.include_temporal = include_temporal
|
self.include_temporal = include_temporal
|
||||||
self.include_provenance = include_provenance
|
self.include_provenance = include_provenance
|
||||||
|
|
||||||
@@ -149,6 +159,114 @@ class TripletExtractor:
|
|||||||
|
|
||||||
self.supported_formats = ["turtle", "ntriples", "jsonld", "xml"]
|
self.supported_formats = ["turtle", "ntriples", "jsonld", "xml"]
|
||||||
|
|
||||||
|
def extract(
|
||||||
|
self,
|
||||||
|
text: Union[str, List[str], List[Dict[str, Any]]],
|
||||||
|
entities: Optional[Union[List[Entity], List[List[Entity]]]] = None,
|
||||||
|
relations: Optional[Union[List[Relation], List[List[Relation]]]] = None,
|
||||||
|
pipeline_id: Optional[str] = None,
|
||||||
|
**kwargs
|
||||||
|
) -> Union[List[Triplet], List[List[Triplet]]]:
|
||||||
|
"""
|
||||||
|
Extract triplets from text or list of documents.
|
||||||
|
Handles batch processing with progress tracking.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: Input text or list of documents
|
||||||
|
entities: Optional pre-extracted entities (single list or list of lists)
|
||||||
|
relations: Optional pre-extracted relations (single list or list of lists)
|
||||||
|
pipeline_id: Optional pipeline ID for progress tracking
|
||||||
|
**kwargs: Extraction options
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Union[List[Triplet], List[List[Triplet]]]: Extracted triplets
|
||||||
|
"""
|
||||||
|
if isinstance(text, list):
|
||||||
|
# Handle batch extraction with progress tracking
|
||||||
|
tracking_id = self.progress_tracker.start_tracking(
|
||||||
|
module="semantic_extract",
|
||||||
|
submodule="TripletExtractor",
|
||||||
|
message=f"Batch extracting triplets from {len(text)} documents",
|
||||||
|
pipeline_id=pipeline_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
results = []
|
||||||
|
total_items = len(text)
|
||||||
|
total_triplets_count = 0
|
||||||
|
|
||||||
|
# Determine update interval
|
||||||
|
if total_items <= 10:
|
||||||
|
update_interval = 1
|
||||||
|
else:
|
||||||
|
update_interval = max(1, min(10, total_items // 100))
|
||||||
|
|
||||||
|
# Initial progress update
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=0,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Starting batch extraction... 0/{total_items} (remaining: {total_items})"
|
||||||
|
)
|
||||||
|
|
||||||
|
for idx, item in enumerate(text):
|
||||||
|
# Prepare arguments for single item
|
||||||
|
doc_text = item["content"] if isinstance(item, dict) and "content" in item else str(item)
|
||||||
|
|
||||||
|
doc_entities = None
|
||||||
|
if entities and isinstance(entities, list) and idx < len(entities):
|
||||||
|
doc_entities = entities[idx]
|
||||||
|
|
||||||
|
doc_relations = None
|
||||||
|
if relations and isinstance(relations, list) and idx < len(relations):
|
||||||
|
doc_relations = relations[idx]
|
||||||
|
|
||||||
|
# Extract
|
||||||
|
current_triplets = self.extract_triplets(
|
||||||
|
doc_text,
|
||||||
|
entities=doc_entities,
|
||||||
|
relations=doc_relations,
|
||||||
|
**kwargs
|
||||||
|
)
|
||||||
|
|
||||||
|
# Add provenance metadata
|
||||||
|
for triplet in current_triplets:
|
||||||
|
if triplet.metadata is None:
|
||||||
|
triplet.metadata = {}
|
||||||
|
triplet.metadata["batch_index"] = idx
|
||||||
|
if isinstance(item, dict) and "id" in item:
|
||||||
|
triplet.metadata["document_id"] = item["id"]
|
||||||
|
|
||||||
|
results.append(current_triplets)
|
||||||
|
total_triplets_count += len(current_triplets)
|
||||||
|
|
||||||
|
# Update progress
|
||||||
|
if (idx + 1) % update_interval == 0 or (idx + 1) == total_items:
|
||||||
|
remaining = total_items - (idx + 1)
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=idx + 1,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Processing... {idx + 1}/{total_items} (remaining: {remaining}) - Extracted {total_triplets_count} triplets"
|
||||||
|
)
|
||||||
|
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id,
|
||||||
|
status="completed",
|
||||||
|
message=f"Batch extraction completed. Processed {len(results)} documents, extracted {total_triplets_count} triplets.",
|
||||||
|
)
|
||||||
|
return results
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id, status="failed", message=str(e)
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
else:
|
||||||
|
# Single item
|
||||||
|
return self.extract_triplets(text, entities=entities, relations=relations, **kwargs)
|
||||||
|
|
||||||
def extract_triplets(
|
def extract_triplets(
|
||||||
self,
|
self,
|
||||||
text: str,
|
text: str,
|
||||||
@@ -201,6 +319,8 @@ class TripletExtractor:
|
|||||||
if isinstance(methods, str):
|
if isinstance(methods, str):
|
||||||
methods = [methods]
|
methods = [methods]
|
||||||
|
|
||||||
|
triplet_types = options.get("triplet_types", self.triplet_types)
|
||||||
|
|
||||||
# Merge config with options
|
# Merge config with options
|
||||||
all_options = {**self.config, **options}
|
all_options = {**self.config, **options}
|
||||||
|
|
||||||
@@ -234,6 +354,11 @@ class TripletExtractor:
|
|||||||
|
|
||||||
# Prepare method-specific options
|
# Prepare method-specific options
|
||||||
method_options = all_options.copy()
|
method_options = all_options.copy()
|
||||||
|
|
||||||
|
# Pass triplet_types to all methods
|
||||||
|
if triplet_types:
|
||||||
|
method_options["triplet_types"] = triplet_types
|
||||||
|
|
||||||
if method_name == "huggingface":
|
if method_name == "huggingface":
|
||||||
method_options["model"] = all_options.get(
|
method_options["model"] = all_options.get(
|
||||||
"huggingface_model", all_options.get("model")
|
"huggingface_model", all_options.get("model")
|
||||||
@@ -265,6 +390,20 @@ class TripletExtractor:
|
|||||||
**method_options,
|
**method_options,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Apply weighted scoring if triplet_types are provided
|
||||||
|
if triplet_types:
|
||||||
|
try:
|
||||||
|
from .methods import calculate_weighted_confidence
|
||||||
|
for t in triplets:
|
||||||
|
t.confidence = calculate_weighted_confidence(
|
||||||
|
item_type=t.predicate,
|
||||||
|
original_confidence=t.confidence,
|
||||||
|
valid_types=triplet_types,
|
||||||
|
item_text=t.predicate # For triplets, predicate is the key text
|
||||||
|
)
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
|
||||||
# Filter by confidence
|
# Filter by confidence
|
||||||
min_conf = options.get("min_confidence", self.min_confidence)
|
min_conf = options.get("min_confidence", self.min_confidence)
|
||||||
filtered = [t for t in triplets if t.confidence >= min_conf]
|
filtered = [t for t in triplets if t.confidence >= min_conf]
|
||||||
@@ -293,20 +432,29 @@ class TripletExtractor:
|
|||||||
triplets = all_triplets[0][1]
|
triplets = all_triplets[0][1]
|
||||||
else:
|
else:
|
||||||
# Fallback: Convert relations to triplets
|
# Fallback: Convert relations to triplets
|
||||||
self.progress_tracker.update_tracking(
|
if relations:
|
||||||
tracking_id,
|
self.progress_tracker.update_tracking(
|
||||||
message=f"Converting {len(relations)} relations to triplets...",
|
tracking_id,
|
||||||
)
|
message=f"Converting {len(relations)} relations to triplets...",
|
||||||
triplets = []
|
|
||||||
for relation in relations:
|
|
||||||
triplet = Triplet(
|
|
||||||
subject=self._format_uri(relation.subject.text),
|
|
||||||
predicate=self._format_uri(relation.predicate),
|
|
||||||
object=self._format_uri(relation.object.text),
|
|
||||||
confidence=relation.confidence,
|
|
||||||
metadata={"context": relation.context, **relation.metadata},
|
|
||||||
)
|
)
|
||||||
triplets.append(triplet)
|
triplets = []
|
||||||
|
for relation in relations:
|
||||||
|
triplet = Triplet(
|
||||||
|
subject=self._format_uri(relation.subject.text),
|
||||||
|
predicate=self._format_uri(relation.predicate),
|
||||||
|
object=self._format_uri(relation.object.text),
|
||||||
|
confidence=relation.confidence,
|
||||||
|
metadata={"context": relation.context, **relation.metadata},
|
||||||
|
)
|
||||||
|
triplets.append(triplet)
|
||||||
|
else:
|
||||||
|
# Last resort: Try rule-based extraction if no relations exist
|
||||||
|
self.progress_tracker.update_tracking(
|
||||||
|
tracking_id,
|
||||||
|
message="No relations found. Trying rule-based triplet extraction...",
|
||||||
|
)
|
||||||
|
method_func = get_triplet_method("rules")
|
||||||
|
triplets = method_func(text, entities=entities, relations=[], **all_options)
|
||||||
|
|
||||||
# Validate triplets
|
# Validate triplets
|
||||||
if options.get("validate", self._should_validate):
|
if options.get("validate", self._should_validate):
|
||||||
@@ -382,7 +530,65 @@ class TripletExtractor:
|
|||||||
Returns:
|
Returns:
|
||||||
list: List of triplet lists for each text
|
list: List of triplet lists for each text
|
||||||
"""
|
"""
|
||||||
return [self.extract_triplets(text, **options) for text in texts]
|
tracking_id = self.progress_tracker.start_tracking(
|
||||||
|
module="semantic_extract",
|
||||||
|
submodule="TripletExtractor",
|
||||||
|
message=f"Batch extracting triplets from {len(texts)} documents",
|
||||||
|
)
|
||||||
|
|
||||||
|
results = []
|
||||||
|
total_triplets_count = 0
|
||||||
|
total_items = len(texts)
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Determine update interval
|
||||||
|
if total_items <= 10:
|
||||||
|
update_interval = 1
|
||||||
|
else:
|
||||||
|
update_interval = max(1, min(10, total_items // 100))
|
||||||
|
|
||||||
|
for idx, text in enumerate(texts, 1):
|
||||||
|
# Extract triplets
|
||||||
|
triplets = self.extract_triplets(text, **options)
|
||||||
|
|
||||||
|
# Add provenance metadata
|
||||||
|
for triplet in triplets:
|
||||||
|
if triplet.metadata is None:
|
||||||
|
triplet.metadata = {}
|
||||||
|
triplet.metadata["batch_index"] = idx - 1
|
||||||
|
|
||||||
|
results.append(triplets)
|
||||||
|
total_triplets_count += len(triplets)
|
||||||
|
|
||||||
|
# Update progress
|
||||||
|
should_update = (
|
||||||
|
idx % update_interval == 0 or
|
||||||
|
idx == total_items or
|
||||||
|
idx == 1 or
|
||||||
|
total_items <= 10
|
||||||
|
)
|
||||||
|
|
||||||
|
if should_update:
|
||||||
|
remaining = total_items - idx
|
||||||
|
self.progress_tracker.update_progress(
|
||||||
|
tracking_id,
|
||||||
|
processed=idx,
|
||||||
|
total=total_items,
|
||||||
|
message=f"Processing documents... {idx}/{total_items} (remaining: {remaining}) - Extracted {total_triplets_count} triplets so far"
|
||||||
|
)
|
||||||
|
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id,
|
||||||
|
status="completed",
|
||||||
|
message=f"Batch extraction completed. Processed {len(results)} documents, extracted {total_triplets_count} triplets.",
|
||||||
|
)
|
||||||
|
return results
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.progress_tracker.stop_tracking(
|
||||||
|
tracking_id, status="failed", message=str(e)
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
class TripletValidator:
|
class TripletValidator:
|
||||||
@@ -428,27 +634,6 @@ class TripletValidator:
|
|||||||
"""
|
"""
|
||||||
return [t for t in triplets if self.validate_triplet(t, **criteria)]
|
return [t for t in triplets if self.validate_triplet(t, **criteria)]
|
||||||
|
|
||||||
def check_triplet_consistency(self, triplets: List[Triplet]) -> Dict[str, Any]:
|
|
||||||
"""
|
|
||||||
Check consistency among triplets.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
triplets: List of triplets
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
dict: Consistency report
|
|
||||||
"""
|
|
||||||
issues = []
|
|
||||||
|
|
||||||
# Check for contradictory triplets
|
|
||||||
# (simplified - would need domain knowledge for full implementation)
|
|
||||||
|
|
||||||
return {
|
|
||||||
"total_triplets": len(triplets),
|
|
||||||
"issues": issues,
|
|
||||||
"consistent": len(issues) == 0,
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
class RDFSerializer:
|
class RDFSerializer:
|
||||||
"""RDF serialization handler."""
|
"""RDF serialization handler."""
|
||||||
|
|||||||
@@ -0,0 +1,151 @@
|
|||||||
|
|
||||||
|
import pytest
|
||||||
|
import sys
|
||||||
|
|
||||||
|
from semantica.semantic_extract.ner_extractor import NERExtractor, Entity
|
||||||
|
from semantica.semantic_extract.relation_extractor import RelationExtractor, Relation
|
||||||
|
from semantica.semantic_extract.triplet_extractor import TripletExtractor
|
||||||
|
|
||||||
|
class TestRobustnessFallback:
|
||||||
|
|
||||||
|
def test_ner_last_resort_fallback(self):
|
||||||
|
"""Test that NER extractor finds entities even in obscure text via last resort."""
|
||||||
|
extractor = NERExtractor()
|
||||||
|
|
||||||
|
# Text with single capitalized word that shouldn't match PERSON pattern (requires 2+ words)
|
||||||
|
text = "Zylophone"
|
||||||
|
|
||||||
|
entities = extractor.extract_entities(text)
|
||||||
|
|
||||||
|
assert len(entities) > 0, "Should have extracted at least one entity via last resort"
|
||||||
|
# Check if they are the capitalized words
|
||||||
|
texts = [e.text for e in entities]
|
||||||
|
assert "Zylophone" in texts
|
||||||
|
|
||||||
|
# Verify metadata
|
||||||
|
for e in entities:
|
||||||
|
assert e.metadata is not None
|
||||||
|
assert "extraction_method" in e.metadata
|
||||||
|
# Should be last_resort_pattern
|
||||||
|
assert e.metadata["extraction_method"] == "last_resort_pattern"
|
||||||
|
|
||||||
|
def test_relation_last_resort_fallback(self):
|
||||||
|
"""Test that Relation extractor creates adjacency relations when no patterns match."""
|
||||||
|
extractor = RelationExtractor()
|
||||||
|
|
||||||
|
# Create entities far apart to avoid "co_occurrence" fallback which triggers < 100 chars
|
||||||
|
padding = " " * 105
|
||||||
|
text = f"Alpha{padding}Beta{padding}Gamma"
|
||||||
|
|
||||||
|
# Alpha at start
|
||||||
|
e1_start = 0
|
||||||
|
e1_end = 5
|
||||||
|
|
||||||
|
# Beta after padding
|
||||||
|
e2_start = e1_end + 105
|
||||||
|
e2_end = e2_start + 4
|
||||||
|
|
||||||
|
# Gamma after padding
|
||||||
|
e3_start = e2_end + 105
|
||||||
|
e3_end = e3_start + 5
|
||||||
|
|
||||||
|
e1 = Entity(text="Alpha", label="UNKNOWN", start_char=e1_start, end_char=e1_end)
|
||||||
|
e2 = Entity(text="Beta", label="UNKNOWN", start_char=e2_start, end_char=e2_end)
|
||||||
|
e3 = Entity(text="Gamma", label="UNKNOWN", start_char=e3_start, end_char=e3_end)
|
||||||
|
|
||||||
|
entities = [e1, e2, e3]
|
||||||
|
|
||||||
|
# This text has no "is a", "works for", etc. patterns.
|
||||||
|
# And entities are too far for co-occurrence (< 100).
|
||||||
|
# It should trigger the last resort adjacency fallback.
|
||||||
|
relations = extractor.extract_relations(text, entities)
|
||||||
|
|
||||||
|
assert len(relations) > 0, "Should have extracted relations via last resort"
|
||||||
|
|
||||||
|
# Expect relations between adjacent entities: Alpha->Beta, Beta->Gamma
|
||||||
|
pairs = [(r.subject.text, r.object.text) for r in relations]
|
||||||
|
assert ("Alpha", "Beta") in pairs
|
||||||
|
assert ("Beta", "Gamma") in pairs
|
||||||
|
|
||||||
|
# Verify metadata
|
||||||
|
for r in relations:
|
||||||
|
assert r.metadata is not None
|
||||||
|
assert "extraction_method" in r.metadata
|
||||||
|
assert r.metadata.get("extraction_method") == "last_resort_adjacency"
|
||||||
|
|
||||||
|
def test_triplet_fallback_conversion(self):
|
||||||
|
"""Test that Triplet extractor falls back to converting relations if extraction fails."""
|
||||||
|
# Setup mocks or use real classes
|
||||||
|
ner = NERExtractor() # We'll just pass entities directly
|
||||||
|
rel_extractor = RelationExtractor()
|
||||||
|
triplet_extractor = TripletExtractor()
|
||||||
|
|
||||||
|
text = "Alpha is connected to Beta."
|
||||||
|
e1 = Entity(text="Alpha", label="Thing", start_char=0, end_char=5)
|
||||||
|
e2 = Entity(text="Beta", label="Thing", start_char=22, end_char=26)
|
||||||
|
entities = [e1, e2]
|
||||||
|
|
||||||
|
# Create a relation manually to ensure we have one to convert
|
||||||
|
relation = Relation(
|
||||||
|
subject=e1,
|
||||||
|
predicate="connected_to",
|
||||||
|
object=e2,
|
||||||
|
confidence=0.9,
|
||||||
|
context=text
|
||||||
|
)
|
||||||
|
|
||||||
|
# We want to test the fallback in extract_triplets.
|
||||||
|
# Since we can't easily force the primary triplet method to return empty without mocking,
|
||||||
|
# we can pass the relations explicitly and rely on the fact that standard triplet extraction
|
||||||
|
# might not support "connected_to" if it relies on strict patterns, or we can use a method that fails.
|
||||||
|
|
||||||
|
# However, the triplet extractor calls relation extractor internally if not provided.
|
||||||
|
# Let's test the flow where we provide relations.
|
||||||
|
|
||||||
|
triplets = triplet_extractor.extract_triplets(text, entities=entities, relations=[relation])
|
||||||
|
|
||||||
|
assert len(triplets) > 0
|
||||||
|
assert triplets[0].subject == "Alpha"
|
||||||
|
assert triplets[0].object == "Beta"
|
||||||
|
assert triplets[0].predicate == "connected_to"
|
||||||
|
|
||||||
|
def test_batch_metadata_propagation(self):
|
||||||
|
"""Verify batch_index and document_id are propagated in batch mode with fallbacks."""
|
||||||
|
ner = NERExtractor()
|
||||||
|
|
||||||
|
docs = [
|
||||||
|
{"content": "First doc", "id": "doc_1"},
|
||||||
|
{"content": "Second doc", "id": "doc_2"}
|
||||||
|
]
|
||||||
|
|
||||||
|
# These docs are simple, might trigger fallback or simple patterns
|
||||||
|
results = ner.extract(docs)
|
||||||
|
|
||||||
|
assert len(results) == 2
|
||||||
|
|
||||||
|
# Check first doc results
|
||||||
|
for e in results[0]:
|
||||||
|
assert e.metadata["batch_index"] == 0
|
||||||
|
assert e.metadata["document_id"] == "doc_1"
|
||||||
|
|
||||||
|
# Check second doc results
|
||||||
|
for e in results[1]:
|
||||||
|
assert e.metadata["batch_index"] == 1
|
||||||
|
assert e.metadata["document_id"] == "doc_2"
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
# Manually run if executed as script
|
||||||
|
t = TestRobustnessFallback()
|
||||||
|
try:
|
||||||
|
t.test_ner_last_resort_fallback()
|
||||||
|
print("NER Fallback Test Passed")
|
||||||
|
t.test_relation_last_resort_fallback()
|
||||||
|
print("Relation Fallback Test Passed")
|
||||||
|
t.test_triplet_fallback_conversion()
|
||||||
|
print("Triplet Fallback Test Passed")
|
||||||
|
t.test_batch_metadata_propagation()
|
||||||
|
print("Batch Metadata Test Passed")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Test Failed: {e}")
|
||||||
|
import traceback
|
||||||
|
traceback.print_exc()
|
||||||
@@ -0,0 +1,177 @@
|
|||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
from typing import Type, List, Optional
|
||||||
|
from pydantic import BaseModel, ValidationError
|
||||||
|
|
||||||
|
from semantica.semantic_extract.providers import BaseProvider
|
||||||
|
from semantica.semantic_extract.methods import (
|
||||||
|
extract_entities_llm,
|
||||||
|
extract_relations_llm,
|
||||||
|
extract_triplets_llm
|
||||||
|
)
|
||||||
|
from semantica.semantic_extract.schemas import EntitiesResponse, RelationsResponse, TripletsResponse
|
||||||
|
from semantica.semantic_extract.ner_extractor import Entity
|
||||||
|
from semantica.semantic_extract.relation_extractor import Relation
|
||||||
|
|
||||||
|
# Mock Pydantic models for responses
|
||||||
|
class MockEntity(BaseModel):
|
||||||
|
text: str
|
||||||
|
label: str
|
||||||
|
start: int = 0
|
||||||
|
end: int = 0
|
||||||
|
confidence: float = 1.0
|
||||||
|
|
||||||
|
class MockEntitiesResponse(BaseModel):
|
||||||
|
entities: List[MockEntity]
|
||||||
|
|
||||||
|
class MockRelation(BaseModel):
|
||||||
|
subject: str
|
||||||
|
predicate: str
|
||||||
|
object: str
|
||||||
|
confidence: float = 1.0
|
||||||
|
|
||||||
|
class MockRelationsResponse(BaseModel):
|
||||||
|
relations: List[MockRelation]
|
||||||
|
|
||||||
|
class MockTriplet(BaseModel):
|
||||||
|
subject: str
|
||||||
|
predicate: str
|
||||||
|
object: str
|
||||||
|
confidence: float = 1.0
|
||||||
|
|
||||||
|
class MockTripletsResponse(BaseModel):
|
||||||
|
triplets: List[MockTriplet]
|
||||||
|
|
||||||
|
class MockProvider(BaseProvider):
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
self.generate_typed_called = False
|
||||||
|
self.generate_structured_called = False
|
||||||
|
self.model = "mock-model"
|
||||||
|
self.is_available_val = True
|
||||||
|
|
||||||
|
def is_available(self) -> bool:
|
||||||
|
return self.is_available_val
|
||||||
|
|
||||||
|
def generate(self, prompt: str, **kwargs) -> str:
|
||||||
|
return "{}"
|
||||||
|
|
||||||
|
def generate_structured(self, prompt: str, **kwargs) -> dict:
|
||||||
|
self.generate_structured_called = True
|
||||||
|
if "entities" in prompt.lower():
|
||||||
|
return [{"text": "Apple", "label": "ORG", "start": 0, "end": 5}]
|
||||||
|
elif "relations" in prompt.lower():
|
||||||
|
return [{"subject": "Steve Jobs", "predicate": "founded", "object": "Apple"}]
|
||||||
|
elif "triplets" in prompt.lower():
|
||||||
|
return [{"subject": "Steve Jobs", "predicate": "founded", "object": "Apple"}]
|
||||||
|
return {}
|
||||||
|
|
||||||
|
def generate_typed(
|
||||||
|
self,
|
||||||
|
prompt: str,
|
||||||
|
schema: Type[BaseModel],
|
||||||
|
max_retries: int = 3,
|
||||||
|
**kwargs
|
||||||
|
) -> BaseModel:
|
||||||
|
self.generate_typed_called = True
|
||||||
|
|
||||||
|
if schema.__name__ == "EntitiesResponse":
|
||||||
|
return EntitiesResponse(entities=[
|
||||||
|
{"text": "Apple", "label": "ORG", "start_char": 0, "end_char": 5, "confidence": 0.99}
|
||||||
|
])
|
||||||
|
elif schema.__name__ == "RelationsResponse":
|
||||||
|
return RelationsResponse(relations=[
|
||||||
|
{"subject": "Steve Jobs", "predicate": "founded", "object": "Apple", "confidence": 0.95}
|
||||||
|
])
|
||||||
|
elif schema.__name__ == "TripletsResponse":
|
||||||
|
return TripletsResponse(triplets=[
|
||||||
|
{"subject": "Steve Jobs", "predicate": "founded", "object": "Apple", "confidence": 0.95}
|
||||||
|
])
|
||||||
|
return schema()
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_provider():
|
||||||
|
return MockProvider()
|
||||||
|
|
||||||
|
@patch("semantica.semantic_extract.methods.create_provider")
|
||||||
|
def test_extract_entities_typed(mock_create_provider, mock_provider):
|
||||||
|
mock_create_provider.return_value = mock_provider
|
||||||
|
|
||||||
|
text = "Apple was founded by Steve Jobs."
|
||||||
|
entities = extract_entities_llm(
|
||||||
|
text,
|
||||||
|
provider="mock",
|
||||||
|
structured_output_mode="typed"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert mock_provider.generate_typed_called
|
||||||
|
assert len(entities) == 1
|
||||||
|
assert entities[0].text == "Apple"
|
||||||
|
assert entities[0].label == "ORG"
|
||||||
|
assert entities[0].metadata["extraction_method"] == "llm_typed"
|
||||||
|
|
||||||
|
@patch("semantica.semantic_extract.methods.create_provider")
|
||||||
|
def test_extract_entities_legacy(mock_create_provider, mock_provider):
|
||||||
|
mock_create_provider.return_value = mock_provider
|
||||||
|
|
||||||
|
text = "Apple was founded by Steve Jobs."
|
||||||
|
entities = extract_entities_llm(
|
||||||
|
text,
|
||||||
|
provider="mock",
|
||||||
|
structured_output_mode="legacy"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Legacy mode now redirects to typed mode
|
||||||
|
assert mock_provider.generate_typed_called
|
||||||
|
assert not mock_provider.generate_structured_called
|
||||||
|
assert len(entities) == 1
|
||||||
|
assert entities[0].text == "Apple"
|
||||||
|
assert entities[0].label == "ORG"
|
||||||
|
assert entities[0].metadata["extraction_method"] == "llm_typed"
|
||||||
|
|
||||||
|
@patch("semantica.semantic_extract.methods.create_provider")
|
||||||
|
def test_extract_relations_typed(mock_create_provider, mock_provider):
|
||||||
|
mock_create_provider.return_value = mock_provider
|
||||||
|
|
||||||
|
text = "Steve Jobs founded Apple."
|
||||||
|
entities = [
|
||||||
|
Entity(text="Steve Jobs", label="PERSON", start_char=0, end_char=10),
|
||||||
|
Entity(text="Apple", label="ORG", start_char=19, end_char=24)
|
||||||
|
]
|
||||||
|
|
||||||
|
relations = extract_relations_llm(
|
||||||
|
text,
|
||||||
|
entities=entities,
|
||||||
|
provider="mock",
|
||||||
|
structured_output_mode="typed"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert mock_provider.generate_typed_called
|
||||||
|
assert len(relations) == 1
|
||||||
|
assert relations[0].subject.text == "Steve Jobs"
|
||||||
|
assert relations[0].object.text == "Apple"
|
||||||
|
assert relations[0].predicate == "founded"
|
||||||
|
assert relations[0].metadata["extraction_method"] == "llm_typed"
|
||||||
|
|
||||||
|
@patch("semantica.semantic_extract.methods.create_provider")
|
||||||
|
def test_extract_triplets_typed(mock_create_provider, mock_provider):
|
||||||
|
mock_create_provider.return_value = mock_provider
|
||||||
|
|
||||||
|
text = "Steve Jobs founded Apple."
|
||||||
|
|
||||||
|
triplets = extract_triplets_llm(
|
||||||
|
text,
|
||||||
|
provider="mock",
|
||||||
|
structured_output_mode="typed"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert mock_provider.generate_typed_called
|
||||||
|
assert len(triplets) == 1
|
||||||
|
assert triplets[0].subject == "Steve Jobs"
|
||||||
|
assert triplets[0].object == "Apple"
|
||||||
|
assert triplets[0].predicate == "founded"
|
||||||
|
assert triplets[0].metadata["extraction_method"] == "llm_typed"
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
pytest.main([__file__])
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -34,7 +34,7 @@ class TestLLMExtractionFixes(unittest.TestCase):
|
|||||||
"""Test that methods raise ProcessingError by default on failure."""
|
"""Test that methods raise ProcessingError by default on failure."""
|
||||||
mock_llm = MagicMock()
|
mock_llm = MagicMock()
|
||||||
mock_llm.is_available.return_value = True
|
mock_llm.is_available.return_value = True
|
||||||
mock_llm.generate_structured.side_effect = ProcessingError("LLM Error")
|
mock_llm.generate_typed.side_effect = ProcessingError("LLM Error")
|
||||||
mock_create.return_value = mock_llm
|
mock_create.return_value = mock_llm
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -48,7 +48,7 @@ class TestLLMExtractionFixes(unittest.TestCase):
|
|||||||
"""Test that silent_fail=True returns empty list instead of raising."""
|
"""Test that silent_fail=True returns empty list instead of raising."""
|
||||||
mock_llm = MagicMock()
|
mock_llm = MagicMock()
|
||||||
mock_llm.is_available.return_value = True
|
mock_llm.is_available.return_value = True
|
||||||
mock_llm.generate_structured.side_effect = Exception("LLM Error")
|
mock_llm.generate_typed.side_effect = Exception("LLM Error")
|
||||||
mock_create.return_value = mock_llm
|
mock_create.return_value = mock_llm
|
||||||
|
|
||||||
entities = extract_entities_llm("test text", provider="openai", silent_fail=True)
|
entities = extract_entities_llm("test text", provider="openai", silent_fail=True)
|
||||||
@@ -83,7 +83,7 @@ class TestLLMExtractionFixes(unittest.TestCase):
|
|||||||
"""Test that long text triggers chunking."""
|
"""Test that long text triggers chunking."""
|
||||||
mock_llm = MagicMock()
|
mock_llm = MagicMock()
|
||||||
mock_llm.is_available.return_value = True
|
mock_llm.is_available.return_value = True
|
||||||
mock_llm.generate_structured.return_value = []
|
mock_llm.generate_typed.return_value = MagicMock(entities=[]) # Mock response
|
||||||
mock_create.return_value = mock_llm
|
mock_create.return_value = mock_llm
|
||||||
|
|
||||||
long_text = "This is a long text that should be chunked into multiple pieces."
|
long_text = "This is a long text that should be chunked into multiple pieces."
|
||||||
|
|||||||
@@ -32,7 +32,8 @@ class TestModelSelection(unittest.TestCase):
|
|||||||
|
|
||||||
def test_generator_switching(self):
|
def test_generator_switching(self):
|
||||||
print("\nTesting EmbeddingGenerator Switching...")
|
print("\nTesting EmbeddingGenerator Switching...")
|
||||||
generator = EmbeddingGenerator()
|
# Initialize with explicit method to ensure consistent starting state for test
|
||||||
|
generator = EmbeddingGenerator(text={"method": "sentence_transformers"})
|
||||||
|
|
||||||
# Default check
|
# Default check
|
||||||
self.assertEqual(generator.get_text_method(), "sentence_transformers")
|
self.assertEqual(generator.get_text_method(), "sentence_transformers")
|
||||||
|
|||||||
@@ -0,0 +1,122 @@
|
|||||||
|
|
||||||
|
print("Starting tests module...")
|
||||||
|
import unittest
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
try:
|
||||||
|
from semantica.semantic_extract.semantic_network_extractor import SemanticNetworkExtractor, SemanticNetwork, SemanticNode, SemanticEdge
|
||||||
|
from semantica.semantic_extract.event_detector import EventDetector, Event
|
||||||
|
from semantica.semantic_extract.semantic_analyzer import SemanticAnalyzer
|
||||||
|
from semantica.semantic_extract.coreference_resolver import CoreferenceResolver, CoreferenceChain, Mention
|
||||||
|
from semantica.semantic_extract.ner_extractor import Entity
|
||||||
|
print("Imports successful")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Import failed: {e}")
|
||||||
|
|
||||||
|
class TestSemanticExtractBatch(unittest.TestCase):
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
# Mock progress tracker to avoid console spam
|
||||||
|
self.tracker_patcher = patch('semantica.utils.progress_tracker.get_progress_tracker')
|
||||||
|
self.mock_tracker_cls = self.tracker_patcher.start()
|
||||||
|
self.mock_tracker = self.mock_tracker_cls.return_value
|
||||||
|
self.mock_tracker.enabled = True
|
||||||
|
self.mock_tracker.start_tracking.return_value = "tracking_id"
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
self.tracker_patcher.stop()
|
||||||
|
|
||||||
|
def test_semantic_network_batch(self):
|
||||||
|
print("Running test_semantic_network_batch")
|
||||||
|
from copy import deepcopy
|
||||||
|
extractor = SemanticNetworkExtractor()
|
||||||
|
|
||||||
|
# Mock extract_network
|
||||||
|
mock_network = SemanticNetwork(
|
||||||
|
nodes=[SemanticNode(id="1", label="test", type="test", metadata={})],
|
||||||
|
edges=[SemanticEdge(source="1", target="1", label="self", metadata={})],
|
||||||
|
metadata={}
|
||||||
|
)
|
||||||
|
# Use side_effect to return a fresh copy each time
|
||||||
|
extractor.extract_network = MagicMock(side_effect=lambda *args, **kwargs: deepcopy(mock_network))
|
||||||
|
|
||||||
|
# Test input
|
||||||
|
docs = [{"content": "doc1", "id": "doc_1"}, {"content": "doc2", "id": "doc_2"}]
|
||||||
|
|
||||||
|
# Run batch
|
||||||
|
results = extractor.extract(docs)
|
||||||
|
|
||||||
|
self.assertEqual(len(results), 2)
|
||||||
|
# Check provenance
|
||||||
|
self.assertEqual(results[0].metadata["batch_index"], 0)
|
||||||
|
self.assertEqual(results[0].metadata["document_id"], "doc_1")
|
||||||
|
self.assertEqual(results[0].nodes[0].metadata["batch_index"], 0)
|
||||||
|
self.assertEqual(results[0].nodes[0].metadata["document_id"], "doc_1")
|
||||||
|
|
||||||
|
self.assertEqual(results[1].metadata["batch_index"], 1)
|
||||||
|
self.assertEqual(results[1].metadata["document_id"], "doc_2")
|
||||||
|
|
||||||
|
def test_event_detector_batch(self):
|
||||||
|
print("Running test_event_detector_batch")
|
||||||
|
detector = EventDetector()
|
||||||
|
|
||||||
|
# Mock detect_events
|
||||||
|
mock_event = Event(
|
||||||
|
text="event", event_type="test", start_char=0, end_char=5
|
||||||
|
)
|
||||||
|
detector.detect_events = MagicMock(return_value=[mock_event])
|
||||||
|
|
||||||
|
docs = [{"content": "doc1", "id": "doc_1"}]
|
||||||
|
|
||||||
|
results = detector.extract(docs)
|
||||||
|
|
||||||
|
self.assertEqual(len(results), 1)
|
||||||
|
self.assertEqual(len(results[0]), 1)
|
||||||
|
self.assertEqual(results[0][0].metadata["batch_index"], 0)
|
||||||
|
self.assertEqual(results[0][0].metadata["document_id"], "doc_1")
|
||||||
|
|
||||||
|
def test_semantic_analyzer_batch(self):
|
||||||
|
print("Running test_semantic_analyzer_batch")
|
||||||
|
analyzer = SemanticAnalyzer()
|
||||||
|
|
||||||
|
# Mock analyze_semantics
|
||||||
|
mock_result = {
|
||||||
|
"text": "test",
|
||||||
|
"semantic_roles": [{"word": "test", "role": "agent"}]
|
||||||
|
}
|
||||||
|
analyzer.analyze_semantics = MagicMock(return_value=mock_result)
|
||||||
|
|
||||||
|
docs = [{"content": "doc1", "id": "doc_1"}]
|
||||||
|
|
||||||
|
results = analyzer.analyze(docs)
|
||||||
|
|
||||||
|
self.assertEqual(len(results), 1)
|
||||||
|
self.assertEqual(results[0]["batch_index"], 0)
|
||||||
|
self.assertEqual(results[0]["document_id"], "doc_1")
|
||||||
|
self.assertEqual(results[0]["semantic_roles"][0]["metadata"]["batch_index"], 0)
|
||||||
|
self.assertEqual(results[0]["semantic_roles"][0]["metadata"]["document_id"], "doc_1")
|
||||||
|
|
||||||
|
def test_coreference_resolver_batch(self):
|
||||||
|
print("Running test_coreference_resolver_batch")
|
||||||
|
resolver = CoreferenceResolver()
|
||||||
|
|
||||||
|
# Mock resolve_coreferences
|
||||||
|
mock_mention = Mention(text="he", start_char=0, end_char=2, mention_type="pronoun")
|
||||||
|
mock_chain = CoreferenceChain(
|
||||||
|
mentions=[mock_mention],
|
||||||
|
representative=mock_mention
|
||||||
|
)
|
||||||
|
resolver.resolve_coreferences = MagicMock(return_value=[mock_chain])
|
||||||
|
|
||||||
|
docs = [{"content": "doc1", "id": "doc_1"}]
|
||||||
|
|
||||||
|
results = resolver.resolve(docs)
|
||||||
|
|
||||||
|
self.assertEqual(len(results), 1)
|
||||||
|
self.assertEqual(results[0][0].mentions[0].metadata["batch_index"], 0)
|
||||||
|
self.assertEqual(results[0][0].mentions[0].metadata["document_id"], "doc_1")
|
||||||
|
self.assertEqual(results[0][0].representative.metadata["batch_index"], 0)
|
||||||
|
self.assertEqual(results[0][0].representative.metadata["document_id"], "doc_1")
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
print("Running main...")
|
||||||
|
unittest.main()
|
||||||
@@ -1,7 +1,7 @@
|
|||||||
import unittest
|
import unittest
|
||||||
from unittest.mock import MagicMock, patch
|
from unittest.mock import MagicMock, patch
|
||||||
from semantica.triplet_store.triplet_manager import TripletManager, TripletStore
|
from semantica.triplet_store.triplet_store import TripletStore
|
||||||
from semantica.triplet_store.query_engine import QueryEngine, QueryResult
|
from semantica.triplet_store.query_engine import QueryEngine
|
||||||
from semantica.semantic_extract.triplet_extractor import Triplet
|
from semantica.semantic_extract.triplet_extractor import Triplet
|
||||||
|
|
||||||
class TestTripletStore(unittest.TestCase):
|
class TestTripletStore(unittest.TestCase):
|
||||||
@@ -10,125 +10,91 @@ class TestTripletStore(unittest.TestCase):
|
|||||||
self.mock_logger = MagicMock()
|
self.mock_logger = MagicMock()
|
||||||
self.mock_tracker = MagicMock()
|
self.mock_tracker = MagicMock()
|
||||||
|
|
||||||
self.logger_patcher = patch('semantica.triplet_store.triplet_manager.get_logger', return_value=self.mock_logger)
|
self.logger_patcher = patch('semantica.triplet_store.triplet_store.get_logger', return_value=self.mock_logger)
|
||||||
self.tracker_patcher = patch('semantica.triplet_store.triplet_manager.get_progress_tracker', return_value=self.mock_tracker)
|
self.tracker_patcher = patch('semantica.triplet_store.triplet_store.get_progress_tracker', return_value=self.mock_tracker)
|
||||||
self.logger_patcher_qe = patch('semantica.triplet_store.query_engine.get_logger', return_value=self.mock_logger)
|
|
||||||
self.tracker_patcher_qe = patch('semantica.triplet_store.query_engine.get_progress_tracker', return_value=self.mock_tracker)
|
|
||||||
|
|
||||||
self.logger_patcher.start()
|
self.logger_patcher.start()
|
||||||
self.tracker_patcher.start()
|
self.tracker_patcher.start()
|
||||||
self.logger_patcher_qe.start()
|
|
||||||
self.tracker_patcher_qe.start()
|
|
||||||
|
|
||||||
def tearDown(self):
|
def tearDown(self):
|
||||||
self.logger_patcher.stop()
|
self.logger_patcher.stop()
|
||||||
self.tracker_patcher.stop()
|
self.tracker_patcher.stop()
|
||||||
self.logger_patcher_qe.stop()
|
|
||||||
self.tracker_patcher_qe.stop()
|
|
||||||
|
|
||||||
def test_triplet_manager_init(self):
|
@patch('semantica.triplet_store.blazegraph_store.BlazegraphStore')
|
||||||
manager = TripletManager(default_store="main")
|
def test_triplet_store_init(self, mock_blazegraph_store):
|
||||||
self.assertEqual(manager.default_store_id, "main")
|
store = TripletStore(backend="blazegraph", endpoint="http://localhost:9999")
|
||||||
self.assertEqual(manager.stores, {})
|
self.assertEqual(store.backend_type, "blazegraph")
|
||||||
|
|
||||||
def test_register_store(self):
|
|
||||||
manager = TripletManager()
|
|
||||||
store = manager.register_store("main", "blazegraph", "http://localhost:9999")
|
|
||||||
self.assertIsInstance(store, TripletStore)
|
|
||||||
self.assertEqual(store.store_id, "main")
|
|
||||||
self.assertEqual(store.store_type, "blazegraph")
|
|
||||||
self.assertEqual(store.endpoint, "http://localhost:9999")
|
self.assertEqual(store.endpoint, "http://localhost:9999")
|
||||||
self.assertIn("main", manager.stores)
|
mock_blazegraph_store.assert_called_once()
|
||||||
|
|
||||||
@patch('semantica.triplet_store.triplet_manager.TripletManager._get_store_backend')
|
@patch('semantica.triplet_store.blazegraph_store.BlazegraphStore')
|
||||||
def test_add_triplet(self, mock_get_store_backend):
|
def test_add_triplet(self, mock_blazegraph_store):
|
||||||
manager = TripletManager()
|
# Setup mock backend
|
||||||
manager.register_store("main", "blazegraph", "http://localhost:9999")
|
mock_backend_instance = MagicMock()
|
||||||
|
mock_blazegraph_store.return_value = mock_backend_instance
|
||||||
mock_store = MagicMock()
|
mock_backend_instance.add_triplet.return_value = {"status": "success"}
|
||||||
mock_get_store_backend.return_value = mock_store
|
|
||||||
mock_store.add_triplet.return_value = {"status": "success"}
|
|
||||||
|
|
||||||
|
store = TripletStore(backend="blazegraph")
|
||||||
triplet = Triplet(subject="s", predicate="p", object="o")
|
triplet = Triplet(subject="s", predicate="p", object="o")
|
||||||
result = manager.add_triplet(triplet, store_id="main")
|
|
||||||
|
|
||||||
self.assertTrue(result["success"])
|
result = store.add_triplet(triplet)
|
||||||
self.assertEqual(result["store_id"], "main")
|
|
||||||
mock_store.add_triplet.assert_called_once_with(triplet)
|
self.assertEqual(result, {"status": "success"})
|
||||||
|
mock_backend_instance.add_triplet.assert_called_once_with(triplet)
|
||||||
|
|
||||||
@patch('semantica.triplet_store.triplet_manager.TripletManager._get_store_backend')
|
@patch('semantica.triplet_store.blazegraph_store.BlazegraphStore')
|
||||||
def test_add_triplets(self, mock_get_store_backend):
|
def test_add_triplets(self, mock_blazegraph_store):
|
||||||
manager = TripletManager()
|
# Setup mock backend and bulk loader
|
||||||
manager.register_store("main", "blazegraph", "http://localhost:9999")
|
mock_backend_instance = MagicMock()
|
||||||
|
mock_blazegraph_store.return_value = mock_backend_instance
|
||||||
|
|
||||||
mock_store = MagicMock()
|
store = TripletStore(backend="blazegraph")
|
||||||
mock_get_store_backend.return_value = mock_store
|
|
||||||
mock_store.add_triplets.return_value = {"status": "success"}
|
# Mock bulk loader
|
||||||
|
mock_loader = MagicMock()
|
||||||
|
store.bulk_loader = mock_loader
|
||||||
|
mock_progress = MagicMock()
|
||||||
|
mock_progress.metadata = {"success": True}
|
||||||
|
mock_progress.total_triplets = 2
|
||||||
|
mock_progress.loaded_triplets = 2
|
||||||
|
mock_progress.failed_triplets = 0
|
||||||
|
mock_progress.total_batches = 1
|
||||||
|
mock_loader.load_triplets.return_value = mock_progress
|
||||||
|
|
||||||
triplets = [
|
triplets = [
|
||||||
Triplet(subject="s1", predicate="p1", object="o1"),
|
Triplet(subject="s1", predicate="p1", object="o1"),
|
||||||
Triplet(subject="s2", predicate="p2", object="o2")
|
Triplet(subject="s2", predicate="p2", object="o2")
|
||||||
]
|
]
|
||||||
|
|
||||||
result = manager.add_triplets(triplets, store_id="main", batch_size=2)
|
result = store.add_triplets(triplets, batch_size=2)
|
||||||
|
|
||||||
self.assertTrue(result["success"])
|
self.assertTrue(result["success"])
|
||||||
self.assertEqual(result["store_id"], "main")
|
self.assertEqual(result["total"], 2)
|
||||||
mock_store.add_triplets.assert_called()
|
mock_loader.load_triplets.assert_called_once()
|
||||||
|
|
||||||
@patch('semantica.triplet_store.triplet_manager.TripletManager._get_store_backend')
|
@patch('semantica.triplet_store.blazegraph_store.BlazegraphStore')
|
||||||
def test_get_triplets(self, mock_get_store_backend):
|
def test_get_triplets(self, mock_blazegraph_store):
|
||||||
manager = TripletManager()
|
mock_backend_instance = MagicMock()
|
||||||
manager.register_store("main", "blazegraph", "http://localhost:9999")
|
mock_blazegraph_store.return_value = mock_backend_instance
|
||||||
|
|
||||||
mock_store = MagicMock()
|
|
||||||
mock_get_store_backend.return_value = mock_store
|
|
||||||
expected_triplets = [Triplet(subject="s", predicate="p", object="o")]
|
expected_triplets = [Triplet(subject="s", predicate="p", object="o")]
|
||||||
mock_store.get_triplets.return_value = expected_triplets
|
mock_backend_instance.get_triplets.return_value = expected_triplets
|
||||||
|
|
||||||
result = manager.get_triplets(subject="s", store_id="main")
|
store = TripletStore(backend="blazegraph")
|
||||||
|
result = store.get_triplets(subject="s")
|
||||||
|
|
||||||
self.assertEqual(result, expected_triplets)
|
self.assertEqual(result, expected_triplets)
|
||||||
mock_store.get_triplets.assert_called_once_with("s", None, None)
|
mock_backend_instance.get_triplets.assert_called_once_with(subject="s", predicate=None, object=None)
|
||||||
|
|
||||||
@patch('semantica.triplet_store.triplet_manager.TripletManager._get_store_backend')
|
@patch('semantica.triplet_store.blazegraph_store.BlazegraphStore')
|
||||||
def test_delete_triplet(self, mock_get_store_backend):
|
def test_delete_triplet(self, mock_blazegraph_store):
|
||||||
manager = TripletManager()
|
mock_backend_instance = MagicMock()
|
||||||
manager.register_store("main", "blazegraph", "http://localhost:9999")
|
mock_blazegraph_store.return_value = mock_backend_instance
|
||||||
|
mock_backend_instance.delete_triplet.return_value = {"success": True}
|
||||||
mock_store = MagicMock()
|
|
||||||
mock_get_store_backend.return_value = mock_store
|
|
||||||
mock_store.delete_triplet.return_value = {"status": "deleted"}
|
|
||||||
|
|
||||||
|
store = TripletStore(backend="blazegraph")
|
||||||
triplet = Triplet(subject="s", predicate="p", object="o")
|
triplet = Triplet(subject="s", predicate="p", object="o")
|
||||||
result = manager.delete_triplet(triplet, store_id="main")
|
|
||||||
|
result = store.delete_triplet(triplet)
|
||||||
|
|
||||||
self.assertTrue(result["success"])
|
self.assertTrue(result["success"])
|
||||||
mock_store.delete_triplet.assert_called_once_with(triplet)
|
mock_backend_instance.delete_triplet.assert_called_once_with(triplet)
|
||||||
|
|
||||||
@patch('semantica.triplet_store.triplet_manager.TripletManager._get_store_backend')
|
|
||||||
def test_update_triplet(self, mock_get_store_backend):
|
|
||||||
manager = TripletManager()
|
|
||||||
manager.register_store("main", "blazegraph", "http://localhost:9999")
|
|
||||||
|
|
||||||
mock_store = MagicMock()
|
|
||||||
mock_get_store_backend.return_value = mock_store
|
|
||||||
mock_store.delete_triplet.return_value = {"status": "deleted"}
|
|
||||||
mock_store.add_triplet.return_value = {"status": "added"}
|
|
||||||
|
|
||||||
old_triplet = Triplet(subject="s", predicate="p", object="o_old")
|
|
||||||
new_triplet = Triplet(subject="s", predicate="p", object="o_new")
|
|
||||||
|
|
||||||
result = manager.update_triplet(old_triplet, new_triplet, store_id="main")
|
|
||||||
|
|
||||||
self.assertTrue(result["success"])
|
|
||||||
mock_store.delete_triplet.assert_called_once_with(old_triplet)
|
|
||||||
mock_store.add_triplet.assert_called_once_with(new_triplet)
|
|
||||||
|
|
||||||
def test_query_engine_init(self):
|
|
||||||
engine = QueryEngine(enable_caching=True)
|
|
||||||
self.assertTrue(engine.enable_caching)
|
|
||||||
self.assertEqual(engine.query_cache, {})
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
unittest.main()
|
|
||||||
|
|||||||
Reference in New Issue
Block a user